Compare commits

..

48 Commits

Author SHA1 Message Date
Alex Cheema 0dde45aa1f fix: keep TRUST_REMOTE_CODE=True for built-in models
The constant is the default for built-in models with known model cards,
which are trusted. Custom models added via API already default to
trust_remote_code=False in ModelCard.fetch_from_hf(). The CLI flag
overrides custom models only.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 15:44:06 +00:00
Alex Cheema 60cbe237ea feat: add --trust-remote-code CLI flag for custom model tokenizers
Some custom models (e.g. Kimi) require trust_remote_code=True to load
their tokenizers. This adds an opt-in CLI flag that sets an env var
read by runner subprocesses, following the same pattern as --fast-synch.
The flag is intentionally CLI-only (not API-accessible) to prevent
remote code execution attacks via the API.

Also changes the default TRUST_REMOTE_CODE constant from True to False,
making remote code execution fully opt-in.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 15:44:06 +00:00
Alex Cheema b65982ddd7 fix: improve text contrast on HOME and DOWNLOADS nav links (#1609)
## Motivation

Follow-up to #1601 (downloads page contrast fix). The HOME and DOWNLOADS
navigation links in the top-right header use `text-exo-light-gray`
(`oklch(0.6 0 0)`) which is too dim against the dark header background.

## Changes

Changed both nav links in `HeaderNav.svelte` from `text-exo-light-gray`
to `text-white/70` for better visibility. Hover state
(`text-exo-yellow`) is unchanged.

## Why It Works

`text-white/70` provides noticeably better contrast against
`bg-exo-dark-gray` while still looking subdued relative to the yellow
accent color on hover. This is consistent with the approach used in
#1601.

## Test Plan

### Manual Testing
- Verified both links are clearly readable on the home page and
downloads page
- Hover state still transitions to yellow as expected

### Automated Testing
- Dashboard builds successfully

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 13:48:40 +00:00
rltakashige 2fe689315b download .model files in exo bench (#1607)
## Motivation

failed again for kimi on a machine that had never downloaded it.

## Test Plan

### Manual Testing
it worked this time
2026-02-24 11:13:04 +00:00
Alex Cheema 644c5573ce fix: improve text contrast on downloads page (#1601)
## Summary
- Bumps opacity on dark-grey text/icons on the downloads page that were
nearly invisible against the dark background
- Informational text (GB sizes, speeds, disk free, model IDs) → full
`text-exo-light-gray`
- Interactive icons (delete, resume, retry) → `/70` at rest
- Hover-only elements (download button) → `/60`

## Test plan
- [ ] Open http://localhost:52415/downloads with models downloaded
- [ ] Verify GB downloaded amounts are clearly visible
- [ ] Verify delete trash icons are visible (not just on hover)
- [ ] Verify download speed text is readable
- [ ] Verify "paused" labels and resume buttons are visible
- [ ] Verify "X GB free" disk labels in column headers are readable

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 10:42:49 +00:00
rltakashige 12c3015f52 fix qwen moe tensor sharding (#1604)
## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
2026-02-23 21:23:27 +00:00
rltakashige 365dd68d9a Final fixes for release (#1603)
## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
2026-02-23 21:10:15 +00:00
Alex Cheema d3d129581e test: verify instance deletion cancels ongoing tasks (#1508)
## Summary
- The cancellation logic for issue #1215 already exists in
`get_transition_events()` (`src/exo/master/placement.py:208-227`) — when
an instance is deleted, `TaskStatusUpdated(Cancelled)` events are
emitted for all Pending/Running tasks on that instance
- Combined with PR #1276's token-boundary cancellation in runners, the
full pipeline works end-to-end
- However, the existing test
`test_get_transition_events_delete_instance` passed `{}` for tasks, so
this path was never exercised
- This PR adds 4 tests covering the cancellation behavior:
  - Running tasks are cancelled on instance deletion
  - Pending tasks are cancelled on instance deletion
  - Completed/Failed/TimedOut/Cancelled tasks are left alone
  - Only tasks matching the deleted instance are cancelled

Closes #1215

## Test plan
- [x] `uv run pytest src/exo/master/tests/test_placement.py -v` — all 15
tests pass
- [x] `uv run basedpyright` — 0 errors
- [x] `uv run ruff check` — all checks passed

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 20:12:23 +00:00
Alex Cheema c90a0cec78 fix: suppress closure errors in runnersupervisor and force spawn start method (#1547)
some errors could be thrown during shutdown - we can dismiss these safely

co-authored by me :)
2026-02-23 18:30:41 +00:00
Alex Cheema e8c1337168 fix: add download/resume buttons to pending downloads (#1581)
## Summary
- Adds a **resume button** (download icon) to paused pending downloads
(those with partial progress)
- Adds a **download button** to not-started pending downloads
- Both buttons call the existing `startDownload()` function which
handles both new downloads and resuming partial ones
- Previously, paused downloads only showed a "paused" label with no
action, and not-started downloads showed "..." with no way to trigger
them

## Test plan
- [ ] Build dashboard (`cd dashboard && npm run build`)
- [ ] Start exo, navigate to Downloads tab
- [ ] Verify paused downloads show a clickable resume (download arrow)
icon below the progress bar
- [ ] Verify not-started pending downloads show a clickable download
icon
- [ ] Click both button types and confirm downloads start/resume

> Note: Screenshot could not be captured because the dashboard requires
the exo backend API to render, and exo has a pre-existing
`Keypair.generate()` startup bug on main.

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 18:20:35 +00:00
Alex Cheema 7024ddcf3e fix: detect completed downloads by checking final file exists (#1582)
## Summary

Split from #1547 per review feedback.

When scanning existing download status, `get_downloaded_size()` returns
bytes from either the final file or its `.partial` counterpart — so a
`.partial` file with all bytes downloaded (e.g. process killed before
hash verification and rename) could be falsely reported as complete.

The previous approach (commit 3b54e7df) added a byte-comparison fallback
in the coordinator (`downloaded >= total > 0`), but this suffered from
the same `.partial` conflation issue.

**Fix:** Check whether the final (non-`.partial`) file actually exists
on disk before marking status as `"complete"` in `download_utils.py`.
This is the only reliable signal that hash verification passed and the
rename from `.partial` succeeded. The coordinator-level byte comparison
is removed since the source now reports correctly.

### Changes
- `download_utils.py`: Add `final_file_exists` check — only report
`"complete"` when the renamed, hash-verified file exists with matching
size
- `coordinator.py`: Revert to simple `progress.status == "complete"`
check, removing the byte-comparison fallback

**Note:** The corresponding byte-comparison workaround in #1547 should
also be removed.

## Test plan
- [x] basedpyright — 0 errors
- [x] ruff check — passes
- [x] pytest — 218 passed, 1 skipped (1 pre-existing Rust bindings
failure)

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 18:12:07 +00:00
Alex Cheema dc89ba662a feat: add info button to model picker variant rows (#1589)
## Motivation

When a model has multiple quantization variants (e.g., 4-bit, 8-bit),
expanding the group shows each variant's quantization and size — but
there's no way to see which specific nodes have each variant downloaded.
The top-level model group has an (i) info button, but the individual
variant rows don't. This makes it hard to tell which nodes have which
quantization.

## Changes

- Added an (i) info button to each expanded variant row in
`ModelPickerGroup.svelte`
- When clicked, opens the existing info modal scoped to that single
variant (showing its quantization, size, and which nodes have it
downloaded)
- Changed the variant row element from `<button>` to `<div
role="button">` to allow nesting the info `<button>` (matching the
pattern used by the top-level model row)

## Why It Works

Reuses the existing `onShowInfo` callback and info modal by constructing
a synthetic single-variant `ModelGroup` from the clicked variant. No
changes needed to `ModelPickerModal.svelte` — the info modal already
handles single-variant groups correctly.

## Test Plan

### Manual Testing
<!-- Hardware: MacBook Pro -->
- Open dashboard, open model picker
- Expand a model with multiple variants — each variant row now has an
(i) icon
- Click a variant's (i) — the info modal shows that single variant's
quantization/size and which nodes have it downloaded

### Automated Testing
- Dashboard builds successfully with `npm run build`
- `nix fmt` passes

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 09:53:36 -08:00
Alex Cheema 5cd96b5060 feat: seamless chat UX with auto model selection and smart recommendations (#1590)
## Motivation

The current chat experience requires users to manually select a model
before they can start chatting. This creates unnecessary friction —
especially for new users who may not know which model to pick.
Additionally, models that are already running or downloading appear
greyed out in the picker if available memory is low, making them
impossible to select even though they're already loaded.

## Changes

### New: ChatModelSelector component
- Category-based model recommendations (Best for Coding, Writing,
Agentic, Biggest) shown on the idle chat screen
- Tiered model ranking system (`AUTO_TIERS`) that picks the best model
fitting available memory
- Fixed-position tooltips that escape overflow-hidden ancestors

### Auto model selection from chat
- Typing in an empty chat auto-picks the best model using tier rankings
- Prefers already-running models over launching smaller new ones (avoids
the "launched a big model but recommends a smaller one" bug)
- If a model is already downloading/loading, attaches to existing
progress instead of launching a duplicate

### Launch-from-chat state machine
- Full `idle → launching → downloading → loading → ready` flow with
inline progress display
- `pendingAutoMessage` queue: messages typed during model launch are
remembered and sent once model is ready
- Fixed `$effect` race condition where restore effect competed with
auto-advance effect, causing lost messages

### Model picker improvements
- Instance status badges: green dot (ready), blue pulsing (downloading),
yellow pulsing (loading) on both group and variant rows
- "Ready" availability filter in ModelFilterPopover
- **Models with active instances (ready, downloading, loading) are never
greyed out** regardless of memory — they're already loaded/in-progress
and must be selectable

### Home page model display
- Model picker button shows the best running model name instead of "—
SELECT MODEL —" when a model is already active
- Uses the same `bestRunningModelId` tier-based derived as the chat
auto-selection

### Chat sidebar & form
- ChatSidebar conversation management with rename, delete, "New Chat"
- `userForcedIdle` guard prevents reactive effects from overriding
explicit user actions (e.g., clicking "New Chat")
- ChatForm simplified: uses ModelPickerModal instead of inline dropdown,
`modelDisplayOverride` prop for showing auto-selected models

## Why It Works

The tier-based auto-selection (`getAutoTierIndex` + `pickAutoModel`)
ensures users always get the best model their cluster can run. By
checking running instances first and comparing tiers, we avoid the bug
where loading a large model reduces available memory and causes a
smaller model to be recommended. The `$effect` ordering fix (moving the
`chatLaunchState === "ready"` guard after the `hasRunningInstance`
check) ensures pending messages are always delivered. The instance
status override in ModelPickerGroup (`anyVariantHasInstance`) ensures
models with active instances bypass memory-based greying regardless of
available RAM.

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Launch exo, navigate to Chat view
- Verify idle chat shows category recommendations (Coding, Writing,
Agentic, Biggest) based on available memory
- Type a message without selecting a model → verify auto-selection picks
the best tier model
- Launch a model, go to New Chat → verify the input bar shows the
running model name (not "SELECT MODEL")
- Launch a model, go to Home → verify model picker button shows the
running model
- While a model is downloading, send a message → verify download
progress appears and message is queued
- Open model picker with a model running → verify it is NOT greyed out
and is selectable
- Click "New Chat" while in a conversation → verify it resets to idle
with model selector

### Automated Testing
- Dashboard builds successfully (`cd dashboard && npm run build`)
- No TypeScript errors in the build output


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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 09:32:55 -08:00
rltakashige 05986f77aa add exo bench protobuf dependency (#1596)
kimi k2.5 requires protobuf
2026-02-23 16:45:22 +00:00
vskiwi dab7ed4821 fix: handle gossipsub MessageTooLarge error to prevent silent crash (#1583)
## Summary

Large prompts (70K+ tokens / ~500KB+ JSON) cause exo to silently crash.
The root cause is an unhandled `PublishError::MessageTooLarge` from
gossipsub when serialized `TextGeneration` commands exceed the 1MB
`max_transmit_size` limit.

The error propagates as a generic Python exception through the PyO3
bindings. Since `_networking_publish` in `router.py` only catches
`NoPeersSubscribedToTopicError` and `AllQueuesFullError`, the unhandled
exception crashes the networking async task, causing exo to shut down
silently — no error message, no API response.

## Changes

- **Rust (PyO3 bindings):** Add `MessageTooLargeError` exception class
and handle `PublishError::MessageTooLarge` explicitly in the gossipsub
publish path, matching the existing pattern for
`NoPeersSubscribedToTopicError` and `AllQueuesFullError`
- **Python (router):** Catch `MessageTooLargeError` in
`_networking_publish` and log a warning with the message size,
preventing the networking task from crashing

## Reproduction

On a multi-node cluster with a large model (e.g., GLM-5 754B tensor
parallel over JACCL RDMA):
1. Send a chat completion request with ~70K+ tokens
2. exo silently shuts down — no error logged, curl gets no response
3. With shorter prompts (< ~50K tokens): works fine

## Test plan

- Verified `cargo check` passes for `networking` and `exo_pyo3_bindings`
crates
- Verified `ruff check` passes for modified Python files
- Manual testing on 4× Mac Studio M3 Ultra cluster: 50K token requests
pass, 70K+ previously caused silent shutdown, now logs a warning and
drops the oversized message gracefully

Co-authored-by: vsm <vsm@nomail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-23 16:21:21 +00:00
Evan Quiney 2261014715 runner process checks (#1592)
partial solve to some of the more mysterious failures

notably, runner now switches to failed after EXO RUNNER MUST OOM
2026-02-23 16:11:18 +00:00
Evan Quiney 61d2a2b6cf add lazy task group (#1569)
in a few places we instantiate task groups early, in others we have an
optional task group.

this standardizes these patterns into a lazy task group, which queues
tasks until it is entered, at which point it enters its inner task group
and starts all its queued tasks. it should be noted that no queued tasks
will be started until the group itself is started.
2026-02-23 16:05:59 +00:00
Evan Quiney 0ff99a2c40 fix isinstance for qwen3Moe (#1595)
we were checking if qwen3 had a transformers Qwen3DecoderLayer rather
than an mlx Qwen3MoeDecoderLayer causing an assertion error on loading
qwen models - this corrects it to the actual layer type
2026-02-23 15:39:13 +00:00
rltakashige fbb80e1cc9 Address ring slowdown by turning on FAST SYNCH (#1594)
## Motivation

Large models + large prompts + pipeline RING = 0.2tps generation speed
Large models + large prompts + pipeline JACCL = 15 tps generation speed

Why? Well, MLX_METAL_FAST_SYNCH is on in pipeline JACCL.

## Changes

Just turn on fast synch everywhere, especially as GPU locks are old news

Also, changed to use mx.device_info as mx.metal.device_info is going to
be deprecated.

## Why It Works

Some magic thing that happens in the mlx backend. I really tried to find
a regression but couldn't. I will probably try again at some point.

## Test Plan

### Manual Testing
Did a bunch, no longer 0.2tps

### Automated Testing
We'll do that today.
2026-02-23 15:14:58 +00:00
Jake Hillion 8d94eab6c6 bench: fix KeyError on DownloadCompleted total field
The bench harness accessed the serialized DownloadCompleted field as
"totalBytes", but the Python field is `total: Memory` which serializes
to "total" (not snake_case, so the camelCase alias generator leaves it
unchanged). This caused a KeyError when --danger-delete-downloads
needed to free disk space by deleting existing models.

Changed the key from "totalBytes" to "total" to match the actual
serialized JSON structure.

Test plan:
- CI
- Reproduced KeyError on unfixed code by filling disk on a test node
  to 14GB free and running exo-bench with a 16GB model
  (Meta-Llama-3.1-8B-Instruct-bf16) with --danger-delete-downloads
- Verified fixed code successfully deletes smaller models to free
  space and completes the benchmark run
2026-02-23 14:17:41 +00:00
Alex Cheema f370452d7e Better onboarding UX (#1533)
## Summary
- **Complete onboarding wizard**: 7-step flow guiding new users from
Welcome → Your Devices (topology) → Add More Devices (animation) →
Choose Model → Download → Load → Chat
- **Native macOS integration**: NSPopover welcome callout anchored to
menu bar icon on first launch, polished DMG installer with
drag-to-Applications arrow
- **Dashboard UX polish**: auto-download on model select, toast
notifications, connection banner, skeleton loading, download progress in
header, recommended model tags, sidebar hidden in home state for cleaner
first impression
- **Settings & menu bar overhaul**: native Settings window with Advanced
tab, onboarding reset, chat sidebar toggle

## Test plan
- [ ] Fresh install: verify onboarding wizard appears and flows Welcome
→ Topology → Animation → Model → Download → Load → Chat
- [ ] Verify topology shows real device data in onboarding step 2
- [ ] Verify selecting a model in the main dashboard picker
auto-triggers download
- [ ] Verify chat sidebar is hidden on home view, appears when chat is
active
- [ ] Verify DMG installer has white background with curved arrow
- [ ] Verify NSPopover appears anchored to menu bar icon on first launch

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
2026-02-23 11:27:28 +00:00
Alex Cheema a4c2aa2b87 fix: raise error when MlxJaccl requested without RDMA cycles (#1585)
## Summary
- When MlxJaccl (RDMA) placement is requested but no RDMA-connected
cycles exist, raise a clear ValueError instead of silently falling back
to non-RDMA cycles
- Split from #1519 for independent review

## Test plan
- [x] basedpyright — 0 errors
- [x] ruff check — passes

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-23 10:09:01 +00:00
Alex Cheema 7312c535b4 feat: add user context prompt and GitHub issue option to macOS bug report (#1544)
## Summary

- When clicking "Send Bug Report" in the macOS app, users are now
prompted with "What's the issue? (optional)" before the diagnostic
upload begins
- The user's description is included in the uploaded report JSON
(`user_description` field)
- After successful upload, a "Create GitHub Issue" button opens the
browser to `github.com/exo-explore/exo/issues/new` pre-filled with the
user's description, macOS version, and EXO version

## Changes

- **`ContentView.swift`**: Replaced simple button with multi-phase
inline UI (idle → prompting → sending → success/failure). Added
`openGitHubIssue()` helper using `URLComponents` for pre-filled GitHub
issue URLs.
- **`BugReportService.swift`**: Added `userDescription` parameter to
`sendReport()` and `makeReportJson()`, included in the JSON payload when
non-empty.
- Removed the previous `BugReportModal.svelte` approach (the feature
belongs in the macOS app, not the dashboard).

## Test plan

- [ ] Build the Xcode project and run the macOS app
- [ ] Click "Send Bug Report" → verify text prompt appears
- [ ] Type a description, click Send → verify upload succeeds and
"Create GitHub Issue" button appears
- [ ] Click "Create GitHub Issue" → verify browser opens with pre-filled
template
- [ ] Test Cancel returns to idle, test empty description still works
- [ ] Test failure case (e.g. with exo stopped) shows error and dismiss

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-22 15:37:43 +00:00
Alex Cheema 18717023ad chore: remove deprecated MlxIbv dashboard references (#1584)
## Summary
- Remove legacy MlxIbvInstance references from ChatSidebar and ModelCard
components
- MlxIbv was replaced by MlxJaccl; these are leftover type checks
- Split from #1519 for independent review

## Test plan
- [x] Visual inspection of dashboard components

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-22 06:56:12 -08:00
Alex Cheema 1780e4ade4 fix: change RDMA AVAILABLE to RDMA NOT ENABLED warning (#1580)
## Summary
- Changed blue info badge "RDMA AVAILABLE" to yellow warning badge "RDMA
NOT ENABLED" — more accurately describes the state
- Added hover tooltip with enable instructions to all views (was missing
in 2 of 4 instances)
- Warning icon instead of info icon, consistent with other cluster
warnings (TB cycle, macOS mismatch)

## Screenshots

**Badge (yellow warning):**
![RDMA warning
badge](https://raw.githubusercontent.com/exo-explore/exo/3f7bdb482c5011d60f140aa84ab21023032e4a57/rdma-warning.png)

**Hover tooltip with instructions:**
![RDMA warning
hover](https://raw.githubusercontent.com/exo-explore/exo/3f7bdb482c5011d60f140aa84ab21023032e4a57/rdma-warning-hover.png)

## Test plan
- [x] Dashboard builds successfully
- [ ] Verify badge appears when 2+ TB5 nodes have RDMA disabled
- [ ] Verify hover tooltip shows in normal layout
- [ ] Verify hover tooltip shows in topology-only mode
- [ ] Verify dismiss button works
- [ ] Verify compact badge in status bar shows yellow warning

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-20 21:40:07 +00:00
Jake Hillion ab9273e723 downloads: add read_only flag to DownloadCompleted for EXO_MODELS_PATH
Models in EXO_MODELS_PATH are pre-downloaded into read-only directories
and must not be deleted. The DownloadCoordinator had no awareness of
these paths, so they never appeared as completed downloads in cluster
state, and the bench harness could attempt to delete them when freeing
disk space.

Added a `read_only: bool` field to `DownloadCompleted` (default False).
The DownloadCoordinator now checks `resolve_model_in_path` in
`_start_download`, proactively scans EXO_MODELS_PATH in
`_emit_existing_download_progress` to emit DownloadCompleted events for
all pre-downloaded models (overriding DownloadPending from the regular
scan), and refuses deletion of read-only models. The bench harness
filters out read-only models from deletion candidates.

Test plan:
- Ran with EXO_MODELS_PATH. Available models now show as downloaded in
  the UI. There isn't good UI for the fact they can't be deleted, but it
  should work with exo_bench.
2026-02-20 20:27:45 +00:00
Jake Hillion 71e48c0f62 model-cards: add missing metadata for Qwen3 Coder Next variants (#1576)
The Qwen3-Coder-Next model card TOML files were missing family,
quantization, base_model, and capabilities fields. This caused them not
to appear under the Qwen family filter in the dashboard's model picker.

Added the missing metadata to all five variants (4bit, 5bit, 6bit, 8bit,
bf16), matching the format used by the existing Qwen3-Coder-480B model
cards.

Test plan:
- Eyeballs
2026-02-20 18:25:49 +00:00
Jake Hillion 42da58c297 worker: add EXO_MODELS_PATH for pre-downloaded model directories
Users with pre-existing model files (e.g. on shared NFS mounts or from
prior downloads) had no way to point exo at those directories without
going through the download coordinator. EXO_MODELS_DIR only moves the
download target directory, it doesn't support read-only search paths.

Added EXO_MODELS_PATH environment variable as a colon-separated list of
directories to search for models. When the worker's plan loop encounters
a DownloadModel task, it checks these directories first and emits a
synthetic DownloadCompleted event if found, bypassing the download
coordinator entirely. The runner's build_model_path also checks these
directories first so the correct path is used during model loading.

This keeps the existing event sourcing state machine unchanged — the
DownloadCompleted event propagates naturally through the system, so
_load_model and all downstream logic work without modification.

Test plan:
- `s1@s1s-Mac-Studio ~ % EXO_LIBP2P_NAMESPACE=jake EXO_MODELS_PATH="/Volumes/Definitely Leo's SSD" nix --extra-experimental-features 'nix-command flakes' run github:exo-explore/exo/f2babbc2f742357d97dc177619fec062ef545be4`
- Started mlx-community/Qwen3-Coder-Next-4bit - it's present on the disk
  and it worked.
- Renamed one safetensor of mlx-community/Qwen3-Coder-Next-4bit on the
  disk. It then started the download locally, as expected.
2026-02-20 18:17:56 +00:00
Mustafa Alp Yılmaz 6b5a705959 fix: immediate cancel check after prefill completes (#1575)
## Problem

When a request is cancelled during prefill, the cancellation is not
detected until `check_for_cancel_every` additional tokens have been
generated. This is because `tokens_since_last_cancel_check` is
initialized to `0`, meaning the first cancel check only happens after
generating `check_for_cancel_every` tokens post-prefill.

For long prefills (which are the most likely to be cancelled), this adds
unnecessary latency before the cancellation is actually honoured.

## Fix

Initialize `tokens_since_last_cancel_check` to `check_for_cancel_every`
instead of `0`, so the very first token generated after prefill triggers
an immediate cancel check.

```diff
- tokens_since_last_cancel_check = 0
+ tokens_since_last_cancel_check = check_for_cancel_every
```

## Impact

- Cancellations issued during prefill are detected immediately when
generation begins
- No change in behaviour for non-cancelled requests (the counter resets
to `0` after each check as before)
- 1 line changed

Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-20 18:00:59 +00:00
Alex Cheema 6b54a27019 fix: add downloaded_bytes to DownloadPending event (#1564)
## Summary
- Add downloaded_bytes field to existing DownloadPending event for
accurate resume progress
- Minimal change per maintainer directive — no new download states
introduced

## Test plan
- [x] 42 tests passed, 1 skipped
- [x] Verified downloaded_bytes populates correctly for partial
downloads

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-20 17:54:18 +00:00
rltakashige e01f50a5cd Update mlx fork (#1565)
## Motivation

Some fixes upstream. This sort of commit will probably be quite common
until GPU locks are resolved.
2026-02-20 17:23:52 +00:00
Evan Quiney 1093080214 cancel active downloads on coordinator shutdown (#1567)
we were seeing some crashes as lost download tasks were trying to push
data toward a deleted coordinator. this cancels download tasks with the
coordinator's shutdown on master election
2026-02-20 17:17:43 +00:00
rltakashige 1a2b8b044a Refactor runner into separate runners (#1570)
## Motivation

We're going to be refactoring the llm inference code, so we should split
the runner up into parts while we can.

## Test Plan

### Manual Testing
Works on single node, at least.

### Automated Testing
Passes CI. Will be tested by our tests today.
2026-02-20 17:11:01 +00:00
Evan Quiney dc8d42b4dc add system ids (#1536)
addresses some election edge cases where a new worker with an old master
would get stuck on the old workers buffer index - we now use new system
ids each time we instantiate a node, and each event-producing system has
a unique system id for its lifespan (until the master moves).
2026-02-20 15:41:59 +00:00
Jake Hillion d484b062e8 bench: add download timing to bench output (#1566)
The bench script downloads models during the planning phase but doesn't
record how long the download took, making it difficult to track download
performance for a given model over time.

Modified `run_planning_phase` to return download metadata: whether a
fresh download occurred, the wall-clock duration, and the model size in
bytes. These fields are included in every JSON output row alongside the
existing per-run metrics, and a summary line is logged to the console.

This allows filtering bench results by `download_occurred` and grouping
by `model_id` to compute average download times across runs.

Test plan:

```
# existing model
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8 --pp 128 --tg 128
...
2026-02-20 15:23:49.081 | INFO     | __main__:main:340 - Planning phase: checking downloads...
2026-02-20 15:23:49.152 | INFO     | harness:run_planning_phase:402 - Started download on 12D3KooWKx41iikn188ozrxSdoG26g88jFCfie9wEA1eQR8csbPm
2026-02-20 15:23:49.184 | INFO     | __main__:main:352 - Download: model already cached
...
Wrote results JSON: bench/results.json
jake@maverick:/data/users/jake/repos/exo/ > cat bench/results.json
[
  {
    "elapsed_s": 2.9446684420108795,
    "output_text_preview": "The user just typed a long series of \"a\". Possibly they are testing. There's no explicit question. Could be they want a response? Might be a test of handling long input. We can respond politely, ask i",
    "stats": {
      "prompt_tps": 117.7872141515621,
      "generation_tps": 85.49598231498028,
      "prompt_tokens": 129,
      "generation_tokens": 128,
      "peak_memory_usage": {
        "inBytes": 68215145744
      }
    },
    "model_short_id": "gpt-oss-120b-MXFP4-Q8",
    "model_id": "mlx-community/gpt-oss-120b-MXFP4-Q8",
    "placement_sharding": "Pipeline",
    "placement_instance_meta": "MlxRing",
    "placement_nodes": 1,
    "instance_id": "68babc2a-6e94-4c70-aa07-7ec681f7c856",
    "pp_tokens": 128,
    "tg": 128,
    "repeat_index": 0
  }
]%
# no change to output
```

```
# missing model
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --host s1 --model mlx-community/Meta-Llama-3.1-8B-Instruct-4bit --pp 128 --tg 128
...
2026-02-20 15:24:42.553 | INFO     | __main__:main:340 - Planning phase: checking downloads...
2026-02-20 15:24:42.625 | INFO     | harness:run_planning_phase:402 - Started download on 12D3KooWKx41iikn188ozrxSdoG26g88jFCfie9wEA1eQR8csbPm
2026-02-20 15:25:37.494 | INFO     | __main__:main:350 - Download: 54.9s (freshly downloaded)
...
Wrote results JSON: bench/results.json
jake@maverick:/data/users/jake/repos/exo/ > cat bench/results.json
[
  {
    "elapsed_s": 1.500349276990164,
    "output_text_preview": "It seems like you've entered a large number of 'a's. If you'd like to discuss something or ask a question, I'm here to help. If not, is there anything else I can assist you with? \n\nIf you're intereste",
    "stats": {
      "prompt_tps": 395.43264952543666,
      "generation_tps": 128.03520443181478,
      "prompt_tokens": 129,
      "generation_tokens": 128,
      "peak_memory_usage": {
        "inBytes": 5116952079
      }
    },
    "model_short_id": "Meta-Llama-3.1-8B-Instruct-4bit",
    "model_id": "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
    "placement_sharding": "Pipeline",
    "placement_instance_meta": "MlxRing",
    "placement_nodes": 1,
    "instance_id": "ccd9bd71-d4cc-4b75-a37f-98090544626a",
    "pp_tokens": 128,
    "tg": 128,
    "repeat_index": 0,
    "download_duration_s": 54.88322358299047
  }
]%
# one new field
```
2026-02-20 15:33:08 +00:00
Alex Cheema e32b649d2f fix: enable psutil fallback for memory monitoring when macmon is missing (#1478)
## Summary
- On macOS, memory monitoring relied exclusively on `macmon` — the
psutil fallback was explicitly disabled (`memory_poll_rate = None`)
- When `macmon` is not installed (e.g., mac-mini-2 through mac-mini-4 in
our cluster), **no memory data was reported**, causing nodes to show 0GB
memory in the cluster state
- This blocked the scheduler from placing shards on those nodes since it
had no memory data to work with
- Fix: when `macmon` is not found on Darwin, fall back to psutil-based
memory polling (`memory_poll_rate = 1`)

## Root cause
`InfoGatherer` has two memory monitoring paths:
1. `macmon` (Darwin-only): provides memory + GPU/CPU/power stats
2. `psutil` (non-Darwin fallback): provides memory via
`MemoryUsage.from_psutil()`

Line 378 disabled psutil on Darwin: `memory_poll_rate = None if
IS_DARWIN else 1`
Line 389 only starts macmon if the binary exists: `if
shutil.which("macmon") is not None`

If macmon is missing on Darwin, **neither path runs** — zero memory
reported.

## Test plan
- [ ] Verify `uv run basedpyright` passes (0 errors confirmed)
- [ ] Verify `uv run ruff check` passes (confirmed)
- [ ] Verify `uv run pytest src/exo/utils/info_gatherer/` passes (2/2
confirmed)
- [ ] Deploy to cluster nodes without macmon and verify memory appears
in `/state`

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-20 13:12:22 +00:00
Alex Cheema bddad7e79c feat: show ETA on prefill progress bar (#1557)
## Summary
- Show estimated time remaining during prefill (prompt processing phase)
- Track prefill start time via performance.now() and extrapolate from
observed token throughput
- Display ~Xs remaining or ~Xm Ys remaining next to the percentage on
the progress bar
- Wait 200ms before showing ETA to ensure a stable sample window

## Changes
**PrefillProgressBar.svelte**: Add etaText derived computation that
calculates remaining time from (remainingTokens / tokensPerMs). Renders
in a new flex row below the progress bar alongside the percentage.

**app.svelte.ts**: Add startedAt: number field to PrefillProgress
interface. Set on first prefill_progress SSE event, preserved across
subsequent updates.

## Test plan
- [ ] Start inference with a long prompt (10k+ tokens) on a multi-node
cluster
- [ ] Verify the progress bar shows ~Xs remaining after ~200ms of
prefill
- [ ] Verify the ETA decreases as prefill progresses
- [ ] Verify short prefills (<200ms) dont flash a briefly-visible ETA
- [ ] Verify ETA disappears when prefill completes and token generation
begins

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-20 12:37:56 +00:00
rltakashige addf73a144 Add support for Ollama API (#1560)
## Motivation

Ollama has a bunch of integrations, such as OpenWebUI, that are very
handy. Let's support it :)

## Test Plan

### Manual Testing
<img width="3426" height="1998" alt="image"
src="https://github.com/user-attachments/assets/44b07f1e-308e-4ff1-9a11-922d8279939f"
/>
2026-02-20 12:03:27 +00:00
Mustafa Alp Yılmaz a16ff2c047 fix: correct misleading docstring in seed_models (#1561)
## Summary
- Fixed stale docstring in `seed_models()` that referenced
`.cache/huggingface/hub` when the function actually moves models to
`EXO_MODELS_DIR` (resolved via `ensure_models_dir()`)
- The old docstring was misleading for AI coding agents analyzing the
codebase, causing incorrect conclusions about model storage paths

## Changes
`src/exo/download/download_utils.py`: Updated docstring from `"Move
model in resources folder of app to .cache/huggingface/hub"` to `"Move
models from resources folder to EXO_MODELS_DIR."`

Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-20 11:57:55 +00:00
rltakashige 3006c8ea4e Ensure coordinator is rank 0 (#1559)
## Motivation

Coordinator can be a random rank. Let's just fix this to rank 0 as
that's what we typically assume.

## Test Plan

### Manual Testing
Works as normal on 2 nodes.


Let's wait for a little more testing to merge this.

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-02-20 11:46:24 +00:00
rltakashige f662c129dd Prioritise tb for ring instances (#1556)
## Motivation

TB has better bandwidth and latency than ethernet. We should prioritise
TB5 where possible. This drastically improves distributed image
generation performance.

## Test Plan

### Manual Testing
Saw on the dashboard that TB (169.254) addresses were prioritised.

Tested that image models scale much better.

### Automated Testing
No regression on Kimi K2.5
2026-02-19 21:32:48 +00:00
Evan Quiney c45ff9ad43 memory tidy (#1558)
add some pythonic extensions to memory, did a bunch of cleanup.
2026-02-19 21:15:33 +00:00
rltakashige 7031901ae5 Prevent common fatal crashes (#1555)
## Motivation
Occasionally, memory does not get released when we shut down. There is
no reason to delay deleting the model.

Also handles can become None during shutdown, causing TypeErrors which
are not handled and bringing down exo.

Similarly, we were closing the event sender in the wrong place.

Also let's not verify the SSL certificate for http connections to local
peers, as this is failing sometimes and crashing.

## Test Plan

### Manual Testing
No more crashes as described.
2026-02-19 20:51:17 +00:00
rltakashige cf648a53b8 Add thinking in thinking blocks, and fix DeepSeek interleaved tool calls (#1548)
## Motivation

OpenCode shows <think> tags and not thinking blocks as we aren't
following the API specs properly.

Claude was also getting horrible prefix cache hits because it sends
headers.

## Changes

Handle thinking tokens properly by placing them in think tags for each
of the API endpoints.
Also support DeepSeekV3.2 tool calling properly as a minor feature.
Strips Claude headers at the API level.

## Test Plan

### Manual Testing
Tested OpenCode manually
Needs testing with Claude.

### Automated Testing
All CI and tests passing - added a new e2e test for DeepSeekV32 tool
parsing.
2026-02-19 18:44:49 +00:00
Alex Cheema 94b2ce6922 feat: Mac Studio en2 RDMA port warning v2 (#1551)
Rebuilt from scratch (replaces PR #1543). Detects when Mac Studio uses
RDMA over en2 (TB5 port next to Ethernet) which does not support RDMA.
Shows dismissible warning banner with hover tooltip showing affected
devices, SVG rear panel illustration, and fix instructions. 205 lines in
+page.svelte.

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-19 18:39:17 +00:00
rltakashige 423ed0f07f Strip Claude headers to improve prefix cache hit rates (#1552)
## Motivation
Our hits are really bad at the moment (0.2%). This PR makes it 98.5% on
average.

## Changes

Also adds an example for how to run Claude using Exo.

## Why It Works
Claude sends some billing and session headers that change with each
message.

## Test Plan

### Manual Testing
Works in manual testing.
2026-02-19 18:29:34 +00:00
Evan Quiney ed001f2409 remove prefillprogress event (#1550)
this should never have been a separate event, but i didnt quite
communicate that well when this was merged. convert PrefillProgress to a
chunk like the rest of the runner responses.

tested with Llama-3.3-70B, prefill progress events still show up in the
dashboard as usual
2026-02-19 18:23:28 +00:00
Evan Quiney 4c4c6ce99f simplify rust ident module
this is partly dead code, partly narrowing the rust-python boundary in
prep for future rewrites. no testing as this is all type safe
refactoring.
2026-02-19 17:19:31 +00:00
112 changed files with 10811 additions and 2011 deletions
+16
View File
@@ -215,6 +215,22 @@ class StreamContext:
traceback: object | None = ...,
) -> None: ...
def device_info() -> dict[str, str | int]:
"""
Get information about the GPU device and system settings.
Currently returns:
* ``architecture``
* ``max_buffer_size``
* ``max_recommended_working_set_size``
* ``memory_size``
* ``resource_limit``
Returns:
dict: A dictionary with string keys and string or integer values.
"""
def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
"""
Element-wise absolute value.
+215 -148
View File
@@ -15,14 +15,23 @@ struct ContentView: View {
@EnvironmentObject private var localNetworkChecker: LocalNetworkChecker
@EnvironmentObject private var updater: SparkleUpdater
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
@EnvironmentObject private var settingsWindowController: SettingsWindowController
@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
@State private var bugReportInFlight = false
@State private var bugReportMessage: String?
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 = ""
@@ -258,139 +267,79 @@ struct ContentView: View {
VStack(alignment: .leading, spacing: 0) {
if controller.status != .stopped {
dashboardButton
baseURLRow
Divider()
.padding(.vertical, 8)
} else {
Divider()
.padding(.vertical, 4)
}
advancedSection
.padding(.bottom, 8)
controlButton(title: "Quit", tint: .secondary) {
HoverButton(
title: "Settings",
tint: .primary,
trailingSystemImage: "gear"
) {
settingsWindowController.open(
controller: controller,
updater: updater,
networkStatusService: networkStatusService,
thunderboltBridgeService: thunderboltBridgeService,
stateService: stateService
)
}
HoverButton(
title: "Check for Updates",
tint: .primary,
trailingSystemImage: "arrow.triangle.2.circlepath"
) {
updater.checkForUpdates()
}
.padding(.bottom, 8)
HoverButton(title: "Quit", tint: .secondary) {
controller.stop()
NSApplication.shared.terminate(nil)
}
}
}
private var advancedSection: some View {
VStack(alignment: .leading, spacing: 6) {
HStack {
Text("Advanced")
.font(.caption)
.foregroundColor(.secondary)
Spacer()
collapseButton(isExpanded: $showAdvanced)
}
.animation(nil, value: showAdvanced)
if showAdvanced {
VStack(alignment: .leading, spacing: 8) {
VStack(alignment: .leading, spacing: 4) {
Text("Cluster Namespace")
.font(.caption2)
.foregroundColor(.secondary)
HStack {
TextField("optional", text: $pendingNamespace)
.textFieldStyle(.roundedBorder)
.font(.caption2)
.onAppear {
pendingNamespace = controller.customNamespace
}
Button("Save & Restart") {
controller.customNamespace = pendingNamespace
if controller.status == .running || controller.status == .starting {
controller.restart()
}
}
.font(.caption2)
.disabled(pendingNamespace == controller.customNamespace)
}
}
VStack(alignment: .leading, spacing: 4) {
Text("HuggingFace Token")
.font(.caption2)
.foregroundColor(.secondary)
HStack {
SecureField("optional", text: $pendingHFToken)
.textFieldStyle(.roundedBorder)
.font(.caption2)
.onAppear {
pendingHFToken = controller.hfToken
}
Button("Save & Restart") {
controller.hfToken = pendingHFToken
if controller.status == .running || controller.status == .starting {
controller.restart()
}
}
.font(.caption2)
.disabled(pendingHFToken == controller.hfToken)
}
}
Divider()
HStack {
Toggle(
"Enable Image Models (experimental)", isOn: $pendingEnableImageModels
)
.toggleStyle(.switch)
.font(.caption2)
.onAppear {
pendingEnableImageModels = controller.enableImageModels
}
Spacer()
Button("Save & Restart") {
controller.enableImageModels = pendingEnableImageModels
if controller.status == .running || controller.status == .starting {
controller.restart()
}
}
.font(.caption2)
.disabled(pendingEnableImageModels == controller.enableImageModels)
}
HoverButton(title: "Check for Updates", small: true) {
updater.checkForUpdates()
}
debugSection
HoverButton(title: "Uninstall", tint: .red, small: true) {
showUninstallConfirmationAlert()
}
.disabled(uninstallInProgress)
}
.transition(.opacity)
}
}
.animation(.easeInOut(duration: 0.25), value: showAdvanced)
}
private func controlButton(title: String, tint: Color = .primary, action: @escaping () -> Void)
-> some View
{
HoverButton(title: title, tint: tint, trailingSystemImage: nil, action: action)
}
private var dashboardButton: some View {
Button {
HoverButton(
title: "Web Dashboard",
tint: .primary,
trailingSystemImage: "arrow.up.right"
) {
guard let url = URL(string: "http://localhost:52415/") else { return }
NSWorkspace.shared.open(url)
} label: {
HStack {
Image(systemName: "arrow.up.right.square")
.imageScale(.small)
Text("Dashboard")
.fontWeight(.medium)
Spacer()
}
.padding(.vertical, 8)
.padding(.horizontal, 10)
.background(
RoundedRectangle(cornerRadius: 8, style: .continuous)
.fill(Color(red: 1.0, green: 0.87, blue: 0.0).opacity(0.2))
)
}
.buttonStyle(.plain)
.padding(.bottom, 4)
}
private var baseURLRow: some View {
HStack(spacing: 6) {
Image(systemName: "link")
.imageScale(.small)
.foregroundColor(.secondary)
Text("localhost:52415/v1")
.font(.system(.caption, design: .monospaced))
.foregroundColor(.primary)
Spacer()
Button {
NSPasteboard.general.clearContents()
NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string)
baseURLCopied = true
DispatchQueue.main.asyncAfter(deadline: .now() + 2) {
baseURLCopied = false
}
} label: {
Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc")
.imageScale(.small)
.foregroundColor(baseURLCopied ? .green : .secondary)
.contentTransition(.symbolEffect(.replace))
}
.buttonStyle(.plain)
.help("Copy API base URL")
}
.padding(.vertical, 4)
.padding(.horizontal, 8)
}
private func collapseButton(isExpanded: Binding<Bool>) -> some View {
@@ -611,39 +560,115 @@ struct ContentView: View {
}
private var sendBugReportButton: some View {
VStack(alignment: .leading, spacing: 4) {
Button {
Task {
await sendBugReport()
}
} label: {
HStack {
if bugReportInFlight {
ProgressView()
.scaleEffect(0.6)
VStack(alignment: .leading, spacing: 6) {
switch bugReportPhase {
case .idle:
Button {
bugReportPhase = .prompting
bugReportUserDescription = ""
} label: {
HStack {
Text("Send Bug Report")
.font(.caption)
.fontWeight(.semibold)
Spacer()
}
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))
)
}
.padding(.vertical, 6)
.padding(.horizontal, 8)
.buttonStyle(.plain)
case .prompting:
VStack(alignment: .leading, spacing: 6) {
Text("What's the issue? (optional)")
.font(.caption2)
.foregroundColor(.secondary)
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.12))
.fill(Color.accentColor.opacity(0.06))
)
}
.buttonStyle(.plain)
.disabled(bugReportInFlight)
if let message = bugReportMessage {
Text(message)
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)
.fixedSize(horizontal: false, vertical: true)
}
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> {
@@ -687,16 +712,58 @@ struct ContentView: View {
}
private func sendBugReport() async {
bugReportInFlight = true
bugReportMessage = "Collecting logs..."
bugReportPhase = .sending("Collecting logs...")
let service = BugReportService()
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
do {
let outcome = try await service.sendReport(isManual: true)
bugReportMessage = outcome.message
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 {
bugReportMessage = error.localizedDescription
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)
}
bugReportInFlight = false
}
private func showUninstallConfirmationAlert() {
+15 -1
View File
@@ -21,7 +21,9 @@ struct EXOApp: App {
@StateObject private var localNetworkChecker: LocalNetworkChecker
@StateObject private var updater: SparkleUpdater
@StateObject private var thunderboltBridgeService: ThunderboltBridgeService
@StateObject private var settingsWindowController: SettingsWindowController
private let terminationObserver: TerminationObserver
private let firstLaunchPopout = FirstLaunchPopout()
private let ciContext = CIContext(options: nil)
init() {
@@ -43,12 +45,13 @@ struct EXOApp: App {
_updater = StateObject(wrappedValue: updater)
let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
_thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
_settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
enableLaunchAtLoginIfNeeded()
// Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
NetworkSetupHelper.promptAndInstallIfNeeded()
// Check local network access periodically (warning disappears when user grants permission)
localNetwork.startPeriodicChecking(interval: 10)
controller.scheduleLaunch(after: 15)
controller.scheduleLaunch(after: 5)
service.startPolling()
networkStatus.startPolling()
}
@@ -62,8 +65,19 @@ struct EXOApp: App {
.environmentObject(localNetworkChecker)
.environmentObject(updater)
.environmentObject(thunderboltBridgeService)
.environmentObject(settingsWindowController)
} label: {
menuBarIcon
.onReceive(controller.$isFirstLaunchReady) { ready in
if ready {
DispatchQueue.main.asyncAfter(deadline: .now() + 3.0) {
self.firstLaunchPopout.onComplete = { [weak controller] in
controller?.markOnboardingCompleted()
}
self.firstLaunchPopout.show()
}
}
}
}
.menuBarExtraStyle(.window)
}
+18
View File
@@ -5,6 +5,7 @@ import Foundation
private let customNamespaceKey = "EXOCustomNamespace"
private let hfTokenKey = "EXOHFToken"
private let enableImageModelsKey = "EXOEnableImageModels"
private let onboardingCompletedKey = "EXOOnboardingCompleted"
@MainActor
final class ExoProcessController: ObservableObject {
@@ -60,6 +61,9 @@ final class ExoProcessController: ObservableObject {
}
}
/// Fires once when EXO transitions to `.running` for the very first time (fresh install).
@Published private(set) var isFirstLaunchReady = false
private var process: Process?
private var runtimeDirectoryURL: URL?
private var pendingLaunchTask: Task<Void, Never>?
@@ -113,6 +117,9 @@ final class ExoProcessController: ObservableObject {
try child.run()
process = child
status = .running
// Show welcome popout on every launch
isFirstLaunchReady = true
} catch {
process = nil
status = .failed(message: "Launch error")
@@ -164,6 +171,17 @@ final class ExoProcessController: ObservableObject {
launch()
}
/// Mark onboarding as completed (user interacted with the welcome popout).
func markOnboardingCompleted() {
UserDefaults.standard.set(true, forKey: onboardingCompletedKey)
}
/// Reset onboarding so the welcome popout appears on next launch.
func resetOnboarding() {
UserDefaults.standard.removeObject(forKey: onboardingCompletedKey)
isFirstLaunchReady = false
}
func scheduleLaunch(after seconds: TimeInterval) {
cancelPendingLaunch()
let start = max(1, Int(ceil(seconds)))
+9 -3
View File
@@ -38,7 +38,8 @@ struct BugReportService {
func sendReport(
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
now: Date = Date(),
isManual: Bool = false
isManual: Bool = false,
userDescription: String? = nil
) async throws -> BugReportOutcome {
let timestamp = Self.runTimestampString(now)
let dayPrefix = Self.dayPrefixString(now)
@@ -60,7 +61,8 @@ struct BugReportService {
ifconfig: ifconfigText,
debugInfo: debugInfo,
isManual: isManual,
clusterTbBridgeStatus: clusterTbBridgeStatus
clusterTbBridgeStatus: clusterTbBridgeStatus,
userDescription: userDescription
)
let eventLogFiles = readAllEventLogs()
@@ -306,7 +308,8 @@ struct BugReportService {
ifconfig: String,
debugInfo: DebugInfo,
isManual: Bool,
clusterTbBridgeStatus: [[String: Any]]?
clusterTbBridgeStatus: [[String: Any]]?,
userDescription: String? = nil
) -> Data? {
let system = readSystemMetadata()
let exo = readExoMetadata()
@@ -323,6 +326,9 @@ struct BugReportService {
if let tbStatus = clusterTbBridgeStatus {
payload["cluster_thunderbolt_bridge"] = tbStatus
}
if let desc = userDescription, !desc.isEmpty {
payload["user_description"] = desc
}
return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
}
+200
View File
@@ -0,0 +1,200 @@
import AppKit
import SwiftUI
/// A popover callout anchored to the menu bar icon on every launch,
/// pointing the user to the web dashboard with an arrow connecting to the icon.
@MainActor
final class FirstLaunchPopout {
private var popover: NSPopover?
private var countdownTask: Task<Void, Never>?
private static let dashboardURL = "http://localhost:52415/"
/// Called when the user completes onboarding (clicks Open Dashboard or dismisses).
var onComplete: (() -> Void)?
func show() {
guard popover == nil else { return }
// The status bar button may not exist yet on first launch; retry generously.
showWithRetry(attemptsRemaining: 15)
}
private func showWithRetry(attemptsRemaining: Int) {
guard attemptsRemaining > 0 else {
// Exhausted retries fall back to just opening the dashboard directly.
openDashboard()
onComplete?()
return
}
guard let button = Self.findStatusItemButton() else {
DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { [weak self] in
self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1)
}
return
}
let pop = NSPopover()
pop.behavior = .applicationDefined
pop.animates = true
pop.contentSize = NSSize(width: 280, height: 120)
pop.contentViewController = NSHostingController(
rootView: WelcomeCalloutView(
countdownDuration: 10,
onDismiss: { [weak self] in
self?.onComplete?()
self?.dismiss()
},
onOpen: { [weak self] in
self?.openDashboard()
self?.onComplete?()
self?.dismiss()
}
)
)
self.popover = pop
pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY)
// Auto-open dashboard after 10s then dismiss
countdownTask = Task {
try? await Task.sleep(nanoseconds: 10_000_000_000)
if !Task.isCancelled {
openDashboard()
onComplete?()
dismiss()
}
}
}
func dismiss() {
countdownTask?.cancel()
countdownTask = nil
guard let pop = popover else { return }
popover = nil
pop.performClose(nil)
}
private func openDashboard() {
guard let url = URL(string: Self.dashboardURL) else { return }
NSWorkspace.shared.open(url)
}
/// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra.
/// Walks the view hierarchy to find the actual button rather than the content view.
private static func findStatusItemButton() -> NSView? {
for window in NSApp.windows {
let className = NSStringFromClass(type(of: window))
// Match NSStatusBarWindow or any internal SwiftUI status bar window
guard className.contains("StatusBar") || className.contains("MenuBarExtra") else {
continue
}
if let content = window.contentView {
if let button = findButton(in: content) {
return button
}
// Fall back to the content view itself if it has a non-zero frame
if content.frame.width > 0 {
return content
}
}
}
return nil
}
/// Recursively searches the view hierarchy for an NSStatusBarButton.
private static func findButton(in view: NSView) -> NSView? {
let className = NSStringFromClass(type(of: view))
if className.contains("StatusBarButton") || className.contains("StatusItem") {
return view
}
for subview in view.subviews {
if let found = findButton(in: subview) {
return found
}
}
return nil
}
}
/// Minimal welcome callout friendly pointer, not a wall of text.
/// Rendered inside the NSPopover which provides its own chrome and arrow.
private struct WelcomeCalloutView: View {
let countdownDuration: Int
let onDismiss: () -> Void
let onOpen: () -> Void
@State private var countdown: Int
@State private var timerTask: Task<Void, Never>?
init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) {
self.countdownDuration = countdownDuration
self.onDismiss = onDismiss
self.onOpen = onOpen
self._countdown = State(initialValue: countdownDuration)
}
var body: some View {
VStack(alignment: .leading, spacing: 10) {
HStack(alignment: .top) {
Text("EXO is running")
.font(.system(.headline, design: .rounded))
.fontWeight(.semibold)
.foregroundColor(.primary)
Spacer()
Button {
onDismiss()
} label: {
Image(systemName: "xmark.circle.fill")
.font(.system(size: 14))
.foregroundStyle(.tertiary)
}
.buttonStyle(.plain)
}
Text("Run your first model here:")
.font(.system(.subheadline, design: .default))
.foregroundColor(.secondary)
HStack {
Button {
onOpen()
} label: {
Label("Open Dashboard", systemImage: "arrow.up.right.square")
.font(.system(.caption, design: .default))
.fontWeight(.medium)
}
.buttonStyle(.borderedProminent)
.tint(.accentColor)
.controlSize(.small)
Spacer()
if countdown > 0 {
Text("Launching in \(countdown) secs...")
.font(.system(.caption2, design: .default))
.foregroundColor(.secondary.opacity(0.6))
.monospacedDigit()
}
}
}
.padding(14)
.onAppear {
startCountdown()
}
.onDisappear {
timerTask?.cancel()
timerTask = nil
}
}
private func startCountdown() {
timerTask = Task {
while countdown > 0 {
try? await Task.sleep(nanoseconds: 1_000_000_000)
if !Task.isCancelled {
countdown -= 1
}
}
}
}
}
+478
View File
@@ -0,0 +1,478 @@
import AppKit
import SwiftUI
/// Native macOS Settings window following Apple HIG.
/// Organized into General, Model, Advanced, and About sections.
struct SettingsView: View {
@EnvironmentObject private var controller: ExoProcessController
@EnvironmentObject private var updater: SparkleUpdater
@EnvironmentObject private var networkStatusService: NetworkStatusService
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
@EnvironmentObject private var stateService: ClusterStateService
@State private var pendingNamespace: String = ""
@State private var pendingHFToken: String = ""
@State private var pendingEnableImageModels = false
@State private var needsRestart = false
@State private var bugReportInFlight = false
@State private var bugReportMessage: String?
@State private var uninstallInProgress = false
var body: some View {
TabView {
generalTab
.tabItem {
Label("General", systemImage: "gear")
}
modelTab
.tabItem {
Label("Model", systemImage: "cube")
}
advancedTab
.tabItem {
Label("Advanced", systemImage: "wrench.and.screwdriver")
}
aboutTab
.tabItem {
Label("About", systemImage: "info.circle")
}
}
.frame(width: 450, height: 400)
.onAppear {
pendingNamespace = controller.customNamespace
pendingHFToken = controller.hfToken
pendingEnableImageModels = controller.enableImageModels
needsRestart = false
}
}
// MARK: - General Tab
private var generalTab: some View {
Form {
Section {
LabeledContent("Cluster Namespace") {
TextField("default", text: $pendingNamespace)
.textFieldStyle(.roundedBorder)
.frame(width: 200)
}
Text("Nodes with the same namespace form a cluster. Leave empty for default.")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
LabeledContent("HuggingFace Token") {
SecureField("optional", text: $pendingHFToken)
.textFieldStyle(.roundedBorder)
.frame(width: 200)
}
Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
HStack {
Spacer()
Button("Save & Restart") {
applyGeneralSettings()
}
.disabled(!hasGeneralChanges)
}
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - Model Tab
private var modelTab: some View {
Form {
Section {
Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels)
Text("Allow text-to-image and image-to-image models in the model picker.")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
HStack {
Spacer()
Button("Save & Restart") {
applyModelSettings()
}
.disabled(!hasModelChanges)
}
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - Advanced Tab
private var advancedTab: some View {
Form {
Section("Onboarding") {
HStack {
VStack(alignment: .leading) {
Text("Reset Onboarding")
Text("Opens the dashboard and resets the onboarding wizard.")
.font(.caption)
.foregroundColor(.secondary)
}
Spacer()
Button("Reset") {
guard let url = URL(string: "http://localhost:52415/?reset-onboarding")
else { return }
NSWorkspace.shared.open(url)
}
}
}
Section("Debug Info") {
LabeledContent("Thunderbolt Bridge") {
Text(thunderboltStatusText)
.foregroundColor(thunderboltStatusColor)
}
VStack(alignment: .leading, spacing: 2) {
clusterThunderboltBridgeView
}
VStack(alignment: .leading, spacing: 2) {
interfaceIpList
}
VStack(alignment: .leading, spacing: 2) {
rdmaStatusView
}
sendBugReportButton
}
Section("Danger Zone") {
Button(role: .destructive) {
showUninstallConfirmationAlert()
} label: {
HStack {
Text("Uninstall EXO")
Spacer()
Image(systemName: "trash")
.imageScale(.small)
}
}
.disabled(uninstallInProgress)
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - About Tab
private var aboutTab: some View {
Form {
Section {
LabeledContent("Version") {
Text(buildTag)
.textSelection(.enabled)
}
LabeledContent("Commit") {
Text(buildCommit)
.font(.system(.body, design: .monospaced))
.textSelection(.enabled)
}
}
Section {
Button("Check for Updates") {
updater.checkForUpdates()
}
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - Debug Info Views (moved from ContentView)
private var thunderboltStatusText: String {
switch networkStatusService.status.thunderboltBridgeState {
case .some(.disabled):
return "Disabled"
case .some(.deleted):
return "Deleted"
case .some(.enabled):
return "Enabled"
case nil:
return "Unknown"
}
}
private var thunderboltStatusColor: Color {
switch networkStatusService.status.thunderboltBridgeState {
case .some(.disabled), .some(.deleted):
return .green
case .some(.enabled):
return .red
case nil:
return .secondary
}
}
private var clusterThunderboltBridgeView: some View {
let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
let localNodeId = stateService.localNodeId
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
return VStack(alignment: .leading, spacing: 1) {
if bridgeStatuses.isEmpty {
Text("Cluster TB Bridge: No data")
.font(.caption2)
.foregroundColor(.secondary)
} else {
Text("Cluster TB Bridge Status:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
if let status = bridgeStatuses[nodeId] {
let nodeName =
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
let isLocal = nodeId == localNodeId
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
let statusText =
!status.exists
? "N/A"
: (status.enabled ? "Enabled" : "Disabled")
let color: Color =
!status.exists
? .secondary
: (status.enabled ? .red : .green)
Text("\(prefix) \(statusText)")
.font(.caption2)
.foregroundColor(color)
}
}
}
}
}
private var interfaceIpList: some View {
let statuses = networkStatusService.status.interfaceStatuses
return VStack(alignment: .leading, spacing: 1) {
Text("Interfaces (en0en7):")
.font(.caption2)
.foregroundColor(.secondary)
if statuses.isEmpty {
Text(" Unknown")
.font(.caption2)
.foregroundColor(.secondary)
} else {
ForEach(statuses, id: \.interfaceName) { status in
let ipText = status.ipAddress ?? "No IP"
Text(" \(status.interfaceName): \(ipText)")
.font(.caption2)
.foregroundColor(status.ipAddress == nil ? .red : .green)
}
}
}
}
private var rdmaStatusView: some View {
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
let localNodeId = stateService.localNodeId
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
let localDevices = networkStatusService.status.localRdmaDevices
let localPorts = networkStatusService.status.localRdmaActivePorts
return VStack(alignment: .leading, spacing: 1) {
if rdmaStatuses.isEmpty {
Text("Cluster RDMA: No data")
.font(.caption2)
.foregroundColor(.secondary)
} else {
Text("Cluster RDMA Status:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
if let status = rdmaStatuses[nodeId] {
let nodeName =
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
let isLocal = nodeId == localNodeId
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
let statusText = status.enabled ? "Enabled" : "Disabled"
let color: Color = status.enabled ? .green : .orange
Text("\(prefix) \(statusText)")
.font(.caption2)
.foregroundColor(color)
}
}
}
if !localDevices.isEmpty {
Text(" Local Devices: \(localDevices.joined(separator: ", "))")
.font(.caption2)
.foregroundColor(.secondary)
}
if !localPorts.isEmpty {
Text(" Local Active Ports:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(localPorts, id: \.device) { port in
Text(" \(port.device) port \(port.port): \(port.state)")
.font(.caption2)
.foregroundColor(.green)
}
}
}
}
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"
alert.informativeText = """
This will remove EXO and all its system components:
• Network configuration daemon
• Launch at login registration
• EXO network location
The app will be moved to Trash.
"""
alert.alertStyle = .warning
alert.addButton(withTitle: "Uninstall")
alert.addButton(withTitle: "Cancel")
if let uninstallButton = alert.buttons.first {
uninstallButton.hasDestructiveAction = true
}
let response = alert.runModal()
if response == .alertFirstButtonReturn {
performUninstall()
}
}
private func performUninstall() {
uninstallInProgress = true
controller.cancelPendingLaunch()
controller.stop()
stateService.stopPolling()
DispatchQueue.global(qos: .utility).async {
do {
try NetworkSetupHelper.uninstall()
DispatchQueue.main.async {
LaunchAtLoginHelper.disable()
self.moveAppToTrash()
DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
NSApplication.shared.terminate(nil)
}
}
} catch {
DispatchQueue.main.async {
let errorAlert = NSAlert()
errorAlert.messageText = "Uninstall Failed"
errorAlert.informativeText = error.localizedDescription
errorAlert.alertStyle = .critical
errorAlert.addButton(withTitle: "OK")
errorAlert.runModal()
self.uninstallInProgress = false
}
}
}
}
private func moveAppToTrash() {
guard let appURL = Bundle.main.bundleURL as URL? else { return }
do {
try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
} catch {
// If we can't trash the app, that's OK - user can do it manually
}
}
// MARK: - Helpers
private var hasGeneralChanges: Bool {
pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
}
private var hasModelChanges: Bool {
pendingEnableImageModels != controller.enableImageModels
}
private func applyGeneralSettings() {
controller.customNamespace = pendingNamespace
controller.hfToken = pendingHFToken
restartIfRunning()
}
private func applyModelSettings() {
controller.enableImageModels = pendingEnableImageModels
restartIfRunning()
}
private func restartIfRunning() {
if controller.status == .running || controller.status == .starting {
controller.restart()
}
}
private var buildTag: String {
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
}
private var buildCommit: String {
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
}
}
@@ -0,0 +1,47 @@
import AppKit
import SwiftUI
/// Manages a standalone native macOS Settings window.
/// Ensures only one instance exists and brings it to front on repeated opens.
@MainActor
final class SettingsWindowController: ObservableObject {
private var window: NSWindow?
func open(
controller: ExoProcessController,
updater: SparkleUpdater,
networkStatusService: NetworkStatusService,
thunderboltBridgeService: ThunderboltBridgeService,
stateService: ClusterStateService
) {
if let existing = window, existing.isVisible {
existing.makeKeyAndOrderFront(nil)
NSApp.activate(ignoringOtherApps: true)
return
}
let settingsView = SettingsView()
.environmentObject(controller)
.environmentObject(updater)
.environmentObject(networkStatusService)
.environmentObject(thunderboltBridgeService)
.environmentObject(stateService)
let hostingView = NSHostingView(rootView: settingsView)
let newWindow = NSWindow(
contentRect: NSRect(x: 0, y: 0, width: 450, height: 400),
styleMask: [.titled, .closable],
backing: .buffered,
defer: false
)
newWindow.title = "EXO Settings"
newWindow.contentView = hostingView
newWindow.center()
newWindow.isReleasedWhenClosed = false
newWindow.makeKeyAndOrderFront(nil)
NSApp.activate(ignoringOtherApps: true)
window = newWindow
}
}
+11 -2
View File
@@ -75,7 +75,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
model_path = Path(
snapshot_download(
model_id,
allow_patterns=["*.json", "*.py", "*.tiktoken"],
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
)
)
@@ -338,7 +338,7 @@ def main() -> int:
)
logger.info("Planning phase: checking downloads...")
run_planning_phase(
download_duration_s = run_planning_phase(
client,
full_model_id,
selected[0],
@@ -346,6 +346,10 @@ def main() -> int:
args.timeout,
settle_deadline,
)
if download_duration_s is not None:
logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
else:
logger.info("Download: model already cached")
all_rows: list[dict[str, Any]] = []
@@ -409,6 +413,11 @@ def main() -> int:
"pp_tokens": actual_pp_tokens,
"tg": tg,
"repeat_index": r,
**(
{"download_duration_s": download_duration_s}
if download_duration_s is not None
else {}
),
}
)
runs.append(row)
+18 -6
View File
@@ -289,8 +289,12 @@ def run_planning_phase(
danger_delete: bool,
timeout: float,
settle_deadline: float | None,
) -> None:
"""Check disk space and ensure model is downloaded before benchmarking."""
) -> float | None:
"""Check disk space and ensure model is downloaded before benchmarking.
Returns the wall-clock download duration in seconds if a fresh download
was needed, or None if the model was already cached on all nodes.
"""
# Get model size from /models
models = client.request_json("GET", "/models") or {}
model_bytes = 0
@@ -303,7 +307,7 @@ def run_planning_phase(
logger.warning(
f"Could not determine size for {full_model_id}, skipping disk check"
)
return
return None
# Get nodes from preview
inner = unwrap_instance(preview["instance"])
@@ -314,6 +318,8 @@ def run_planning_phase(
downloads = state.get("downloads", {})
node_disk = state.get("nodeDisk", {})
needs_download = False
for node_id in node_ids:
node_downloads = downloads.get(node_id, [])
@@ -329,6 +335,8 @@ def run_planning_phase(
if already_downloaded:
continue
needs_download = True
# Wait for disk info if settle_deadline is set
disk_info = node_disk.get(node_id, {})
backoff = _SETTLE_INITIAL_BACKOFF_S
@@ -357,16 +365,17 @@ def run_planning_phase(
f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
)
# Delete from smallest to largest
# Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
completed = [
(
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
],
p["DownloadCompleted"]["totalBytes"]["inBytes"],
p["DownloadCompleted"]["total"]["inBytes"],
)
for p in node_downloads
if "DownloadCompleted" in p
and not p["DownloadCompleted"].get("readOnly", False)
]
for del_model, size in sorted(completed, key=lambda x: x[1]):
logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
@@ -379,6 +388,7 @@ def run_planning_phase(
raise RuntimeError(f"Could not free enough space on {node_id}")
# Start downloads (idempotent)
download_t0 = time.perf_counter() if needs_download else None
for node_id in node_ids:
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
shard = runner_to_shard[runner_id]
@@ -421,7 +431,9 @@ def run_planning_phase(
if not done:
all_done = False
if all_done:
return
if download_t0 is not None:
return time.perf_counter() - download_t0
return None
time.sleep(1)
raise TimeoutError("Downloads did not complete in time")
+1
View File
@@ -10,6 +10,7 @@ dependencies = [
"huggingface-hub>=0.33.4",
"tiktoken>=0.12.0",
"jinja2>=3.1.0",
"protobuf>=5.29.0",
]
[build-system]
+65
View File
@@ -202,6 +202,23 @@
filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5));
}
/* Onboarding step 2: connection line between devices */
.onboarding-connection-line {
stroke: oklch(0.85 0.18 85 / 0.5);
stroke-width: 1.5px;
stroke-dasharray: 6, 6;
animation: flowAnimation 1s linear infinite;
filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4));
}
/* Onboarding step 4: red connection line for disconnect */
.onboarding-connection-line-red {
stroke: rgba(220, 38, 38, 0.7);
stroke-width: 1.5px;
stroke-dasharray: 6, 6;
filter: drop-shadow(0 0 2px rgba(220, 38, 38, 0.3));
}
.graph-link-active {
stroke: oklch(0.85 0.18 85 / 0.8);
stroke-width: 2px;
@@ -320,3 +337,51 @@ input:focus, textarea:focus {
transform: translate(400px, 400px);
}
}
/* Onboarding smooth fade-in */
@keyframes onb-fade-in {
from { opacity: 0; transform: translateY(8px); }
to { opacity: 1; transform: none; }
}
/* Opacity-only fade-in — no transform, safe for containers with fixed-position children */
@keyframes onb-fade-opacity {
from { opacity: 0; }
to { opacity: 1; }
}
/* Respect reduced motion preference */
@media (prefers-reduced-motion: reduce) {
.shooting-star,
.shooting-star::before {
animation: none !important;
opacity: 0 !important;
}
.graph-link {
animation: none;
}
.status-pulse {
animation: none;
}
.cursor-blink {
animation: none;
}
.onboarding-connection-line {
animation: none;
}
.onboarding-connection-line-red {
animation: none;
}
[style*="onb-fade-in"],
[style*="onb-fade-opacity"] {
animation: none !important;
opacity: 1 !important;
}
*,
*::before,
*::after {
transition-duration: 0.01ms !important;
animation-duration: 0.01ms !important;
animation-iteration-count: 1 !important;
}
}
+43 -174
View File
@@ -7,8 +7,6 @@
editingImage,
clearEditingImage,
selectedChatModel,
setSelectedChatModel,
instances,
ttftMs,
tps,
totalTokens,
@@ -29,6 +27,19 @@
showModelSelector?: boolean;
modelTasks?: Record<string, string[]>;
modelCapabilities?: Record<string, string[]>;
onSend?: () => void;
onAutoSend?: (
content: string,
files?: {
id: string;
name: string;
type: string;
textContent?: string;
preview?: string;
}[],
) => void;
onOpenModelPicker?: () => void;
modelDisplayOverride?: string;
}
let {
@@ -39,6 +50,10 @@
showModelSelector = false,
modelTasks = {},
modelCapabilities = {},
onSend,
onAutoSend,
onOpenModelPicker,
modelDisplayOverride,
}: Props = $props();
let message = $state("");
@@ -49,27 +64,12 @@
const thinkingEnabled = $derived(thinkingEnabledStore());
let loading = $derived(isLoading());
const currentModel = $derived(selectedChatModel());
const instanceData = $derived(instances());
const currentTtft = $derived(ttftMs());
const currentTps = $derived(tps());
const currentTokens = $derived(totalTokens());
const currentEditingImage = $derived(editingImage());
const isEditMode = $derived(currentEditingImage !== null);
// Custom dropdown state
let isModelDropdownOpen = $state(false);
let dropdownButtonRef: HTMLButtonElement | undefined = $state();
let dropdownPosition = $derived(() => {
if (!dropdownButtonRef || !isModelDropdownOpen)
return { top: 0, left: 0, width: 0 };
const rect = dropdownButtonRef.getBoundingClientRect();
return {
top: rect.top,
left: rect.left,
width: rect.width,
};
});
// Accept all supported file types
const acceptString = getAcceptString(["image", "text", "pdf"]);
@@ -122,73 +122,14 @@
uploadedFiles.length > 0),
);
// Extract available models from running instances
const availableModels = $derived(() => {
const models: Array<{ id: string; label: string; isImageModel: boolean }> =
[];
for (const [, instance] of Object.entries(instanceData)) {
const modelId = getInstanceModelId(instance);
if (
modelId &&
modelId !== "Unknown" &&
!models.some((m) => m.id === modelId)
) {
models.push({
id: modelId,
label: modelId.split("/").pop() || modelId,
isImageModel: modelSupportsImageGeneration(modelId),
});
}
}
return models;
});
// Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
let previousModelIds: Set<string> = new Set();
// Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
$effect(() => {
const models = availableModels();
const currentModelIds = new Set(models.map((m) => m.id));
if (models.length > 0) {
// Find newly added models (in current but not in previous)
const newModels = models.filter((m) => !previousModelIds.has(m.id));
// If no model selected, select the first available
if (!currentModel) {
setSelectedChatModel(models[0].id);
}
// If current model is stale (no longer has a running instance), reset to first available
else if (!models.some((m) => m.id === currentModel)) {
setSelectedChatModel(models[0].id);
}
// If a new model was just added, select it
else if (newModels.length > 0 && previousModelIds.size > 0) {
setSelectedChatModel(newModels[0].id);
}
} else {
// No instances running - clear the selected model
if (currentModel) {
setSelectedChatModel("");
}
}
// Update previous model IDs for next comparison
previousModelIds = currentModelIds;
});
function getInstanceModelId(instanceWrapped: unknown): string {
if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
const keys = Object.keys(instanceWrapped as Record<string, unknown>);
if (keys.length === 1) {
const instance = (instanceWrapped as Record<string, unknown>)[
keys[0]
] as { shardAssignments?: { modelId?: string } };
return instance?.shardAssignments?.modelId || "";
}
return "";
}
// Short label for the currently selected model
const currentModelLabel = $derived(
currentModel
? currentModel.split("/").pop() || currentModel
: modelDisplayOverride
? modelDisplayOverride.split("/").pop() || modelDisplayOverride
: "",
);
async function handleFiles(files: File[]) {
if (files.length === 0) return;
@@ -275,6 +216,15 @@
uploadedFiles = [];
resetTextareaHeight();
// When onAutoSend is provided, the parent controls all send logic
// (including launching non-running models before sending)
if (onAutoSend) {
onAutoSend(content, files);
onSend?.();
setTimeout(() => textareaRef?.focus(), 10);
return;
}
// Use image editing if in edit mode
if (isEditMode && currentEditingImage && content) {
editImage(content, currentEditingImage.imageDataUrl);
@@ -306,6 +256,8 @@
);
}
onSend?.();
// Refocus the textarea after sending
setTimeout(() => textareaRef?.focus(), 10);
}
@@ -416,7 +368,7 @@
{/if}
<!-- Model selector (when enabled) -->
{#if showModelSelector && availableModels().length > 0}
{#if showModelSelector}
<div
class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
>
@@ -425,33 +377,22 @@
class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
>MODEL:</span
>
<!-- Custom dropdown -->
<!-- Model button — opens the full model picker -->
<div class="relative flex-1 max-w-xs">
<button
bind:this={dropdownButtonRef}
type="button"
onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
? 'border-exo-yellow/70'
: ''}"
onclick={() => onOpenModelPicker?.()}
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
>
{#if availableModels().find((m) => m.id === currentModel)}
<span class="text-exo-yellow truncate"
>{availableModels().find((m) => m.id === currentModel)
?.label}</span
>
{:else if availableModels().length > 0}
<span class="text-exo-yellow truncate"
>{availableModels()[0].label}</span
{#if currentModelLabel}
<span class="text-exo-yellow truncate">{currentModelLabel}</span
>
{:else}
<span class="text-exo-light-gray/50">— SELECT MODEL —</span>
{/if}
</button>
<div
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
? 'rotate-180'
: ''}"
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
>
<svg
class="w-3 h-3 text-exo-yellow/60"
@@ -468,78 +409,6 @@
</svg>
</div>
</div>
{#if isModelDropdownOpen}
<!-- Backdrop to close dropdown -->
<button
type="button"
class="fixed inset-0 z-[9998] cursor-default"
onclick={() => (isModelDropdownOpen = false)}
aria-label="Close dropdown"
></button>
<!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
<div
class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
style="bottom: calc(100vh - {dropdownPosition()
.top}px + 4px); left: {dropdownPosition()
.left}px; width: {dropdownPosition().width}px;"
>
<div class="py-1">
{#each availableModels() as model}
<button
type="button"
onclick={() => {
setSelectedChatModel(model.id);
isModelDropdownOpen = false;
}}
class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
model.id
? 'bg-transparent text-exo-yellow'
: 'text-exo-light-gray hover:text-exo-yellow'}"
>
{#if currentModel === model.id}
<svg
class="w-3 h-3 flex-shrink-0"
fill="currentColor"
viewBox="0 0 20 20"
>
<path
fill-rule="evenodd"
d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
clip-rule="evenodd"
/>
</svg>
{:else}
<span class="w-3"></span>
{/if}
{#if model.isImageModel}
<svg
class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
aria-label="Image generation model"
>
<rect
x="3"
y="3"
width="18"
height="18"
rx="2"
ry="2"
/>
<circle cx="8.5" cy="8.5" r="1.5" />
<polyline points="21 15 16 10 5 21" />
</svg>
{/if}
<span class="truncate flex-1">{model.label}</span>
</button>
{/each}
</div>
</div>
{/if}
</div>
<!-- Thinking toggle -->
{#if modelSupportsThinking()}
@@ -807,8 +807,8 @@
>
AWAITING INPUT
</p>
<p class="text-sm sm:text-xs text-exo-light-gray tracking-wider mt-1">
ENTER A QUERY TO BEGIN
<p class="text-xs text-white/30 tracking-wider mt-1.5 font-mono">
Type a message below &middot; Shift+Enter for newline
</p>
</div>
{/if}
@@ -823,6 +823,7 @@
onclick={scrollToBottom}
class="sticky bottom-4 left-1/2 -translate-x-1/2 w-10 h-10 rounded-full bg-exo-dark-gray/90 border border-exo-medium-gray/50 flex items-center justify-center text-exo-light-gray hover:text-exo-yellow hover:border-exo-yellow/50 transition-all shadow-lg cursor-pointer z-10"
title="Scroll to bottom"
aria-label="Scroll to bottom of messages"
>
<svg
class="w-5 h-5"
@@ -0,0 +1,401 @@
<script lang="ts" module>
export interface ChatModelInfo {
id: string;
name: string;
base_model: string;
storage_size_megabytes: number;
capabilities: string[];
family: string;
quantization: string;
}
// Auto mode tier list (for when user just starts typing)
export const AUTO_TIERS: string[][] = [
// Tier 1 (frontier)
["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
// Tier 2 (excellent)
[
"Kimi K2",
"Qwen3 235B",
"MiniMax M2.5",
"Step 3.5 Flash",
"Qwen3 Next 80B",
],
// Tier 3 (great)
[
"GLM 4.7",
"MiniMax M2.1",
"Qwen3 Coder 480B",
"GLM 4.5 Air",
"Llama 3.3 70B",
],
// Tier 4 (good)
["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
// Tier 5 (small/fast)
[
"Llama 3.1 8B",
"GPT-OSS 20B",
"Llama 3.2 3B",
"Qwen3 0.6B",
"Llama 3.2 1B",
],
];
/** Return the tier index (0 = best) for a base_model name. */
export function getAutoTierIndex(baseModel: string): number {
for (let i = 0; i < AUTO_TIERS.length; i++) {
if (AUTO_TIERS[i].includes(baseModel)) return i;
}
return AUTO_TIERS.length; // not in any tier → lowest priority
}
/** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
export function pickAutoModel(
modelList: ChatModelInfo[],
memoryGB: number,
): ChatModelInfo | null {
for (const tier of AUTO_TIERS) {
const candidates: ChatModelInfo[] = [];
for (const baseModel of tier) {
const variants = modelList
.filter(
(m) =>
m.base_model === baseModel &&
(m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
(m.storage_size_megabytes || 0) > 0,
)
.sort(
(a, b) =>
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
);
if (variants[0]) candidates.push(variants[0]);
}
if (candidates.length > 0) {
candidates.sort(
(a, b) =>
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
);
return candidates[0];
}
}
return null;
}
</script>
<script lang="ts">
interface CategoryRecommendation {
category: "coding" | "writing" | "agentic" | "biggest";
label: string;
model: ChatModelInfo | null;
tooltip: string;
}
interface Props {
models: ChatModelInfo[];
clusterLabel: string;
totalMemoryGB: number;
onSelect: (modelId: string, category: string) => void;
onAddModel: () => void;
class?: string;
}
let {
models,
clusterLabel,
totalMemoryGB,
onSelect,
onAddModel,
class: className = "",
}: Props = $props();
// --- Hardcoded Rankings ---
const CODING_RANKING = [
"Qwen3 Coder Next",
"Qwen3 Coder 480B",
"Qwen3 30B",
"GPT-OSS 20B",
"Llama 3.1 8B",
"Llama 3.2 3B",
"Qwen3 0.6B",
];
const WRITING_RANKING = [
"Kimi K2.5",
"Kimi K2",
"Qwen3 Next 80B",
"Llama 3.3 70B",
"MiniMax M2.5",
"GLM 4.5 Air",
"GLM 4.7 Flash",
"GPT-OSS 20B",
"Llama 3.1 8B",
"Llama 3.2 3B",
"Qwen3 0.6B",
];
const AGENTIC_RANKING = [
"DeepSeek V3.1",
"GLM-5",
"Qwen3 235B",
"Step 3.5 Flash",
"GLM 4.7",
"MiniMax M2.1",
"GPT-OSS 120B",
"Llama 3.3 70B",
"Llama 3.1 70B",
"GLM 4.7 Flash",
"GPT-OSS 20B",
"Qwen3 30B",
"Llama 3.1 8B",
"Llama 3.2 3B",
"Qwen3 0.6B",
];
function getModelSizeGB(m: ChatModelInfo): number {
return (m.storage_size_megabytes || 0) / 1024;
}
function fitsInMemory(m: ChatModelInfo): boolean {
return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
}
/** For a given base_model name, find the biggest quant variant that fits in memory. */
function pickBestVariant(baseModel: string): ChatModelInfo | null {
const variants = models
.filter((m) => m.base_model === baseModel && fitsInMemory(m))
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
return variants[0] ?? null;
}
/** Walk a ranked list of base_model names, return the first that has a fitting variant. */
function pickFromRanking(ranking: string[]): ChatModelInfo | null {
for (const baseModel of ranking) {
const pick = pickBestVariant(baseModel);
if (pick) return pick;
}
return null;
}
/** Pick the single biggest model that fits. */
function pickBiggest(): ChatModelInfo | null {
const fitting = models
.filter((m) => fitsInMemory(m))
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
return fitting[0] ?? null;
}
const recommendations = $derived.by((): CategoryRecommendation[] => {
return [
{
category: "coding",
label: "Best for Coding",
model: pickFromRanking(CODING_RANKING),
tooltip:
"Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
},
{
category: "writing",
label: "Best for Writing",
model: pickFromRanking(WRITING_RANKING),
tooltip: "Ranked by creative writing quality and instruction following",
},
{
category: "agentic",
label: "Best Agentic",
model: pickFromRanking(AGENTIC_RANKING),
tooltip: "Ranked by reasoning, planning, and tool-use capability",
},
{
category: "biggest",
label: "Biggest",
model: pickBiggest(),
tooltip: "Largest model that fits in your available memory",
},
];
});
function formatSize(mb: number): string {
const gb = mb / 1024;
if (gb >= 100) return `${Math.round(gb)} GB`;
return `${gb.toFixed(1)} GB`;
}
// Category icons (SVG paths)
const categoryIcons: Record<string, string> = {
coding:
"M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
writing:
"M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
agentic:
"M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
biggest:
"M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
};
let hoveredTooltip = $state<string | null>(null);
let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
function showTooltip(category: string, e: MouseEvent | FocusEvent) {
hoveredTooltip = category;
const target = e.currentTarget as HTMLElement;
const rect = target.getBoundingClientRect();
tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
}
function hideTooltip() {
hoveredTooltip = null;
tooltipAnchor = null;
}
</script>
<div class="flex flex-col items-center justify-center gap-6 {className}">
<!-- Header -->
<div class="text-center">
<p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
Recommended for your
</p>
<p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
</div>
<!-- Category Cards Grid -->
<div class="grid grid-cols-2 gap-3 w-full max-w-md">
{#each recommendations as rec}
{#if rec.model}
<button
type="button"
onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
>
<!-- Category icon + label -->
<div class="flex items-center gap-2 w-full">
<svg
class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="1.5"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={categoryIcons[rec.category]}
/>
</svg>
<span
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
>
{rec.label}
</span>
<!-- Info tooltip -->
<div class="ml-auto flex-shrink-0">
<span
role="button"
tabindex="-1"
class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
onmouseleave={() => hideTooltip()}
onclick={(e: MouseEvent) => {
e.stopPropagation();
if (hoveredTooltip === rec.category) {
hideTooltip();
} else {
showTooltip(rec.category, e);
}
}}
>
<svg
class="w-3.5 h-3.5"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<circle cx="12" cy="12" r="10" />
<path d="M12 16v-4m0-4h.01" />
</svg>
</span>
</div>
</div>
<!-- Model name + size -->
<div class="w-full">
<p class="text-sm text-white font-mono truncate">
{rec.model.base_model}
</p>
<p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
{formatSize(rec.model.storage_size_megabytes)}
{#if rec.model.quantization}
<span class="text-exo-light-gray/40"
>&middot; {rec.model.quantization}</span
>
{/if}
</p>
</div>
</button>
{:else}
<!-- No model fits for this category -->
<div
class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
>
<div class="flex items-center gap-2">
<svg
class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="1.5"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={categoryIcons[rec.category]}
/>
</svg>
<span
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
>{rec.label}</span
>
</div>
<p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
</div>
{/if}
{/each}
</div>
<!-- Add Model Button -->
<button
type="button"
onclick={onAddModel}
class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
>
<svg
class="w-3.5 h-3.5"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
</svg>
Add Model
</button>
<!-- Auto hint -->
<p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
Or just start typing &mdash; we'll pick the best model automatically
</p>
</div>
<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
{#if hoveredTooltip && tooltipAnchor}
{@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
{#if rec}
<div
class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
8}px; transform: translate(-50%, -100%);"
>
{rec.tooltip}
</div>
{/if}
{/if}
+29 -29
View File
@@ -2,7 +2,6 @@
import {
conversations,
activeConversationId,
createConversation,
loadConversation,
deleteConversation,
deleteAllConversations,
@@ -17,9 +16,15 @@
interface Props {
class?: string;
onNewChat?: () => void;
onSelectConversation?: () => void;
}
let { class: className = "" }: Props = $props();
let {
class: className = "",
onNewChat,
onSelectConversation,
}: Props = $props();
const conversationList = $derived(conversations());
const activeId = $derived(activeConversationId());
@@ -42,10 +47,11 @@
);
function handleNewChat() {
createConversation();
onNewChat?.();
}
function handleSelectConversation(id: string) {
onSelectConversation?.();
loadConversation(id);
}
@@ -185,11 +191,7 @@
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (
instanceTag === "MlxIbvInstance" ||
instanceTag === "MlxJacclInstance"
)
instanceType = "MLX RDMA";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
let sharding: string | null = null;
const inst = instance as {
@@ -307,7 +309,7 @@
<div class="py-2">
<div class="px-4 py-2">
<span
class="text-sm text-white/70 font-mono tracking-wider uppercase"
class="text-xs text-exo-light-gray font-mono tracking-wider uppercase"
>
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"}
</span>
@@ -376,39 +378,37 @@
onkeydown={(e) =>
e.key === "Enter" &&
handleSelectConversation(conversation.id)}
class="group w-full flex items-center justify-between p-2 rounded mb-1 transition-all text-left cursor-pointer
class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer
{activeId === conversation.id
? 'bg-transparent border border-exo-yellow/30'
: 'hover:border-exo-yellow/20 border border-transparent'}"
? 'bg-exo-yellow/5 border border-exo-yellow/30'
: 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}"
>
<div class="flex-1 min-w-0 pr-2">
<div
class="text-sm truncate {activeId === conversation.id
class="text-sm font-medium truncate {activeId ===
conversation.id
? 'text-exo-yellow'
: 'text-white/90'}"
: 'text-white'}"
>
{conversation.name}
</div>
<div class="text-sm text-white/50 mt-0.5">
<div class="text-xs text-white/60 mt-0.5">
{formatDate(conversation.updatedAt)}
</div>
<div class="text-sm text-white/70 truncate">
<div class="text-xs text-exo-light-gray truncate">
{info.modelLabel}
</div>
<div class="text-xs text-white/60 font-mono">
Strategy: <span class="text-white/80"
>{info.strategyLabel}</span
>
</div>
{#if stats}
<div class="text-xs text-white/60 font-mono mt-1">
{#if stats.ttftMs}<span class="text-white/40">TTFT</span>
{stats.ttftMs.toFixed(
0,
)}ms{/if}{#if stats.ttftMs && stats.tps}<span
class="text-white/30 mx-1.5"></span
>{/if}{#if stats.tps}{stats.tps.toFixed(1)}
<span class="text-white/40">tok/s</span>{/if}
<div class="text-xs text-white/70 font-mono mt-1">
{#if stats.ttftMs}<span class="text-white/50">TTFT</span>
<span class="text-exo-yellow/80"
>{stats.ttftMs.toFixed(0)}ms</span
>{/if}{#if stats.ttftMs && stats.tps}<span
class="text-white/30 mx-1.5">·</span
>{/if}{#if stats.tps}<span class="text-exo-yellow/80"
>{stats.tps.toFixed(1)}</span
>
<span class="text-white/50">tok/s</span>{/if}
</div>
{/if}
</div>
@@ -0,0 +1,20 @@
<script lang="ts">
import { isConnected } from "$lib/stores/app.svelte";
import { slide } from "svelte/transition";
const connected = $derived(isConnected());
</script>
{#if !connected}
<div
transition:slide={{ duration: 200 }}
class="relative z-50 flex items-center justify-center gap-2 px-4 py-2 bg-red-950/80 border-b border-red-500/30"
role="alert"
aria-live="assertive"
>
<div class="w-2 h-2 bg-red-500 rounded-full animate-pulse"></div>
<span class="text-xs font-mono text-red-300 tracking-wider uppercase">
Connection lost &mdash; Reconnecting to backend&hellip;
</span>
</div>
{/if}
@@ -0,0 +1,277 @@
<script lang="ts">
/**
* DeviceIcon — renders a device icon as an SVG <g> element.
* Uses the exact same proportional math as TopologyGraph.svelte
* so that devices look identical in both the topology view and
* the onboarding animation.
*
* Must be placed inside an <svg> element.
*/
interface Props {
/** "macbook pro" | "mac studio" | "mac mini" etc. */
deviceType: string;
/** Center X coordinate in SVG space */
cx: number;
/** Center Y coordinate in SVG space */
cy: number;
/** Base sizing factor (equivalent to TopologyGraph's nodeRadius) */
size?: number;
/** RAM usage 0100 */
ramPercent?: number;
/** Unique id suffix for clip-path ids */
uid?: string;
}
let {
deviceType,
cx,
cy,
size = 60,
ramPercent = 60,
uid = "dev",
}: Props = $props();
// Apple logo path — same constant used by TopologyGraph
const APPLE_LOGO_PATH =
"M788.1 340.9c-5.8 4.5-108.2 62.2-108.2 190.5 0 148.4 130.3 200.9 134.2 202.2-.6 3.2-20.7 71.9-68.7 141.9-42.8 61.6-87.5 123.1-155.5 123.1s-85.5-39.5-164-39.5c-76.5 0-103.7 40.8-165.9 40.8s-105.6-57-155.5-127C46.7 790.7 0 663 0 541.8c0-194.4 126.4-297.5 250.8-297.5 66.1 0 121.2 43.4 162.7 43.4 39.5 0 101.1-46 176.3-46 28.5 0 130.9 2.6 198.3 99.2zm-234-181.5c31.1-36.9 53.1-88.1 53.1-139.3 0-7.1-.6-14.3-1.9-20.1-50.6 1.9-110.8 33.7-147.1 75.8-28.5 32.4-55.1 83.6-55.1 135.5 0 7.8 1.3 15.6 1.9 18.1 3.2.6 8.4 1.3 13.6 1.3 45.4 0 102.5-30.4 135.5-71.3z";
const LOGO_NATIVE_WIDTH = 814;
const LOGO_NATIVE_HEIGHT = 1000;
const wireColor = "rgba(179,179,179,0.8)";
const strokeWidth = 1.5;
const modelLower = $derived(deviceType.toLowerCase());
// ── Mac Studio dimensions (same ratios as TopologyGraph) ──
const studioW = $derived(size * 1.25);
const studioH = $derived(size * 0.85);
const studioX = $derived(cx - studioW / 2);
const studioY = $derived(cy - studioH / 2);
const studioCorner = 4;
const studioTopH = $derived(studioH * 0.15);
// Studio front panel details
const studioSlotH = $derived(studioH * 0.14);
const studioVSlotW = $derived(studioW * 0.05);
const studioVSlotY = $derived(
studioY + studioTopH + (studioH - studioTopH) * 0.6,
);
const studioVSlot1X = $derived(studioX + studioW * 0.18);
const studioVSlot2X = $derived(studioX + studioW * 0.28);
const studioHSlotW = $derived(studioW * 0.2);
const studioHSlotX = $derived(studioX + studioW * 0.5 - studioHSlotW / 2);
// Studio memory fill
const studioMemTotalH = $derived(studioH - studioTopH);
const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
// ── MacBook dimensions (same ratios as TopologyGraph) ──
const mbW = $derived((size * 1.6 * 0.85) / 1.15);
const mbH = $derived(size * 0.85);
const mbX = $derived(cx - mbW / 2);
const mbY = $derived(cy - mbH / 2);
const mbScreenH = $derived(mbH * 0.7);
const mbBaseH = $derived(mbH * 0.3);
const mbScreenW = $derived(mbW * 0.85);
const mbScreenX = $derived(cx - mbScreenW / 2);
const mbBezel = 3;
// MacBook memory fill
const mbMemTotalH = $derived(mbScreenH - mbBezel * 2);
const mbMemH = $derived((ramPercent / 100) * mbMemTotalH);
// Apple logo sizing
const mbLogoTargetH = $derived(mbScreenH * 0.22);
const mbLogoScale = $derived(mbLogoTargetH / LOGO_NATIVE_HEIGHT);
const mbLogoX = $derived(cx - (LOGO_NATIVE_WIDTH * mbLogoScale) / 2);
const mbLogoY = $derived(
mbY + mbScreenH / 2 - (LOGO_NATIVE_HEIGHT * mbLogoScale) / 2,
);
// MacBook base (trapezoidal)
const mbBaseY = $derived(mbY + mbScreenH);
const mbBaseTopW = $derived(mbScreenW);
const mbBaseBottomW = $derived(mbW);
const mbBaseTopX = $derived(cx - mbBaseTopW / 2);
const mbBaseBottomX = $derived(cx - mbBaseBottomW / 2);
// Keyboard
const mbKbX = $derived(mbBaseTopX + 6);
const mbKbY = $derived(mbBaseY + 3);
const mbKbW = $derived(mbBaseTopW - 12);
const mbKbH = $derived(mbBaseH * 0.55);
// Trackpad
const mbTpW = $derived(mbBaseTopW * 0.4);
const mbTpX = $derived(cx - mbTpW / 2);
const mbTpY = $derived(mbBaseY + mbKbH + 5);
const mbTpH = $derived(mbBaseH * 0.3);
// Clip IDs
const screenClipId = $derived(`di-screen-${uid}`);
const studioClipId = $derived(`di-studio-${uid}`);
</script>
{#if modelLower === "mac studio" || modelLower === "mac mini"}
<!-- Mac Studio / Mac Mini -->
<defs>
<clipPath id={studioClipId}>
<rect
x={studioX}
y={studioY + studioTopH}
width={studioW}
height={studioH - studioTopH}
rx={studioCorner - 1}
/>
</clipPath>
</defs>
<!-- Main body -->
<rect
x={studioX}
y={studioY}
width={studioW}
height={studioH}
rx={studioCorner}
fill="#1a1a1a"
stroke={wireColor}
stroke-width={strokeWidth}
/>
<!-- Memory fill -->
{#if ramPercent > 0}
<rect
x={studioX}
y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
width={studioW}
height={studioMemH}
fill="rgba(255,215,0,0.75)"
clip-path="url(#{studioClipId})"
/>
{/if}
<!-- Top surface divider -->
<line
x1={studioX}
y1={studioY + studioTopH}
x2={studioX + studioW}
y2={studioY + studioTopH}
stroke="rgba(179,179,179,0.3)"
stroke-width="0.5"
/>
<!-- Front panel: vertical slots -->
<rect
x={studioVSlot1X - studioVSlotW / 2}
y={studioVSlotY}
width={studioVSlotW}
height={studioSlotH}
fill="rgba(0,0,0,0.35)"
rx="1.5"
/>
<rect
x={studioVSlot2X - studioVSlotW / 2}
y={studioVSlotY}
width={studioVSlotW}
height={studioSlotH}
fill="rgba(0,0,0,0.35)"
rx="1.5"
/>
<!-- Horizontal slot (SD card) -->
<rect
x={studioHSlotX}
y={studioVSlotY}
width={studioHSlotW}
height={studioSlotH * 0.6}
fill="rgba(0,0,0,0.35)"
rx="1"
/>
{:else}
<!-- MacBook Pro -->
<defs>
<clipPath id={screenClipId}>
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel}
width={mbScreenW - mbBezel * 2}
height={mbScreenH - mbBezel * 2}
rx="2"
/>
</clipPath>
</defs>
<!-- Screen outer frame -->
<rect
x={mbScreenX}
y={mbY}
width={mbScreenW}
height={mbScreenH}
rx="3"
fill="#1a1a1a"
stroke={wireColor}
stroke-width={strokeWidth}
/>
<!-- Screen inner (dark) -->
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel}
width={mbScreenW - mbBezel * 2}
height={mbScreenH - mbBezel * 2}
rx="2"
fill="#0a0a12"
/>
<!-- Memory fill on screen -->
{#if ramPercent > 0}
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
width={mbScreenW - mbBezel * 2}
height={mbMemH}
fill="rgba(255,215,0,0.85)"
clip-path="url(#{screenClipId})"
/>
{/if}
<!-- Apple logo -->
<path
d={APPLE_LOGO_PATH}
transform="translate({mbLogoX}, {mbLogoY}) scale({mbLogoScale})"
fill="#FFFFFF"
opacity="0.9"
/>
<!-- Keyboard base (trapezoidal) -->
<path
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
fill="#2c2c2c"
stroke={wireColor}
stroke-width="1"
/>
<!-- Keyboard area -->
<rect
x={mbKbX}
y={mbKbY}
width={mbKbW}
height={mbKbH}
fill="rgba(0,0,0,0.2)"
rx="2"
/>
<!-- Trackpad -->
<rect
x={mbTpX}
y={mbTpY}
width={mbTpW}
height={mbTpH}
fill="rgba(255,255,255,0.08)"
rx="2"
/>
{/if}
+65 -20
View File
@@ -6,6 +6,10 @@
export let showSidebarToggle = false;
export let sidebarVisible = true;
export let onToggleSidebar: (() => void) | null = null;
export let downloadProgress: {
count: number;
percentage: number;
} | null = null;
function handleHome(): void {
if (onHome) {
@@ -33,13 +37,17 @@
<div class="absolute left-6 top-1/2 -translate-y-1/2">
<button
onclick={handleToggleSidebar}
class="p-2 rounded border border-exo-medium-gray/40 hover:border-exo-yellow/50 transition-colors cursor-pointer"
class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer"
title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
aria-label={sidebarVisible
? "Hide conversation sidebar"
: "Show conversation sidebar"}
aria-pressed={sidebarVisible}
>
<svg
class="w-5 h-5 {sidebarVisible
? 'text-exo-yellow'
: 'text-exo-medium-gray'}"
: 'text-exo-light-gray'}"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
@@ -75,18 +83,19 @@
<img
src="/exo-logo.png"
alt="EXO"
class="h-18 drop-shadow-[0_0_20px_rgba(255,215,0,0.5)]"
class="h-18 drop-shadow-[0_0_4px_rgba(255,215,0,0.3)]"
/>
</button>
<!-- Right: Home + Downloads -->
<div
<nav
class="absolute right-6 top-1/2 -translate-y-1/2 flex items-center gap-4"
aria-label="Main navigation"
>
{#if showHome}
<button
onclick={handleHome}
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
title="Back to topology view"
>
<svg
@@ -107,23 +116,59 @@
{/if}
<a
href="/#/downloads"
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
title="View downloads overview"
>
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M12 3v12" />
<path d="M7 12l5 5 5-5" />
<path d="M5 21h14" />
</svg>
{#if downloadProgress}
<!-- Compact download progress indicator -->
<div class="relative w-4 h-4 flex-shrink-0">
<svg class="w-4 h-4 -rotate-90" viewBox="0 0 20 20">
<circle
cx="10"
cy="10"
r="8"
fill="none"
stroke="currentColor"
stroke-width="2"
opacity="0.2"
/>
<circle
cx="10"
cy="10"
r="8"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-dasharray={2 * Math.PI * 8}
stroke-dashoffset={2 *
Math.PI *
8 *
(1 - downloadProgress.percentage / 100)}
class="text-blue-400 transition-all duration-300"
/>
</svg>
<div
class="absolute inset-0 flex items-center justify-center text-[6px] font-mono text-blue-400"
>
{downloadProgress.count}
</div>
</div>
{:else}
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M12 3v12" />
<path d="M7 12l5 5 5-5" />
<path d="M5 21h14" />
</svg>
{/if}
Downloads
</a>
</div>
</nav>
</header>
+29 -3
View File
@@ -21,7 +21,7 @@
} | null;
nodes?: Record<string, NodeInfo>;
sharding?: "Pipeline" | "Tensor";
runtime?: "MlxRing" | "MlxIbv" | "MlxJaccl";
runtime?: "MlxRing" | "MlxJaccl";
onLaunch?: () => void;
tags?: string[];
apiPreview?: PlacementPreview | null;
@@ -348,7 +348,7 @@
// Debug mode state
const isDebugMode = $derived(debugMode());
const topology = $derived(topologyData());
const isRdma = $derived(runtime === "MlxIbv" || runtime === "MlxJaccl");
const isRdma = $derived(runtime === "MlxJaccl");
// Get interface name for an IP from node data
function getInterfaceForIp(nodeId: string, ip?: string): string | null {
@@ -567,20 +567,46 @@
<div class="flex items-center gap-1.5 mb-2">
<span
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
title={sharding === "Pipeline"
? "Pipeline: splits model into sequential stages across devices. Lower network overhead."
: "Tensor: splits each layer across devices. Best with high-bandwidth connections (Thunderbolt)."}
>
{sharding}
</span>
<span
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
title={runtime === "MlxRing"
? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
>
{runtime === "MlxRing"
? "MLX Ring"
: runtime === "MlxIbv" || runtime === "MlxJaccl"
: runtime === "MlxJaccl"
? "MLX RDMA"
: runtime}
</span>
</div>
<!-- Download Status -->
{#if isDownloading && progress}
<div class="mb-2 space-y-1">
<div class="flex items-center justify-between text-xs font-mono">
<span class="text-blue-400 tracking-wider uppercase">Downloading</span
>
<span class="text-white/60"
>{percentage.toFixed(1)}% &middot; {formatSpeed(progress.speed)}
&middot; {formatEta(progress.etaMs)}</span
>
</div>
<div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
<div
class="h-full bg-blue-500/70 transition-all duration-300"
style="width: {percentage}%"
></div>
</div>
</div>
{/if}
<!-- Mini Topology Preview -->
{#if placementPreview().nodes.length > 0}
{@const preview = placementPreview()}
@@ -6,6 +6,7 @@
capabilities: string[];
sizeRange: { min: number; max: number } | null;
downloadedOnly: boolean;
readyOnly: boolean;
}
type ModelFilterPopoverProps = {
@@ -190,34 +191,58 @@
</div>
</div>
<!-- Downloaded only -->
<!-- Availability filters -->
<div>
<h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
<div class="flex flex-wrap gap-1.5">
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
<span class="ml-1">Downloaded</span>
</button>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
<span class="ml-1">Downloaded</span>
</button>
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, readyOnly: !filters.readyOnly })}
>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<circle cx="12" cy="12" r="10" />
<path d="m9 12 2 2 4-4" />
</svg>
<span class="ml-1">Ready</span>
</button>
</div>
</div>
<!-- Size range -->
@@ -41,6 +41,7 @@
onShowInfo: (group: ModelGroup) => void;
downloadStatusMap?: Map<string, DownloadAvailability>;
launchedAt?: number;
instanceStatuses?: Record<string, { status: string; statusClass: string }>;
};
let {
@@ -56,6 +57,7 @@
onShowInfo,
downloadStatusMap,
launchedAt,
instanceStatuses = {},
}: ModelPickerGroupProps = $props();
// Group-level download status: show if any variant is downloaded
@@ -92,6 +94,12 @@
const anyVariantFits = $derived(
group.variants.some((v) => canModelFit(v.id)),
);
// Check if any variant has an active instance (ready, loading, downloading)
const anyVariantHasInstance = $derived(
instanceStatuses
? group.variants.some((v) => instanceStatuses[v.id] != null)
: false,
);
const groupFitStatus = $derived.by((): ModelFitStatus => {
let hasClusterCapacityOnly = false;
for (const variant of group.variants) {
@@ -122,16 +130,35 @@
!group.hasMultipleVariants &&
group.variants.some((v) => v.id === selectedModelId),
);
// Group-level instance status: show the "best" status across all variants
const groupInstanceStatus = $derived.by(() => {
if (!instanceStatuses) return null;
const readyStatuses = ["READY", "LOADED", "RUNNING"];
const loadingStatuses = ["LOADING", "WARMING UP"];
let bestStatus: { status: string; statusClass: string } | null = null;
for (const variant of group.variants) {
const s = instanceStatuses[variant.id];
if (!s) continue;
if (readyStatuses.includes(s.status)) return s; // Ready is best
if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
bestStatus = s;
}
}
return bestStatus;
});
</script>
<div
class="border-b border-white/5 last:border-b-0 {!anyVariantFits
class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
!anyVariantHasInstance
? 'opacity-50'
: ''}"
>
<!-- Main row -->
<div
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
anyVariantHasInstance
? 'hover:bg-white/5 cursor-pointer'
: 'cursor-not-allowed'} {isMainSelected
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
@@ -141,7 +168,7 @@
onToggleExpand();
} else {
const modelId = group.variants[0]?.id;
if (modelId && canModelFit(modelId)) {
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
onSelectModel(modelId);
}
}
@@ -155,7 +182,7 @@
onToggleExpand();
} else {
const modelId = group.variants[0]?.id;
if (modelId && canModelFit(modelId)) {
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
onSelectModel(modelId);
}
}
@@ -346,6 +373,47 @@
</span>
{/if}
<!-- Instance status badge -->
{#if groupInstanceStatus}
{#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
<span class="flex-shrink-0" title="Running">
<svg
class="w-3 h-3 text-green-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{:else if groupInstanceStatus.status === "DOWNLOADING"}
<span class="flex-shrink-0 animate-pulse" title="Downloading">
<svg
class="w-3.5 h-3.5 text-blue-400"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
<polyline points="7 10 12 15 17 10" />
<line x1="12" y1="15" x2="12" y2="3" />
</svg>
</span>
{:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
<span class="flex-shrink-0 animate-pulse" title="Loading">
<svg
class="w-3 h-3 text-yellow-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{/if}
{/if}
<!-- Check mark if selected (single-variant) -->
{#if isMainSelected}
<svg
@@ -420,20 +488,30 @@
{#each group.variants as variant}
{@const fitStatus = getModelFitStatus(variant.id)}
{@const modelCanFit = canModelFit(variant.id)}
{@const variantHasInstance = instanceStatuses[variant.id] != null}
{@const isSelected = selectedModelId === variant.id}
<button
type="button"
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
<div
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
!variantHasInstance
? 'opacity-50 cursor-not-allowed'
: 'cursor-pointer'} {isSelected
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
: 'border-l-2 border-transparent'}"
disabled={!modelCanFit}
role="button"
tabindex="0"
onclick={() => {
if (modelCanFit) {
if (modelCanFit || variantHasInstance) {
onSelectModel(variant.id);
}
}}
onkeydown={(e) => {
if (e.key === "Enter" || e.key === " ") {
e.preventDefault();
if (modelCanFit) {
onSelectModel(variant.id);
}
}
}}
>
<!-- Quantization badge -->
<span
@@ -478,6 +556,48 @@
{/if}
{/if}
<!-- Instance status badge -->
{#if instanceStatuses[variant.id]}
{@const instStatus = instanceStatuses[variant.id]}
{#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
<span class="flex-shrink-0" title="Running">
<svg
class="w-3 h-3 text-green-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{:else if instStatus.status === "DOWNLOADING"}
<span class="flex-shrink-0 animate-pulse" title="Downloading">
<svg
class="w-3.5 h-3.5 text-blue-400"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
<polyline points="7 10 12 15 17 10" />
<line x1="12" y1="15" x2="12" y2="3" />
</svg>
</span>
{:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
<span class="flex-shrink-0 animate-pulse" title="Loading">
<svg
class="w-3 h-3 text-yellow-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{/if}
{/if}
<!-- Check mark if selected -->
{#if isSelected}
<svg
@@ -490,7 +610,36 @@
/>
</svg>
{/if}
</button>
<!-- Info button -->
<button
type="button"
class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
onclick={(e) => {
e.stopPropagation();
onShowInfo({
id: variant.id,
name: variant.name || variant.id,
capabilities: group.capabilities,
family: group.family,
variants: [variant],
smallestVariant: variant,
hasMultipleVariants: false,
});
}}
title="View variant details"
>
<svg
class="w-4 h-4 text-white/30 hover:text-white/50"
viewBox="0 0 24 24"
fill="currentColor"
>
<path
d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
/>
</svg>
</button>
</div>
{/each}
</div>
{/if}
@@ -36,6 +36,7 @@
capabilities: string[];
sizeRange: { min: number; max: number } | null;
downloadedOnly: boolean;
readyOnly: boolean;
}
interface HuggingFaceModel {
@@ -49,6 +50,11 @@
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
export type InstanceStatus = {
status: string;
statusClass: string;
};
type ModelPickerModalProps = {
isOpen: boolean;
models: ModelInfo[];
@@ -75,6 +81,7 @@
macmon_info?: { memory?: { ram_total?: number } };
}
>;
instanceStatuses?: Record<string, InstanceStatus>;
};
let {
@@ -96,6 +103,7 @@
usedMemoryGB,
downloadsData,
topologyNodes,
instanceStatuses = {},
}: ModelPickerModalProps = $props();
// Local state
@@ -107,6 +115,7 @@
capabilities: [],
sizeRange: null,
downloadedOnly: false,
readyOnly: false,
});
let infoGroup = $state<ModelGroup | null>(null);
@@ -440,6 +449,16 @@
);
}
// Filter to ready/running models only
if (filters.readyOnly) {
result = result.filter((g) =>
g.variants.some((v) => {
const s = instanceStatuses[v.id];
return s && s.statusClass === "ready";
}),
);
}
// Sort: fits-now first, then fits-cluster-capacity, then too-large
result.sort((a, b) => {
const getGroupFitRank = (group: ModelGroup): number => {
@@ -512,6 +531,18 @@
);
});
// Split filtered groups into recommended (fits_now) and others for visual separation
const recommendedGroups = $derived(
filteredGroups.filter((g) =>
g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
),
);
const otherGroups = $derived(
filteredGroups.filter(
(g) => !g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
),
);
function toggleGroupExpanded(groupId: string) {
const next = new Set(expandedGroups);
if (next.has(groupId)) {
@@ -538,13 +569,19 @@
}
function clearFilters() {
filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
filters = {
capabilities: [],
sizeRange: null,
downloadedOnly: false,
readyOnly: false,
};
}
const hasActiveFilters = $derived(
filters.capabilities.length > 0 ||
filters.sizeRange !== null ||
filters.downloadedOnly,
filters.downloadedOnly ||
filters.readyOnly,
);
</script>
@@ -813,6 +850,7 @@
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
{instanceStatuses}
/>
{/each}
{/if}
@@ -840,7 +878,34 @@
{/if}
</div>
{:else}
{#each filteredGroups as group}
<!-- Recommended for your cluster -->
{#if recommendedGroups.length > 0 && otherGroups.length > 0 && !searchQuery.trim()}
<div
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-green-950/60 border-b border-green-500/20 backdrop-blur-sm"
>
<svg
class="w-3.5 h-3.5 text-green-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<span
class="text-xs font-mono text-green-400 tracking-wider uppercase"
>Recommended for your cluster</span
>
<span class="text-xs font-mono text-green-400/50"
>— fits in available memory</span
>
</div>
{/if}
{#each recommendedGroups as group}
<ModelPickerGroup
{group}
isExpanded={expandedGroups.has(group.id)}
@@ -853,6 +918,34 @@
{onToggleFavorite}
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
{instanceStatuses}
/>
{/each}
<!-- Other models -->
{#if otherGroups.length > 0 && recommendedGroups.length > 0 && !searchQuery.trim()}
<div
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-exo-dark-gray/80 border-y border-exo-medium-gray/20 backdrop-blur-sm"
>
<span
class="text-xs font-mono text-white/40 tracking-wider uppercase"
>Other models</span
>
</div>
{/if}
{#each otherGroups as group}
<ModelPickerGroup
{group}
isExpanded={expandedGroups.has(group.id)}
isFavorite={favorites.has(group.id)}
{selectedModelId}
{canModelFit}
{getModelFitStatus}
onToggleExpand={() => toggleGroupExpanded(group.id)}
onSelectModel={handleSelect}
{onToggleFavorite}
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
{instanceStatuses}
/>
{/each}
{/if}
@@ -875,6 +968,11 @@
>Downloaded</span
>
{/if}
{#if filters.readyOnly}
<span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
>Ready</span
>
{/if}
{#if filters.sizeRange}
<span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
{filters.sizeRange.min}GB - {filters.sizeRange.max}GB
@@ -14,6 +14,21 @@
: 0,
);
const etaText = $derived.by(() => {
if (progress.processed <= 0 || progress.total <= 0) return null;
const elapsedMs = performance.now() - progress.startedAt;
if (elapsedMs < 200) return null; // need a minimum sample window
const tokensPerMs = progress.processed / elapsedMs;
const remainingTokens = progress.total - progress.processed;
const remainingMs = remainingTokens / tokensPerMs;
const remainingSec = Math.ceil(remainingMs / 1000);
if (remainingSec <= 0) return null;
if (remainingSec < 60) return `~${remainingSec}s remaining`;
const mins = Math.floor(remainingSec / 60);
const secs = remainingSec % 60;
return `~${mins}m ${secs}s remaining`;
});
function formatTokenCount(count: number | undefined): string {
if (count == null) return "0";
if (count >= 1000) {
@@ -40,8 +55,11 @@
style="width: {percentage}%"
></div>
</div>
<div class="text-right text-xs text-exo-light-gray/70 mt-0.5 font-mono">
{percentage}%
<div
class="flex items-center justify-between text-xs text-exo-light-gray/70 mt-0.5 font-mono"
>
<span>{etaText ?? ""}</span>
<span>{percentage}%</span>
</div>
</div>
@@ -0,0 +1,117 @@
<script lang="ts">
import { toasts, dismissToast, type Toast } from "$lib/stores/toast.svelte";
import { fly, fade } from "svelte/transition";
import { flip } from "svelte/animate";
const items = $derived(toasts());
const typeStyles: Record<
Toast["type"],
{ border: string; icon: string; iconColor: string }
> = {
success: {
border: "border-l-green-500",
icon: "M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z",
iconColor: "text-green-400",
},
error: {
border: "border-l-red-500",
icon: "M12 9v3.75m9-.75a9 9 0 11-18 0 9 9 0 0118 0zm-9 3.75h.008v.008H12v-.008z",
iconColor: "text-red-400",
},
warning: {
border: "border-l-yellow-500",
icon: "M12 9v3.75m-9.303 3.376c-.866 1.5.217 3.374 1.948 3.374h14.71c1.73 0 2.813-1.874 1.948-3.374L13.949 3.378c-.866-1.5-3.032-1.5-3.898 0L2.697 16.126z",
iconColor: "text-yellow-400",
},
info: {
border: "border-l-blue-500",
icon: "M11.25 11.25l.041-.02a.75.75 0 011.063.852l-.708 2.836a.75.75 0 001.063.853l.041-.021M21 12a9 9 0 11-18 0 9 9 0 0118 0zm-9-3.75h.008v.008H12V8.25z",
iconColor: "text-blue-400",
},
};
</script>
{#if items.length > 0}
<div
class="fixed bottom-6 right-6 z-[9999] flex flex-col gap-2 pointer-events-none"
role="log"
aria-live="polite"
aria-label="Notifications"
>
{#each items as toast (toast.id)}
{@const style = typeStyles[toast.type]}
<div
class="pointer-events-auto max-w-sm w-80 bg-exo-dark-gray/95 backdrop-blur-sm border border-exo-medium-gray/60 border-l-[3px] {style.border} rounded shadow-lg shadow-black/40"
in:fly={{ x: 80, duration: 250 }}
out:fade={{ duration: 150 }}
animate:flip={{ duration: 200 }}
role="alert"
>
<div class="flex items-start gap-3 px-4 py-3">
<!-- Icon -->
<svg
class="w-5 h-5 flex-shrink-0 mt-0.5 {style.iconColor}"
fill="none"
viewBox="0 0 24 24"
stroke-width="1.5"
stroke="currentColor"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={style.icon}
/>
</svg>
<!-- Message -->
<p class="flex-1 text-sm text-white/90 font-mono leading-snug">
{toast.message}
</p>
<!-- Dismiss button -->
<button
onclick={() => dismissToast(toast.id)}
class="flex-shrink-0 p-0.5 text-white/40 hover:text-white/80 transition-colors cursor-pointer"
aria-label="Dismiss notification"
>
<svg
class="w-4 h-4"
fill="none"
viewBox="0 0 24 24"
stroke-width="2"
stroke="currentColor"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M6 18L18 6M6 6l12 12"
/>
</svg>
</button>
</div>
<!-- Auto-dismiss progress bar -->
{#if toast.duration > 0}
<div class="h-0.5 bg-white/5 rounded-b overflow-hidden">
<div
class="h-full {style.border.replace('border-l-', 'bg-')}/60"
style="animation: shrink {toast.duration}ms linear forwards"
></div>
</div>
{/if}
</div>
{/each}
</div>
{/if}
<style>
@keyframes shrink {
from {
width: 100%;
}
to {
width: 0%;
}
}
</style>
+1
View File
@@ -12,3 +12,4 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
export { default as ChatModelSelector } from "./ChatModelSelector.svelte";
+101 -31
View File
@@ -168,7 +168,7 @@ export interface ModelDownloadStatus {
export interface PlacementPreview {
model_id: string;
sharding: "Pipeline" | "Tensor";
instance_meta: "MlxRing" | "MlxIbv" | "MlxJaccl";
instance_meta: "MlxRing" | "MlxJaccl";
instance: unknown | null;
memory_delta_by_node: Record<string, number> | null;
error: string | null;
@@ -219,7 +219,6 @@ interface RawStateResponse {
string,
{
MlxRingInstance?: Instance;
MlxIbvInstance?: Instance;
MlxJacclInstance?: Instance;
}
>;
@@ -250,6 +249,11 @@ interface RawStateResponse {
>;
// Thunderbolt bridge cycles (nodes with bridge enabled forming loops)
thunderboltBridgeCycles?: string[][];
// Disk usage per node
nodeDisk?: Record<
string,
{ total: { inBytes: number }; available: { inBytes: number } }
>;
}
export interface MessageAttachment {
@@ -276,6 +280,8 @@ export interface TokenData {
export interface PrefillProgress {
processed: number;
total: number;
/** Timestamp (performance.now()) when prefill started. */
startedAt: number;
}
export interface Message {
@@ -583,6 +589,12 @@ class AppStore {
// Image editing state
editingImage = $state<EditingImage | null>(null);
/** True when the backend is reachable. */
isConnected = $state<boolean>(true);
/** Number of consecutive fetch failures. */
private consecutiveFailures = 0;
private static readonly CONNECTION_LOST_THRESHOLD = 3;
private fetchInterval: ReturnType<typeof setInterval> | null = null;
private previewsInterval: ReturnType<typeof setInterval> | null = null;
private lastConversationPersistTs = 0;
@@ -835,6 +847,10 @@ class AppStore {
this.thinkingEnabled = conversation.enableThinking ?? true;
this.refreshConversationModelFromInstances();
// Sync global selection to the loaded conversation's model so reactive
// effects in +page.svelte can determine the correct chat launch state.
this.selectedChatModel = conversation.modelId || "";
return true;
}
@@ -905,11 +921,7 @@ class AppStore {
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (
instanceTag === "MlxIbvInstance" ||
instanceTag === "MlxJacclInstance"
)
instanceType = "MLX RDMA";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
let sharding: string | null = null;
const inst = instance as {
@@ -1288,7 +1300,19 @@ class AppStore {
// Thunderbolt bridge status per node
this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
this.lastUpdate = Date.now();
// Connection recovered
if (!this.isConnected) {
this.isConnected = true;
}
this.consecutiveFailures = 0;
} catch (error) {
this.consecutiveFailures++;
if (
this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD &&
this.isConnected
) {
this.isConnected = false;
}
console.error("Error fetching state:", error);
}
}
@@ -1652,11 +1676,12 @@ class AppStore {
if (!reader) throw new Error("No response body");
let fullContent = prefixText;
let streamedThinking = "";
const collectedTokens: TokenData[] = [...tokensToKeep];
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string };
delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -1677,6 +1702,7 @@ class AppStore {
(parsed) => {
const choice = parsed.choices?.[0];
const delta = choice?.delta?.content;
const thinkingDelta = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -1695,7 +1721,11 @@ class AppStore {
}
}
if (delta) {
if (thinkingDelta) {
streamedThinking += thinkingDelta;
}
if (delta || thinkingDelta) {
if (firstTokenTime === null) {
firstTokenTime = performance.now();
this.ttftMs = firstTokenTime - requestStartTime;
@@ -1709,9 +1739,14 @@ class AppStore {
this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000;
}
fullContent += delta;
const { displayContent, thinkingContent } =
if (delta) {
fullContent += delta;
}
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(fullContent);
const combinedThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
if (this.activeConversationId === targetConversationId) {
this.currentResponse = displayContent;
@@ -1723,7 +1758,7 @@ class AppStore {
messageId,
(m) => {
m.content = displayContent;
m.thinking = thinkingContent || undefined;
m.thinking = combinedThinking || undefined;
m.tokens = [...collectedTokens];
},
);
@@ -1735,11 +1770,14 @@ class AppStore {
// Final update
if (this.conversationExists(targetConversationId)) {
const { displayContent, thinkingContent } =
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(fullContent);
const finalThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
this.updateConversationMessage(targetConversationId, messageId, (m) => {
m.content = displayContent;
m.thinking = thinkingContent || undefined;
m.thinking = finalThinking || undefined;
m.tokens = [...collectedTokens];
if (this.ttftMs !== null) m.ttftMs = this.ttftMs;
if (this.tps !== null) m.tps = this.tps;
@@ -1815,7 +1853,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content =
"Error: No model available. Please launch an instance first.";
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.";
},
);
this.syncActiveMessagesIfNeeded(targetConversationId);
@@ -1847,11 +1885,12 @@ class AppStore {
}
let streamedContent = "";
let streamedThinking = "";
const collectedTokens: TokenData[] = [];
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string };
delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -1872,6 +1911,7 @@ class AppStore {
(parsed) => {
const choice = parsed.choices?.[0];
const delta = choice?.delta?.content;
const thinkingDelta = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -1890,10 +1930,19 @@ class AppStore {
}
}
if (delta) {
streamedContent += delta;
const { displayContent, thinkingContent } =
if (thinkingDelta) {
streamedThinking += thinkingDelta;
}
if (delta || thinkingDelta) {
if (delta) {
streamedContent += delta;
}
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const combinedThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
// Only update currentResponse if target conversation is active
if (this.activeConversationId === targetConversationId) {
@@ -1906,7 +1955,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = combinedThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -1918,14 +1967,17 @@ class AppStore {
// Final cleanup of the message (if conversation still exists)
if (this.conversationExists(targetConversationId)) {
const { displayContent, thinkingContent } =
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const finalThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
this.updateConversationMessage(
targetConversationId,
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = finalThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -2272,7 +2324,7 @@ class AppStore {
const modelToUse = this.getModelForRequest();
if (!modelToUse) {
throw new Error(
"No model selected and no running instances available. Please launch an instance first.",
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.",
);
}
@@ -2317,10 +2369,11 @@ class AppStore {
}
let streamedContent = "";
let streamedThinking = "";
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string };
delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -2348,6 +2401,7 @@ class AppStore {
const choice = parsed.choices?.[0];
const tokenContent = choice?.delta?.content;
const thinkingContent = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -2366,7 +2420,11 @@ class AppStore {
}
}
if (tokenContent) {
if (thinkingContent) {
streamedThinking += thinkingContent;
}
if (tokenContent || thinkingContent) {
// Track first token for TTFT
if (firstTokenTime === null) {
firstTokenTime = performance.now();
@@ -2383,11 +2441,16 @@ class AppStore {
this.tps = (tokenCount / elapsed) * 1000;
}
streamedContent += tokenContent;
if (tokenContent) {
streamedContent += tokenContent;
}
// Strip thinking tags for display and extract thinking content
const { displayContent, thinkingContent } =
// Use stripThinkingTags as fallback for any <think> tags still in content
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const combinedThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
// Only update currentResponse if target conversation is active
if (this.activeConversationId === targetConversationId) {
@@ -2400,7 +2463,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = combinedThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -2420,6 +2483,7 @@ class AppStore {
this.prefillProgress = {
processed: inner.processed_tokens,
total: inner.total_tokens,
startedAt: this.prefillProgress?.startedAt ?? performance.now(),
};
},
},
@@ -2436,14 +2500,17 @@ class AppStore {
// Final cleanup of the message (if conversation still exists)
if (this.conversationExists(targetConversationId)) {
const { displayContent, thinkingContent } =
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const finalThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
this.updateConversationMessage(
targetConversationId,
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = finalThinking || undefined;
msg.tokens = [...collectedTokens];
// Store performance metrics on the message
if (this.ttftMs !== null) {
@@ -3200,6 +3267,9 @@ export const setChatSidebarVisible = (visible: boolean) =>
appStore.setChatSidebarVisible(visible);
export const refreshState = () => appStore.fetchState();
// Connection status
export const isConnected = () => appStore.isConnected;
// Node identities (for OS version mismatch detection)
export const nodeIdentities = () => appStore.nodeIdentities;
+87
View File
@@ -0,0 +1,87 @@
/**
* Toast notification store - Global notification system for the EXO dashboard.
*
* Usage:
* import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte";
* addToast({ type: "success", message: "Model launched" });
* addToast({ type: "error", message: "Connection lost", persistent: true });
*/
type ToastType = "success" | "error" | "warning" | "info";
export interface Toast {
id: string;
type: ToastType;
message: string;
/** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */
duration: number;
createdAt: number;
}
interface ToastInput {
type: ToastType;
message: string;
/** If true, toast stays until manually dismissed. Default: false. */
persistent?: boolean;
/** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */
duration?: number;
}
const DEFAULT_DURATIONS: Record<ToastType, number> = {
success: 4000,
info: 4000,
warning: 6000,
error: 6000,
};
let toastList = $state<Toast[]>([]);
const timers = new Map<string, ReturnType<typeof setTimeout>>();
function generateId(): string {
return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
}
export function addToast(input: ToastInput): string {
const id = generateId();
const duration = input.persistent
? 0
: (input.duration ?? DEFAULT_DURATIONS[input.type]);
const toast: Toast = {
id,
type: input.type,
message: input.message,
duration,
createdAt: Date.now(),
};
toastList = [...toastList, toast];
if (duration > 0) {
const timer = setTimeout(() => dismissToast(id), duration);
timers.set(id, timer);
}
return id;
}
export function dismissToast(id: string): void {
const timer = timers.get(id);
if (timer) {
clearTimeout(timer);
timers.delete(id);
}
toastList = toastList.filter((t) => t.id !== id);
}
/** Dismiss all toasts matching a message (useful for dedup). */
export function dismissByMessage(message: string): void {
const matching = toastList.filter((t) => t.message === message);
for (const t of matching) {
dismissToast(t.id);
}
}
export function toasts(): Toast[] {
return toastList;
}
+4
View File
@@ -1,5 +1,7 @@
<script lang="ts">
import "../app.css";
import ToastContainer from "$lib/components/ToastContainer.svelte";
import ConnectionBanner from "$lib/components/ConnectionBanner.svelte";
let { children } = $props();
</script>
@@ -10,5 +12,7 @@
</svelte:head>
<div class="min-h-screen bg-background text-foreground">
<ConnectionBanner />
{@render children?.()}
<ToastContainer />
</div>
File diff suppressed because it is too large Load Diff
+156 -42
View File
@@ -29,7 +29,12 @@
etaMs: number;
modelDirectory?: string;
}
| { kind: "pending"; modelDirectory?: string }
| {
kind: "pending";
downloaded: number;
total: number;
modelDirectory?: string;
}
| { kind: "failed"; modelDirectory?: string }
| { kind: "not_present" };
@@ -74,7 +79,6 @@
if (typeof value === "number") return value;
if (value && typeof value === "object") {
const v = value as Record<string, unknown>;
if (typeof v.in_bytes === "number") return v.in_bytes;
if (typeof v.inBytes === "number") return v.inBytes;
}
return 0;
@@ -231,23 +235,14 @@
undefined;
let cell: CellStatus;
if (tag === "DownloadCompleted") {
const totalBytes = getBytes(
payload.total_bytes ?? payload.totalBytes,
);
const totalBytes = getBytes(payload.total);
cell = { kind: "completed", totalBytes, modelDirectory };
} else if (tag === "DownloadOngoing") {
const rawProgress =
payload.download_progress ?? payload.downloadProgress ?? {};
const prog = rawProgress as Record<string, unknown>;
const totalBytes = getBytes(
prog.total_bytes ??
prog.totalBytes ??
payload.total_bytes ??
payload.totalBytes,
);
const downloadedBytes = getBytes(
prog.downloaded_bytes ?? prog.downloadedBytes,
);
const totalBytes = getBytes(prog.total ?? payload.total);
const downloadedBytes = getBytes(prog.downloaded);
const speed = (prog.speed as number) ?? 0;
const etaMs =
(prog.eta_ms as number) ?? (prog.etaMs as number) ?? 0;
@@ -265,7 +260,20 @@
} else if (tag === "DownloadFailed") {
cell = { kind: "failed", modelDirectory };
} else {
cell = { kind: "pending", modelDirectory };
const downloaded = getBytes(
payload.downloaded ??
payload.downloaded_bytes ??
payload.downloadedBytes,
);
const total = getBytes(
payload.total ?? payload.total_bytes ?? payload.totalBytes,
);
cell = {
kind: "pending",
downloaded,
total,
modelDirectory,
};
}
const existing = row.cells[nodeId];
@@ -275,14 +283,51 @@
}
}
function rowSortKey(row: ModelRow): number {
// in progress (4) -> completed (3) -> paused (2) -> not started (1) -> not present (0)
let best = 0;
for (const cell of Object.values(row.cells)) {
let score = 0;
if (cell.kind === "downloading") score = 4;
else if (cell.kind === "completed") score = 3;
else if (cell.kind === "pending" && cell.downloaded > 0)
score = 2; // paused
else if (cell.kind === "pending" || cell.kind === "failed") score = 1; // not started
if (score > best) best = score;
}
return best;
}
function totalCompletedBytes(row: ModelRow): number {
let total = 0;
for (const cell of Object.values(row.cells)) {
if (cell.kind === "completed") total += cell.totalBytes;
}
return total;
}
const rows = Array.from(rowMap.values()).sort((a, b) => {
const aCompleted = Object.values(a.cells).filter(
(c) => c.kind === "completed",
).length;
const bCompleted = Object.values(b.cells).filter(
(c) => c.kind === "completed",
).length;
if (aCompleted !== bCompleted) return bCompleted - aCompleted;
const aPriority = rowSortKey(a);
const bPriority = rowSortKey(b);
if (aPriority !== bPriority) return bPriority - aPriority;
// Within completed or paused, sort by biggest size first
if (aPriority === 3 && bPriority === 3) {
const sizeDiff = totalCompletedBytes(b) - totalCompletedBytes(a);
if (sizeDiff !== 0) return sizeDiff;
}
if (aPriority === 2 && bPriority === 2) {
const aSize = Math.max(
...Object.values(a.cells).map((c) =>
c.kind === "pending" ? c.total : 0,
),
);
const bSize = Math.max(
...Object.values(b.cells).map((c) =>
c.kind === "pending" ? c.total : 0,
),
);
if (aSize !== bSize) return bSize - aSize;
}
return a.modelId.localeCompare(b.modelId);
});
@@ -367,7 +412,7 @@
<div>{col.label}</div>
{#if col.diskAvailable != null}
<div
class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
>
{formatBytes(col.diskAvailable)} free
</div>
@@ -391,7 +436,7 @@
</div>
{#if row.prettyName}
<div
class="text-[10px] text-exo-light-gray/60"
class="text-[10px] text-white/60"
title={row.modelId}
>
{row.modelId}
@@ -405,7 +450,7 @@
title="View model details"
>
<svg
class="w-4 h-4 text-white/30 hover:text-white/60"
class="w-4 h-4 text-white/60 hover:text-white/80"
viewBox="0 0 24 24"
fill="currentColor"
>
@@ -424,11 +469,11 @@
<td class="px-4 py-3 text-center align-middle">
{#if cell.kind === "completed"}
<div
class="flex flex-col items-center gap-0.5"
class="flex flex-col items-center gap-1"
title="Completed ({formatBytes(cell.totalBytes)})"
>
<svg
class="w-5 h-5 text-green-400"
class="w-7 h-7 text-green-400"
viewBox="0 0 20 20"
fill="currentColor"
>
@@ -438,18 +483,18 @@
clip-rule="evenodd"
></path>
</svg>
<span class="text-[10px] text-exo-light-gray/70"
<span class="text-xs text-white/70"
>{formatBytes(cell.totalBytes)}</span
>
<button
type="button"
class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
onclick={() =>
deleteDownload(col.nodeId, row.modelId)}
title="Delete from this node"
>
<svg
class="w-3.5 h-3.5"
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
@@ -472,11 +517,11 @@
cell.speed,
)} - ETA {formatEta(cell.etaMs)}"
>
<span class="text-exo-yellow text-xs font-medium"
<span class="text-exo-yellow text-sm font-medium"
>{clampPercent(cell.percentage).toFixed(1)}%</span
>
<div
class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
>
<div
class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
@@ -485,24 +530,93 @@
).toFixed(1)}%"
></div>
</div>
<span class="text-[9px] text-exo-light-gray/60"
<span class="text-[10px] text-white/70"
>{formatSpeed(cell.speed)}</span
>
</div>
{:else if cell.kind === "pending"}
<div
class="flex flex-col items-center gap-0.5"
title="Download pending"
class="flex flex-col items-center gap-1"
title={cell.downloaded > 0
? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
: "Download pending"}
>
<span class="text-exo-light-gray/50 text-sm">...</span>
{#if cell.downloaded > 0 && cell.total > 0}
<span class="text-white/70 text-xs"
>{formatBytes(cell.downloaded)} / {formatBytes(
cell.total,
)}</span
>
<div
class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
>
<div
class="h-full bg-exo-light-gray/40 rounded-full"
style="width: {(
(cell.downloaded / cell.total) *
100
).toFixed(1)}%"
></div>
</div>
{#if row.shardMetadata}
<button
type="button"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Resume download on this node"
>
<svg
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
stroke-width="2"
>
<path
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
stroke-linecap="round"
stroke-linejoin="round"
></path>
</svg>
</button>
{:else}
<span class="text-white/50 text-[10px]">paused</span
>
{/if}
{:else if row.shardMetadata}
<button
type="button"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Start download on this node"
>
<svg
class="w-6 h-6"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
stroke-width="2"
>
<path
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
stroke-linecap="round"
stroke-linejoin="round"
></path>
</svg>
</button>
{:else}
<span class="text-white/40 text-sm">...</span>
{/if}
</div>
{:else if cell.kind === "failed"}
<div
class="flex flex-col items-center gap-0.5"
class="flex flex-col items-center gap-1"
title="Download failed"
>
<svg
class="w-5 h-5 text-red-400"
class="w-7 h-7 text-red-400"
viewBox="0 0 20 20"
fill="currentColor"
>
@@ -515,13 +629,13 @@
{#if row.shardMetadata}
<button
type="button"
class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Retry download on this node"
>
<svg
class="w-3.5 h-3.5"
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
@@ -547,13 +661,13 @@
{#if row.shardMetadata}
<button
type="button"
class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Download to this node"
>
<svg
class="w-3.5 h-3.5"
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
+3 -3
View File
@@ -41,7 +41,7 @@ let
mlx = stdenv.mkDerivation rec {
pname = "mlx";
version = let v = "0.30.7.dev20260218+14841977"; in
version = let v = "0.30.7.dev20260220+13998a05"; in
assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
v;
pyproject = true;
@@ -49,8 +49,8 @@ let
src = fetchFromGitHub {
owner = "rltakashige";
repo = "mlx-jaccl-fix-small-recv";
rev = "1484197707f35186ad3bd614357c7c47fdf86ebc";
hash = "sha256-FupCMoK/SF/ldfKuvMSAKECcOP8c+ANgkQlPZttDsLk=";
rev = "13998a054715edcdc93618fb1496c79c7c25ff7c";
hash = "sha256-fAqA3hFwNBx7FcoGnhQsIFpAIRbC2EerACm4Fvne0Cc=";
};
patches = [
Binary file not shown.

After

Width:  |  Height:  |  Size: 6.8 KiB

+112
View File
@@ -0,0 +1,112 @@
#!/usr/bin/env bash
# create-dmg.sh — Build a polished macOS DMG installer for EXO
#
# Usage:
# ./packaging/dmg/create-dmg.sh <app-path> <output-dmg> [volume-name]
#
# Example:
# ./packaging/dmg/create-dmg.sh output/EXO.app EXO-1.0.0.dmg "EXO"
#
# Creates a DMG with:
# - Custom background image with drag-to-Applications arrow
# - App icon on left, Applications alias on right
# - Proper window size and icon positioning
set -euo pipefail
APP_PATH="${1:?Usage: create-dmg.sh <app-path> <output-dmg> [volume-name]}"
OUTPUT_DMG="${2:?Usage: create-dmg.sh <app-path> <output-dmg> [volume-name]}"
VOLUME_NAME="${3:-EXO}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
BACKGROUND_SCRIPT="${SCRIPT_DIR}/generate-background.py"
TEMP_DIR="$(mktemp -d)"
DMG_STAGING="${TEMP_DIR}/dmg-root"
TEMP_DMG="${TEMP_DIR}/temp.dmg"
BACKGROUND_PNG="${TEMP_DIR}/background.png"
cleanup() { rm -rf "$TEMP_DIR"; }
trap cleanup EXIT
echo "==> Creating DMG installer for ${VOLUME_NAME}"
# ── Step 1: Generate background image ────────────────────────────────────────
if command -v python3 &>/dev/null; then
python3 "$BACKGROUND_SCRIPT" "$BACKGROUND_PNG"
echo " Background image generated"
else
echo " Warning: python3 not found, skipping custom background"
BACKGROUND_PNG=""
fi
# ── Step 2: Prepare staging directory ─────────────────────────────────────────
mkdir -p "$DMG_STAGING"
cp -R "$APP_PATH" "$DMG_STAGING/"
ln -s /Applications "$DMG_STAGING/Applications"
# ── Step 3: Create writable DMG ──────────────────────────────────────────────
# Calculate required size (app size + 20MB headroom)
APP_SIZE_KB=$(du -sk "$APP_PATH" | cut -f1)
DMG_SIZE_KB=$((APP_SIZE_KB + 20480))
hdiutil create \
-volname "$VOLUME_NAME" \
-size "${DMG_SIZE_KB}k" \
-fs HFS+ \
-layout SPUD \
"$TEMP_DMG"
# ── Step 4: Mount and configure ──────────────────────────────────────────────
MOUNT_DIR=$(hdiutil attach "$TEMP_DMG" -readwrite -noverify | awk -F'\t' '/Apple_HFS/ {gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF); print $NF}')
echo " Mounted at: $MOUNT_DIR"
# Copy contents
cp -R "$DMG_STAGING/"* "$MOUNT_DIR/"
# Add background image
if [[ -n $BACKGROUND_PNG && -f $BACKGROUND_PNG ]]; then
mkdir -p "$MOUNT_DIR/.background"
cp "$BACKGROUND_PNG" "$MOUNT_DIR/.background/background.png"
fi
# ── Step 5: Configure window appearance via AppleScript ──────────────────────
# Window: 800×400, app icon on left, Applications on right (matches Ollama layout)
# Background image is 1600×740 (2× retina for 800×400 logical window).
APP_NAME="$(basename "$APP_PATH")"
osascript <<APPLESCRIPT
tell application "Finder"
tell disk "$VOLUME_NAME"
open
set current view of container window to icon view
set toolbar visible of container window to false
set statusbar visible of container window to false
set bounds of container window to {200, 120, 1000, 520}
set opts to icon view options of container window
set icon size of opts to 128
set text size of opts to 12
set arrangement of opts to not arranged
if exists file ".background:background.png" then
set background picture of opts to file ".background:background.png"
end if
set position of item "$APP_NAME" of container window to {200, 190}
set position of item "Applications" of container window to {600, 190}
close
open
update without registering applications
delay 1
close
end tell
end tell
APPLESCRIPT
echo " Window layout configured"
# Ensure Finder updates are flushed
sync
# ── Step 6: Finalise ─────────────────────────────────────────────────────────
hdiutil detach "$MOUNT_DIR" -quiet
hdiutil convert "$TEMP_DMG" -format UDZO -imagekey zlib-level=9 -o "$OUTPUT_DMG"
echo "==> DMG created: $OUTPUT_DMG"
echo " Size: $(du -h "$OUTPUT_DMG" | cut -f1)"
+91
View File
@@ -0,0 +1,91 @@
#!/usr/bin/env python3
"""Generate the DMG background image with a centered drag-to-Applications arrow.
The output is a 1600×740 retina PNG (2× for 800×400 logical window).
Icons are positioned at (200, 190) and (600, 190) in logical coordinates;
the arrow is drawn centered between them.
Usage:
python3 generate-background.py [output.png]
If no output path is given, overwrites the bundled background.png in-place.
"""
from __future__ import annotations
import math
import sys
from pathlib import Path
from PIL import Image, ImageDraw
# Retina dimensions (2× logical 800×400)
WIDTH = 1600
HEIGHT = 740
# Icon positions in logical coords → retina coords
# App icon at (200, 190), Applications at (600, 190)
APP_X = 200 * 2 # 400
APPS_X = 600 * 2 # 1200
ICON_Y = 190 * 2 # 380
# Arrow drawn between icons, slightly above icon center
ARROW_START_X = APP_X + 160 # past the icon
ARROW_END_X = APPS_X - 160 # before the Applications icon
ARROW_Y = ICON_Y # same height as icons
ARROW_RISE = 120 # upward arc height
def draw_arrow(draw: ImageDraw.ImageDraw) -> None:
"""Draw a hand-drawn-style curved arrow from app icon toward Applications."""
color = (30, 30, 30)
line_width = 8
# Compute bezier curve points for a gentle upward arc
points: list[tuple[float, float]] = []
steps = 80
for i in range(steps + 1):
t = i / steps
# Quadratic bezier: start → control → end
cx = (ARROW_START_X + ARROW_END_X) / 2
cy = ARROW_Y - ARROW_RISE
x = (1 - t) ** 2 * ARROW_START_X + 2 * (1 - t) * t * cx + t**2 * ARROW_END_X
y = (1 - t) ** 2 * ARROW_Y + 2 * (1 - t) * t * cy + t**2 * ARROW_Y
points.append((x, y))
# Draw the curve as connected line segments
for i in range(len(points) - 1):
draw.line([points[i], points[i + 1]], fill=color, width=line_width)
# Arrowhead at the end
end_x, end_y = points[-1]
# Direction from second-to-last to last point
prev_x, prev_y = points[-3]
angle = math.atan2(end_y - prev_y, end_x - prev_x)
head_len = 36
head_angle = math.radians(25)
left_x = end_x - head_len * math.cos(angle - head_angle)
left_y = end_y - head_len * math.sin(angle - head_angle)
right_x = end_x - head_len * math.cos(angle + head_angle)
right_y = end_y - head_len * math.sin(angle + head_angle)
draw.polygon(
[(end_x, end_y), (left_x, left_y), (right_x, right_y)],
fill=color,
)
def generate_background(output_path: str) -> None:
"""Generate a white DMG background with a centered arrow."""
img = Image.new("RGBA", (WIDTH, HEIGHT), (255, 255, 255, 255))
draw = ImageDraw.Draw(img)
draw_arrow(draw)
img.save(output_path, "PNG")
if __name__ == "__main__":
default_output = str(Path(__file__).parent / "background.png")
out = sys.argv[1] if len(sys.argv) >= 2 else default_output
generate_background(out)
print(f"Background image written to {out}")
@@ -3,6 +3,10 @@ n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "4bit"
base_model = "Qwen3 Coder Next"
capabilities = ["text", "code"]
[storage_size]
in_bytes = 45644286500
@@ -3,6 +3,10 @@ n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "5bit"
base_model = "Qwen3 Coder Next"
capabilities = ["text", "code"]
[storage_size]
in_bytes = 57657697020
@@ -3,6 +3,10 @@ n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "6bit"
base_model = "Qwen3 Coder Next"
capabilities = ["text", "code"]
[storage_size]
in_bytes = 68899327465
@@ -3,6 +3,10 @@ n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3 Coder Next"
capabilities = ["text", "code"]
[storage_size]
in_bytes = 89357758772
@@ -3,6 +3,10 @@ n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "bf16"
base_model = "Qwen3 Coder Next"
capabilities = ["text", "code"]
[storage_size]
in_bytes = 157548627945
+11 -106
View File
@@ -19,7 +19,7 @@ class ConnectionUpdate:
Whether this is a connection or disconnection event
"""
@property
def peer_id(self) -> PeerId:
def peer_id(self) -> builtins.str:
r"""
Identity of the peer that we have connected to or disconnected from.
"""
@@ -40,92 +40,22 @@ class Keypair:
Identity keypair of a node.
"""
@staticmethod
def generate_ed25519() -> Keypair:
def generate() -> Keypair:
r"""
Generate a new Ed25519 keypair.
"""
@staticmethod
def generate_ecdsa() -> Keypair:
def from_bytes(bytes: bytes) -> Keypair:
r"""
Generate a new ECDSA keypair.
"""
@staticmethod
def generate_secp256k1() -> Keypair:
r"""
Generate a new Secp256k1 keypair.
"""
@staticmethod
def from_protobuf_encoding(bytes: bytes) -> Keypair:
r"""
Decode a private key from a protobuf structure and parse it as a `Keypair`.
"""
@staticmethod
def rsa_from_pkcs8(bytes: bytes) -> Keypair:
r"""
Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
format (i.e. unencrypted) as defined in [RFC5208].
[RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
"""
@staticmethod
def secp256k1_from_der(bytes: bytes) -> Keypair:
r"""
Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
structure as defined in [RFC5915].
[RFC5915]: https://tools.ietf.org/html/rfc5915
"""
@staticmethod
def ed25519_from_bytes(bytes: bytes) -> Keypair: ...
def to_protobuf_encoding(self) -> bytes:
r"""
Encode a private key as protobuf structure.
"""
def to_peer_id(self) -> PeerId:
r"""
Convert the `Keypair` into the corresponding `PeerId`.
"""
@typing.final
class Multiaddr:
r"""
Representation of a Multiaddr.
"""
@staticmethod
def empty() -> Multiaddr:
r"""
Create a new, empty multiaddress.
"""
@staticmethod
def with_capacity(n: builtins.int) -> Multiaddr:
r"""
Create a new, empty multiaddress with the given capacity.
"""
@staticmethod
def from_bytes(bytes: bytes) -> Multiaddr:
r"""
Parse a `Multiaddr` value from its byte slice representation.
"""
@staticmethod
def from_string(string: builtins.str) -> Multiaddr:
r"""
Parse a `Multiaddr` value from its string representation.
"""
def len(self) -> builtins.int:
r"""
Return the length in bytes of this multiaddress.
"""
def is_empty(self) -> builtins.bool:
r"""
Returns true if the length of this multiaddress is 0.
Construct an Ed25519 keypair from secret key bytes
"""
def to_bytes(self) -> bytes:
r"""
Return a copy of this [`Multiaddr`]'s byte representation.
Get the secret key bytes underlying the keypair
"""
def to_string(self) -> builtins.str:
def to_node_id(self) -> builtins.str:
r"""
Convert a Multiaddr to a string.
Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
"""
@typing.final
@@ -175,39 +105,14 @@ class NetworkingHandle:
"""
@typing.final
class NoPeersSubscribedToTopicError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
class MessageTooLargeError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class PeerId:
r"""
Identifier of a peer of the network.
The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
"""
@staticmethod
def random() -> PeerId:
r"""
Generates a random peer ID from a cryptographically secure PRNG.
This is useful for randomly walking on a DHT, or for testing purposes.
"""
@staticmethod
def from_bytes(bytes: bytes) -> PeerId:
r"""
Parses a `PeerId` from bytes.
"""
def to_bytes(self) -> bytes:
r"""
Returns a raw bytes representation of this `PeerId`.
"""
def to_base58(self) -> builtins.str:
r"""
Returns a base-58 encoded string of this `PeerId`.
"""
class NoPeersSubscribedToTopicError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
+17 -129
View File
@@ -1,8 +1,6 @@
use crate::ext::ResultExt as _;
use libp2p::PeerId;
use libp2p::identity::Keypair;
use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
use pyo3::types::PyBytes;
use pyo3::types::{PyBytes, PyBytesMethods as _};
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
@@ -18,142 +16,32 @@ pub struct PyKeypair(pub Keypair);
impl PyKeypair {
/// Generate a new Ed25519 keypair.
#[staticmethod]
fn generate_ed25519() -> Self {
fn generate() -> Self {
Self(Keypair::generate_ed25519())
}
/// Generate a new ECDSA keypair.
/// Construct an Ed25519 keypair from secret key bytes
#[staticmethod]
fn generate_ecdsa() -> Self {
Self(Keypair::generate_ecdsa())
}
/// Generate a new Secp256k1 keypair.
#[staticmethod]
fn generate_secp256k1() -> Self {
Self(Keypair::generate_secp256k1())
}
/// Decode a private key from a protobuf structure and parse it as a `Keypair`.
#[staticmethod]
fn from_protobuf_encoding(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::from_protobuf_encoding(&bytes).pyerr()?))
}
/// Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
/// format (i.e. unencrypted) as defined in [RFC5208].
///
/// [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
#[staticmethod]
fn rsa_from_pkcs8(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::rsa_from_pkcs8(&mut bytes).pyerr()?))
}
/// Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
/// structure as defined in [RFC5915].
///
/// [RFC5915]: https://tools.ietf.org/html/rfc5915
#[staticmethod]
fn secp256k1_from_der(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::secp256k1_from_der(&mut bytes).pyerr()?))
}
#[staticmethod]
fn ed25519_from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
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()?))
}
/// Encode a private key as protobuf structure.
fn to_protobuf_encoding<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
let bytes = self.0.to_protobuf_encoding().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`.
fn to_peer_id(&self) -> PyPeerId {
PyPeerId(self.0.public().to_peer_id())
}
// /// Hidden constructor for pickling support. TODO: figure out how to do pickling...
// #[gen_stub(skip)]
// #[new]
// fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
// Self::from_protobuf_encoding(bytes)
// }
//
// #[gen_stub(skip)]
// fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
// *self = Self::from_protobuf_encoding(state)?;
// Ok(())
// }
//
// #[gen_stub(skip)]
// fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
// self.to_protobuf_encoding(py)
// }
//
// #[gen_stub(skip)]
// pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> {
// Ok((self.to_protobuf_encoding(py)?,))
// }
}
/// Identifier of a peer of the network.
///
/// The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
/// as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
#[gen_stub_pyclass]
#[pyclass(name = "PeerId", frozen)]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct PyPeerId(pub PeerId);
#[gen_stub_pymethods]
#[pymethods]
#[allow(clippy::needless_pass_by_value)]
impl PyPeerId {
/// Generates a random peer ID from a cryptographically secure PRNG.
///
/// This is useful for randomly walking on a DHT, or for testing purposes.
#[staticmethod]
fn random() -> Self {
Self(PeerId::random())
}
/// Parses a `PeerId` from bytes.
#[staticmethod]
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let bytes = Vec::from(bytes.as_bytes());
Ok(Self(PeerId::from_bytes(&bytes).pyerr()?))
}
/// Returns a raw bytes representation of this `PeerId`.
fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
let bytes = self.0.to_bytes();
PyBytes::new(py, &bytes)
}
/// Returns a base-58 encoded string of this `PeerId`.
fn to_base58(&self) -> String {
self.0.to_base58()
}
fn __repr__(&self) -> String {
format!("PeerId({})", self.to_base58())
}
fn __str__(&self) -> String {
self.to_base58()
/// 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()
}
}
pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyKeypair>()?;
m.add_class::<PyPeerId>()?;
Ok(())
}
+3 -2
View File
@@ -8,9 +8,10 @@ mod allow_threading;
mod ident;
mod networking;
use crate::ident::ident_submodule;
use crate::ident::PyKeypair;
use crate::networking::networking_submodule;
use pyo3::prelude::PyModule;
use pyo3::types::PyModuleMethods;
use pyo3::{Bound, PyResult, pyclass, pymodule};
use pyo3_stub_gen::define_stub_info_gatherer;
@@ -158,7 +159,7 @@ 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...
ident_submodule(m)?;
m.add_class::<PyKeypair>()?;
networking_submodule(m)?;
// top-level constructs
+38 -4
View File
@@ -8,7 +8,7 @@
use crate::r#const::MPSC_CHANNEL_SIZE;
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
use crate::ident::{PyKeypair, PyPeerId};
use crate::ident::PyKeypair;
use crate::pyclass;
use libp2p::futures::StreamExt as _;
use libp2p::gossipsub;
@@ -98,6 +98,37 @@ mod exception {
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()
}
}
}
/// Connection or disconnection event discriminant type.
@@ -119,7 +150,7 @@ struct PyConnectionUpdate {
/// Identity of the peer that we have connected to or disconnected from.
#[pyo3(get)]
peer_id: PyPeerId,
peer_id: String,
/// Remote connection's IPv4 address.
#[pyo3(get)]
@@ -200,6 +231,8 @@ async fn networking_task(
Err(exception::PyNoPeersSubscribedToTopicError::new_err())
} else if let Err(PublishError::AllQueuesFull(_)) = result {
Err(exception::PyAllQueuesFullError::new_err())
} else if let Err(PublishError::MessageTooLarge) = result {
Err(exception::PyMessageTooLargeError::new_err())
} else {
result.pyerr()
};
@@ -251,7 +284,7 @@ async fn networking_task(
// send connection event to channel (or exit if connection closed)
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
update_type: PyConnectionUpdateType::Connected,
peer_id: PyPeerId(peer_id),
peer_id: peer_id.to_base58(),
remote_ipv4,
remote_tcp_port,
}).await {
@@ -272,7 +305,7 @@ async fn networking_task(
// send disconnection event to channel (or exit if connection closed)
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
update_type: PyConnectionUpdateType::Disconnected,
peer_id: PyPeerId(peer_id),
peer_id: peer_id.to_base58(),
remote_ipv4,
remote_tcp_port,
}).await {
@@ -561,6 +594,7 @@ impl PyNetworkingHandle {
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::<PyConnectionUpdateType>()?;
m.add_class::<PyConnectionUpdate>()?;
+119 -24
View File
@@ -1,31 +1,34 @@
import asyncio
import socket
from dataclasses import dataclass, field
from typing import Iterator
from random import random
import anyio
from anyio import current_time
from anyio.abc import TaskGroup
from loguru import logger
from exo.download.download_utils import (
RepoDownloadProgress,
delete_model,
map_repo_download_progress_to_download_progress_data,
resolve_model_in_path,
)
from exo.download.shard_downloader import ShardDownloader
from exo.shared.constants import EXO_MODELS_DIR
from exo.shared.models.model_cards import ModelId
from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
from exo.shared.models.model_cards import ModelId, get_model_cards
from exo.shared.types.commands import (
CancelDownload,
DeleteDownload,
ForwarderDownloadCommand,
StartDownload,
)
from exo.shared.types.common import NodeId, SessionId
from exo.shared.types.common import NodeId, SessionId, SystemId
from exo.shared.types.events import (
Event,
ForwarderEvent,
EventId,
# TODO(evan): just for acks, should delete this ASAP
GlobalForwarderEvent,
LocalForwarderEvent,
NodeDownloadProgress,
)
from exo.shared.types.worker.downloads import (
@@ -35,8 +38,9 @@ from exo.shared.types.worker.downloads import (
DownloadPending,
DownloadProgress,
)
from exo.shared.types.worker.shards import ShardMetadata
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.task_group import TaskGroup
@dataclass
@@ -45,10 +49,16 @@ class DownloadCoordinator:
session_id: SessionId
shard_downloader: ShardDownloader
download_command_receiver: Receiver[ForwarderDownloadCommand]
local_event_sender: Sender[ForwarderEvent]
event_index_counter: Iterator[int]
local_event_sender: Sender[LocalForwarderEvent]
# ack stuff
_global_event_receiver: Receiver[GlobalForwarderEvent]
_out_for_delivery: dict[EventId, LocalForwarderEvent] = field(default_factory=dict)
offline: bool = False
_system_id: SystemId = field(default_factory=SystemId)
# Local state
download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
@@ -56,7 +66,7 @@ class DownloadCoordinator:
# Internal event channel for forwarding (initialized in __post_init__)
event_sender: Sender[Event] = field(init=False)
event_receiver: Receiver[Event] = field(init=False)
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
# Per-model throttle for download progress events
_last_progress_time: dict[ModelId, float] = field(default_factory=dict)
@@ -80,7 +90,7 @@ class DownloadCoordinator:
completed = DownloadCompleted(
shard_metadata=callback_shard,
node_id=self.node_id,
total_bytes=progress.total_bytes,
total=progress.total,
model_directory=self._model_dir(model_id),
)
self.download_status[model_id] = completed
@@ -115,12 +125,18 @@ class DownloadCoordinator:
)
if not self.offline:
self._test_internet_connection()
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._forward_events)
tg.start_soon(self._emit_existing_download_progress)
if not self.offline:
tg.start_soon(self._check_internet_connection)
try:
async with self._tg as tg:
tg.start_soon(self._command_processor)
tg.start_soon(self._forward_events)
tg.start_soon(self._emit_existing_download_progress)
tg.start_soon(self._resend_out_for_delivery)
tg.start_soon(self._clear_ofd)
if not self.offline:
tg.start_soon(self._check_internet_connection)
finally:
for task in self.active_downloads.values():
task.cancel()
def _test_internet_connection(self) -> None:
# Try multiple endpoints since some ISPs/networks block specific IPs
@@ -151,7 +167,21 @@ class DownloadCoordinator:
self._tg.start_soon(self._emit_existing_download_progress)
def shutdown(self) -> None:
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
# directly copied from worker
async def _resend_out_for_delivery(self) -> None:
# This can also be massively tightened, we should check events are at least a certain age before resending.
# Exponential backoff would also certainly help here.
while True:
await anyio.sleep(1 + random())
for event in self._out_for_delivery.copy().values():
await self.local_event_sender.send(event)
async def _clear_ofd(self) -> None:
with self._global_event_receiver as events:
async for event in events:
self._out_for_delivery.pop(event.event.event_id, None)
async def _command_processor(self) -> None:
with self.download_command_receiver as commands:
@@ -185,6 +215,25 @@ class DownloadCoordinator:
)
return
# Check EXO_MODELS_PATH for pre-downloaded models
found_path = resolve_model_in_path(model_id)
if found_path is not None:
logger.info(
f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}"
)
completed = DownloadCompleted(
shard_metadata=shard,
node_id=self.node_id,
total=shard.model_card.storage_size,
model_directory=str(found_path),
read_only=True,
)
self.download_status[model_id] = completed
await self.event_sender.send(
NodeDownloadProgress(download_progress=completed)
)
return
# Emit pending status
progress = DownloadPending(
shard_metadata=shard,
@@ -203,7 +252,7 @@ class DownloadCoordinator:
completed = DownloadCompleted(
shard_metadata=shard,
node_id=self.node_id,
total_bytes=initial_progress.total_bytes,
total=initial_progress.total,
model_directory=self._model_dir(model_id),
)
self.download_status[model_id] = completed
@@ -269,6 +318,15 @@ class DownloadCoordinator:
self.active_downloads[model_id] = task
async def _delete_download(self, model_id: ModelId) -> None:
# Protect read-only models (from EXO_MODELS_PATH) from deletion
if model_id in self.download_status:
current = self.download_status[model_id]
if isinstance(current, DownloadCompleted) and current.read_only:
logger.warning(
f"Refusing to delete read-only model {model_id} (from EXO_MODELS_PATH)"
)
return
# Cancel if active
if model_id in self.active_downloads:
logger.info(f"Cancelling active download for {model_id} before deletion")
@@ -298,19 +356,21 @@ class DownloadCoordinator:
del self.download_status[model_id]
async def _forward_events(self) -> None:
idx = 0
with self.event_receiver as events:
async for event in events:
idx = next(self.event_index_counter)
fe = ForwarderEvent(
fe = LocalForwarderEvent(
origin_idx=idx,
origin=self.node_id,
origin=self._system_id,
session=self.session_id,
event=event,
)
idx += 1
logger.debug(
f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
)
await self.local_event_sender.send(fe)
self._out_for_delivery[event.event_id] = fe
async def _emit_existing_download_progress(self) -> None:
try:
@@ -332,19 +392,21 @@ class DownloadCoordinator:
status: DownloadProgress = DownloadCompleted(
node_id=self.node_id,
shard_metadata=progress.shard,
total_bytes=progress.total_bytes,
total=progress.total,
model_directory=self._model_dir(
progress.shard.model_card.model_id
),
)
elif progress.status in ["in_progress", "not_started"]:
if progress.downloaded_bytes_this_session.in_bytes == 0:
if progress.downloaded_this_session.in_bytes == 0:
status = DownloadPending(
node_id=self.node_id,
shard_metadata=progress.shard,
model_directory=self._model_dir(
progress.shard.model_card.model_id
),
downloaded=progress.downloaded,
total=progress.total,
)
else:
status = DownloadOngoing(
@@ -364,6 +426,39 @@ class DownloadCoordinator:
await self.event_sender.send(
NodeDownloadProgress(download_progress=status)
)
# Scan EXO_MODELS_PATH for pre-downloaded models
if EXO_MODELS_PATH is not None:
for card in await get_model_cards():
mid = card.model_id
if mid in self.active_downloads:
continue
if isinstance(
self.download_status.get(mid),
(DownloadCompleted, DownloadOngoing, DownloadFailed),
):
continue
found = resolve_model_in_path(mid)
if found is not None:
path_shard = PipelineShardMetadata(
model_card=card,
device_rank=0,
world_size=1,
start_layer=0,
end_layer=card.n_layers,
n_layers=card.n_layers,
)
path_completed: DownloadProgress = DownloadCompleted(
node_id=self.node_id,
shard_metadata=path_shard,
total=card.storage_size,
model_directory=str(found),
read_only=True,
)
self.download_status[mid] = path_completed
await self.event_sender.send(
NodeDownloadProgress(download_progress=path_completed)
)
logger.debug(
"DownloadCoordinator: Done emitting existing download progress."
)
+111 -76
View File
@@ -20,7 +20,6 @@ from huggingface_hub import (
)
from loguru import logger
from pydantic import (
DirectoryPath,
TypeAdapter,
)
@@ -31,7 +30,7 @@ from exo.download.huggingface_utils import (
get_hf_endpoint,
get_hf_token,
)
from exo.shared.constants import EXO_MODELS_DIR
from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
from exo.shared.models.model_cards import ModelTask
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
@@ -80,9 +79,9 @@ def map_repo_file_download_progress_to_download_progress_data(
repo_file_download_progress: RepoFileDownloadProgress,
) -> DownloadProgressData:
return DownloadProgressData(
downloaded_bytes=repo_file_download_progress.downloaded,
downloaded_bytes_this_session=repo_file_download_progress.downloaded_this_session,
total_bytes=repo_file_download_progress.total,
downloaded=repo_file_download_progress.downloaded,
downloaded_this_session=repo_file_download_progress.downloaded_this_session,
total=repo_file_download_progress.total,
completed_files=1 if repo_file_download_progress.status == "complete" else 0,
total_files=1,
speed=repo_file_download_progress.speed,
@@ -95,9 +94,9 @@ def map_repo_download_progress_to_download_progress_data(
repo_download_progress: RepoDownloadProgress,
) -> DownloadProgressData:
return DownloadProgressData(
total_bytes=repo_download_progress.total_bytes,
downloaded_bytes=repo_download_progress.downloaded_bytes,
downloaded_bytes_this_session=repo_download_progress.downloaded_bytes_this_session,
total=repo_download_progress.total,
downloaded=repo_download_progress.downloaded,
downloaded_this_session=repo_download_progress.downloaded_this_session,
completed_files=repo_download_progress.completed_files,
total_files=repo_download_progress.total_files,
speed=repo_download_progress.overall_speed,
@@ -111,7 +110,27 @@ def map_repo_download_progress_to_download_progress_data(
)
def build_model_path(model_id: ModelId) -> DirectoryPath:
def resolve_model_in_path(model_id: ModelId) -> Path | None:
"""Search EXO_MODELS_PATH directories for a pre-existing model.
Checks each directory for the normalized name (org--model). A candidate
is only returned if ``is_model_directory_complete`` confirms all weight
files are present.
"""
if EXO_MODELS_PATH is None:
return None
normalized = model_id.normalize()
for search_dir in EXO_MODELS_PATH:
candidate = search_dir / normalized
if candidate.is_dir() and is_model_directory_complete(candidate):
return candidate
return None
def build_model_path(model_id: ModelId) -> Path:
found = resolve_model_in_path(model_id)
if found is not None:
return found
return EXO_MODELS_DIR / model_id.normalize()
@@ -142,7 +161,7 @@ async def delete_model(model_id: ModelId) -> bool:
async def seed_models(seed_dir: str | Path):
"""Move model in resources folder of app to .cache/huggingface/hub"""
"""Move models from resources folder to EXO_MODELS_DIR."""
source_dir = Path(seed_dir)
dest_dir = await ensure_models_dir()
for path in source_dir.iterdir():
@@ -158,6 +177,72 @@ async def seed_models(seed_dir: str | Path):
logger.error(traceback.format_exc())
def _scan_model_directory(
model_dir: Path, recursive: bool = False
) -> list[FileListEntry] | None:
"""Scan a local model directory and build a file list.
Requires at least one ``*.safetensors.index.json``. Every weight file
referenced by the index that is missing on disk gets ``size=None``.
"""
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
if not index_files:
return None
entries_by_path: dict[str, FileListEntry] = {}
if recursive:
for dirpath, _, filenames in os.walk(model_dir):
for filename in filenames:
if filename.endswith(".partial"):
continue
full_path = Path(dirpath) / filename
rel_path = str(full_path.relative_to(model_dir))
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=full_path.stat().st_size,
)
else:
for item in model_dir.iterdir():
if item.is_file() and not item.name.endswith(".partial"):
entries_by_path[item.name] = FileListEntry(
type="file",
path=item.name,
size=item.stat().st_size,
)
# Add expected weight files from index that haven't been downloaded yet
for index_file in index_files:
try:
index_data = ModelSafetensorsIndex.model_validate_json(
index_file.read_text()
)
relative_dir = index_file.parent.relative_to(model_dir)
for filename in set(index_data.weight_map.values()):
rel_path = (
str(relative_dir / filename)
if relative_dir != Path(".")
else filename
)
if rel_path not in entries_by_path:
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=None,
)
except Exception:
continue
return list(entries_by_path.values())
def is_model_directory_complete(model_dir: Path) -> bool:
"""Check if a model directory contains all required weight files."""
file_list = _scan_model_directory(model_dir, recursive=True)
return file_list is not None and all(f.size is not None for f in file_list)
async def _build_file_list_from_local_directory(
model_id: ModelId,
recursive: bool = False,
@@ -172,59 +257,7 @@ async def _build_file_list_from_local_directory(
if not await aios.path.exists(model_dir):
return None
def _scan() -> list[FileListEntry] | None:
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
if not index_files:
return None
entries_by_path: dict[str, FileListEntry] = {}
if recursive:
for dirpath, _, filenames in os.walk(model_dir):
for filename in filenames:
if filename.endswith(".partial"):
continue
full_path = Path(dirpath) / filename
rel_path = str(full_path.relative_to(model_dir))
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=full_path.stat().st_size,
)
else:
for item in model_dir.iterdir():
if item.is_file() and not item.name.endswith(".partial"):
entries_by_path[item.name] = FileListEntry(
type="file",
path=item.name,
size=item.stat().st_size,
)
# Add expected weight files from index that haven't been downloaded yet
for index_file in index_files:
try:
index_data = ModelSafetensorsIndex.model_validate_json(
index_file.read_text()
)
relative_dir = index_file.parent.relative_to(model_dir)
for filename in set(index_data.weight_map.values()):
rel_path = (
str(relative_dir / filename)
if relative_dir != Path(".")
else filename
)
if rel_path not in entries_by_path:
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=None,
)
except Exception:
continue
return list(entries_by_path.values())
file_list = await asyncio.to_thread(_scan)
file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive)
if not file_list:
return None
return file_list
@@ -578,19 +611,20 @@ def calculate_repo_progress(
file_progress: dict[str, RepoFileDownloadProgress],
all_start_time: float,
) -> RepoDownloadProgress:
all_total_bytes = sum((p.total.in_bytes for p in file_progress.values()), 0)
all_downloaded_bytes = sum(
(p.downloaded.in_bytes for p in file_progress.values()), 0
all_total = sum((p.total for p in file_progress.values()), Memory.from_bytes(0))
all_downloaded = sum(
(p.downloaded for p in file_progress.values()), Memory.from_bytes(0)
)
all_downloaded_bytes_this_session = sum(
(p.downloaded_this_session.in_bytes for p in file_progress.values()), 0
all_downloaded_this_session = sum(
(p.downloaded_this_session for p in file_progress.values()),
Memory.from_bytes(0),
)
elapsed_time = time.time() - all_start_time
all_speed = (
all_downloaded_bytes_this_session / elapsed_time if elapsed_time > 0 else 0
all_downloaded_this_session.in_bytes / elapsed_time if elapsed_time > 0 else 0
)
all_eta = (
timedelta(seconds=(all_total_bytes - all_downloaded_bytes) / all_speed)
timedelta(seconds=(all_total - all_downloaded).in_bytes / all_speed)
if all_speed > 0
else timedelta(seconds=0)
)
@@ -609,11 +643,9 @@ def calculate_repo_progress(
[p for p in file_progress.values() if p.downloaded == p.total]
),
total_files=len(file_progress),
downloaded_bytes=Memory.from_bytes(all_downloaded_bytes),
downloaded_bytes_this_session=Memory.from_bytes(
all_downloaded_bytes_this_session
),
total_bytes=Memory.from_bytes(all_total_bytes),
downloaded=all_downloaded,
downloaded_this_session=all_downloaded_this_session,
total=all_total,
overall_speed=all_speed,
overall_eta=all_eta,
status=status,
@@ -791,6 +823,7 @@ async def download_shard(
for file in filtered_file_list:
downloaded_bytes = await get_downloaded_size(target_dir / file.path)
final_file_exists = await aios.path.exists(target_dir / file.path)
file_progress[file.path] = RepoFileDownloadProgress(
repo_id=shard.model_card.model_id,
repo_revision=revision,
@@ -800,7 +833,9 @@ async def download_shard(
total=Memory.from_bytes(file.size or 0),
speed=0,
eta=timedelta(0),
status="complete" if downloaded_bytes == file.size else "not_started",
status="complete"
if final_file_exists and downloaded_bytes == file.size
else "not_started",
start_time=time.time(),
)
+3 -3
View File
@@ -107,9 +107,9 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
),
completed_files=0,
total_files=0,
downloaded_bytes=Memory.from_bytes(0),
downloaded_bytes_this_session=Memory.from_bytes(0),
total_bytes=Memory.from_bytes(0),
downloaded=Memory.from_bytes(0),
downloaded_this_session=Memory.from_bytes(0),
total=Memory.from_bytes(0),
overall_speed=0,
overall_eta=timedelta(seconds=0),
status="complete",
@@ -0,0 +1,98 @@
from typing import Any
import anyio
import pytest
from exo.download.coordinator import DownloadCoordinator
from exo.download.shard_downloader import NoopShardDownloader
from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.common import ModelId, NodeId, SessionId
from exo.shared.types.events import (
GlobalForwarderEvent,
LocalForwarderEvent,
NodeDownloadProgress,
)
from exo.shared.types.memory import Memory
from exo.shared.types.worker.downloads import (
DownloadPending,
)
from exo.shared.types.worker.shards import PipelineShardMetadata
from exo.utils.channels import channel
# Use the builtin NoopShardDownloader directly it already implements the required abstract interface.
# No additional subclass is needed for this test.
@pytest.mark.anyio
async def test_ack_behaviour():
# Create channels (type Any for simplicity)
_, command_receiver = channel[Any]()
local_sender, _ = channel[Any]()
global_sender, global_receiver = channel[Any]()
# Minimal identifiers
node_id = NodeId()
session_id = SessionId(master_node_id=node_id, election_clock=0)
# Create a dummy model card and shard metadata
model_id = ModelId("test/model")
model_card = ModelCard(
model_id=model_id,
storage_size=Memory.from_bytes(0),
n_layers=1,
hidden_size=1,
supports_tensor=True,
tasks=[ModelTask.TextGeneration],
)
shard = PipelineShardMetadata(
model_card=model_card,
device_rank=0,
world_size=1,
start_layer=0,
end_layer=1,
n_layers=1,
)
# Instantiate the coordinator with the dummy downloader
coord = DownloadCoordinator(
node_id=node_id,
session_id=session_id,
shard_downloader=NoopShardDownloader(),
download_command_receiver=command_receiver,
local_event_sender=local_sender,
_global_event_receiver=global_receiver,
)
async with anyio.create_task_group() as tg:
# Start the forwarding and ackclearing loops
tg.start_soon(coord._forward_events) # pyright: ignore[reportPrivateUsage]
tg.start_soon(coord._clear_ofd) # pyright: ignore[reportPrivateUsage]
# Send a pending download progress event via the internal event sender
pending = DownloadPending(
node_id=node_id,
shard_metadata=shard,
model_directory="/tmp/model",
)
await coord.event_sender.send(NodeDownloadProgress(download_progress=pending))
# Allow the forwarder to process the event
await anyio.sleep(0.1)
# There should be exactly one entry awaiting ACK
assert len(coord._out_for_delivery) == 1 # pyright: ignore[reportPrivateUsage]
# Retrieve the stored LocalForwarderEvent
stored_fe: LocalForwarderEvent = next(iter(coord._out_for_delivery.values())) # pyright: ignore[reportPrivateUsage]
# Simulate receiving a global ack for this event
ack = GlobalForwarderEvent(
origin_idx=0,
origin=node_id,
session=session_id,
event=stored_fe.event,
)
await global_sender.send(ack)
# Give the clearofd task a moment to process the ack
await anyio.sleep(0.1)
# The outfordelivery map should now be empty
assert len(coord._out_for_delivery) == 0 # pyright: ignore[reportPrivateUsage]
# Cancel background tasks
tg.cancel_scope.cancel()
+27 -20
View File
@@ -1,14 +1,12 @@
import argparse
import itertools
import multiprocessing as mp
import os
import resource
import signal
from dataclasses import dataclass, field
from typing import Iterator, Self
from typing import Self
import anyio
from anyio.abc import TaskGroup
from loguru import logger
from pydantic import PositiveInt
@@ -24,6 +22,7 @@ from exo.shared.logging import logger_cleanup, logger_setup
from exo.shared.types.common import NodeId, SessionId
from exo.utils.channels import Receiver, channel
from exo.utils.pydantic_ext import CamelCaseModel
from exo.utils.task_group import TaskGroup
from exo.worker.main import Worker
@@ -38,14 +37,13 @@ class Node:
api: API | None
node_id: NodeId
event_index_counter: Iterator[int]
offline: bool
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
@classmethod
async def create(cls, args: "Args") -> "Self":
async def create(cls, args: "Args") -> Self:
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_peer_id().to_base58())
node_id = NodeId(keypair.to_node_id())
session_id = SessionId(master_node_id=node_id, election_clock=0)
router = Router.create(keypair)
await router.register_topic(topics.GLOBAL_EVENTS)
@@ -57,9 +55,6 @@ class Node:
logger.info(f"Starting node {node_id}")
# Create shared event index counter for Worker and DownloadCoordinator
event_index_counter = itertools.count()
# Create DownloadCoordinator (unless --no-downloads)
if not args.no_downloads:
download_coordinator = DownloadCoordinator(
@@ -68,8 +63,9 @@ class Node:
exo_shard_downloader(),
download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
local_event_sender=router.sender(topics.LOCAL_EVENTS),
event_index_counter=event_index_counter,
offline=args.offline,
# TODO(evan): remove
_global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
)
else:
download_coordinator = None
@@ -95,7 +91,6 @@ class Node:
local_event_sender=router.sender(topics.LOCAL_EVENTS),
command_sender=router.sender(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
event_index_counter=event_index_counter,
)
else:
worker = None
@@ -133,7 +128,6 @@ class Node:
master,
api,
node_id,
event_index_counter,
args.offline,
)
@@ -155,11 +149,11 @@ class Node:
def shutdown(self):
# if this is our second call to shutdown, just sys.exit
if self._tg.cancel_scope.cancel_called:
if self._tg.cancel_called():
import sys
sys.exit(1)
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _elect_loop(self):
with self.election_result_receiver as results:
@@ -212,8 +206,6 @@ class Node:
)
if result.is_new_master:
await anyio.sleep(0)
# Fresh counter for new session (buffer expects indices from 0)
self.event_index_counter = itertools.count()
if self.download_coordinator:
self.download_coordinator.shutdown()
self.download_coordinator = DownloadCoordinator(
@@ -224,8 +216,11 @@ class Node:
topics.DOWNLOAD_COMMANDS
),
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
event_index_counter=self.event_index_counter,
offline=self.offline,
# TODO(evan): remove
_global_event_receiver=self.router.receiver(
topics.GLOBAL_EVENTS
),
)
self._tg.start_soon(self.download_coordinator.run)
if self.worker:
@@ -242,7 +237,6 @@ class Node:
download_command_sender=self.router.sender(
topics.DOWNLOAD_COMMANDS
),
event_index_counter=self.event_index_counter,
)
self._tg.start_soon(self.worker.run)
if self.api:
@@ -258,7 +252,7 @@ def main():
target = min(max(soft, 65535), hard)
resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
mp.set_start_method("spawn")
mp.set_start_method("spawn", force=True)
# TODO: Refactor the current verbosity system
logger_setup(EXO_LOG, args.verbosity)
logger.info("Starting EXO")
@@ -267,6 +261,13 @@ def main():
if args.offline:
logger.info("Running in OFFLINE mode — no internet checks, local models only")
# Set trust_remote_code override env var for runner subprocesses
if args.trust_remote_code:
os.environ["EXO_TRUST_REMOTE_CODE"] = "1"
logger.warning(
"--trust-remote-code enabled: models may execute arbitrary code during loading"
)
# Set FAST_SYNCH override env var for runner subprocesses
if args.fast_synch is True:
os.environ["EXO_FAST_SYNCH"] = "on"
@@ -291,6 +292,7 @@ class Args(CamelCaseModel):
no_downloads: bool = False
offline: bool = False
fast_synch: bool | None = None # None = auto, True = force on, False = force off
trust_remote_code: bool = False
@classmethod
def parse(cls) -> Self:
@@ -342,6 +344,11 @@ class Args(CamelCaseModel):
action="store_true",
help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
)
parser.add_argument(
"--trust-remote-code",
action="store_true",
help="Allow models to execute custom code during tokenizer loading (security-sensitive, CLI-only)",
)
fast_synch_group = parser.add_mutually_exclusive_group()
fast_synch_group.add_argument(
"--fast-synch",
+18 -3
View File
@@ -59,7 +59,11 @@ def chat_request_to_text_generation(
chat_template_messages.append({"role": "system", "content": content})
else:
# Skip messages with no meaningful content
if msg.content is None and msg.thinking is None and msg.tool_calls is None:
if (
msg.content is None
and msg.reasoning_content is None
and msg.tool_calls is None
):
continue
if msg.role in ("user", "assistant", "developer"):
@@ -111,6 +115,11 @@ def chunk_to_response(
]
)
if chunk.is_thinking:
delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text)
else:
delta = ChatCompletionMessage(role="assistant", content=chunk.text)
return ChatCompletionResponse(
id=command_id,
created=int(time.time()),
@@ -118,7 +127,7 @@ def chunk_to_response(
choices=[
StreamingChoiceResponse(
index=0,
delta=ChatCompletionMessage(role="assistant", content=chunk.text),
delta=delta,
logprobs=logprobs,
finish_reason=chunk.finish_reason,
)
@@ -208,6 +217,7 @@ async def collect_chat_response(
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ChatCompletionResponse."""
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_calls: list[ToolCall] = []
logprobs_content: list[LogprobsContentItem] = []
model: str | None = None
@@ -228,7 +238,10 @@ async def collect_chat_response(
if model is None:
model = chunk.model
last_usage = chunk.usage or last_usage
text_parts.append(chunk.text)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.logprob is not None:
logprobs_content.append(
LogprobsContentItem(
@@ -258,6 +271,7 @@ async def collect_chat_response(
raise ValueError(error_message)
combined_text = "".join(text_parts)
combined_thinking = "".join(thinking_parts) if thinking_parts else None
assert model is not None
yield ChatCompletionResponse(
@@ -270,6 +284,7 @@ async def collect_chat_response(
message=ChatCompletionMessage(
role="assistant",
content=combined_text,
reasoning_content=combined_thinking,
tool_calls=tool_calls if tool_calls else None,
),
logprobs=Logprobs(content=logprobs_content)
+112 -21
View File
@@ -1,6 +1,7 @@
"""Claude Messages API adapter for converting requests/responses."""
import json
import re
from collections.abc import AsyncGenerator
from typing import Any
@@ -28,6 +29,8 @@ from exo.shared.types.claude_api import (
ClaudeStopReason,
ClaudeTextBlock,
ClaudeTextDelta,
ClaudeThinkingBlock,
ClaudeThinkingDelta,
ClaudeToolResultBlock,
ClaudeToolUseBlock,
ClaudeUsage,
@@ -61,6 +64,22 @@ def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str:
return "".join(sub_block.text for sub_block in block.content)
# Matches "x-anthropic-billing-header: ...;" (with optional trailing newline)
# or similar telemetry headers that change every request and break KV prefix caching.
_VOLATILE_HEADER_RE = re.compile(r"^x-anthropic-[^\n]*;\n?", re.MULTILINE)
def _strip_volatile_headers(text: str) -> str:
"""Remove Anthropic billing/telemetry headers from system prompt text.
Claude Code prepends headers like 'x-anthropic-billing-header: cc_version=...;
cc_entrypoint=...; cch=...;' that contain per-request content hashes. These
change every request and break KV prefix caching (the prefix diverges at ~20
tokens instead of matching thousands of conversation tokens).
"""
return _VOLATILE_HEADER_RE.sub("", text)
def claude_request_to_text_generation(
request: ClaudeMessagesRequest,
) -> TextGenerationTaskParams:
@@ -73,6 +92,8 @@ def claude_request_to_text_generation(
instructions = request.system
else:
instructions = "".join(block.text for block in request.system)
instructions = _strip_volatile_headers(instructions)
chat_template_messages.append({"role": "system", "content": instructions})
# Convert messages to input
@@ -85,12 +106,15 @@ def claude_request_to_text_generation(
# Process structured content blocks
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_calls: list[dict[str, Any]] = []
tool_results: list[ClaudeToolResultBlock] = []
for block in msg.content:
if isinstance(block, ClaudeTextBlock):
text_parts.append(block.text)
elif isinstance(block, ClaudeThinkingBlock):
thinking_parts.append(block.thinking)
elif isinstance(block, ClaudeToolUseBlock):
tool_calls.append(
{
@@ -106,6 +130,7 @@ def claude_request_to_text_generation(
tool_results.append(block)
content = "".join(text_parts)
reasoning_content = "".join(thinking_parts) if thinking_parts else None
# Build InputMessage from text content
if msg.role in ("user", "assistant"):
@@ -113,9 +138,14 @@ def claude_request_to_text_generation(
# Build chat_template_messages preserving tool structure
if tool_calls:
chat_template_messages.append(
{"role": "assistant", "content": content, "tool_calls": tool_calls}
)
chat_msg: dict[str, Any] = {
"role": "assistant",
"content": content,
"tool_calls": tool_calls,
}
if reasoning_content:
chat_msg["reasoning_content"] = reasoning_content
chat_template_messages.append(chat_msg)
elif tool_results:
for tr in tool_results:
chat_template_messages.append(
@@ -126,7 +156,10 @@ def claude_request_to_text_generation(
}
)
else:
chat_template_messages.append({"role": msg.role, "content": content})
chat_msg = {"role": msg.role, "content": content}
if reasoning_content:
chat_msg["reasoning_content"] = reasoning_content
chat_template_messages.append(chat_msg)
# Convert Claude tool definitions to OpenAI-style function tools
tools: list[dict[str, Any]] | None = None
@@ -143,6 +176,10 @@ def claude_request_to_text_generation(
for tool in request.tools
]
enable_thinking: bool | None = None
if request.thinking is not None:
enable_thinking = request.thinking.type in ("enabled", "adaptive")
return TextGenerationTaskParams(
model=request.model,
input=input_messages
@@ -156,6 +193,7 @@ def claude_request_to_text_generation(
stop=request.stop_sequences,
stream=request.stream,
tools=tools,
enable_thinking=enable_thinking,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
@@ -173,6 +211,7 @@ async def collect_claude_response(
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ClaudeMessagesResponse."""
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_use_blocks: list[ClaudeToolUseBlock] = []
stop_reason: ClaudeStopReason | None = None
last_usage: Usage | None = None
@@ -200,7 +239,10 @@ async def collect_claude_response(
stop_reason = "tool_use"
continue
text_parts.append(chunk.text)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -209,9 +251,12 @@ async def collect_claude_response(
raise ValueError(error_message)
combined_text = "".join(text_parts)
combined_thinking = "".join(thinking_parts)
# Build content blocks
content: list[ClaudeContentBlock] = []
if combined_thinking:
content.append(ClaudeThinkingBlock(thinking=combined_thinking))
if combined_text:
content.append(ClaudeTextBlock(text=combined_text))
content.extend(tool_use_blocks)
@@ -256,16 +301,16 @@ async def generate_claude_stream(
start_event = ClaudeMessageStartEvent(message=initial_message)
yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
# content_block_start for text block at index 0
block_start = ClaudeContentBlockStartEvent(
index=0, content_block=ClaudeTextBlock(text="")
)
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
output_tokens = 0
stop_reason: ClaudeStopReason | None = None
last_usage: Usage | None = None
next_block_index = 1 # text block is 0, tool blocks start at 1
next_block_index = 0
# Track whether we've started thinking/text blocks
thinking_block_started = False
thinking_block_index = -1
text_block_started = False
text_block_index = -1
async for chunk in chunk_stream:
if isinstance(chunk, PrefillProgressChunk):
@@ -310,12 +355,45 @@ async def generate_claude_stream(
output_tokens += 1 # Count each chunk as one token
# content_block_delta
delta_event = ClaudeContentBlockDeltaEvent(
index=0,
delta=ClaudeTextDelta(text=chunk.text),
)
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
if chunk.is_thinking:
# Start thinking block on first thinking token
if not thinking_block_started:
thinking_block_started = True
thinking_block_index = next_block_index
next_block_index += 1
block_start = ClaudeContentBlockStartEvent(
index=thinking_block_index,
content_block=ClaudeThinkingBlock(thinking=""),
)
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
delta_event = ClaudeContentBlockDeltaEvent(
index=thinking_block_index,
delta=ClaudeThinkingDelta(thinking=chunk.text),
)
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
else:
# Close thinking block when transitioning to text
if thinking_block_started and text_block_index == -1:
block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
# Start text block on first text token
if not text_block_started:
text_block_started = True
text_block_index = next_block_index
next_block_index += 1
block_start = ClaudeContentBlockStartEvent(
index=text_block_index,
content_block=ClaudeTextBlock(text=""),
)
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
delta_event = ClaudeContentBlockDeltaEvent(
index=text_block_index,
delta=ClaudeTextDelta(text=chunk.text),
)
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -324,9 +402,22 @@ async def generate_claude_stream(
if last_usage is not None:
output_tokens = last_usage.completion_tokens
# content_block_stop for text block
block_stop = ClaudeContentBlockStopEvent(index=0)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
# Close any open blocks
if thinking_block_started and text_block_index == -1:
block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
if text_block_started:
block_stop = ClaudeContentBlockStopEvent(index=text_block_index)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
if not thinking_block_started and not text_block_started:
empty_start = ClaudeContentBlockStartEvent(
index=0, content_block=ClaudeTextBlock(text="")
)
yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n"
empty_stop = ClaudeContentBlockStopEvent(index=0)
yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n"
# message_delta
message_delta = ClaudeMessageDeltaEvent(
+456
View File
@@ -0,0 +1,456 @@
from __future__ import annotations
import json
from collections.abc import AsyncGenerator
from typing import Any
from exo.shared.types.chunks import (
ErrorChunk,
PrefillProgressChunk,
TokenChunk,
ToolCallChunk,
)
from exo.shared.types.common import CommandId
from exo.shared.types.ollama_api import (
OllamaChatRequest,
OllamaChatResponse,
OllamaDoneReason,
OllamaGenerateRequest,
OllamaGenerateResponse,
OllamaMessage,
OllamaToolCall,
OllamaToolFunction,
)
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
def _map_done_reason(
finish_reason: str | None,
) -> OllamaDoneReason | None:
if finish_reason is None:
return None
if finish_reason == "stop":
return "stop"
if finish_reason == "length":
return "length"
if finish_reason in ("tool_calls", "function_call"):
return "tool_call"
if finish_reason == "error":
return "error"
return "stop"
def _try_parse_json(value: str) -> dict[str, Any] | str:
try:
return json.loads(value) # type: ignore
except json.JSONDecodeError:
return value
def _build_tool_calls(chunk: ToolCallChunk) -> list[OllamaToolCall]:
tool_calls: list[OllamaToolCall] = []
for index, tool in enumerate(chunk.tool_calls):
# tool.arguments is always str; try to parse as JSON dict for Ollama format
arguments: dict[str, Any] | str = _try_parse_json(tool.arguments)
tool_calls.append(
OllamaToolCall(
id=tool.id,
type="function",
function=OllamaToolFunction(
name=tool.name, arguments=arguments, index=index
),
)
)
return tool_calls
def _get_usage(
chunk: TokenChunk | ToolCallChunk,
) -> tuple[int | None, int | None]:
"""Extract (prompt_eval_count, eval_count) from a chunk."""
if chunk.usage is not None:
return (chunk.usage.prompt_tokens, chunk.usage.completion_tokens)
if chunk.stats is not None:
return (chunk.stats.prompt_tokens, chunk.stats.generation_tokens)
return (None, None)
def ollama_request_to_text_generation(
request: OllamaChatRequest,
) -> TextGenerationTaskParams:
"""Convert Ollama chat request to exo's internal text generation format."""
instructions: str | None = None
input_messages: list[InputMessage] = []
chat_template_messages: list[dict[str, Any]] = []
tool_message_index = 0
for msg in request.messages:
content = msg.content or ""
if msg.role == "system":
if instructions is None:
instructions = content
else:
instructions = f"{instructions}\n{content}"
chat_template_messages.append({"role": "system", "content": content})
continue
if msg.role in ("user", "assistant") and (
msg.content is not None or msg.thinking is not None or msg.tool_calls
):
input_messages.append(InputMessage(role=msg.role, content=content))
dumped: dict[str, Any] = {"role": msg.role, "content": content}
if msg.thinking is not None:
dumped["thinking"] = msg.thinking
if msg.tool_calls is not None:
tool_calls_list: list[dict[str, Any]] = []
for tc in msg.tool_calls:
function: dict[str, Any] = {
"name": tc.function.name,
"arguments": (
json.dumps(tc.function.arguments)
if isinstance(tc.function.arguments, dict)
else tc.function.arguments
),
}
if tc.function.index is not None:
function["index"] = tc.function.index
tool_call: dict[str, Any] = {"function": function}
if tc.id is not None:
tool_call["id"] = tc.id
if tc.type is not None:
tool_call["type"] = tc.type
tool_calls_list.append(tool_call)
dumped["tool_calls"] = tool_calls_list
if msg.name is not None:
dumped["name"] = msg.name
if msg.role == "tool":
tool_message_index += 1
tool_call_id = msg.tool_name or msg.name or f"tool_{tool_message_index}"
dumped["tool_call_id"] = tool_call_id
if msg.tool_name is not None:
dumped["tool_name"] = msg.tool_name
chat_template_messages.append(dumped)
options = request.options
return TextGenerationTaskParams(
model=request.model,
input=input_messages
if input_messages
else [InputMessage(role="user", content="")],
instructions=instructions,
max_output_tokens=options.num_predict if options else None,
temperature=options.temperature if options else None,
top_p=options.top_p if options else None,
top_k=options.top_k if options else None,
stop=options.stop if options else None,
seed=options.seed if options else None,
stream=request.stream,
tools=request.tools,
enable_thinking=request.think,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
)
async def generate_ollama_chat_stream(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str, None]:
"""Generate streaming responses in Ollama format (newline-delimited JSON)."""
thinking_parts: list[str] = []
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
error_response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant", content=chunk.error_message
),
done=True,
done_reason="error",
)
yield f"{error_response.model_dump_json(exclude_none=True)}\n"
return
case ToolCallChunk():
prompt_eval, eval_count = _get_usage(chunk)
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant",
content="",
tool_calls=_build_tool_calls(chunk),
thinking="".join(thinking_parts) if thinking_parts else None,
),
done=True,
done_reason="tool_call",
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
return
case TokenChunk():
done = chunk.finish_reason is not None
if chunk.is_thinking:
thinking_parts.append(chunk.text)
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant", content="", thinking=chunk.text
),
done=False,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
elif done:
prompt_eval, eval_count = _get_usage(chunk)
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant",
content=chunk.text,
),
done=True,
done_reason=_map_done_reason(chunk.finish_reason),
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
else:
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(role="assistant", content=chunk.text),
done=False,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
if done:
return
async def collect_ollama_chat_response(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str]:
"""Collect streaming chunks into a single non-streaming Ollama response.
Returns an AsyncGenerator[str] (single yield) for consistency with FastAPI
StreamingResponse cancellation handling.
"""
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_calls: list[OllamaToolCall] = []
model: str | None = None
finish_reason: str | None = None
prompt_eval_count: int | None = None
eval_count: int | None = None
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
raise ValueError(chunk.error_message or "Internal server error")
case TokenChunk():
if model is None:
model = str(chunk.model)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.finish_reason is not None:
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
case ToolCallChunk():
if model is None:
model = str(chunk.model)
tool_calls.extend(_build_tool_calls(chunk))
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
combined_text = "".join(text_parts)
combined_thinking = "".join(thinking_parts) if thinking_parts else None
assert model is not None
yield OllamaChatResponse(
model=model,
message=OllamaMessage(
role="assistant",
content=combined_text,
thinking=combined_thinking,
tool_calls=tool_calls if tool_calls else None,
),
done=True,
done_reason=_map_done_reason(finish_reason),
prompt_eval_count=prompt_eval_count,
eval_count=eval_count,
).model_dump_json(exclude_none=True)
return
# ── /api/generate ──
def ollama_generate_request_to_text_generation(
request: OllamaGenerateRequest,
) -> TextGenerationTaskParams:
"""Convert Ollama generate request to exo's internal text generation format."""
chat_template_messages: list[dict[str, Any]] = []
if request.system:
chat_template_messages.append({"role": "system", "content": request.system})
chat_template_messages.append({"role": "user", "content": request.prompt})
options = request.options
return TextGenerationTaskParams(
model=request.model,
input=[InputMessage(role="user", content=request.prompt)],
instructions=request.system,
max_output_tokens=options.num_predict if options else None,
temperature=options.temperature if options else None,
top_p=options.top_p if options else None,
top_k=options.top_k if options else None,
stop=options.stop if options else None,
seed=options.seed if options else None,
stream=request.stream,
enable_thinking=request.think,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
)
async def generate_ollama_generate_stream(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str, None]:
"""Generate streaming responses for /api/generate in Ollama NDJSON format."""
thinking_parts: list[str] = []
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
resp = OllamaGenerateResponse(
model=str(chunk.model),
response="",
done=True,
done_reason="error",
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
return
case ToolCallChunk():
# generate endpoint doesn't support tools; emit as done
prompt_eval, eval_count = _get_usage(chunk)
resp = OllamaGenerateResponse(
model=str(chunk.model),
response="",
done=True,
done_reason="stop",
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
return
case TokenChunk():
done = chunk.finish_reason is not None
if chunk.is_thinking:
thinking_parts.append(chunk.text)
resp = OllamaGenerateResponse(
model=str(chunk.model),
response="",
thinking=chunk.text,
done=False,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
elif done:
prompt_eval, eval_count = _get_usage(chunk)
resp = OllamaGenerateResponse(
model=str(chunk.model),
response=chunk.text,
done=True,
done_reason=_map_done_reason(chunk.finish_reason),
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
else:
resp = OllamaGenerateResponse(
model=str(chunk.model),
response=chunk.text,
done=False,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
if done:
return
async def collect_ollama_generate_response(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str]:
"""Collect chunks into a single non-streaming /api/generate response."""
text_parts: list[str] = []
thinking_parts: list[str] = []
model: str | None = None
finish_reason: str | None = None
prompt_eval_count: int | None = None
eval_count: int | None = None
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
raise ValueError(chunk.error_message or "Internal server error")
case TokenChunk():
if model is None:
model = str(chunk.model)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.finish_reason is not None:
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
case ToolCallChunk():
if model is None:
model = str(chunk.model)
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
assert model is not None
yield OllamaGenerateResponse(
model=model,
response="".join(text_parts),
thinking="".join(thinking_parts) if thinking_parts else None,
done=True,
done_reason=_map_done_reason(finish_reason),
prompt_eval_count=prompt_eval_count,
eval_count=eval_count,
).model_dump_json(exclude_none=True)
return
+207 -30
View File
@@ -29,6 +29,12 @@ from exo.shared.types.openai_responses import (
ResponseOutputItemAddedEvent,
ResponseOutputItemDoneEvent,
ResponseOutputText,
ResponseReasoningItem,
ResponseReasoningSummaryPartAddedEvent,
ResponseReasoningSummaryPartDoneEvent,
ResponseReasoningSummaryText,
ResponseReasoningSummaryTextDeltaEvent,
ResponseReasoningSummaryTextDoneEvent,
ResponsesRequest,
ResponsesResponse,
ResponsesStreamEvent,
@@ -141,7 +147,9 @@ async def collect_responses_response(
"""Collect all token chunks and return a single ResponsesResponse."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
reasoning_id = f"rs_{command_id}"
accumulated_text = ""
thinking_parts: list[str] = []
function_call_items: list[ResponseFunctionCallItem] = []
last_usage: Usage | None = None
error_message: str | None = None
@@ -168,6 +176,10 @@ async def collect_responses_response(
)
continue
if chunk.is_thinking:
thinking_parts.append(chunk.text)
continue
accumulated_text += chunk.text
if error_message is not None:
@@ -182,13 +194,21 @@ async def collect_responses_response(
total_tokens=last_usage.total_tokens,
)
output: list[ResponseItem] = [
output: list[ResponseItem] = []
if thinking_parts:
output.append(
ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))],
)
)
output.append(
ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text=accumulated_text)],
status="completed",
)
]
)
output.extend(function_call_items)
yield ResponsesResponse(
@@ -212,6 +232,7 @@ async def generate_responses_stream(
"""Generate OpenAI Responses API streaming events from TokenChunks."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
reasoning_id = f"rs_{command_id}"
seq = count(1)
# response.created
@@ -233,32 +254,17 @@ async def generate_responses_stream(
)
yield _format_sse(in_progress_event)
# response.output_item.added
initial_item = ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text="")],
status="in_progress",
)
item_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq), output_index=0, item=initial_item
)
yield _format_sse(item_added)
# response.content_part.added
initial_part = ResponseOutputText(text="")
part_added = ResponseContentPartAddedEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
content_index=0,
part=initial_part,
)
yield _format_sse(part_added)
accumulated_text = ""
accumulated_thinking = ""
function_call_items: list[ResponseFunctionCallItem] = []
last_usage: Usage | None = None
next_output_index = 1 # message item is at 0
next_output_index = 0
# Track dynamic block creation
reasoning_started = False
reasoning_output_index = -1
message_started = False
message_output_index = -1
async for chunk in chunk_stream:
if isinstance(chunk, PrefillProgressChunk):
@@ -327,23 +333,184 @@ async def generate_responses_stream(
next_output_index += 1
continue
if chunk.is_thinking:
# Start reasoning block on first thinking token
if not reasoning_started:
reasoning_started = True
reasoning_output_index = next_output_index
next_output_index += 1
# response.output_item.added for reasoning
reasoning_item = ResponseReasoningItem(
id=reasoning_id,
summary=[],
status="in_progress",
)
rs_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq),
output_index=reasoning_output_index,
item=reasoning_item,
)
yield _format_sse(rs_added)
# response.reasoning_summary_part.added
part_added = ResponseReasoningSummaryPartAddedEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
part=ResponseReasoningSummaryText(text=""),
)
yield _format_sse(part_added)
accumulated_thinking += chunk.text
# response.reasoning_summary_text.delta
rs_delta = ResponseReasoningSummaryTextDeltaEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
delta=chunk.text,
)
yield _format_sse(rs_delta)
continue
# Close reasoning block when transitioning to text
if reasoning_started and not message_started:
# response.reasoning_summary_text.done
rs_text_done = ResponseReasoningSummaryTextDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
text=accumulated_thinking,
)
yield _format_sse(rs_text_done)
# response.reasoning_summary_part.done
rs_part_done = ResponseReasoningSummaryPartDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
part=ResponseReasoningSummaryText(text=accumulated_thinking),
)
yield _format_sse(rs_part_done)
# response.output_item.done for reasoning
rs_item_done = ResponseOutputItemDoneEvent(
sequence_number=next(seq),
output_index=reasoning_output_index,
item=ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
),
)
yield _format_sse(rs_item_done)
# Start message block on first text token
if not message_started:
message_started = True
message_output_index = next_output_index
next_output_index += 1
initial_item = ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text="")],
status="in_progress",
)
item_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq),
output_index=message_output_index,
item=initial_item,
)
yield _format_sse(item_added)
initial_part = ResponseOutputText(text="")
part_added = ResponseContentPartAddedEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=message_output_index,
content_index=0,
part=initial_part,
)
yield _format_sse(part_added)
accumulated_text += chunk.text
# response.output_text.delta
delta_event = ResponseTextDeltaEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
output_index=message_output_index,
content_index=0,
delta=chunk.text,
)
yield _format_sse(delta_event)
# Close reasoning block if it was never followed by text
if reasoning_started and not message_started:
rs_text_done = ResponseReasoningSummaryTextDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
text=accumulated_thinking,
)
yield _format_sse(rs_text_done)
rs_part_done = ResponseReasoningSummaryPartDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
part=ResponseReasoningSummaryText(text=accumulated_thinking),
)
yield _format_sse(rs_part_done)
rs_item_done = ResponseOutputItemDoneEvent(
sequence_number=next(seq),
output_index=reasoning_output_index,
item=ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
),
)
yield _format_sse(rs_item_done)
# If no message block was started, create one now (empty text)
if not message_started:
message_output_index = next_output_index
next_output_index += 1
initial_item = ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text="")],
status="in_progress",
)
item_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq),
output_index=message_output_index,
item=initial_item,
)
yield _format_sse(item_added)
initial_part = ResponseOutputText(text="")
part_added_evt = ResponseContentPartAddedEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=message_output_index,
content_index=0,
part=initial_part,
)
yield _format_sse(part_added_evt)
# response.output_text.done
text_done = ResponseTextDoneEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
output_index=message_output_index,
content_index=0,
text=accumulated_text,
)
@@ -354,7 +521,7 @@ async def generate_responses_stream(
part_done = ResponseContentPartDoneEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
output_index=message_output_index,
content_index=0,
part=final_part,
)
@@ -367,7 +534,9 @@ async def generate_responses_stream(
status="completed",
)
item_done = ResponseOutputItemDoneEvent(
sequence_number=next(seq), output_index=0, item=final_message_item
sequence_number=next(seq),
output_index=message_output_index,
item=final_message_item,
)
yield _format_sse(item_done)
@@ -381,7 +550,15 @@ async def generate_responses_stream(
)
# response.completed
output: list[ResponseItem] = [final_message_item]
output: list[ResponseItem] = []
if reasoning_started:
output.append(
ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
)
)
output.append(final_message_item)
output.extend(function_call_items)
final_response = ResponsesResponse(
id=response_id,
+222 -32
View File
@@ -11,8 +11,7 @@ from typing import Annotated, Literal, cast
from uuid import uuid4
import anyio
from anyio import BrokenResourceError, create_task_group
from anyio.abc import TaskGroup
from anyio import BrokenResourceError
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -32,6 +31,14 @@ from exo.master.adapters.claude import (
collect_claude_response,
generate_claude_stream,
)
from exo.master.adapters.ollama import (
collect_ollama_chat_response,
collect_ollama_generate_response,
generate_ollama_chat_stream,
generate_ollama_generate_stream,
ollama_generate_request_to_text_generation,
ollama_request_to_text_generation,
)
from exo.master.adapters.responses import (
collect_responses_response,
generate_responses_stream,
@@ -43,6 +50,7 @@ from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
from exo.shared.constants import (
DASHBOARD_DIR,
EXO_CACHE_HOME,
EXO_EVENT_LOG_DIR,
EXO_IMAGE_CACHE_DIR,
EXO_MAX_CHUNK_SIZE,
@@ -132,16 +140,28 @@ from exo.shared.types.commands import (
TaskFinished,
TextGeneration,
)
from exo.shared.types.common import CommandId, Id, NodeId, SessionId
from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
from exo.shared.types.events import (
ChunkGenerated,
Event,
ForwarderEvent,
GlobalForwarderEvent,
IndexedEvent,
PrefillProgress,
TracesMerged,
)
from exo.shared.types.memory import Memory
from exo.shared.types.ollama_api import (
OllamaChatRequest,
OllamaChatResponse,
OllamaGenerateRequest,
OllamaGenerateResponse,
OllamaModelDetails,
OllamaModelTag,
OllamaPsModel,
OllamaPsResponse,
OllamaShowRequest,
OllamaShowResponse,
OllamaTagsResponse,
)
from exo.shared.types.openai_responses import (
ResponsesRequest,
ResponsesResponse,
@@ -153,8 +173,10 @@ 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.event_buffer import OrderedBuffer
from exo.utils.task_group import TaskGroup
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
@@ -177,8 +199,7 @@ class API:
session_id: SessionId,
*,
port: int,
# Ideally this would be a MasterForwarderEvent but type system says no :(
global_event_receiver: Receiver[ForwarderEvent],
global_event_receiver: Receiver[GlobalForwarderEvent],
command_sender: Sender[ForwarderCommand],
download_command_sender: Sender[ForwarderDownloadCommand],
# This lets us pause the API if an election is running
@@ -186,6 +207,7 @@ class API:
) -> 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.global_event_receiver = global_event_receiver
@@ -230,13 +252,14 @@ class API:
CommandId, Sender[ImageChunk | ErrorChunk]
] = {}
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
self._tg: TaskGroup | None = None
self._tg: TaskGroup = TaskGroup()
def reset(self, new_session_id: SessionId, result_clock: int):
logger.info("Resetting API State")
self._event_log.close()
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
self.state = State()
self._system_id = SystemId()
self.session_id = new_session_id
self.event_buffer = OrderedBuffer[Event]()
self._text_generation_queues = {}
@@ -301,6 +324,21 @@ class API:
self.app.get("/images/{image_id}")(self.get_image)
self.app.post("/v1/messages", response_model=None)(self.claude_messages)
self.app.post("/v1/responses", response_model=None)(self.openai_responses)
# Ollama API
self.app.head("/ollama/")(self.ollama_version)
self.app.head("/ollama/api/version")(self.ollama_version)
self.app.post("/ollama/api/chat", response_model=None)(self.ollama_chat)
self.app.post("/ollama/api/api/chat", response_model=None)(self.ollama_chat)
self.app.post("/ollama/api/v1/chat", response_model=None)(self.ollama_chat)
self.app.post("/ollama/api/generate", response_model=None)(self.ollama_generate)
self.app.get("/ollama/api/tags")(self.ollama_tags)
self.app.get("/ollama/api/api/tags")(self.ollama_tags)
self.app.get("/ollama/api/v1/tags")(self.ollama_tags)
self.app.post("/ollama/api/show")(self.ollama_show)
self.app.get("/ollama/api/ps")(self.ollama_ps)
self.app.get("/ollama/api/version")(self.ollama_version)
self.app.get("/state")(lambda: self.state)
self.app.get("/events")(self.stream_events)
self.app.post("/download/start")(self.start_download)
@@ -309,6 +347,8 @@ class API:
self.app.get("/v1/traces/{task_id}")(self.get_trace)
self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
self.app.get("/onboarding")(self.get_onboarding)
self.app.post("/onboarding")(self.complete_onboarding)
async def place_instance(self, payload: PlaceInstanceParams):
command = PlaceInstance(
@@ -554,7 +594,7 @@ class API:
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
@@ -902,7 +942,7 @@ class API:
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
@@ -988,7 +1028,7 @@ class API:
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
@@ -1294,6 +1334,163 @@ class API:
media_type="application/json",
)
async def _ollama_root(self) -> JSONResponse:
"""Respond to HEAD / from Ollama CLI connectivity checks."""
return JSONResponse(content="Ollama is running")
async def ollama_chat(
self, request: Request
) -> OllamaChatResponse | StreamingResponse:
"""Ollama Chat API — accepts JSON regardless of Content-Type."""
body = await request.body()
payload = OllamaChatRequest.model_validate_json(body)
task_params = ollama_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
command = TextGeneration(task_params=task_params)
await self._send(command)
if payload.stream:
return StreamingResponse(
generate_ollama_chat_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/x-ndjson",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
else:
return StreamingResponse(
collect_ollama_chat_response(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
async def ollama_generate(
self, request: Request
) -> OllamaGenerateResponse | StreamingResponse:
"""Ollama Generate API — accepts JSON regardless of Content-Type."""
body = await request.body()
payload = OllamaGenerateRequest.model_validate_json(body)
task_params = ollama_generate_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
command = TextGeneration(task_params=task_params)
await self._send(command)
if payload.stream:
return StreamingResponse(
generate_ollama_generate_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/x-ndjson",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
else:
return StreamingResponse(
collect_ollama_generate_response(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
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()
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
]
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
return OllamaTagsResponse(
models=[
OllamaModelTag(
name=str(card.model_id),
model=str(card.model_id),
modified_at=now,
size=card.storage_size.in_bytes,
digest="sha256:000000000000",
details=OllamaModelDetails(
family=none_if_empty(card.family),
quantization_level=none_if_empty(card.quantization),
),
)
for card in cards
]
)
async def ollama_show(self, request: Request) -> OllamaShowResponse:
"""Returns model information in Ollama show format."""
body = await request.body()
payload = OllamaShowRequest.model_validate_json(body)
model_name = payload.name or payload.model
if not model_name:
raise HTTPException(status_code=400, detail="name or model is required")
try:
card = await ModelCard.load(ModelId(model_name))
except Exception as exc:
raise HTTPException(
status_code=404, detail=f"Model not found: {model_name}"
) from exc
return OllamaShowResponse(
modelfile=f"FROM {card.model_id}",
template="{{ .Prompt }}",
details=OllamaModelDetails(
family=card.family or None,
quantization_level=card.quantization or None,
),
)
async def ollama_ps(self) -> OllamaPsResponse:
"""Returns list of running models (active instances)."""
models: list[OllamaPsModel] = []
seen: set[str] = set()
for instance in self.state.instances.values():
model_id = str(instance.shard_assignments.model_id)
if model_id in seen:
continue
seen.add(model_id)
models.append(
OllamaPsModel(
name=model_id,
model=model_id,
size=0,
)
)
return OllamaPsResponse(models=models)
async def ollama_version(self) -> dict[str, str]:
"""Returns version information for Ollama API compatibility."""
return {"version": "exo v1.0"}
def _calculate_total_available_memory(self) -> Memory:
"""Calculate total available memory across all nodes in bytes."""
total_available = Memory()
@@ -1323,7 +1520,7 @@ class API:
name=card.model_id.short(),
description="",
tags=[],
storage_size_megabytes=int(card.storage_size.in_mb),
storage_size_megabytes=card.storage_size.in_mb,
supports_tensor=card.supports_tensor,
tasks=[task.value for task in card.tasks],
is_custom=is_custom_card(card.model_id),
@@ -1394,8 +1591,7 @@ class API:
shutdown_ev = anyio.Event()
try:
async with create_task_group() as tg:
self._tg = tg
async with self._tg as tg:
logger.info("Starting API")
tg.start_soon(self._apply_state)
tg.start_soon(self._pause_on_new_election)
@@ -1429,6 +1625,8 @@ class API:
async def _apply_state(self):
with self.global_event_receiver as events:
async for f_event in events:
if f_event.session != self.session_id:
continue
if f_event.origin != self.session_id.master_node_id:
continue
self.event_buffer.ingest(f_event.origin_idx, f_event.event)
@@ -1455,22 +1653,6 @@ class API:
await queue.send(event.chunk)
except BrokenResourceError:
self._text_generation_queues.pop(event.command_id, None)
elif isinstance(event, PrefillProgress):
if queue := self._text_generation_queues.get(
event.command_id, None
):
try:
await queue.send(
PrefillProgressChunk(
model=event.model,
processed_tokens=event.processed_tokens,
total_tokens=event.total_tokens,
)
)
except BrokenResourceError:
self._text_generation_queues.pop(event.command_id, None)
if isinstance(event, TracesMerged):
self._save_merged_trace(event)
@@ -1508,12 +1690,12 @@ class API:
while self.paused:
await self.paused_ev.wait()
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
async def _send_download(self, command: DownloadCommand):
await self.download_command_sender.send(
ForwarderDownloadCommand(origin=self.node_id, command=command)
ForwarderDownloadCommand(origin=self._system_id, command=command)
)
async def start_download(
@@ -1635,3 +1817,11 @@ class API:
media_type="application/json",
filename=f"trace_{task_id}.json",
)
async def get_onboarding(self) -> JSONResponse:
return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
async def complete_onboarding(self) -> JSONResponse:
ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True)
ONBOARDING_COMPLETE_FILE.write_text("true")
return JSONResponse({"completed": True})
+15 -13
View File
@@ -1,7 +1,6 @@
from datetime import datetime, timedelta, timezone
import anyio
from anyio.abc import TaskGroup
from loguru import logger
from exo.master.event_log import DiskEventLog
@@ -29,13 +28,14 @@ from exo.shared.types.commands import (
TestCommand,
TextGeneration,
)
from exo.shared.types.common import CommandId, NodeId, SessionId
from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
from exo.shared.types.events import (
Event,
ForwarderEvent,
GlobalForwarderEvent,
IndexedEvent,
InputChunkReceived,
InstanceDeleted,
LocalForwarderEvent,
NodeGatheredInfo,
NodeTimedOut,
TaskCreated,
@@ -62,6 +62,7 @@ from exo.shared.types.tasks import (
from exo.shared.types.worker.instances import InstanceId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import MultiSourceBuffer
from exo.utils.task_group import TaskGroup
class Master:
@@ -71,12 +72,12 @@ class Master:
session_id: SessionId,
*,
command_receiver: Receiver[ForwarderCommand],
local_event_receiver: Receiver[ForwarderEvent],
global_event_sender: Sender[ForwarderEvent],
local_event_receiver: Receiver[LocalForwarderEvent],
global_event_sender: Sender[GlobalForwarderEvent],
download_command_sender: Sender[ForwarderDownloadCommand],
):
self.state = State()
self._tg: TaskGroup = anyio.create_task_group()
self._tg: TaskGroup = TaskGroup()
self.node_id = node_id
self.session_id = session_id
self.command_task_mapping: dict[CommandId, TaskId] = {}
@@ -87,10 +88,11 @@ class Master:
send, recv = channel[Event]()
self.event_sender: Sender[Event] = send
self._loopback_event_receiver: Receiver[Event] = recv
self._loopback_event_sender: Sender[ForwarderEvent] = (
self._loopback_event_sender: Sender[LocalForwarderEvent] = (
local_event_receiver.clone_sender()
)
self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
self._system_id = SystemId()
self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
self._expected_ranks: dict[TaskId, set[int]] = {}
@@ -114,7 +116,7 @@ class Master:
async def shutdown(self):
logger.info("Stopping Master")
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _command_processor(self) -> None:
with self.command_receiver as commands:
@@ -288,7 +290,7 @@ class Master:
):
await self.download_command_sender.send(
ForwarderDownloadCommand(
origin=self.node_id, command=cmd
origin=self._system_id, command=cmd
)
)
generated_events.extend(transition_events)
@@ -414,8 +416,8 @@ class Master:
with self._loopback_event_receiver as events:
async for event in events:
await self._loopback_event_sender.send(
ForwarderEvent(
origin=NodeId(f"master_{self.node_id}"),
LocalForwarderEvent(
origin=self._system_id,
origin_idx=local_index,
session=self.session_id,
event=event,
@@ -427,7 +429,7 @@ class Master:
async def _send_event(self, event: IndexedEvent):
# Convenience method since this line is ugly
await self.global_event_sender.send(
ForwarderEvent(
GlobalForwarderEvent(
origin=self.node_id,
origin_idx=event.idx,
session=self.session_id,
+26 -6
View File
@@ -112,7 +112,11 @@ def place_instance(
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
]
if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
if command.instance_meta == InstanceMeta.MlxJaccl:
if not smallest_rdma_cycles:
raise ValueError(
"Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
)
smallest_cycles = smallest_rdma_cycles
cycles_with_leaf_nodes: list[Cycle] = [
@@ -129,6 +133,11 @@ def place_instance(
),
)
# Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
if len(selected_cycle) == 1:
command.instance_meta = InstanceMeta.MlxRing
command.sharding = Sharding.Pipeline
shard_assignments = get_shard_assignments(
command.model_card, selected_cycle, command.sharding, node_memory
)
@@ -138,18 +147,29 @@ def place_instance(
instance_id = InstanceId()
target_instances = dict(deepcopy(current_instances))
if len(selected_cycle) == 1:
command.instance_meta = InstanceMeta.MlxRing
# TODO: Single node instances
match command.instance_meta:
case InstanceMeta.MlxJaccl:
# TODO(evan): shard assignments should contain information about ranks, this is ugly
def get_device_rank(node_id: NodeId) -> int:
runner_id = shard_assignments.node_to_runner[node_id]
shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
assert shard_metadata is not None
return shard_metadata.device_rank
zero_node_ids = [
node_id
for node_id in selected_cycle.node_ids
if get_device_rank(node_id) == 0
]
assert len(zero_node_ids) == 1
coordinator_node_id = zero_node_ids[0]
mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
[node_id for node_id in selected_cycle],
cycle_digraph,
)
mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
coordinator=selected_cycle.node_ids[0],
coordinator=coordinator_node_id,
coordinator_port=random_ephemeral_port(),
cycle_digraph=cycle_digraph,
node_network=node_network,
+32 -16
View File
@@ -102,22 +102,21 @@ def _allocate_and_validate_layers(
layer_allocations = allocate_layers_proportionally(
total_layers=model_card.n_layers,
memory_fractions=[
node_memory[node_id].ram_available.in_bytes / total_memory.in_bytes
for node_id in node_ids
node_memory[node_id].ram_available / total_memory for node_id in node_ids
],
)
total_storage_bytes = model_card.storage_size.in_bytes
total_storage = model_card.storage_size
total_layers = model_card.n_layers
for i, node_id in enumerate(node_ids):
node_layers = layer_allocations[i]
required_memory = (total_storage_bytes * node_layers) // total_layers
available_memory = node_memory[node_id].ram_available.in_bytes
required_memory = (total_storage * node_layers) // total_layers
available_memory = node_memory[node_id].ram_available
if required_memory > available_memory:
raise ValueError(
f"Node {i} ({node_id}) has insufficient memory: "
f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
f"but only has {available_memory / (1024**3):.2f} GB available"
f"requires {required_memory.in_gb:.2f} GB for {node_layers} layers, "
f"but only has {available_memory.in_gb:.2f} GB available"
)
return layer_allocations
@@ -342,6 +341,7 @@ def _find_ip_prioritised(
other_node_id: NodeId,
cycle_digraph: Topology,
node_network: Mapping[NodeId, NodeNetworkInfo],
ring: bool,
) -> str | None:
"""Find an IP address between nodes with prioritization.
@@ -354,13 +354,27 @@ def _find_ip_prioritised(
ip_to_type = {
iface.ip_address: iface.interface_type for iface in other_network.interfaces
}
priority = {
"ethernet": 0,
"wifi": 1,
"unknown": 2,
"maybe_ethernet": 3,
"thunderbolt": 4,
}
# Ring should prioritise fastest connection. As a best-effort, we prioritise TB.
# TODO: Profile and get actual connection speeds.
if ring:
priority = {
"thunderbolt": 0,
"maybe_ethernet": 1,
"ethernet": 2,
"wifi": 3,
"unknown": 4,
}
# RDMA prefers ethernet coordinator
else:
priority = {
"ethernet": 0,
"wifi": 1,
"unknown": 2,
"maybe_ethernet": 3,
"thunderbolt": 4,
}
return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
@@ -400,7 +414,7 @@ def get_mlx_ring_hosts_by_node(
continue
connection_ip = _find_ip_prioritised(
node_id, other_node_id, cycle_digraph, node_network
node_id, other_node_id, cycle_digraph, node_network, ring=True
)
if connection_ip is None:
raise ValueError(
@@ -431,7 +445,9 @@ def get_mlx_jaccl_coordinators(
if n == coordinator:
return "0.0.0.0"
ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network)
ip = _find_ip_prioritised(
n, coordinator, cycle_digraph, node_network, ring=False
)
if ip is not None:
return ip
+4 -5
View File
@@ -261,7 +261,7 @@ class TestGenerateClaudeStreamToolUse:
parsed = _parse_sse_events(events)
# Two tool block starts (at indices 1 and 2)
# Two tool block starts (at indices 0 and 1 — no text block when only tools)
tool_starts = [
e
for e in parsed
@@ -270,12 +270,11 @@ class TestGenerateClaudeStreamToolUse:
== "tool_use"
]
assert len(tool_starts) == 2
assert tool_starts[0]["index"] == 1
assert tool_starts[1]["index"] == 2
assert tool_starts[0]["index"] == 0
assert tool_starts[1]["index"] == 1
# Two tool block stops (at indices 1 and 2), plus text block stop at 0
# Two tool block stops (at indices 0 and 1)
block_stops = [e for e in parsed if e.get("type") == "content_block_stop"]
stop_indices = [e["index"] for e in block_stops]
assert 0 in stop_indices
assert 1 in stop_indices
assert 2 in stop_indices
+10 -10
View File
@@ -15,11 +15,12 @@ from exo.shared.types.commands import (
PlaceInstance,
TextGeneration,
)
from exo.shared.types.common import ModelId, NodeId, SessionId
from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
from exo.shared.types.events import (
ForwarderEvent,
GlobalForwarderEvent,
IndexedEvent,
InstanceCreated,
LocalForwarderEvent,
NodeGatheredInfo,
TaskCreated,
)
@@ -42,12 +43,12 @@ from exo.utils.channels import channel
@pytest.mark.asyncio
async def test_master():
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_peer_id().to_base58())
node_id = NodeId(keypair.to_node_id())
session_id = SessionId(master_node_id=node_id, election_clock=0)
ge_sender, global_event_receiver = channel[ForwarderEvent]()
ge_sender, global_event_receiver = channel[GlobalForwarderEvent]()
command_sender, co_receiver = channel[ForwarderCommand]()
local_event_sender, le_receiver = channel[ForwarderEvent]()
local_event_sender, le_receiver = channel[LocalForwarderEvent]()
fcds, _fcdr = channel[ForwarderDownloadCommand]()
all_events: list[IndexedEvent] = []
@@ -75,13 +76,12 @@ async def test_master():
async with anyio.create_task_group() as tg:
tg.start_soon(master.run)
sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender")
# inject a NodeGatheredInfo event
logger.info("inject a NodeGatheredInfo event")
await local_event_sender.send(
ForwarderEvent(
LocalForwarderEvent(
origin_idx=0,
origin=sender_node_id,
origin=SystemId("Worker"),
session=session_id,
event=(
NodeGatheredInfo(
@@ -108,7 +108,7 @@ async def test_master():
logger.info("inject a CreateInstance Command")
await command_sender.send(
ForwarderCommand(
origin=node_id,
origin=SystemId("API"),
command=(
PlaceInstance(
command_id=CommandId(),
@@ -133,7 +133,7 @@ async def test_master():
logger.info("inject a TextGeneration Command")
await command_sender.send(
ForwarderCommand(
origin=node_id,
origin=SystemId("API"),
command=(
TextGeneration(
command_id=CommandId(),
+120 -4
View File
@@ -14,10 +14,12 @@ from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.topology import Topology
from exo.shared.types.commands import PlaceInstance
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.events import InstanceCreated, InstanceDeleted
from exo.shared.types.events import InstanceCreated, InstanceDeleted, TaskStatusUpdated
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.tasks import TaskId, TaskStatus, TextGeneration
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.instances import (
Instance,
@@ -80,8 +82,8 @@ def test_get_instance_placements_create_instance(
):
# arrange
model_card.n_layers = total_layers
model_card.storage_size.in_bytes = sum(
available_memory
model_card.storage_size = Memory.from_bytes(
sum(available_memory)
) # make it exactly fit across all nodes
topology = Topology()
@@ -349,7 +351,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
# arrange
topology = Topology()
model_card.n_layers = 12
model_card.storage_size.in_bytes = 1500
model_card.storage_size = Memory.from_bytes(1500)
node_a = NodeId()
node_b = NodeId()
@@ -456,3 +458,117 @@ def test_tensor_rdma_backend_connectivity_matrix(
else:
ip_part = coordinator.split(":")[0]
assert len(ip_part.split(".")) == 4
def _make_task(
instance_id: InstanceId,
status: TaskStatus = TaskStatus.Running,
) -> TextGeneration:
return TextGeneration(
task_id=TaskId(),
task_status=status,
instance_id=instance_id,
command_id=CommandId(),
task_params=TextGenerationTaskParams(
model=ModelId("test-model"),
input=[InputMessage(role="user", content="hello")],
),
)
def test_get_transition_events_delete_instance_cancels_running_tasks(
instance: Instance,
):
# arrange
instance_id = InstanceId()
current_instances: dict[InstanceId, Instance] = {instance_id: instance}
target_instances: dict[InstanceId, Instance] = {}
task = _make_task(instance_id, TaskStatus.Running)
tasks = {task.task_id: task}
# act
events = get_transition_events(current_instances, target_instances, tasks)
# assert cancellation event should come before the deletion event
assert len(events) == 2
assert isinstance(events[0], TaskStatusUpdated)
assert events[0].task_id == task.task_id
assert events[0].task_status == TaskStatus.Cancelled
assert isinstance(events[1], InstanceDeleted)
assert events[1].instance_id == instance_id
def test_get_transition_events_delete_instance_cancels_pending_tasks(
instance: Instance,
):
# arrange
instance_id = InstanceId()
current_instances: dict[InstanceId, Instance] = {instance_id: instance}
target_instances: dict[InstanceId, Instance] = {}
task = _make_task(instance_id, TaskStatus.Pending)
tasks = {task.task_id: task}
# act
events = get_transition_events(current_instances, target_instances, tasks)
# assert
assert len(events) == 2
assert isinstance(events[0], TaskStatusUpdated)
assert events[0].task_id == task.task_id
assert events[0].task_status == TaskStatus.Cancelled
assert isinstance(events[1], InstanceDeleted)
def test_get_transition_events_delete_instance_ignores_completed_tasks(
instance: Instance,
):
# arrange
instance_id = InstanceId()
current_instances: dict[InstanceId, Instance] = {instance_id: instance}
target_instances: dict[InstanceId, Instance] = {}
tasks = {
t.task_id: t
for t in [
_make_task(instance_id, TaskStatus.Complete),
_make_task(instance_id, TaskStatus.Failed),
_make_task(instance_id, TaskStatus.TimedOut),
_make_task(instance_id, TaskStatus.Cancelled),
]
}
# act
events = get_transition_events(current_instances, target_instances, tasks)
# assert only the InstanceDeleted event, no cancellations
assert len(events) == 1
assert isinstance(events[0], InstanceDeleted)
def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
instance: Instance,
):
# arrange
instance_id_a = InstanceId()
instance_id_b = InstanceId()
current_instances: dict[InstanceId, Instance] = {
instance_id_a: instance,
instance_id_b: instance,
}
# only delete instance A, keep instance B
target_instances: dict[InstanceId, Instance] = {instance_id_b: instance}
task_a = _make_task(instance_id_a, TaskStatus.Running)
task_b = _make_task(instance_id_b, TaskStatus.Running)
tasks = {task_a.task_id: task_a, task_b.task_id: task_b}
# act
events = get_transition_events(current_instances, target_instances, tasks)
# assert only task_a should be cancelled
cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)]
delete_events = [e for e in events if isinstance(e, InstanceDeleted)]
assert len(cancel_events) == 1
assert cancel_events[0].task_id == task_a.task_id
assert cancel_events[0].task_status == TaskStatus.Cancelled
assert len(delete_events) == 1
assert delete_events[0].instance_id == instance_id_a
+1 -1
View File
@@ -30,7 +30,7 @@ class ConnectionMessage(CamelCaseModel):
@classmethod
def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
return cls(
node_id=NodeId(update.peer_id.to_base58()),
node_id=NodeId(update.peer_id),
connection_type=ConnectionMessageType.from_update_type(update.update_type),
remote_ipv4=update.remote_ipv4,
remote_tcp_port=update.remote_tcp_port,
+12 -12
View File
@@ -8,14 +8,13 @@ from typing import cast
from anyio import (
BrokenResourceError,
ClosedResourceError,
create_task_group,
move_on_after,
sleep_forever,
)
from anyio.abc import TaskGroup
from exo_pyo3_bindings import (
AllQueuesFullError,
Keypair,
MessageTooLargeError,
NetworkingHandle,
NoPeersSubscribedToTopicError,
)
@@ -25,6 +24,7 @@ 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 CamelCaseModel
from exo.utils.task_group import TaskGroup
from .connection_message import ConnectionMessage
from .topics import CONNECTION_MESSAGES, PublishPolicy, TypedTopic
@@ -111,10 +111,9 @@ class Router:
self._net: NetworkingHandle = handle
self._tmp_networking_sender: Sender[tuple[str, bytes]] | None = send
self._id_count = count()
self._tg: TaskGroup | None = None
self._tg: TaskGroup = TaskGroup()
async def register_topic[T: CamelCaseModel](self, topic: TypedTopic[T]):
assert self._tg is None, "Attempted to register topic after setup time"
send = self._tmp_networking_sender
if send:
self._tmp_networking_sender = None
@@ -148,8 +147,7 @@ class Router:
async def run(self):
logger.debug("Starting Router")
try:
async with create_task_group() as tg:
self._tg = tg
async with self._tg as tg:
for topic in self.topic_routers:
router = self.topic_routers[topic]
tg.start_soon(router.run)
@@ -165,9 +163,7 @@ class Router:
async def shutdown(self):
logger.debug("Shutting down Router")
if not self._tg:
return
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _networking_subscribe(self, topic: str):
await self._net.gossipsub_subscribe(topic)
@@ -211,6 +207,10 @@ class Router:
pass
except AllQueuesFullError:
logger.warning(f"All peer queues full, dropping message on {topic}")
except MessageTooLargeError:
logger.warning(
f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
)
def get_node_id_keypair(
@@ -221,7 +221,7 @@ def get_node_id_keypair(
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_ed25519()
return Keypair.generate()
def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path:
return Path(str(path) + ".lock")
@@ -235,12 +235,12 @@ def get_node_id_keypair(
protobuf_encoded = f.read()
try: # if decoded successfully, save & return
return Keypair.from_protobuf_encoding(protobuf_encoded)
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_ed25519()
f.write(keypair.to_protobuf_encoding())
f.write(keypair.to_bytes())
return keypair
+4 -3
View File
@@ -5,7 +5,8 @@ from exo.routing.connection_message import ConnectionMessage
from exo.shared.election import ElectionMessage
from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand
from exo.shared.types.events import (
ForwarderEvent,
GlobalForwarderEvent,
LocalForwarderEvent,
)
from exo.utils.pydantic_ext import CamelCaseModel
@@ -36,8 +37,8 @@ class TypedTopic[T: CamelCaseModel]:
return self.model_type.model_validate_json(b.decode("utf-8"))
GLOBAL_EVENTS = TypedTopic("global_events", PublishPolicy.Always, ForwarderEvent)
LOCAL_EVENTS = TypedTopic("local_events", PublishPolicy.Always, ForwarderEvent)
GLOBAL_EVENTS = TypedTopic("global_events", PublishPolicy.Always, GlobalForwarderEvent)
LOCAL_EVENTS = TypedTopic("local_events", PublishPolicy.Always, LocalForwarderEvent)
COMMANDS = TypedTopic("commands", PublishPolicy.Always, ForwarderCommand)
ELECTION_MESSAGES = TypedTopic(
"election_messages", PublishPolicy.Always, ElectionMessage
-2
View File
@@ -15,7 +15,6 @@ from exo.shared.types.events import (
NodeDownloadProgress,
NodeGatheredInfo,
NodeTimedOut,
PrefillProgress,
RunnerDeleted,
RunnerStatusUpdated,
TaskAcknowledged,
@@ -65,7 +64,6 @@ def event_apply(event: Event, state: State) -> State:
| ChunkGenerated()
| TaskAcknowledged()
| InputChunkReceived()
| PrefillProgress()
| TracesCollected()
| TracesMerged()
): # Pass-through events that don't modify state
+9
View File
@@ -33,6 +33,15 @@ EXO_MODELS_DIR = (
if _EXO_MODELS_DIR_ENV is None
else Path.home() / _EXO_MODELS_DIR_ENV
)
# Read-only search path for pre-downloaded models (colon-separated directories)
_EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
EXO_MODELS_PATH: tuple[Path, ...] | None = (
tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
if _EXO_MODELS_PATH_ENV is not None
else None
)
_RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
RESOURCES_DIR = (
find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
+4 -13
View File
@@ -4,10 +4,8 @@ import anyio
from anyio import (
CancelScope,
Event,
create_task_group,
get_cancelled_exc_class,
)
from anyio.abc import TaskGroup
from loguru import logger
from exo.routing.connection_message import ConnectionMessage
@@ -15,6 +13,7 @@ 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 CamelCaseModel
from exo.utils.task_group import TaskGroup
DEFAULT_ELECTION_TIMEOUT = 3.0
@@ -82,13 +81,12 @@ class Election:
self._candidates: list[ElectionMessage] = []
self._campaign_cancel_scope: CancelScope | None = None
self._campaign_done: Event | None = None
self._tg: TaskGroup | None = None
self._tg = TaskGroup()
async def run(self):
logger.info("Starting Election")
try:
async with create_task_group() as tg:
self._tg = tg
async with self._tg as tg:
tg.start_soon(self._election_receiver)
tg.start_soon(self._connection_receiver)
tg.start_soon(self._command_counter)
@@ -124,12 +122,7 @@ class Election:
)
async def shutdown(self) -> None:
if not self._tg:
logger.warning(
"Attempted to shutdown election service that was not running"
)
return
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _election_receiver(self) -> None:
with self._em_receiver as election_messages:
@@ -143,7 +136,6 @@ class Election:
if message.clock > self.clock:
self.clock = message.clock
logger.debug(f"New clock: {self.clock}")
assert self._tg is not None
logger.debug("Starting new campaign")
candidates: list[ElectionMessage] = [message]
logger.debug(f"Candidates: {candidates}")
@@ -178,7 +170,6 @@ class Election:
# These messages are strictly peer to peer
self.clock += 1
logger.debug(f"New clock: {self.clock}")
assert self._tg is not None
candidates: list[ElectionMessage] = []
self._candidates = candidates
logger.debug("Starting new campaign")
+2
View File
@@ -90,6 +90,7 @@ class ModelCard(CamelCaseModel):
base_model: str = ""
capabilities: list[str] = []
uses_cfg: bool = False
trust_remote_code: bool = True
@field_validator("tasks", mode="before")
@classmethod
@@ -137,6 +138,7 @@ class ModelCard(CamelCaseModel):
hidden_size=config_data.hidden_size or 0,
supports_tensor=config_data.supports_tensor,
tasks=[ModelTask.TextGeneration],
trust_remote_code=False,
)
await mc.save_to_custom_dir()
_card_cache[model_id] = mc
@@ -14,7 +14,7 @@ def test_apply_node_download_progress():
event = DownloadCompleted(
node_id=NodeId("node-1"),
shard_metadata=shard1,
total_bytes=Memory(),
total=Memory(),
)
new_state = apply_node_download_progress(
@@ -30,12 +30,12 @@ def test_apply_two_node_download_progress():
event1 = DownloadCompleted(
node_id=NodeId("node-1"),
shard_metadata=shard1,
total_bytes=Memory(),
total=Memory(),
)
event2 = DownloadCompleted(
node_id=NodeId("node-1"),
shard_metadata=shard2,
total_bytes=Memory(),
total=Memory(),
)
state = State(downloads={NodeId("node-1"): [event1]})
+2 -2
View File
@@ -4,7 +4,7 @@ from anyio import create_task_group, fail_after, move_on_after
from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
from exo.shared.election import Election, ElectionMessage, ElectionResult
from exo.shared.types.commands import ForwarderCommand, TestCommand
from exo.shared.types.common import NodeId, SessionId
from exo.shared.types.common import NodeId, SessionId, SystemId
from exo.utils.channels import channel
# ======= #
@@ -384,7 +384,7 @@ async def test_tie_breaker_prefers_node_with_more_commands_seen() -> None:
# Pump local commands so our commands_seen is high before the round starts
for _ in range(50):
await co_tx.send(
ForwarderCommand(origin=NodeId("SOMEONE"), command=TestCommand())
ForwarderCommand(origin=SystemId("SOMEONE"), command=TestCommand())
)
# Trigger a round at clock=1 with a peer of equal seniority but fewer commands
@@ -23,7 +23,7 @@ def _get_keypair_concurrent_subprocess_task(
sem.release()
# wait to be told to begin simultaneous read
ev.wait()
queue.put(get_node_id_keypair().to_protobuf_encoding())
queue.put(get_node_id_keypair().to_bytes())
def _get_keypair_concurrent(num_procs: int) -> bytes:
+1 -1
View File
@@ -77,7 +77,7 @@ class ChatCompletionMessage(BaseModel):
content: (
str | ChatCompletionMessageText | list[ChatCompletionMessageText] | None
) = None
thinking: str | None = None # Added for GPT-OSS harmony format support
reasoning_content: str | None = None
name: str | None = None
tool_calls: list[ToolCall] | None = None
tool_call_id: str | None = None
+1
View File
@@ -27,6 +27,7 @@ class TokenChunk(BaseChunk):
stats: GenerationStats | None = None
logprob: float | None = None
top_logprobs: list[TopLogprobItem] | None = None
is_thinking: bool = False
class ErrorChunk(BaseChunk):
+31 -4
View File
@@ -47,6 +47,14 @@ class ClaudeImageBlock(BaseModel, frozen=True):
source: ClaudeImageSource
class ClaudeThinkingBlock(BaseModel, frozen=True):
"""Thinking content block in Claude Messages API."""
type: Literal["thinking"] = "thinking"
thinking: str
signature: str | None = None
class ClaudeToolUseBlock(BaseModel, frozen=True):
"""Tool use content block in Claude Messages API."""
@@ -66,11 +74,17 @@ class ClaudeToolResultBlock(BaseModel, frozen=True):
cache_control: dict[str, str] | None = None
ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock
ClaudeContentBlock = (
ClaudeTextBlock | ClaudeImageBlock | ClaudeThinkingBlock | ClaudeToolUseBlock
)
# Input content blocks can also include tool_result (sent by user after tool_use)
ClaudeInputContentBlock = (
ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock | ClaudeToolResultBlock
ClaudeTextBlock
| ClaudeImageBlock
| ClaudeThinkingBlock
| ClaudeToolUseBlock
| ClaudeToolResultBlock
)
@@ -82,6 +96,11 @@ class ClaudeMessage(BaseModel, frozen=True):
content: str | list[ClaudeInputContentBlock]
class ClaudeThinkingConfig(BaseModel, frozen=True):
type: Literal["enabled", "disabled", "adaptive"]
budget_tokens: int | None = None
class ClaudeMessagesRequest(BaseModel):
"""Request body for Claude Messages API."""
@@ -96,6 +115,7 @@ class ClaudeMessagesRequest(BaseModel):
top_k: int | None = None
tools: list[ClaudeToolDefinition] | None = None
metadata: dict[str, str] | None = None
thinking: ClaudeThinkingConfig | None = None
# Response types
@@ -145,7 +165,7 @@ class ClaudeContentBlockStartEvent(BaseModel, frozen=True):
type: Literal["content_block_start"] = "content_block_start"
index: int
content_block: ClaudeTextBlock | ClaudeToolUseBlock
content_block: ClaudeTextBlock | ClaudeThinkingBlock | ClaudeToolUseBlock
class ClaudeTextDelta(BaseModel, frozen=True):
@@ -155,6 +175,13 @@ class ClaudeTextDelta(BaseModel, frozen=True):
text: str
class ClaudeThinkingDelta(BaseModel, frozen=True):
"""Delta for thinking content block."""
type: Literal["thinking_delta"] = "thinking_delta"
thinking: str
class ClaudeInputJsonDelta(BaseModel, frozen=True):
"""Delta for tool use input JSON content block."""
@@ -167,7 +194,7 @@ class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True):
type: Literal["content_block_delta"] = "content_block_delta"
index: int
delta: ClaudeTextDelta | ClaudeInputJsonDelta
delta: ClaudeTextDelta | ClaudeThinkingDelta | ClaudeInputJsonDelta
class ClaudeContentBlockStopEvent(BaseModel, frozen=True):
+3 -3
View File
@@ -6,7 +6,7 @@ from exo.shared.types.api import (
ImageGenerationTaskParams,
)
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.common import CommandId, NodeId, SystemId
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -100,10 +100,10 @@ Command = (
class ForwarderCommand(CamelCaseModel):
origin: NodeId
origin: SystemId
command: Command
class ForwarderDownloadCommand(CamelCaseModel):
origin: NodeId
origin: SystemId
command: DownloadCommand
+4
View File
@@ -25,6 +25,10 @@ class NodeId(Id):
pass
class SystemId(Id):
pass
class ModelId(Id):
def normalize(self) -> str:
return self.replace("/", "--")
+11 -10
View File
@@ -5,7 +5,7 @@ from pydantic import Field
from exo.shared.topology import Connection
from exo.shared.types.chunks import GenerationChunk, InputImageChunk
from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId
from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import DownloadProgress
from exo.shared.types.worker.instances import Instance, InstanceId
@@ -102,13 +102,6 @@ class InputChunkReceived(BaseEvent):
chunk: InputImageChunk
class PrefillProgress(BaseEvent):
command_id: CommandId
model: ModelId
processed_tokens: int
total_tokens: int
class TopologyEdgeCreated(BaseEvent):
conn: Connection
@@ -155,7 +148,6 @@ Event = (
| NodeDownloadProgress
| ChunkGenerated
| InputChunkReceived
| PrefillProgress
| TopologyEdgeCreated
| TopologyEdgeDeleted
| TracesCollected
@@ -170,10 +162,19 @@ class IndexedEvent(CamelCaseModel):
event: Event
class ForwarderEvent(CamelCaseModel):
class GlobalForwarderEvent(CamelCaseModel):
"""An event the forwarder will serialize and send over the network"""
origin_idx: int = Field(ge=0)
origin: NodeId
session: SessionId
event: Event
class LocalForwarderEvent(CamelCaseModel):
"""An event the forwarder will serialize and send over the network"""
origin_idx: int = Field(ge=0)
origin: SystemId
session: SessionId
event: Event
+95 -17
View File
@@ -1,10 +1,10 @@
from math import ceil
from typing import Self
from typing import Self, overload
from exo.utils.pydantic_ext import CamelCaseModel
from exo.utils.pydantic_ext import FrozenModel
class Memory(CamelCaseModel):
class Memory(FrozenModel):
in_bytes: int = 0
@classmethod
@@ -33,12 +33,22 @@ class Memory(CamelCaseModel):
return cls(in_bytes=round(val * 1024))
@property
def in_mb(self) -> float:
"""The approximate megabytes this memory represents. Setting this property rounds to the nearest byte."""
return self.in_bytes / (1024**2)
def in_mb(self) -> int:
"""The approximate megabytes this memory represents, rounded to nearest MB. Setting this property rounds to the nearest byte."""
return round(self.in_bytes / (1024**2))
@in_mb.setter
def in_mb(self, val: float):
def in_mb(self, val: int):
"""Set the megabytes for this memory."""
self.in_bytes = val * (1024**2)
@property
def in_float_mb(self) -> float:
"""The megabytes this memory represents as a float. Setting this property rounds to the nearest byte."""
return self.in_bytes / (1024**2)
@in_float_mb.setter
def in_float_mb(self, val: float):
"""Set the megabytes for this memory, rounded to the nearest byte."""
self.in_bytes = round(val * (1024**2))
@@ -57,17 +67,85 @@ class Memory(CamelCaseModel):
"""The approximate gigabytes this memory represents."""
return self.in_bytes / (1024**3)
def __add__(self, other: "Memory") -> "Memory":
return Memory.from_bytes(self.in_bytes + other.in_bytes)
def __add__(self, other: object) -> "Memory":
if isinstance(other, Memory):
return Memory.from_bytes(self.in_bytes + other.in_bytes)
return NotImplemented
def __lt__(self, other: Self) -> bool:
return self.in_bytes < other.in_bytes
def __radd__(self, other: object) -> "Memory":
if other == 0:
return self
return NotImplemented
def __le__(self, other: Self) -> bool:
return self.in_bytes <= other.in_bytes
def __sub__(self, other: object) -> "Memory":
if isinstance(other, Memory):
return Memory.from_bytes(self.in_bytes - other.in_bytes)
return NotImplemented
def __gt__(self, other: Self) -> bool:
return self.in_bytes > other.in_bytes
def __mul__(self, other: int | float):
return Memory.from_bytes(round(self.in_bytes * other))
def __ge__(self, other: Self) -> bool:
return self.in_bytes >= other.in_bytes
def __rmul__(self, other: int | float):
return self * other
@overload
def __truediv__(self, other: "Memory") -> float: ...
@overload
def __truediv__(self, other: int) -> "Memory": ...
@overload
def __truediv__(self, other: float) -> "Memory": ...
def __truediv__(self, other: object) -> "Memory | float":
if isinstance(other, Memory):
return self.in_bytes / other.in_bytes
if isinstance(other, (int, float)):
return Memory.from_bytes(round(self.in_bytes / other))
return NotImplemented
def __floordiv__(self, other: object) -> "Memory":
if isinstance(other, (int, float)):
return Memory.from_bytes(int(self.in_bytes // other))
return NotImplemented
def __lt__(self, other: object) -> bool:
if isinstance(other, Memory):
return self.in_bytes < other.in_bytes
return NotImplemented
def __le__(self, other: object) -> bool:
if isinstance(other, Memory):
return self.in_bytes <= other.in_bytes
return NotImplemented
def __gt__(self, other: object) -> bool:
if isinstance(other, Memory):
return self.in_bytes > other.in_bytes
return NotImplemented
def __ge__(self, other: object) -> bool:
if isinstance(other, Memory):
return self.in_bytes >= other.in_bytes
return NotImplemented
def __eq__(self, other: object) -> bool:
if isinstance(other, Memory):
return self.in_bytes == other.in_bytes
return NotImplemented
def __repr__(self) -> str:
return f"Memory.from_bytes({self.in_bytes})"
def __str__(self) -> str:
if self.in_gb > 2:
val = self.in_gb
unit = "GiB"
elif self.in_mb > 2:
val = self.in_mb
unit = "MiB"
elif self.in_kb > 3:
val = self.in_kb
unit = "KiB"
else:
val = self.in_bytes
unit = "B"
return f"{val:.2f} {unit}".rstrip("0").rstrip(".") + f" {unit}"
+148
View File
@@ -0,0 +1,148 @@
from __future__ import annotations
import time
from typing import Any, Literal
from pydantic import BaseModel, Field
from exo.shared.models.model_cards import ModelId
# https://github.com/ollama/ollama/blob/main/docs/api.md
OllamaRole = Literal["system", "user", "assistant", "tool"]
OllamaDoneReason = Literal["stop", "length", "tool_call", "error"]
class OllamaToolFunction(BaseModel, frozen=True):
name: str
arguments: dict[str, Any] | str
index: int | None = None
class OllamaToolCall(BaseModel, frozen=True):
id: str | None = None
type: Literal["function"] | None = None
function: OllamaToolFunction
class OllamaMessage(BaseModel, frozen=True):
role: OllamaRole
content: str | None = None
thinking: str | None = None
tool_calls: list[OllamaToolCall] | None = None
name: str | None = None
tool_name: str | None = None
images: list[str] | None = None
class OllamaOptions(BaseModel, frozen=True):
num_predict: int | None = None
temperature: float | None = None
top_p: float | None = None
top_k: int | None = None
stop: str | list[str] | None = None
seed: int | None = None
class OllamaChatRequest(BaseModel, frozen=True):
model: ModelId
messages: list[OllamaMessage]
stream: bool = True
options: OllamaOptions | None = None
tools: list[dict[str, Any]] | None = None
format: Literal["json"] | dict[str, Any] | None = None
keep_alive: str | int | None = None
think: bool | None = None
class OllamaGenerateRequest(BaseModel, frozen=True):
model: ModelId
prompt: str = ""
system: str | None = None
stream: bool = True
options: OllamaOptions | None = None
format: Literal["json"] | dict[str, Any] | None = None
keep_alive: str | int | None = None
think: bool | None = None
raw: bool = False
class OllamaGenerateResponse(BaseModel, frozen=True, strict=True):
model: str
created_at: str = Field(
default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
)
response: str
thinking: str | None = None
done: bool
done_reason: OllamaDoneReason | None = None
total_duration: int | None = None
load_duration: int | None = None
prompt_eval_count: int | None = None
prompt_eval_duration: int | None = None
eval_count: int | None = None
eval_duration: int | None = None
class OllamaShowRequest(BaseModel, frozen=True):
name: str | None = None
model: str | None = None
verbose: bool | None = None
class OllamaChatResponse(BaseModel, frozen=True, strict=True):
model: str
created_at: str = Field(
default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
)
message: OllamaMessage
done: bool
done_reason: OllamaDoneReason | None = None
total_duration: int | None = None
load_duration: int | None = None
prompt_eval_count: int | None = None
prompt_eval_duration: int | None = None
eval_count: int | None = None
eval_duration: int | None = None
class OllamaModelDetails(BaseModel, frozen=True, strict=True):
format: str | None = None
family: str | None = None
parameter_size: str | None = None
quantization_level: str | None = None
class OllamaModelTag(BaseModel, frozen=True, strict=True):
name: str
model: str | None = None
modified_at: str | None = None
size: int | None = None
digest: str | None = None
details: OllamaModelDetails | None = None
class OllamaTagsResponse(BaseModel, frozen=True, strict=True):
models: list[OllamaModelTag]
class OllamaShowResponse(BaseModel, frozen=True, strict=True):
modelfile: str | None = None
parameters: str | None = None
template: str | None = None
details: OllamaModelDetails | None = None
model_info: dict[str, Any] | None = None
class OllamaPsModel(BaseModel, frozen=True, strict=True):
name: str
model: str
size: int
digest: str | None = None
details: OllamaModelDetails | None = None
expires_at: str | None = None
size_vram: int | None = None
class OllamaPsResponse(BaseModel, frozen=True, strict=True):
models: list[OllamaPsModel]
+73 -1
View File
@@ -145,7 +145,23 @@ class ResponseFunctionCallItem(BaseModel, frozen=True):
status: ResponseStatus = "completed"
ResponseItem = ResponseMessageItem | ResponseFunctionCallItem
class ResponseReasoningSummaryText(BaseModel, frozen=True):
"""Summary text part in a reasoning output item."""
type: Literal["summary_text"] = "summary_text"
text: str
class ResponseReasoningItem(BaseModel, frozen=True):
"""Reasoning output item in response output array."""
type: Literal["reasoning"] = "reasoning"
id: str
summary: list[ResponseReasoningSummaryText] = Field(default_factory=list)
status: ResponseStatus = "completed"
ResponseItem = ResponseMessageItem | ResponseFunctionCallItem | ResponseReasoningItem
class ResponseUsage(BaseModel, frozen=True):
@@ -273,6 +289,58 @@ class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True):
arguments: str
class ResponseReasoningSummaryPartAddedEvent(BaseModel, frozen=True):
"""Event sent when a reasoning summary part is added."""
type: Literal["response.reasoning_summary_part.added"] = (
"response.reasoning_summary_part.added"
)
sequence_number: int
item_id: str
output_index: int
summary_index: int
part: ResponseReasoningSummaryText
class ResponseReasoningSummaryTextDeltaEvent(BaseModel, frozen=True):
"""Event sent for reasoning summary text delta during streaming."""
type: Literal["response.reasoning_summary_text.delta"] = (
"response.reasoning_summary_text.delta"
)
sequence_number: int
item_id: str
output_index: int
summary_index: int
delta: str
class ResponseReasoningSummaryTextDoneEvent(BaseModel, frozen=True):
"""Event sent when reasoning summary text is done."""
type: Literal["response.reasoning_summary_text.done"] = (
"response.reasoning_summary_text.done"
)
sequence_number: int
item_id: str
output_index: int
summary_index: int
text: str
class ResponseReasoningSummaryPartDoneEvent(BaseModel, frozen=True):
"""Event sent when a reasoning summary part is done."""
type: Literal["response.reasoning_summary_part.done"] = (
"response.reasoning_summary_part.done"
)
sequence_number: int
item_id: str
output_index: int
summary_index: int
part: ResponseReasoningSummaryText
class ResponseCompletedEvent(BaseModel, frozen=True):
"""Event sent when response is completed."""
@@ -292,5 +360,9 @@ ResponsesStreamEvent = (
| ResponseOutputItemDoneEvent
| ResponseFunctionCallArgumentsDeltaEvent
| ResponseFunctionCallArgumentsDoneEvent
| ResponseReasoningSummaryPartAddedEvent
| ResponseReasoningSummaryTextDeltaEvent
| ResponseReasoningSummaryTextDoneEvent
| ResponseReasoningSummaryPartDoneEvent
| ResponseCompletedEvent
)
+10 -8
View File
@@ -10,9 +10,9 @@ from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
class DownloadProgressData(CamelCaseModel):
total_bytes: Memory
downloaded_bytes: Memory
downloaded_bytes_this_session: Memory
total: Memory
downloaded: Memory
downloaded_this_session: Memory
completed_files: int
total_files: int
@@ -30,11 +30,13 @@ class BaseDownloadProgress(TaggedModel):
class DownloadPending(BaseDownloadProgress):
pass
downloaded: Memory = Memory()
total: Memory = Memory()
class DownloadCompleted(BaseDownloadProgress):
total_bytes: Memory
total: Memory
read_only: bool = False
class DownloadFailed(BaseDownloadProgress):
@@ -86,9 +88,9 @@ class RepoDownloadProgress(BaseModel):
shard: ShardMetadata
completed_files: int
total_files: int
downloaded_bytes: Memory
downloaded_bytes_this_session: Memory
total_bytes: Memory
downloaded: Memory
downloaded_this_session: Memory
total: Memory
overall_speed: float
overall_eta: timedelta
status: Literal["not_started", "in_progress", "complete"]
+8
View File
@@ -2,6 +2,7 @@ from enum import Enum
from pydantic import model_validator
from exo.shared.models.model_cards import ModelTask
from exo.shared.types.common import Host, Id, NodeId
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -49,6 +50,13 @@ class BoundInstance(CamelCaseModel):
assert shard is not None
return shard
@property
def is_image_model(self) -> bool:
return (
ModelTask.TextToImage in self.bound_shard.model_card.tasks
or ModelTask.ImageToImage in self.bound_shard.model_card.tasks
)
@model_validator(mode="after")
def validate_shard_exists(self) -> "BoundInstance":
assert (
@@ -28,6 +28,7 @@ class GenerationResponse(BaseRunnerResponse):
finish_reason: FinishReason | None = None
stats: GenerationStats | None = None
usage: Usage | None
is_thinking: bool = False
class ImageGenerationResponse(BaseRunnerResponse):
+2 -1
View File
@@ -34,7 +34,8 @@ class RunnerConnected(BaseRunnerStatus):
class RunnerLoading(BaseRunnerStatus):
pass
layers_loaded: int = 0
total_layers: int = 0
class RunnerLoaded(BaseRunnerStatus):
+30
View File
@@ -1,8 +1,27 @@
import logging
import os
import sys
import webbrowser
from exo.shared.constants import EXO_CONFIG_HOME
logger = logging.getLogger(__name__)
_FIRST_RUN_MARKER = EXO_CONFIG_HOME / ".dashboard_opened"
def _is_first_run() -> bool:
return not _FIRST_RUN_MARKER.exists()
def _mark_first_run_done() -> None:
_FIRST_RUN_MARKER.parent.mkdir(parents=True, exist_ok=True)
_FIRST_RUN_MARKER.touch()
def print_startup_banner(port: int) -> None:
dashboard_url = f"http://localhost:{port}"
first_run = _is_first_run()
banner = f"""
@@ -30,3 +49,14 @@ def print_startup_banner(port: int) -> None:
"""
print(banner, file=sys.stderr)
if first_run:
# Skip browser open when running inside the native macOS app —
# FirstLaunchPopout.swift handles the auto-open with a countdown.
if not os.environ.get("EXO_RUNTIME_DIR"):
try:
webbrowser.open(dashboard_url)
logger.info("First run detected — opening dashboard in browser")
except Exception:
logger.debug("Could not auto-open browser", exc_info=True)
_mark_first_run_done()
+9 -3
View File
@@ -5,7 +5,7 @@ from math import inf
from multiprocessing.synchronize import Event
from queue import Empty, Full
from types import TracebackType
from typing import Self
from typing import Any, Self
from anyio import (
CapacityLimiter,
@@ -157,7 +157,7 @@ class MpSender[T]:
) -> None:
self.close()
def __getstate__(self):
def __getstate__(self) -> dict[str, Any]:
d = self.__dict__.copy()
d.pop("__orig_class__", None)
return d
@@ -192,7 +192,13 @@ class MpReceiver[T]:
try:
return self.receive_nowait()
except WouldBlock:
item = self._state.buffer.get()
try:
item = self._state.buffer.get()
except (TypeError, OSError):
# Queue pipe can get closed while we are blocked on get().
# The underlying connection._handle becomes None, causing
# TypeError in read(handle, remaining).
raise ClosedResourceError from None
if isinstance(item, _MpEndOfStream):
self.close()
raise EndOfStream from None
+10 -4
View File
@@ -8,8 +8,7 @@ from subprocess import CalledProcessError
from typing import Self, cast
import anyio
from anyio import create_task_group, fail_after, open_process, to_thread
from anyio.abc import TaskGroup
from anyio import fail_after, open_process, to_thread
from anyio.streams.buffered import BufferedByteReceiveStream
from anyio.streams.text import TextReceiveStream
from loguru import logger
@@ -30,6 +29,7 @@ from exo.shared.types.thunderbolt import (
)
from exo.utils.channels import Sender
from exo.utils.pydantic_ext import TaggedModel
from exo.utils.task_group import TaskGroup
from .macmon import MacmonMetrics
from .system_info import (
@@ -381,13 +381,19 @@ class InfoGatherer:
static_info_poll_interval: float | None = 60
rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
disk_poll_interval: float | None = 30
_tg: TaskGroup = field(init=False, default_factory=create_task_group)
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
async def run(self):
async with self._tg as tg:
if IS_DARWIN:
if (macmon_path := shutil.which("macmon")) is not None:
tg.start_soon(self._monitor_macmon, macmon_path)
else:
# macmon not installed — fall back to psutil for memory
logger.warning(
"macmon not found, falling back to psutil for memory monitoring"
)
self.memory_poll_rate = 1
tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
tg.start_soon(self._monitor_thunderbolt_bridge_status)
tg.start_soon(self._monitor_rdma_ctl_status)
@@ -402,7 +408,7 @@ class InfoGatherer:
await self.info_sender.send(nc)
def shutdown(self):
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _monitor_static_info(self):
if self.static_info_poll_interval is None:
+1 -1
View File
@@ -108,7 +108,7 @@ async def check_reachable(
await send.send((target_ip, expected_node_id))
async with (
httpx.AsyncClient(timeout=timeout, limits=limits) as client,
httpx.AsyncClient(timeout=timeout, limits=limits, verify=False) as client,
create_task_group() as tg,
):
for node_id in topology.list_nodes():
+65
View File
@@ -0,0 +1,65 @@
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from types import TracebackType
from typing import Any, Unpack
from anyio import create_task_group
from anyio.abc import TaskGroup as TaskGroupABC
@dataclass
class TaskGroup:
_tg: TaskGroupABC | None = field(default=None, init=False)
_queued: list[tuple[Any, Any, Any]] | None = field(default_factory=list, init=False)
def is_running(self) -> bool:
return self._tg is not None
def cancel_tasks(self):
assert self._tg
self._tg.cancel_scope.cancel()
def cancel_called(self) -> bool:
assert self._tg
return self._tg.cancel_scope.cancel_called
def start_soon[*T](
self,
func: Callable[[Unpack[T]], Awaitable[Any]],
*args: Unpack[T],
name: object = None,
) -> None:
assert self._tg is not None
assert self._queued is None
self._tg.start_soon(func, *args, name=name)
def queue[*T](
self,
func: Callable[[Unpack[T]], Awaitable[Any]],
*args: Unpack[T],
name: object = None,
) -> None:
assert self._tg is None
assert self._queued is not None
self._queued.append((func, args, name))
async def __aenter__(self) -> TaskGroupABC:
assert self._tg is None
assert self._queued is not None
self._tg = create_task_group()
r = await self._tg.__aenter__()
for func, args, name in self._queued: # pyright: ignore[reportAny]
self._tg.start_soon(func, *args, name=name) # pyright: ignore[reportAny]
self._queued = None
return r
async def __aexit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
exc_tb: TracebackType | None,
) -> bool:
"""Exit the task group context waiting for all tasks to finish."""
assert self._tg is not None, "aenter sets self.lazy, so it exists when we aexit"
assert self._queued is None
return await self._tg.__aexit__(exc_type, exc_val, exc_tb)
+2 -2
View File
@@ -166,7 +166,7 @@ def generate_image(
else 0.0
)
peak_memory_gb = mx.get_peak_memory() / (1024**3)
peak_memory = Memory.from_bytes(mx.get_peak_memory())
stats = ImageGenerationStats(
seconds_per_step=seconds_per_step,
@@ -175,7 +175,7 @@ def generate_image(
num_images=num_images,
image_width=width,
image_height=height,
peak_memory_usage=Memory.from_gb(peak_memory_gb),
peak_memory_usage=peak_memory,
)
buffer = io.BytesIO()
+68 -42
View File
@@ -32,13 +32,12 @@ from mlx_lm.models.minimax import MiniMaxAttention
from mlx_lm.models.minimax import Model as MiniMaxModel
from mlx_lm.models.ministral3 import Model as Ministral3Model
from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock
from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock
from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
from mlx_lm.models.step3p5 import Model as Step35Model
from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
from exo.shared.logging import logger
from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -47,6 +46,7 @@ if TYPE_CHECKING:
from mlx_lm.models.cache import Cache
TimeoutCallback = Callable[[], None]
LayerLoadedCallback = Callable[[int, int], None] # (layers_loaded, total_layers)
def eval_with_timeout(
@@ -186,7 +186,7 @@ def set_pipeline_prefill(model: nn.Module, is_prefill: bool) -> None:
layer.is_prefill = is_prefill
def _inner_model(model: nn.Module) -> nn.Module:
def get_inner_model(model: nn.Module) -> nn.Module:
inner = getattr(model, "model", None)
if isinstance(inner, nn.Module):
return inner
@@ -204,7 +204,7 @@ def _inner_model(model: nn.Module) -> nn.Module:
raise ValueError("Model must either have a 'model' or 'transformer' attribute")
def _get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
def get_layers(inner_model_instance: nn.Module) -> list[_LayerCallable]:
# Handle both model.layers and model.h cases
layers: list[_LayerCallable]
if hasattr(inner_model_instance, "layers"):
@@ -221,6 +221,7 @@ def pipeline_auto_parallel(
model: nn.Module,
group: mx.distributed.Group,
model_shard_meta: PipelineShardMetadata,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
"""
Automatically parallelize a model across multiple devices.
@@ -230,16 +231,19 @@ def pipeline_auto_parallel(
Returns:
The parallelized model
"""
inner_model_instance: nn.Module = _inner_model(model)
inner_model_instance: nn.Module = get_inner_model(model)
layers = _get_layers(inner_model_instance)
layers = get_layers(inner_model_instance)
start_layer, end_layer = model_shard_meta.start_layer, model_shard_meta.end_layer
device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size
layers = layers[start_layer:end_layer]
for layer in layers:
total = len(layers)
for i, layer in enumerate(layers):
mx.eval(layer) # type: ignore
if on_layer_loaded is not None:
on_layer_loaded(i, total)
layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
layers[-1] = PipelineLastLayer(
@@ -351,8 +355,9 @@ def patch_tensor_model[T](model: T) -> T:
def tensor_auto_parallel(
model: nn.Module,
group: mx.distributed.Group,
timeout_seconds: float = 60.0,
on_timeout: TimeoutCallback | None = None,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
all_to_sharded_linear = partial(
shard_linear,
@@ -462,7 +467,7 @@ def tensor_auto_parallel(
raise ValueError(f"Unsupported model type: {type(model)}")
model = tensor_parallel_sharding_strategy.shard_model(
model, timeout_seconds, on_timeout
model, timeout_seconds, on_timeout, on_layer_loaded
)
return patch_tensor_model(model)
@@ -489,6 +494,7 @@ class TensorParallelShardingStrategy(ABC):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module: ...
@@ -498,13 +504,13 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(LlamaModel, model)
for layer in model.layers:
total = len(model.layers)
for i, layer in enumerate(model.layers):
# Force load weights before sharding to avoid FAST_SYNCH deadlock
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -517,11 +523,13 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
inner_model_instance = _inner_model(model)
inner_model_instance = get_inner_model(model)
if hasattr(inner_model_instance, "layers"):
inner_model_instance.layers = layers
@@ -552,13 +560,13 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(DeepseekV3Model, model)
total = len(model.layers)
for layer in model.layers:
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
if layer.self_attn.q_lora_rank is None:
@@ -609,6 +617,8 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.sharding_group = self.group
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
@@ -635,13 +645,15 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(GLM4MoeLiteModel, model)
for layer in model.layers: # type: ignore
total = len(model.layers) # type: ignore
for i, layer in enumerate(model.layers): # type: ignore
layer = cast(Glm4MoeLiteDecoderLayer, layer)
eval_with_timeout(
layer.parameters(),
timeout_seconds / len(model.layers), # type: ignore
timeout_seconds / total,
on_timeout,
)
if layer.self_attn.q_lora_rank is None: # type: ignore
@@ -689,6 +701,8 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
layer.mlp = ShardedMoE(layer.mlp) # type: ignore
layer.mlp.sharding_group = self.group # type: ignore
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
@@ -777,12 +791,12 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(MiniMaxModel, model)
for layer in model.layers:
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -807,6 +821,8 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
layer.block_sparse_moe.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
@@ -816,14 +832,14 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(Qwen3MoeModel | Qwen3NextModel, model)
for layer in model.layers:
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
if isinstance(layer, Qwen3DecoderLayer):
if isinstance(layer, Qwen3MoeDecoderLayer):
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
)
@@ -836,6 +852,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj
)
layer.self_attn.n_heads //= self.N
layer.self_attn.n_kv_heads //= self.N
else:
assert isinstance(layer, Qwen3NextDecoderLayer)
if hasattr(layer, "linear_attn"):
@@ -930,6 +948,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
@@ -939,12 +959,12 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(Glm4MoeModel, model)
for layer in model.layers:
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
@@ -976,6 +996,8 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
@@ -985,13 +1007,13 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(GptOssMoeModel, model)
total = len(model.layers)
for layer in model.layers:
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -1017,6 +1039,8 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
layer.mlp = ShardedMoE(layer.mlp) # type: ignore
layer.mlp.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
@@ -1026,13 +1050,13 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module:
model = cast(Step35Model, model)
total = len(model.layers)
for layer in model.layers:
eval_with_timeout(
layer.parameters(), timeout_seconds / len(model.layers), on_timeout
)
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
@@ -1060,4 +1084,6 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
mx.eval(layer)
if on_layer_loaded is not None:
on_layer_loaded(i, total)
return model
+1 -1
View File
@@ -22,7 +22,7 @@ from exo.worker.runner.bootstrap import logger
# Fraction of device memory above which LRU eviction kicks in.
# Smaller machines need more aggressive eviction.
def _default_memory_threshold() -> float:
total_gb = psutil.virtual_memory().total / (1024**3)
total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb
if total_gb >= 128:
return 0.85
if total_gb >= 64:
+2 -1
View File
@@ -13,5 +13,6 @@ KV_CACHE_BITS: int | None = None
DEFAULT_TOP_LOGPROBS: int = 5
# TODO: We should really make this opt-in, but Kimi requires trust_remote_code=True
# True for built-in models with known model cards; custom models added via API default to False
# and can be overridden with the --trust-remote-code CLI flag.
TRUST_REMOTE_CODE: bool = True
@@ -0,0 +1,72 @@
import json
import re
from typing import Any
from mlx_lm.chat_templates import deepseek_v32
from exo.shared.types.api import ToolCallItem
BOS_TOKEN: str = deepseek_v32.bos_token
EOS_TOKEN: str = deepseek_v32.eos_token
DSML_TOKEN: str = deepseek_v32.dsml_token
THINKING_START: str = deepseek_v32.thinking_start_token
THINKING_END: str = deepseek_v32.thinking_end_token
USER_TOKEN = "<\uff5cUser\uff5c>"
ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
encode_messages = deepseek_v32.encode_messages
_INVOKE_PATTERN = re.compile(
rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"
rf"(.*?)"
rf"</{re.escape(DSML_TOKEN)}invoke>",
re.DOTALL,
)
_PARAM_PATTERN = re.compile(
rf"<{re.escape(DSML_TOKEN)}parameter\s+name=\"([^\"]+)\"\s+string=\"(true|false)\">"
rf"(.*?)"
rf"</{re.escape(DSML_TOKEN)}parameter>",
re.DOTALL,
)
def parse_dsml_output(text: str) -> list[ToolCallItem] | None:
"""Parse DSML function_calls block from model output text.
Args:
text: The text containing the DSML function_calls block
(including the start/end markers).
Returns:
List of ToolCallItem, or None if parsing fails.
"""
tool_calls: list[ToolCallItem] = []
for invoke_match in _INVOKE_PATTERN.finditer(text):
func_name = invoke_match.group(1)
invoke_body = invoke_match.group(2)
args: dict[str, Any] = {}
for param_match in _PARAM_PATTERN.finditer(invoke_body):
param_name = param_match.group(1)
is_string = param_match.group(2) == "true"
param_value = param_match.group(3)
if is_string:
args[param_name] = param_value
else:
try:
args[param_name] = json.loads(param_value)
except (json.JSONDecodeError, ValueError):
args[param_name] = param_value
tool_calls.append(
ToolCallItem(
name=func_name,
arguments=json.dumps(args),
)
)
return tool_calls if tool_calls else None
+81 -25
View File
@@ -23,9 +23,7 @@ from mlx_lm.models.deepseek_v3 import DeepseekV3Model
from mlx_lm.tokenizer_utils import TokenizerWrapper
from exo.shared.models.model_cards import ModelId
from exo.worker.engines.mlx.constants import (
TRUST_REMOTE_CODE,
)
from exo.worker.engines.mlx.constants import TRUST_REMOTE_CODE
try:
from mlx_lm.tokenizer_utils import load_tokenizer
@@ -55,8 +53,11 @@ from exo.shared.types.worker.shards import (
)
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.auto_parallel import (
LayerLoadedCallback,
TimeoutCallback,
eval_with_timeout,
get_inner_model,
get_layers,
pipeline_auto_parallel,
tensor_auto_parallel,
)
@@ -171,13 +172,28 @@ def initialize_mlx(
def load_mlx_items(
bound_instance: BoundInstance,
group: Group | None,
on_timeout: TimeoutCallback | None = None,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> tuple[Model, TokenizerWrapper]:
if group is None:
logger.info(f"Single device used for {bound_instance.instance}")
model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
start_time = time.perf_counter()
model, _ = load_model(model_path, strict=True)
model, _ = load_model(model_path, lazy=True, strict=False)
# Eval layers one by one for progress reporting
try:
inner = get_inner_model(model)
layers = get_layers(inner)
total = len(layers)
for i, layer in enumerate(layers):
mx.eval(layer) # type: ignore
if on_layer_loaded is not None:
on_layer_loaded(i, total)
except ValueError as e:
logger.opt(exception=e).debug(
"Model architecture doesn't support layer-by-layer progress tracking",
)
mx.eval(model)
end_time = time.perf_counter()
logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s")
tokenizer = get_tokenizer(model_path, bound_instance.bound_shard)
@@ -186,7 +202,10 @@ def load_mlx_items(
logger.info("Starting distributed init")
start_time = time.perf_counter()
model, tokenizer = shard_and_load(
bound_instance.bound_shard, group=group, on_timeout=on_timeout
bound_instance.bound_shard,
group=group,
on_timeout=on_timeout,
on_layer_loaded=on_layer_loaded,
)
end_time = time.perf_counter()
logger.info(
@@ -201,7 +220,8 @@ def load_mlx_items(
def shard_and_load(
shard_metadata: ShardMetadata,
group: Group,
on_timeout: TimeoutCallback | None = None,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> tuple[nn.Module, TokenizerWrapper]:
model_path = build_model_path(shard_metadata.model_card.model_id)
@@ -232,20 +252,24 @@ def shard_and_load(
# Estimate timeout based on model size (5x default for large queued workloads)
base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300"))
model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3)
timeout_seconds = base_timeout + model_size_gb
model_size = get_weights_size(shard_metadata)
timeout_seconds = base_timeout + model_size.in_gb
logger.info(
f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
f"(model size: {model_size_gb:.1f}GB)"
f"(model size: {model_size.in_gb:.1f}GB)"
)
match shard_metadata:
case TensorShardMetadata():
logger.info(f"loading model from {model_path} with tensor parallelism")
model = tensor_auto_parallel(model, group, timeout_seconds, on_timeout)
model = tensor_auto_parallel(
model, group, timeout_seconds, on_timeout, on_layer_loaded
)
case PipelineShardMetadata():
logger.info(f"loading model from {model_path} with pipeline parallelism")
model = pipeline_auto_parallel(model, group, shard_metadata)
model = pipeline_auto_parallel(
model, group, shard_metadata, on_layer_loaded=on_layer_loaded
)
eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
case CfgShardMetadata():
raise ValueError(
@@ -267,7 +291,15 @@ def shard_and_load(
def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
"""Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
return load_tokenizer_for_model_id(shard_metadata.model_card.model_id, model_path)
trust_remote_code = (
shard_metadata.model_card.trust_remote_code
or os.environ.get("EXO_TRUST_REMOTE_CODE") == "1"
)
return load_tokenizer_for_model_id(
shard_metadata.model_card.model_id,
model_path,
trust_remote_code=trust_remote_code,
)
def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
@@ -299,7 +331,7 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
def load_tokenizer_for_model_id(
model_id: ModelId, model_path: Path
model_id: ModelId, model_path: Path, *, trust_remote_code: bool = TRUST_REMOTE_CODE
) -> TokenizerWrapper:
"""
Load tokenizer for a model given its ID and local path.
@@ -368,7 +400,7 @@ def load_tokenizer_for_model_id(
tokenizer = load_tokenizer(
model_path,
tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
tokenizer_config_extra={"trust_remote_code": trust_remote_code},
eos_token_ids=eos_token_ids,
)
@@ -458,6 +490,19 @@ def _patch_lossy_chat_template(template: str) -> str | None:
return patched if n > 0 else None
def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
if "deepseek-v3.2" not in task_params.model.lower():
return False
# Use DSML encoding when tools are provided or tool results are in the conversation
if task_params.tools:
return True
if task_params.chat_template_messages:
return any(
msg.get("role") == "tool" for msg in task_params.chat_template_messages
)
return False
def apply_chat_template(
tokenizer: TokenizerWrapper,
task_params: TextGenerationTaskParams,
@@ -469,7 +514,6 @@ def apply_chat_template(
When chat_template_messages is available (from Chat Completions API),
uses those directly to preserve tool_calls, thinking, and other fields.
Otherwise builds messages from the task params input/instructions.
"""
formatted_messages: list[dict[str, Any]] = []
if task_params.chat_template_messages is not None:
@@ -497,6 +541,19 @@ def apply_chat_template(
partial_assistant_content = cast(str, formatted_messages[-1].get("content", ""))
formatted_messages = formatted_messages[:-1]
if _needs_dsml_encoding(task_params):
from exo.worker.engines.mlx.dsml_encoding import encode_messages
prompt = encode_messages(
messages=formatted_messages,
thinking_mode="thinking" if task_params.enable_thinking else "chat",
tools=task_params.tools,
)
if partial_assistant_content:
prompt += partial_assistant_content
logger.info(prompt)
return prompt
extra_kwargs: dict[str, Any] = {}
if task_params.enable_thinking is not None:
# Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking".
@@ -591,7 +648,7 @@ class NullKVCache(KVCache):
raise NotImplementedError("We should not be setting a NullKVCache.")
def mlx_force_oom(size: int = 40000) -> None:
def mlx_force_oom(size: int = 200000) -> None:
"""
Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
"""
@@ -617,18 +674,17 @@ def set_wired_limit_for_model(model_size: Memory):
if not mx.metal.is_available():
return
model_bytes = model_size.in_bytes
max_rec_size = int(mx.metal.device_info()["max_recommended_working_set_size"])
if model_bytes > 0.9 * max_rec_size:
model_mb = model_bytes // 2**20
max_rec_mb = max_rec_size // 2**20
max_rec_size = Memory.from_bytes(
int(mx.device_info()["max_recommended_working_set_size"])
)
if model_size > 0.9 * max_rec_size:
logger.warning(
f"Generating with a model that requires {model_mb} MB "
f"which is close to the maximum recommended size of {max_rec_mb} "
f"Generating with a model that requires {model_size.in_float_mb:.1f} MB "
f"which is close to the maximum recommended size of {max_rec_size.in_float_mb:.1f} "
"MB. This can be slow. See the documentation for possible work-arounds: "
"https://github.com/ml-explore/mlx-lm/tree/main#large-models"
)
mx.set_wired_limit(max_rec_size)
mx.set_wired_limit(max_rec_size.in_bytes)
logger.info(f"Wired limit set to {max_rec_size}.")
+58 -28
View File
@@ -1,13 +1,12 @@
from collections import defaultdict
from datetime import datetime, timezone
from random import random
from typing import Iterator
import anyio
from anyio import CancelScope, create_task_group, fail_after
from anyio.abc import TaskGroup
from anyio import CancelScope, fail_after
from loguru import logger
from exo.download.download_utils import resolve_model_in_path
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelId
from exo.shared.types.api import ImageEditsTaskParams
@@ -17,13 +16,15 @@ from exo.shared.types.commands import (
RequestEventLog,
StartDownload,
)
from exo.shared.types.common import CommandId, NodeId, SessionId
from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
from exo.shared.types.events import (
Event,
EventId,
ForwarderEvent,
GlobalForwarderEvent,
IndexedEvent,
InputChunkReceived,
LocalForwarderEvent,
NodeDownloadProgress,
NodeGatheredInfo,
TaskCreated,
TaskStatusUpdated,
@@ -42,12 +43,14 @@ from exo.shared.types.tasks import (
TaskStatus,
)
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.runners import RunnerId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
from exo.utils.info_gatherer.net_profile import check_reachable
from exo.utils.keyed_backoff import KeyedBackoff
from exo.utils.task_group import TaskGroup
from exo.worker.plan import plan
from exo.worker.runner.runner_supervisor import RunnerSupervisor
@@ -58,34 +61,34 @@ class Worker:
node_id: NodeId,
session_id: SessionId,
*,
global_event_receiver: Receiver[ForwarderEvent],
local_event_sender: Sender[ForwarderEvent],
global_event_receiver: Receiver[GlobalForwarderEvent],
local_event_sender: Sender[LocalForwarderEvent],
# 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],
event_index_counter: Iterator[int],
):
self.node_id: NodeId = node_id
self.session_id: SessionId = session_id
self.global_event_receiver = global_event_receiver
self.local_event_sender = local_event_sender
self.event_index_counter = event_index_counter
self.command_sender = command_sender
self.download_command_sender = download_command_sender
self.event_buffer = OrderedBuffer[Event]()
self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
self.out_for_delivery: dict[EventId, LocalForwarderEvent] = {}
self.state: State = State()
self.runners: dict[RunnerId, RunnerSupervisor] = {}
self._tg: TaskGroup = create_task_group()
self._tg: TaskGroup = TaskGroup()
self._nack_cancel_scope: CancelScope | None = None
self._nack_attempts: int = 0
self._nack_base_seconds: float = 0.5
self._nack_cap_seconds: float = 10.0
self._system_id = SystemId()
self.event_sender, self.event_receiver = channel[Event]()
# Buffer for input image chunks (for image editing)
@@ -132,6 +135,8 @@ class Worker:
async def _event_applier(self):
with self.global_event_receiver as events:
async for f_event in events:
if f_event.session != self.session_id:
continue
if f_event.origin != self.session_id.master_node_id:
continue
self.event_buffer.ingest(f_event.origin_idx, f_event.event)
@@ -210,20 +215,44 @@ class Worker:
model_id = shard.model_card.model_id
self._download_backoff.record_attempt(model_id)
await self.download_command_sender.send(
ForwarderDownloadCommand(
origin=self.node_id,
command=StartDownload(
target_node_id=self.node_id,
shard_metadata=shard,
),
found_path = resolve_model_in_path(model_id)
if found_path is not None:
logger.info(
f"Model {model_id} found in EXO_MODELS_PATH at {found_path}"
)
)
await self.event_sender.send(
TaskStatusUpdated(
task_id=task.task_id, task_status=TaskStatus.Running
await self.event_sender.send(
NodeDownloadProgress(
download_progress=DownloadCompleted(
node_id=self.node_id,
shard_metadata=shard,
model_directory=str(found_path),
total=shard.model_card.storage_size,
read_only=True,
)
)
)
await self.event_sender.send(
TaskStatusUpdated(
task_id=task.task_id,
task_status=TaskStatus.Complete,
)
)
else:
await self.download_command_sender.send(
ForwarderDownloadCommand(
origin=self._system_id,
command=StartDownload(
target_node_id=self.node_id,
shard_metadata=shard,
),
)
)
await self.event_sender.send(
TaskStatusUpdated(
task_id=task.task_id,
task_status=TaskStatus.Running,
)
)
)
case Shutdown(runner_id=runner_id):
runner = self.runners.pop(runner_id)
try:
@@ -288,7 +317,7 @@ class Worker:
await self._start_runner_task(task)
def shutdown(self):
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _start_runner_task(self, task: Task):
if (instance := self.state.instances.get(task.instance_id)) is not None:
@@ -317,7 +346,7 @@ class Worker:
)
await self.command_sender.send(
ForwarderCommand(
origin=self.node_id,
origin=self._system_id,
command=RequestEventLog(since_idx=since_idx),
)
)
@@ -344,15 +373,16 @@ class Worker:
return runner
async def _forward_events(self) -> None:
idx = 0
with self.event_receiver as events:
async for event in events:
idx = next(self.event_index_counter)
fe = ForwarderEvent(
fe = LocalForwarderEvent(
origin_idx=idx,
origin=self.node_id,
origin=self._system_id,
session=self.session_id,
event=event,
)
idx += 1
logger.debug(f"Worker published event {idx}: {str(event)[:100]}")
await self.local_event_sender.send(fe)
self.out_for_delivery[event.event_id] = fe
View File
+10 -12
View File
@@ -4,7 +4,7 @@ import loguru
from exo.shared.types.events import Event, RunnerStatusUpdated
from exo.shared.types.tasks import Task, TaskId
from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runners import RunnerFailed
from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
@@ -18,28 +18,26 @@ def entrypoint(
cancel_receiver: MpReceiver[TaskId],
_logger: "loguru.Logger",
) -> None:
global logger
logger = _logger
fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
if fast_synch_override == "on" or (
fast_synch_override != "off"
and (
isinstance(bound_instance.instance, MlxJacclInstance)
and len(bound_instance.instance.jaccl_devices) >= 2
)
):
if fast_synch_override != "off":
os.environ["MLX_METAL_FAST_SYNCH"] = "1"
else:
os.environ["MLX_METAL_FAST_SYNCH"] = "0"
global logger
logger = _logger
logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
# Import main after setting global logger - this lets us just import logger from this module
try:
from exo.worker.runner.runner import main
if bound_instance.is_image_model:
from exo.worker.runner.image_models.runner import main
else:
from exo.worker.runner.llm_inference.runner import main
main(bound_instance, event_sender, task_receiver, cancel_receiver)
except ClosedResourceError:
logger.warning("Runner communication closed unexpectedly")
except Exception as e:
@@ -0,0 +1,453 @@
import base64
import resource
import time
from typing import TYPE_CHECKING, Literal
import mlx.core as mx
from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
from exo.shared.models.model_cards import ModelTask
from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
from exo.shared.types.api import ImageGenerationStats
from exo.shared.types.chunks import ErrorChunk, ImageChunk
from exo.shared.types.common import CommandId, ModelId
from exo.shared.types.events import (
ChunkGenerated,
Event,
RunnerStatusUpdated,
TaskAcknowledged,
TaskStatusUpdated,
TraceEventData,
TracesCollected,
)
from exo.shared.types.tasks import (
ConnectToGroup,
ImageEdits,
ImageGeneration,
LoadModel,
Shutdown,
StartWarmup,
Task,
TaskId,
TaskStatus,
)
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runner_response import (
ImageGenerationResponse,
PartialImageResponse,
)
from exo.shared.types.worker.runners import (
RunnerConnected,
RunnerConnecting,
RunnerFailed,
RunnerIdle,
RunnerLoaded,
RunnerLoading,
RunnerReady,
RunnerRunning,
RunnerShutdown,
RunnerShuttingDown,
RunnerStatus,
RunnerWarmingUp,
)
from exo.shared.types.worker.shards import (
CfgShardMetadata,
PipelineShardMetadata,
ShardMetadata,
)
from exo.utils.channels import MpReceiver, MpSender
from exo.worker.engines.image import (
DistributedImageModel,
generate_image,
initialize_image_model,
warmup_image_generator,
)
from exo.worker.engines.mlx.utils_mlx import (
initialize_mlx,
)
from exo.worker.runner.bootstrap import logger
def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool:
"""Check if this node is the primary output node for image generation.
For CFG models: the last pipeline stage in CFG group 0 (positive prompt).
For non-CFG models: the last pipeline stage.
"""
if isinstance(shard_metadata, CfgShardMetadata):
is_pipeline_last = (
shard_metadata.pipeline_rank == shard_metadata.pipeline_world_size - 1
)
return is_pipeline_last and shard_metadata.cfg_rank == 0
elif isinstance(shard_metadata, PipelineShardMetadata):
return shard_metadata.device_rank == shard_metadata.world_size - 1
return False
def _process_image_response(
response: ImageGenerationResponse | PartialImageResponse,
command_id: CommandId,
shard_metadata: ShardMetadata,
event_sender: MpSender[Event],
image_index: int,
) -> None:
"""Process a single image response and send chunks."""
encoded_data = base64.b64encode(response.image_data).decode("utf-8")
is_partial = isinstance(response, PartialImageResponse)
# Extract stats from final ImageGenerationResponse if available
stats = response.stats if isinstance(response, ImageGenerationResponse) else None
_send_image_chunk(
encoded_data=encoded_data,
command_id=command_id,
model_id=shard_metadata.model_card.model_id,
event_sender=event_sender,
image_index=response.image_index,
is_partial=is_partial,
partial_index=response.partial_index if is_partial else None,
total_partials=response.total_partials if is_partial else None,
stats=stats,
image_format=response.format,
)
def _send_traces_if_enabled(
event_sender: MpSender[Event],
task_id: TaskId,
rank: int,
) -> None:
if not EXO_TRACING_ENABLED:
return
traces = get_trace_buffer()
if traces:
trace_data = [
TraceEventData(
name=t.name,
start_us=t.start_us,
duration_us=t.duration_us,
rank=t.rank,
category=t.category,
)
for t in traces
]
event_sender.send(
TracesCollected(
task_id=task_id,
rank=rank,
traces=trace_data,
)
)
clear_trace_buffer()
def _send_image_chunk(
encoded_data: str,
command_id: CommandId,
model_id: ModelId,
event_sender: MpSender[Event],
image_index: int,
is_partial: bool,
partial_index: int | None = None,
total_partials: int | None = None,
stats: ImageGenerationStats | None = None,
image_format: Literal["png", "jpeg", "webp"] | None = None,
) -> None:
"""Send base64-encoded image data as chunks via events."""
data_chunks = [
encoded_data[i : i + EXO_MAX_CHUNK_SIZE]
for i in range(0, len(encoded_data), EXO_MAX_CHUNK_SIZE)
]
total_chunks = len(data_chunks)
for chunk_index, chunk_data in enumerate(data_chunks):
# Only include stats on the last chunk of the final image
chunk_stats = (
stats if chunk_index == total_chunks - 1 and not is_partial else None
)
event_sender.send(
ChunkGenerated(
command_id=command_id,
chunk=ImageChunk(
model=model_id,
data=chunk_data,
chunk_index=chunk_index,
total_chunks=total_chunks,
image_index=image_index,
is_partial=is_partial,
partial_index=partial_index,
total_partials=total_partials,
stats=chunk_stats,
format=image_format,
),
)
)
def main(
bound_instance: BoundInstance,
event_sender: MpSender[Event],
task_receiver: MpReceiver[Task],
cancel_receiver: MpReceiver[TaskId],
):
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (min(max(soft, 2048), hard), hard))
instance, runner_id, shard_metadata = (
bound_instance.instance,
bound_instance.bound_runner_id,
bound_instance.bound_shard,
)
device_rank = shard_metadata.device_rank
logger.info("hello from the runner")
if getattr(shard_metadata, "immediate_exception", False):
raise Exception("Fake exception - runner failed to spin up.")
if timeout := getattr(shard_metadata, "should_timeout", 0):
time.sleep(timeout)
setup_start_time = time.time()
cancelled_tasks = set[TaskId]()
image_model: DistributedImageModel | None = None
group = None
current_status: RunnerStatus = RunnerIdle()
logger.info("runner created")
event_sender.send(
RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
)
seen = set[TaskId]()
with task_receiver as tasks:
for task in tasks:
if task.task_id in seen:
logger.warning("repeat task - potential error")
seen.add(task.task_id)
cancelled_tasks.discard(TaskId("CANCEL_CURRENT_TASK"))
event_sender.send(
TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
)
match task:
case ConnectToGroup() if isinstance(
current_status, (RunnerIdle, RunnerFailed)
):
logger.info("runner connecting")
current_status = RunnerConnecting()
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
)
)
event_sender.send(TaskAcknowledged(task_id=task.task_id))
group = initialize_mlx(bound_instance)
logger.info("runner connected")
current_status = RunnerConnected()
# we load the model if it's connected with a group, or idle without a group. we should never tell a model to connect if it doesn't need to
case LoadModel() if (
isinstance(current_status, RunnerConnected) and group is not None
) or (isinstance(current_status, RunnerIdle) and group is None):
current_status = RunnerLoading()
logger.info("runner loading")
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
)
)
event_sender.send(TaskAcknowledged(task_id=task.task_id))
assert (
ModelTask.TextToImage in shard_metadata.model_card.tasks
or ModelTask.ImageToImage in shard_metadata.model_card.tasks
), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
image_model = initialize_image_model(bound_instance)
current_status = RunnerLoaded()
logger.info("runner loaded")
case StartWarmup() if isinstance(current_status, RunnerLoaded):
current_status = RunnerWarmingUp()
logger.info("runner warming up")
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
)
)
event_sender.send(TaskAcknowledged(task_id=task.task_id))
logger.info(f"warming up inference for instance: {instance}")
assert image_model
image = warmup_image_generator(model=image_model)
if image is not None:
logger.info(f"warmed up by generating {image.size} image")
else:
logger.info("warmup completed (non-primary node)")
logger.info(
f"runner initialized in {time.time() - setup_start_time} seconds"
)
current_status = RunnerReady()
logger.info("runner ready")
case ImageGeneration(
task_params=task_params, command_id=command_id
) if isinstance(current_status, RunnerReady):
assert image_model
logger.info(f"received image generation request: {str(task)[:500]}")
current_status = RunnerRunning()
logger.info("runner running")
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
)
)
event_sender.send(TaskAcknowledged(task_id=task.task_id))
try:
image_index = 0
for response in generate_image(
model=image_model, task=task_params
):
is_primary_output = _is_primary_output_node(shard_metadata)
if is_primary_output:
match response:
case PartialImageResponse():
logger.info(
f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
)
_process_image_response(
response,
command_id,
shard_metadata,
event_sender,
image_index,
)
case ImageGenerationResponse():
logger.info("sending final ImageChunk")
_process_image_response(
response,
command_id,
shard_metadata,
event_sender,
image_index,
)
image_index += 1
# can we make this more explicit?
except Exception as e:
if _is_primary_output_node(shard_metadata):
event_sender.send(
ChunkGenerated(
command_id=command_id,
chunk=ErrorChunk(
model=shard_metadata.model_card.model_id,
finish_reason="error",
error_message=str(e),
),
)
)
raise
finally:
_send_traces_if_enabled(event_sender, task.task_id, device_rank)
current_status = RunnerReady()
logger.info("runner ready")
case ImageEdits(task_params=task_params, command_id=command_id) if (
isinstance(current_status, RunnerReady)
):
assert image_model
logger.info(f"received image edits request: {str(task)[:500]}")
current_status = RunnerRunning()
logger.info("runner running")
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
)
)
event_sender.send(TaskAcknowledged(task_id=task.task_id))
try:
image_index = 0
for response in generate_image(
model=image_model, task=task_params
):
if _is_primary_output_node(shard_metadata):
match response:
case PartialImageResponse():
logger.info(
f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
)
_process_image_response(
response,
command_id,
shard_metadata,
event_sender,
image_index,
)
case ImageGenerationResponse():
logger.info("sending final ImageChunk")
_process_image_response(
response,
command_id,
shard_metadata,
event_sender,
image_index,
)
image_index += 1
except Exception as e:
if _is_primary_output_node(shard_metadata):
event_sender.send(
ChunkGenerated(
command_id=command_id,
chunk=ErrorChunk(
model=shard_metadata.model_card.model_id,
finish_reason="error",
error_message=str(e),
),
)
)
raise
finally:
_send_traces_if_enabled(event_sender, task.task_id, device_rank)
current_status = RunnerReady()
logger.info("runner ready")
case Shutdown():
current_status = RunnerShuttingDown()
logger.info("runner shutting down")
if not TYPE_CHECKING:
del image_model, group
mx.clear_cache()
import gc
gc.collect()
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=current_status
)
)
event_sender.send(TaskAcknowledged(task_id=task.task_id))
current_status = RunnerShutdown()
case _:
raise ValueError(
f"Received {task.__class__.__name__} outside of state machine in {current_status=}"
)
was_cancelled = (task.task_id in cancelled_tasks) or (
TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
)
if not was_cancelled:
event_sender.send(
TaskStatusUpdated(
task_id=task.task_id, task_status=TaskStatus.Complete
)
)
event_sender.send(
RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
)
if isinstance(current_status, RunnerShutdown):
break

Some files were not shown because too many files have changed in this diff Show More