Compare commits

...

46 Commits

Author SHA1 Message Date
Evan 525aaab808 fix 2026-02-26 13:42:07 +00:00
ciaranbor eaed92952c Use tmpdir for coordination file (#1624)
## Motivation

Coordination files for MLX distributed init were written to the current
working directory (./hosts_*.json)

## Changes

- Move coordination file creation to a tempfile.TemporaryDirectory(),
which auto-cleans on context manager exit
2026-02-26 10:59:36 +00:00
Evan Quiney ba611f9cd0 Revert "report macmon failures more aggressively (#1618)" (#1625)
this pr broke macmon - revert it
2026-02-25 19:15:55 +00:00
Evan Quiney eab3e0b456 report macmon failures more aggressively (#1618)
macmon appears to be going silent. time it out and restart it after
3*interval (Default 3s) and report warning.

Co-authored-by: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
2026-02-25 17:49:10 +00:00
Evan Quiney c4e874e97d skip nan logprobs on tokens (#1622)
sometimes we generate NaN logprobs for tokens, this causes pydantic
validation errors on the receiving end. in this case we should just not
send the logprob items
2026-02-25 17:44:15 +00:00
rltakashige e23c3a3026 Address Mac Mini pipeline GPU timeouts (#1620)
## Motivation
Users were reporting GPU timeout errors on Mac Minis, which we never saw
on testing with Mac Studios. It also seems to only happen with large
models.

## Changes
Eval specific distributed operations.

## Why It Works

As I wrote in a Slack message:
Basically, prefill is too slow for pipeline communications. If there are
both communications and GPU operations as part of an mlx graph, the
communications become subject to the GPU's 5 second command buffer
timeout.

For normal generation, I added evals to the communications (only during
prefill, as it slows down decode) to do this, fixing GPU timeouts.

But we don't do this during warmup, as the prompt is absolutely tiny.
This is still too slow on an M4 Pro on some models that it causes a GPU
timeout during warmup...


----------------------
This was one of the issues. However, there is another issue:

mx.all_gather sometimes reads stale data with FAST_SYNCH enabled. I'm
still investigating the root cause, but the code as it is now works on
Mac Minis.



## Test Plan

### Manual Testing
<img width="2762" height="1808" alt="image"
src="https://github.com/user-attachments/assets/27c88542-606c-4551-8f7c-bd2c0471f54e"
/>

<img width="2820" height="1898" alt="image"
src="https://github.com/user-attachments/assets/0ba3478c-ee39-438d-902c-92893db23d05"
/>


### Automated Testing
Needs a bunch on mac minis
2026-02-25 17:37:32 +00:00
Alex Cheema 190e63e56d fix: log exceptions causing silent node shutdown (#1621)
## Motivation

Nodes silently shut down mid-inference with no exception logged. The
logs show a clean-looking shutdown cascade (unsubscribe all topics →
stop worker → runner communication closed) but no error explaining
*why*. This makes debugging cluster issues extremely difficult.

## Changes

**`src/exo/routing/router.py`** — Added `try/except Exception` around
`_networking_recv()` and `_networking_recv_connection_messages()` loops.
Logs the root cause at ERROR level via
`logger.opt(exception=...).error(...)` before re-raising. Uses
`Exception` (not `BaseException`) so clean SIGTERM cancellation is
unaffected.

**`src/exo/main.py`** — Wrapped `anyio.run(node.run)` in
`try/except/finally`:
- `except BaseException`: logs the fatal exception at CRITICAL level
through loguru
- `finally`: ensures `logger.info("EXO Shutdown complete")` and
`logger_cleanup()` always run, guaranteeing the async log queue is
flushed before process exit

## Why It Works

The Router's recv loops (`_networking_recv`,
`_networking_recv_connection_messages`) are infinite `while True` loops
calling Rust bindings with **no exception handling**. When the Rust
networking layer raises `ConnectionError` (e.g. channel closed), the
exception cascades silently through anyio task groups: Router → Node →
all components shut down. The exception was never logged because (1) no
try-except anywhere in the cascade, and (2) `logger_cleanup()` was
skipped when `anyio.run()` raised, so loguru's async queue was never
flushed.

## Test Plan

### Manual Testing
- Run `uv run exo`, send SIGTERM → should see clean shutdown with "EXO
Shutdown complete", no error/critical logs
- Next time the silent shutdown reproduces, logs should now show the
actual exception with full traceback at ERROR level from the recv loop,
plus CRITICAL level from main()

### Automated Testing
- All 222 existing tests pass (1 pre-existing failure in
`test_python.py::test_sleep_on_multiple_items` unrelated to this change)
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — all checks passed
- `nix fmt` — applied

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-25 08:24:18 -08:00
Alex Cheema 7660893538 fix: replace flaky internet checks with explicit offline mode (#1615)
## Motivation

The socket-based internet connectivity probes (connecting to
1.1.1.1/8.8.8.8/1.0.0.1 every 10 seconds) are flaky — they produce false
negatives on some networks/ISPs, causing downloads to silently fail or
get stuck. Instead of dynamically detecting connectivity, this replaces
it with an explicit offline mode that the user opts into.

## Changes

- **Removed internet check infrastructure**: Deleted
`_test_internet_connection()` (socket probes) and
`_check_internet_connection()` (10s polling loop) from
`DownloadCoordinator`
- **Renamed `internet_connection` → `offline`** on `ShardDownloader`
base class and all subclasses (`SingletonShardDownloader`,
`CachedShardDownloader`, `ResumableShardDownloader`)
- **Removed `on_connection_lost` callbacks** from download calls — no
longer needed without dynamic connectivity detection
- **Added `EXO_OFFLINE` env var support** in `constants.py` and `Args`
default in `main.py`
- **Added Offline Mode toggle** to macOS app Settings → General tab,
which sets `EXO_OFFLINE=true` in the process environment and restarts

## Why It Works

The download behavior (`skip_internet` conditionals in
`download_utils.py`) is unchanged — we just changed what drives it.
Instead of a flaky socket probe setting
`internet_connection=True/False`, the user explicitly sets offline mode
via:
1. `--offline` CLI flag
2. `EXO_OFFLINE=true` environment variable
3. macOS app Settings → General → Offline Mode toggle

This is more reliable and predictable. Users on flaky networks no longer
get intermittent download failures.

## Test Plan

### Manual Testing
- `uv run exo --offline` starts without internet checks, rejects
downloads for unavailable models
- `EXO_OFFLINE=true uv run exo` behaves identically
- macOS app Settings toggle persists across restarts and passes env var
to the exo process

### Automated Testing
- All existing tests pass (222 passed), including 8 offline-mode
specific tests
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — all checks passed
- `nix fmt` — 0 files changed

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 13:36:19 -08:00
Evan Quiney 9a2d2a4a7c bump (#1608)
should be release for 1.0.68 - let's synch our py version with our app
version - 0.3.68.
next minor release should be 1.4.0 and 0.4.0 respectively.

Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-24 19:14:15 +00:00
Alex Cheema bea64fe85e fix: prevent stale loading state and conversation loss when switching chats (#1613)
## Motivation

When navigating back to a previous conversation that used a different
model than the one currently loading, the UI incorrectly displayed the
other model's loading/download progress bar instead of the conversation
messages. Additionally, attempting to continue that old conversation by
sending a message would create an entirely new chat rather than
continuing the existing one.

## Changes

All changes in `dashboard/src/routes/+page.svelte`:

1. **Added fallthrough reset in the `chatLaunchState` restore
`$effect`**: When switching to a conversation whose model has no active
instance (not running, not downloading, not loading), the effect now
resets `chatLaunchState` to `"idle"` instead of leaving stale state from
a different model.

2. **Added `skipCreate` parameter to `launchModelForChat()`**: When
continuing an existing conversation (has messages),
`createConversation()` is skipped so the old conversation is preserved
rather than replaced with a new empty one.

3. **Reordered view conditional**: Progress views
(downloading/loading/launching) take priority, then conversation
messages (when idle with messages or model ready), then model selector
(idle with no messages). This ensures:
   - Old conversations display normally when their model isn't running
- Download progress shows when relaunching a model for any conversation
   - Model selector only appears when there are no messages

## Why It Works

- The `$effect` fallthrough prevents `chatLaunchState` from retaining a
stale value (e.g., `"downloading"`) from Model A when the user switches
to a conversation using Model B that has no active state.
- The `skipCreate` flag ensures `launchModelForChat` can be reused for
both new conversations (from model picker) and continuing existing ones
(from chat input) without always creating a new conversation.
- The reordered view logic ensures each state maps to the correct UI:
active launch → progress, existing messages → chat view, nothing → model
selector.

## Test Plan

### Manual Testing
- Load an old conversation while a different model is downloading →
should see conversation messages, not the other model's loading bar
- Send a message from that old conversation → model should
launch/download with progress shown → conversation continues with the
response appended
- Select a new model from the picker → should still create a new
conversation as before
- New conversation with model download → progress bar shows correctly

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 19:08:31 +00:00
rltakashige 14526d281a update mlx 2 (#1611)
## Motivation

GPU locks because of prompt progress callbacks taking time. Current
solution: Don't fix it, make the symptom better

## Changes
Shortened timeout by 2x
Get event leak fixes from latest upstream
2026-02-24 18:30:48 +00:00
Evan Quiney 73e50df827 fix glm5 tool calling (#1612)
glm5 is a deepseekv32 model, so was parsing dsml style tool calls
instead of glm style tool calls. fix it!!
2026-02-24 18:21:18 +00:00
Alex Cheema 2b417f28ff fix: sync model selectors between sidebar and chat input (#1610)
## Motivation

The "Load Model" dropdown in the right sidebar and the model selector
above the chat input were operating on independent state
(`selectedPreviewModelId` vs `selectedChatModel`). This caused several
UX issues:
- Selecting a model in one selector didn't update the other
- After going home from a chat, the chat selector showed "SELECT MODEL"
while the sidebar still showed the model
- On page refresh, only the sidebar restored the saved model
- Sending a chat message with a running model briefly showed the
recommended models view instead of starting the chat

## Changes

- **`handleModelPickerSelect`**: Also sets `selectedChatModel` when a
model is picked from the sidebar dropdown
- **`handleChatPickerSelect`**: Also sets `selectedPreviewModelId` when
a model is picked from the chat selector
- **`handleGoHome`**: Restores `selectedChatModel` from the sidebar's
`selectedModelId` instead of clearing it
- **`applyLaunchDefaults`**: Syncs `selectedChatModel` when restoring
saved defaults on page load
- **`handleChatSend`**: Sets `chatLaunchState = "ready"` before calling
`sendMessage` when the model is already running, ensuring the chat view
renders correctly on view transition

## Why It Works

The two selectors were backed by different store properties that were
never kept in sync. By updating both properties in every selection path
(sidebar pick, chat pick, go-home, page restore), they always reflect
the same model. The `chatLaunchState = "ready"` fix mirrors what
`launchModelForChat` already does (line 2649) and prevents the chat
state from being "idle" when transitioning from the welcome view to the
chat view.

## Test Plan

### Manual Testing
<!-- Hardware: MacBook Pro M4 Max 48GB -->
- Select a model from sidebar "Load Model" dropdown → chat input
selector updates to match
- Select a model from chat input selector → sidebar dropdown updates to
match
- Launch a model, chat with it, click "Go Home" → both selectors show
the same model
- Refresh the page → both selectors show the previously selected model
- With a running model, type and send a message from the welcome view →
chat starts directly without flashing the recommended models view

### Automated Testing
No new automated tests — this is a UI state synchronization fix in the
Svelte dashboard with no backend changes.

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-24 17:31:20 +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
89 changed files with 7734 additions and 1608 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)
}
+30
View File
@@ -5,6 +5,8 @@ import Foundation
private let customNamespaceKey = "EXOCustomNamespace"
private let hfTokenKey = "EXOHFToken"
private let enableImageModelsKey = "EXOEnableImageModels"
private let offlineModeKey = "EXOOfflineMode"
private let onboardingCompletedKey = "EXOOnboardingCompleted"
@MainActor
final class ExoProcessController: ObservableObject {
@@ -59,6 +61,17 @@ final class ExoProcessController: ObservableObject {
UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
}
}
@Published var offlineMode: Bool = {
return UserDefaults.standard.bool(forKey: offlineModeKey)
}()
{
didSet {
UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
}
}
/// 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?
@@ -113,6 +126,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 +180,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)))
@@ -249,6 +276,9 @@ final class ExoProcessController: ObservableObject {
if enableImageModels {
environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
}
if offlineMode {
environment["EXO_OFFLINE"] = "true"
}
var paths: [String] = []
if let existing = environment["PATH"], !existing.isEmpty {
+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
}
}
}
}
}
+489
View File
@@ -0,0 +1,489 @@
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 pendingOfflineMode = 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
pendingOfflineMode = controller.offlineMode
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 {
Toggle("Offline Mode", isOn: $pendingOfflineMode)
Text("Skip internet checks and use only locally available models.")
.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
|| pendingOfflineMode != controller.offlineMode
}
private var hasModelChanges: Bool {
pendingEnableImageModels != controller.enableImageModels
}
private func applyGeneralSettings() {
controller.customNamespace = pendingNamespace
controller.hfToken = pendingHFToken
controller.offlineMode = pendingOfflineMode
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
@@ -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";
+29 -9
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;
}
>;
@@ -590,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;
@@ -842,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;
}
@@ -912,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 {
@@ -1295,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);
}
}
@@ -1836,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);
@@ -2307,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.",
);
}
@@ -3250,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
+153 -29
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" };
@@ -255,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];
@@ -265,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);
});
@@ -357,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>
@@ -381,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}
@@ -395,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"
>
@@ -414,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"
>
@@ -428,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"
@@ -462,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"
@@ -475,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"
>
@@ -505,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"
@@ -537,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.dev20260225+257d5692"; 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 = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
};
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}")
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "exo"
version = "0.3.0"
version = "0.3.68"
description = "Exo"
readme = "README.md"
requires-python = ">=3.13"
@@ -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
@@ -104,6 +104,12 @@ class NetworkingHandle:
will sleep until a message is sent.
"""
@typing.final
class MessageTooLargeError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class NoPeersSubscribedToTopicError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...
+34
View File
@@ -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.
@@ -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()
};
@@ -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>()?;
+113 -53
View File
@@ -1,31 +1,33 @@
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 +37,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 +48,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,15 +65,13 @@ 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)
def __post_init__(self) -> None:
self.event_sender, self.event_receiver = channel[Event]()
if self.offline:
self.shard_downloader.set_internet_connection(False)
self.shard_downloader.on_progress(self._download_progress_callback)
def _model_dir(self, model_id: ModelId) -> str:
@@ -113,45 +120,33 @@ class DownloadCoordinator:
logger.info(
f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
)
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)
def _test_internet_connection(self) -> None:
# Try multiple endpoints since some ISPs/networks block specific IPs
for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
try:
socket.create_connection((host, 443), timeout=3).close()
self.shard_downloader.set_internet_connection(True)
logger.debug(f"Internet connectivity: True (via {host})")
return
except OSError:
continue
self.shard_downloader.set_internet_connection(False)
logger.debug("Internet connectivity: False")
async def _check_internet_connection(self) -> None:
first_connection = True
while True:
await asyncio.sleep(10)
# Assume that internet connection is set to False on 443 errors.
if self.shard_downloader.internet_connection:
continue
self._test_internet_connection()
if first_connection and self.shard_downloader.internet_connection:
first_connection = False
self._tg.start_soon(self._emit_existing_download_progress)
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)
finally:
for task in self.active_downloads.values():
task.cancel()
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 +180,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,
@@ -269,6 +283,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 +321,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:
@@ -345,6 +370,8 @@ class DownloadCoordinator:
model_directory=self._model_dir(
progress.shard.model_card.model_id
),
downloaded=progress.downloaded,
total=progress.total,
)
else:
status = DownloadOngoing(
@@ -364,6 +391,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."
)
+93 -57
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
@@ -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()
@@ -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
@@ -790,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,
@@ -799,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(),
)
+11 -17
View File
@@ -15,9 +15,13 @@ from exo.shared.types.worker.shards import (
)
def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
def exo_shard_downloader(
max_parallel_downloads: int = 8, offline: bool = False
) -> ShardDownloader:
return SingletonShardDownloader(
CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
CachedShardDownloader(
ResumableShardDownloader(max_parallel_downloads, offline=offline)
)
)
@@ -50,10 +54,6 @@ class SingletonShardDownloader(ShardDownloader):
self.shard_downloader = shard_downloader
self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
def set_internet_connection(self, value: bool) -> None:
self.internet_connection = value
self.shard_downloader.set_internet_connection(value)
def on_progress(
self,
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -90,10 +90,6 @@ class CachedShardDownloader(ShardDownloader):
self.shard_downloader = shard_downloader
self.cache: dict[tuple[str, ShardMetadata], Path] = {}
def set_internet_connection(self, value: bool) -> None:
self.internet_connection = value
self.shard_downloader.set_internet_connection(value)
def on_progress(
self,
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -123,8 +119,9 @@ class CachedShardDownloader(ShardDownloader):
class ResumableShardDownloader(ShardDownloader):
def __init__(self, max_parallel_downloads: int = 8):
def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
self.max_parallel_downloads = max_parallel_downloads
self.offline = offline
self.on_progress_callbacks: list[
Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
] = []
@@ -151,8 +148,7 @@ class ResumableShardDownloader(ShardDownloader):
self.on_progress_wrapper,
max_parallel_downloads=self.max_parallel_downloads,
allow_patterns=allow_patterns,
skip_internet=not self.internet_connection,
on_connection_lost=lambda: self.set_internet_connection(False),
skip_internet=self.offline,
)
return target_dir
@@ -168,8 +164,7 @@ class ResumableShardDownloader(ShardDownloader):
shard,
self.on_progress_wrapper,
skip_download=True,
skip_internet=not self.internet_connection,
on_connection_lost=lambda: self.set_internet_connection(False),
skip_internet=self.offline,
)
semaphore = asyncio.Semaphore(self.max_parallel_downloads)
@@ -198,7 +193,6 @@ class ResumableShardDownloader(ShardDownloader):
shard,
self.on_progress_wrapper,
skip_download=True,
skip_internet=not self.internet_connection,
on_connection_lost=lambda: self.set_internet_connection(False),
skip_internet=self.offline,
)
return progress
-5
View File
@@ -16,11 +16,6 @@ from exo.shared.types.worker.shards import (
# TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
class ShardDownloader(ABC):
internet_connection: bool = False
def set_internet_connection(self, value: bool) -> None:
self.internet_connection = value
@abstractmethod
async def ensure_shard(
self, shard: ShardMetadata, config_only: bool = False
@@ -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 -25
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,12 +37,11 @@ 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_node_id())
session_id = SessionId(master_node_id=node_id, election_clock=0)
@@ -57,19 +55,17 @@ 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(
node_id,
session_id,
exo_shard_downloader(),
exo_shard_downloader(offline=args.offline),
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,20 +206,21 @@ 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(
self.node_id,
result.session_id,
exo_shard_downloader(),
exo_shard_downloader(offline=self.offline),
download_command_receiver=self.router.receiver(
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")
@@ -276,9 +270,16 @@ def main():
logger.info("FAST_SYNCH forced OFF")
node = anyio.run(Node.create, args)
anyio.run(node.run)
logger.info("EXO Shutdown complete")
logger_cleanup()
try:
anyio.run(node.run)
except BaseException as exception:
logger.opt(exception=exception).critical(
"EXO terminated due to unhandled exception"
)
raise
finally:
logger.info("EXO Shutdown complete")
logger_cleanup()
class Args(CamelCaseModel):
@@ -289,7 +290,7 @@ class Args(CamelCaseModel):
tb_only: bool = False
no_worker: bool = False
no_downloads: bool = False
offline: bool = False
offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
fast_synch: bool | None = None # None = auto, True = force on, False = force off
@classmethod
@@ -340,6 +341,7 @@ class Args(CamelCaseModel):
parser.add_argument(
"--offline",
action="store_true",
default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
)
fast_synch_group = parser.add_mutually_exclusive_group()
+34 -20
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
@@ -51,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,
@@ -140,11 +140,11 @@ 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,
TracesMerged,
)
@@ -173,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:
@@ -197,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
@@ -206,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
@@ -250,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 = {}
@@ -344,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(
@@ -519,15 +524,15 @@ class API:
if (
model_card.model_id,
sharding,
instance_meta,
instance.sharding(),
instance.instance_meta(),
len(placement_node_ids),
) not in seen:
previews.append(
PlacementPreview(
model_id=model_card.model_id,
sharding=sharding,
instance_meta=instance_meta,
sharding=instance.sharding(),
instance_meta=instance.instance_meta(),
instance=instance,
memory_delta_by_node=memory_delta_by_node or None,
error=None,
@@ -536,8 +541,8 @@ class API:
seen.add(
(
model_card.model_id,
sharding,
instance_meta,
instance.sharding(),
instance.instance_meta(),
len(placement_node_ids),
)
)
@@ -589,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:
@@ -937,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:
@@ -1023,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:
@@ -1586,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)
@@ -1621,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)
@@ -1684,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(
@@ -1811,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,
+10 -4
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,9 +147,6 @@ def place_instance(
instance_id = InstanceId()
target_instances = dict(deepcopy(current_instances))
if len(selected_cycle) == 1:
command.instance_meta = InstanceMeta.MlxRing
match command.instance_meta:
case InstanceMeta.MlxJaccl:
# TODO(evan): shard assignments should contain information about ranks, this is ugly
+9 -9
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,
)
@@ -45,9 +46,9 @@ async def test_master():
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_node_id()}_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(),
+117 -1
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,
@@ -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
+41 -27
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)
@@ -178,28 +174,42 @@ class Router:
logger.info(f"Unsubscribed from {topic}")
async def _networking_recv(self):
while True:
topic, data = await self._net.gossipsub_recv()
logger.trace(f"Received message on {topic} with payload {data}")
if topic not in self.topic_routers:
logger.warning(f"Received message on unknown or inactive topic {topic}")
continue
try:
while True:
topic, data = await self._net.gossipsub_recv()
logger.trace(f"Received message on {topic} with payload {data}")
if topic not in self.topic_routers:
logger.warning(
f"Received message on unknown or inactive topic {topic}"
)
continue
router = self.topic_routers[topic]
await router.publish_bytes(data)
router = self.topic_routers[topic]
await router.publish_bytes(data)
except Exception as exception:
logger.opt(exception=exception).error(
"Gossipsub receive loop terminated unexpectedly"
)
raise
async def _networking_recv_connection_messages(self):
while True:
update = await self._net.connection_update_recv()
message = ConnectionMessage.from_update(update)
logger.trace(
f"Received message on connection_messages with payload {message}"
try:
while True:
update = await self._net.connection_update_recv()
message = ConnectionMessage.from_update(update)
logger.trace(
f"Received message on connection_messages with payload {message}"
)
if CONNECTION_MESSAGES.topic in self.topic_routers:
router = self.topic_routers[CONNECTION_MESSAGES.topic]
assert router.topic.model_type == ConnectionMessage
router = cast(TopicRouter[ConnectionMessage], router)
await router.publish(message)
except Exception as exception:
logger.opt(exception=exception).error(
"Connection update receive loop terminated unexpectedly"
)
if CONNECTION_MESSAGES.topic in self.topic_routers:
router = self.topic_routers[CONNECTION_MESSAGES.topic]
assert router.topic.model_type == ConnectionMessage
router = cast(TopicRouter[ConnectionMessage], router)
await router.publish(message)
raise
async def _networking_publish(self):
with self.networking_receiver as networked_items:
@@ -211,6 +221,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(
+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
+11
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
@@ -69,4 +78,6 @@ EXO_ENABLE_IMAGE_MODELS = (
os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
)
EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
+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
+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
+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 -2
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, 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
@@ -162,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
+3 -1
View File
@@ -30,11 +30,13 @@ class BaseDownloadProgress(TaggedModel):
class DownloadPending(BaseDownloadProgress):
pass
downloaded: Memory = Memory()
total: Memory = Memory()
class DownloadCompleted(BaseDownloadProgress):
total: Memory
read_only: bool = False
class DownloadFailed(BaseDownloadProgress):
+39 -1
View File
@@ -2,8 +2,15 @@ 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.shared.types.worker.runners import RunnerId, ShardAssignments
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
Sharding,
ShardMetadata,
TensorShardMetadata,
)
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -23,16 +30,40 @@ class BaseInstance(TaggedModel):
def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
return self.shard_assignments.runner_to_shard.get(runner_id, None)
@staticmethod
def instance_meta() -> InstanceMeta: ...
def sharding(self) -> Sharding:
if all(
isinstance(sm, PipelineShardMetadata)
for sm in self.shard_assignments.runner_to_shard.values()
):
return Sharding.Pipeline
if all(
isinstance(sm, TensorShardMetadata)
for sm in self.shard_assignments.runner_to_shard.values()
):
return Sharding.Tensor
raise ValueError("shard metadata malformed")
class MlxRingInstance(BaseInstance):
hosts_by_node: dict[NodeId, list[Host]]
ephemeral_port: int
@staticmethod
def instance_meta() -> InstanceMeta:
return InstanceMeta.MlxRing
class MlxJacclInstance(BaseInstance):
jaccl_devices: list[list[str | None]]
jaccl_coordinators: dict[NodeId, str]
@staticmethod
def instance_meta() -> InstanceMeta:
return InstanceMeta.MlxJaccl
# TODO: Single node instance
Instance = MlxRingInstance | MlxJacclInstance
@@ -49,6 +80,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 (
+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()
+2 -2
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
+4 -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,7 +381,7 @@ 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:
@@ -408,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:
+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)
+80 -50
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(
@@ -128,11 +128,11 @@ class PipelineFirstLayer(CustomMlxLayer):
def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
if self.r != 0:
# We want to avoid GPU timeout errors by evalling the distributed operation
# so that it stays on CPU, which does not have a timeout.
mx.eval(x)
x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
if self.is_prefill:
# We want to avoid GPU timeout errors by evalling the distributed operation
# so that it stays on CPU, which does not have a timeout.
mx.eval(x)
mx.eval(x)
return self.original_layer(x, *args, **kwargs)
@@ -158,6 +158,10 @@ class PipelineLastLayer(CustomMlxLayer):
output: mx.array = self.original_layer(x, *args, **kwargs)
# Eval layer output to materialize it before send — this splits the graph
# so the send is isolated and the receiving rank's recv can complete.
mx.eval(output)
if self.r != self.s - 1:
output = mx.distributed.send(
output, (self.r + 1) % self.s, group=self.group
@@ -167,15 +171,15 @@ class PipelineLastLayer(CustomMlxLayer):
# doesn't have .keys directly; access via first sub-cache.
_cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore
_cache.keys = mx.depends(_cache.keys, output) # type: ignore
if self.is_prefill:
mx.eval(output)
if cache is not None:
mx.eval(_cache.keys) # type: ignore
mx.eval(output)
if cache is not None:
mx.eval(_cache.keys) # type: ignore
if not self.is_prefill:
output = mx.distributed.all_gather(output, group=self.group)[
-output.shape[0] :
]
mx.eval(output)
return output
@@ -186,7 +190,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 +208,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 +225,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 +235,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 +359,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 +471,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 +498,7 @@ class TensorParallelShardingStrategy(ABC):
model: nn.Module,
timeout_seconds: float,
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
) -> nn.Module: ...
@@ -498,13 +508,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 +527,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 +564,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 +621,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 +649,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 +705,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 +795,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 +825,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 +836,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 +856,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 +952,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 +963,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 +1000,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 +1011,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 +1043,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 +1054,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 +1088,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,3 +1,4 @@
import math
import time
from copy import deepcopy
from typing import Callable, Generator, cast, get_args
@@ -248,6 +249,9 @@ def extract_top_logprobs(
for i in range(top_logprobs):
token_id = int(top_indices[i].item())
token_logprob = float(top_values[i].item())
if math.isnan(token_logprob):
continue
# Decode token ID to string
token_str = tokenizer.decode([token_id])
# Get byte representation
+49 -26
View File
@@ -2,6 +2,7 @@ import json
import os
import re
import sys
import tempfile
import time
from pathlib import Path
from typing import Any, cast
@@ -23,9 +24,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 +54,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,
)
@@ -97,14 +99,13 @@ def mlx_distributed_init(
rank = bound_instance.bound_shard.device_rank
logger.info(f"Starting initialization for rank {rank}")
coordination_file = None
try:
with tempfile.TemporaryDirectory() as tmpdir:
coordination_file = str(
Path(tmpdir) / f"hosts_{bound_instance.instance.instance_id}_{rank}.json"
)
# TODO: singleton instances
match bound_instance.instance:
case MlxRingInstance(hosts_by_node=hosts_by_node, ephemeral_port=_):
coordination_file = (
f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
)
hosts_for_node = hosts_by_node[bound_instance.bound_node_id]
hosts_json = HostList.from_hosts(hosts_for_node).model_dump_json()
@@ -127,9 +128,6 @@ def mlx_distributed_init(
jaccl_devices[i][i] is None for i in range(len(jaccl_devices))
)
# Use RDMA connectivity matrix
coordination_file = (
f"./hosts_{bound_instance.instance.instance_id}_{rank}.json"
)
jaccl_devices_json = json.dumps(jaccl_devices)
with open(coordination_file, "w") as f:
@@ -149,10 +147,6 @@ def mlx_distributed_init(
logger.info(f"Rank {rank} mlx distributed initialization complete")
return group
finally:
with contextlib.suppress(FileNotFoundError):
if coordination_file:
os.remove(coordination_file)
def initialize_mlx(
@@ -171,13 +165,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 +195,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(
@@ -195,13 +207,16 @@ def load_mlx_items(
set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
mx.clear_cache()
return cast(Model, model), tokenizer
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)
@@ -242,10 +257,14 @@ def shard_and_load(
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 +286,11 @@ 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)
return load_tokenizer_for_model_id(
shard_metadata.model_card.model_id,
model_path,
trust_remote_code=shard_metadata.model_card.trust_remote_code,
)
def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
@@ -299,7 +322,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 +391,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,
)
@@ -616,7 +639,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.
"""
@@ -643,7 +666,7 @@ def set_wired_limit_for_model(model_size: Memory):
return
max_rec_size = Memory.from_bytes(
int(mx.metal.device_info()["max_recommended_working_set_size"])
int(mx.device_info()["max_recommended_working_set_size"])
)
if model_size > 0.9 * max_rec_size:
logger.warning(
+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
@@ -1,10 +1,9 @@
import base64
import math
import resource
import time
from collections.abc import Generator
from functools import cache
from typing import TYPE_CHECKING, Literal
from typing import TYPE_CHECKING, cast
import mlx.core as mx
from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
@@ -18,31 +17,22 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs]
load_harmony_encoding,
)
from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
from exo.shared.models.model_cards import ModelId, ModelTask
from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
from exo.shared.types.api import ImageGenerationStats
from exo.shared.models.model_cards import ModelTask
from exo.shared.types.chunks import (
ErrorChunk,
ImageChunk,
PrefillProgressChunk,
TokenChunk,
ToolCallChunk,
)
from exo.shared.types.common import CommandId
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,
@@ -55,8 +45,6 @@ from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.runner_response import (
GenerationResponse,
ImageGenerationResponse,
PartialImageResponse,
ToolCallItem,
ToolCallResponse,
)
@@ -74,18 +62,7 @@ from exo.shared.types.worker.runners import (
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 import Model
from exo.worker.engines.mlx.cache import KVPrefixCache
from exo.worker.engines.mlx.generator.generate import (
@@ -106,22 +83,6 @@ from exo.worker.runner.bootstrap import logger
from .tool_parsers import ToolParser, make_mlx_parser
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 main(
bound_instance: BoundInstance,
event_sender: MpSender[Event],
@@ -136,6 +97,7 @@ def main(
bound_instance.bound_runner_id,
bound_instance.bound_shard,
)
model_id = shard_metadata.model_card.model_id
device_rank = shard_metadata.device_rank
logger.info("hello from the runner")
if getattr(shard_metadata, "immediate_exception", False):
@@ -146,9 +108,7 @@ def main(
setup_start_time = time.time()
cancelled_tasks = set[TaskId]()
# type checker was unhappy with me - splitting these fixed it
inference_model: Model | None = None
image_model: DistributedImageModel | None = None
tokenizer = None
tool_parser: ToolParser | None = None
group = None
@@ -191,7 +151,10 @@ def main(
case LoadModel() if (
isinstance(current_status, RunnerConnected) and group is not None
) or (isinstance(current_status, RunnerIdle) and group is None):
current_status = RunnerLoading()
total_layers = shard_metadata.end_layer - shard_metadata.start_layer
current_status = RunnerLoading(
layers_loaded=0, total_layers=total_layers
)
logger.info("runner loading")
event_sender.send(
RunnerStatusUpdated(
@@ -211,33 +174,40 @@ def main(
)
time.sleep(0.5)
if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
inference_model, tokenizer = load_mlx_items(
bound_instance, group, on_timeout=on_model_load_timeout
def on_layer_loaded(layers_loaded: int, total: int) -> None:
nonlocal current_status
current_status = RunnerLoading(
layers_loaded=layers_loaded, total_layers=total
)
logger.info(
f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}"
)
if tokenizer.has_tool_calling:
assert tokenizer.tool_call_start
assert tokenizer.tool_call_end
assert tokenizer.tool_parser # pyright: ignore[reportAny]
tool_parser = make_mlx_parser(
tokenizer.tool_call_start,
tokenizer.tool_call_end,
tokenizer.tool_parser, # pyright: ignore[reportAny]
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id,
runner_status=current_status,
)
kv_prefix_cache = KVPrefixCache(group)
elif (
ModelTask.TextToImage in shard_metadata.model_card.tasks
or ModelTask.ImageToImage in shard_metadata.model_card.tasks
):
image_model = initialize_image_model(bound_instance)
else:
raise ValueError(
f"Unknown model task(s): {shard_metadata.model_card.tasks}"
)
assert (
ModelTask.TextGeneration in shard_metadata.model_card.tasks
), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
inference_model, tokenizer = load_mlx_items(
bound_instance,
group,
on_timeout=on_model_load_timeout,
on_layer_loaded=on_layer_loaded,
)
logger.info(
f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}"
)
if tokenizer.has_tool_calling:
assert tokenizer.tool_call_start
assert tokenizer.tool_call_end
assert tokenizer.tool_parser # pyright: ignore[reportAny]
tool_parser = make_mlx_parser(
tokenizer.tool_call_start,
tokenizer.tool_call_end,
tokenizer.tool_parser, # pyright: ignore[reportAny]
)
kv_prefix_cache = KVPrefixCache(group)
current_status = RunnerLoaded()
logger.info("runner loaded")
case StartWarmup() if isinstance(current_status, RunnerLoaded):
@@ -251,46 +221,34 @@ def main(
event_sender.send(TaskAcknowledged(task_id=task.task_id))
logger.info(f"warming up inference for instance: {instance}")
if ModelTask.TextGeneration in shard_metadata.model_card.tasks:
assert inference_model
assert tokenizer
assert inference_model
assert tokenizer
t = time.monotonic()
toks = warmup_inference(
model=inference_model,
tokenizer=tokenizer,
group=group,
t = time.monotonic()
toks = warmup_inference(
model=cast(Model, inference_model),
tokenizer=tokenizer,
group=group,
)
logger.info(f"warmed up by generating {toks} tokens")
check_for_cancel_every = min(
math.ceil(toks / min(time.monotonic() - t, 0.001)), 100
)
if group is not None:
check_for_cancel_every = int(
mx.max(
mx.distributed.all_gather(
mx.array([check_for_cancel_every]), group=group
)
).item()
)
logger.info(f"warmed up by generating {toks} tokens")
check_for_cancel_every = min(
math.ceil(toks / min(time.monotonic() - t, 0.001)), 100
)
if group is not None:
check_for_cancel_every = int(
mx.max(
mx.distributed.all_gather(
mx.array([check_for_cancel_every]), group=group
)
).item()
)
logger.info(
f"runner checking for cancellation every {check_for_cancel_every} tokens"
)
logger.info(
f"runner initialized in {time.time() - setup_start_time} seconds"
)
elif (
ModelTask.TextToImage in shard_metadata.model_card.tasks
or ModelTask.ImageToImage in shard_metadata.model_card.tasks
):
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 checking for cancellation every {check_for_cancel_every} tokens"
)
logger.info(
f"runner initialized in {time.time() - setup_start_time} seconds"
)
current_status = RunnerReady()
logger.info("runner ready")
case TextGeneration(task_params=task_params, command_id=command_id) if (
@@ -324,7 +282,7 @@ def main(
ChunkGenerated(
command_id=command_id,
chunk=PrefillProgressChunk(
model=shard_metadata.model_card.model_id,
model=model_id,
processed_tokens=processed,
total_tokens=total,
),
@@ -345,7 +303,7 @@ def main(
# Generate responses using the actual MLX generation
mlx_generator = mlx_generate(
model=inference_model,
model=cast(Model, inference_model),
tokenizer=tokenizer,
task=task_params,
prompt=prompt,
@@ -368,13 +326,16 @@ def main(
# Model-specific output parsing for tool calls.
if isinstance(inference_model, GptOssModel):
mlx_generator = parse_gpt_oss(mlx_generator)
elif isinstance(inference_model, DeepseekV32Model):
elif (
isinstance(inference_model, DeepseekV32Model)
and "deepseek" in model_id.normalize().lower()
):
mlx_generator = parse_deepseek_v32(mlx_generator)
elif tool_parser:
mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
completion_tokens = 0
tokens_since_last_cancel_check = 0
tokens_since_last_cancel_check = check_for_cancel_every
for response in mlx_generator:
tokens_since_last_cancel_check += 1
if tokens_since_last_cancel_check >= check_for_cancel_every:
@@ -398,7 +359,7 @@ def main(
command_id=command_id,
chunk=ErrorChunk(
error_message=response.text,
model=shard_metadata.model_card.model_id,
model=model_id,
),
)
)
@@ -413,7 +374,7 @@ def main(
ChunkGenerated(
command_id=command_id,
chunk=TokenChunk(
model=shard_metadata.model_card.model_id,
model=model_id,
text=response.text,
token_id=response.token,
usage=response.usage,
@@ -432,7 +393,7 @@ def main(
command_id=command_id,
chunk=ToolCallChunk(
tool_calls=response.tool_calls,
model=shard_metadata.model_card.model_id,
model=model_id,
usage=response.usage,
stats=response.stats,
),
@@ -448,7 +409,7 @@ def main(
ChunkGenerated(
command_id=command_id,
chunk=ErrorChunk(
model=shard_metadata.model_card.model_id,
model=model_id,
finish_reason="error",
error_message=str(e),
),
@@ -458,138 +419,12 @@ def main(
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, shard_metadata.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, shard_metadata.device_rank
)
current_status = RunnerReady()
logger.info("runner ready")
case Shutdown():
current_status = RunnerShuttingDown()
logger.info("runner shutting down")
if not TYPE_CHECKING:
del inference_model, image_model, tokenizer, group
del inference_model, tokenizer, group
mx.clear_cache()
import gc
@@ -890,111 +725,13 @@ def parse_thinking_models(
yield response.model_copy(update={"is_thinking": in_thinking})
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 _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 _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 parse_tool_calls(
responses: Generator[GenerationResponse], tool_parser: ToolParser
) -> Generator[GenerationResponse | ToolCallResponse]:
in_tool_call = False
tool_call_text_parts: list[str] = []
for response in responses:
if response.text.startswith(tool_parser.start_parsing):
if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
in_tool_call = True
if in_tool_call:
@@ -1032,10 +769,9 @@ def parse_tool_calls(
)
yield response
continue
# fallthrough
yield response
else:
# fallthrough
yield response
EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
@@ -58,15 +58,16 @@ def _flatten(p: dict[str, Any]) -> dict[str, str]:
}
json_tool_parser = ToolParser(
start_parsing="<tool_call>",
end_parsing="</tool_call>",
parse_tool_calls=_parse_json_calls,
)
def make_json_parser() -> ToolParser:
return ToolParser(
start_parsing="<tool_call>",
end_parsing="</tool_call>",
parse_tool_calls=_parse_json_calls,
)
def infer_tool_parser(chat_template: str) -> ToolParser | None:
"""Attempt to auto-infer a tool parser from the chat template."""
if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
return json_tool_parser
return make_json_parser()
return None
+50 -22
View File
@@ -1,7 +1,7 @@
import contextlib
import multiprocessing as mp
import signal
from dataclasses import dataclass, field
from multiprocessing import Process
from typing import Self
import anyio
@@ -32,6 +32,7 @@ from exo.shared.types.worker.runners import (
)
from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
from exo.utils.task_group import TaskGroup
from exo.worker.runner.bootstrap import entrypoint
PREFILL_TIMEOUT_SECONDS = 60
@@ -42,16 +43,20 @@ DECODE_TIMEOUT_SECONDS = 5
class RunnerSupervisor:
shard_metadata: ShardMetadata
bound_instance: BoundInstance
runner_process: Process
runner_process: mp.Process
initialize_timeout: float
_ev_recv: MpReceiver[Event]
_task_sender: MpSender[Task]
_event_sender: Sender[Event]
_cancel_sender: MpSender[TaskId]
_tg: TaskGroup = field(default_factory=TaskGroup, init=False)
status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
completed: set[TaskId] = field(default_factory=set, init=False)
cancelled: set[TaskId] = field(default_factory=set, init=False)
_cancel_watch_runner: anyio.CancelScope = field(
default_factory=anyio.CancelScope, init=False
)
@classmethod
def create(
@@ -65,7 +70,7 @@ class RunnerSupervisor:
task_sender, task_recv = mp_channel[Task]()
cancel_sender, cancel_recv = mp_channel[TaskId]()
runner_process = Process(
runner_process = mp.Process(
target=entrypoint,
args=(
bound_instance,
@@ -94,15 +99,25 @@ class RunnerSupervisor:
async def run(self):
self.runner_process.start()
await self._forward_events()
async with self._tg as tg:
tg.start_soon(self._watch_runner)
tg.start_soon(self._forward_events)
def shutdown(self):
logger.info("Runner supervisor shutting down")
self._ev_recv.close()
self._task_sender.close()
self._tg.cancel_tasks()
if not self._cancel_watch_runner.cancel_called:
self._cancel_watch_runner.cancel()
with contextlib.suppress(ClosedResourceError):
self._ev_recv.close()
with contextlib.suppress(ClosedResourceError):
self._task_sender.close()
with contextlib.suppress(ClosedResourceError):
self._event_sender.close()
with contextlib.suppress(ClosedResourceError):
self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
self._cancel_sender.close()
with contextlib.suppress(ClosedResourceError):
self._cancel_sender.close()
self.runner_process.join(5)
if not self.runner_process.is_alive():
logger.info("Runner process succesfully terminated")
@@ -151,8 +166,8 @@ class RunnerSupervisor:
await self._check_runner(TimeoutError("cancel pipe blocked"))
async def _forward_events(self):
with self._ev_recv as events:
try:
try:
with self._ev_recv as events:
async for event in events:
if isinstance(event, RunnerStatusUpdated):
self.status = event.runner_status
@@ -176,25 +191,34 @@ class RunnerSupervisor:
)
self.completed.add(event.task_id)
await self._event_sender.send(event)
except (ClosedResourceError, BrokenResourceError) as e:
await self._check_runner(e)
for tid in self.pending:
self.pending[tid].set()
self._event_sender.close()
except (ClosedResourceError, BrokenResourceError) as e:
await self._check_runner(e)
finally:
for tid in self.pending:
self.pending[tid].set()
def __del__(self) -> None:
if self.runner_process.is_alive():
logger.warning("RunnerSupervisor was not stopped cleanly.")
logger.critical("RunnerSupervisor was not stopped cleanly.")
with contextlib.suppress(ValueError):
self.runner_process.kill()
async def _watch_runner(self) -> None:
with self._cancel_watch_runner:
while True:
await anyio.sleep(5)
if not self.runner_process.is_alive():
await self._check_runner(RuntimeError("Runner found to be dead"))
async def _check_runner(self, e: Exception) -> None:
if not self._cancel_watch_runner.cancel_called:
self._cancel_watch_runner.cancel()
logger.info("Checking runner's status")
if self.runner_process.is_alive():
logger.info("Runner was found to be alive, attempting to join process")
await to_thread.run_sync(self.runner_process.join, 5)
rc = self.runner_process.exitcode
logger.info(f"RunnerSupervisor exited with exit code {rc}")
logger.info(f"Runner exited with exit code {rc}")
if rc == 0:
return
@@ -207,15 +231,19 @@ class RunnerSupervisor:
else:
cause = f"exitcode={rc}"
logger.opt(exception=e).error(f"Runner terminated ({cause})")
logger.opt(exception=e).error(f"Runner terminated with {cause}")
try:
await self._event_sender.send(
RunnerStatusUpdated(
runner_id=self.bound_instance.bound_runner_id,
runner_status=RunnerFailed(error_message=f"Terminated ({cause})"),
self.status = RunnerFailed(error_message=f"Terminated ({cause})")
with anyio.CancelScope(shield=True):
await self._event_sender.send(
RunnerStatusUpdated(
runner_id=self.bound_instance.bound_runner_id,
runner_status=RunnerFailed(
error_message=f"Terminated ({cause})"
),
)
)
)
except (ClosedResourceError, BrokenResourceError):
logger.warning(
"Event sender already closed, unable to report runner failure"
@@ -96,7 +96,9 @@ def run_gpt_oss_pipeline_device(
n_layers=24,
)
model, tokenizer = shard_and_load(shard_meta, group)
model, tokenizer = shard_and_load(
shard_meta, group, on_timeout=None, on_layer_loaded=None
)
model = cast(Model, model)
# Generate a prompt of exact token length
@@ -172,7 +174,9 @@ def run_gpt_oss_tensor_parallel_device(
n_layers=24,
)
model, tokenizer = shard_and_load(shard_meta, group)
model, tokenizer = shard_and_load(
shard_meta, group, on_timeout=None, on_layer_loaded=None
)
model = cast(Model, model)
base_text = "The quick brown fox jumps over the lazy dog. "
@@ -19,7 +19,7 @@ from exo.worker.engines.mlx.dsml_encoding import (
encode_messages,
parse_dsml_output,
)
from exo.worker.runner.runner import parse_deepseek_v32
from exo.worker.runner.llm_inference.runner import parse_deepseek_v32
# ── Shared fixtures ──────────────────────────────────────────────
@@ -6,7 +6,7 @@ from typing import Callable
import mlx.core as mx
import pytest
import exo.worker.runner.runner as mlx_runner
import exo.worker.runner.llm_inference.runner as mlx_runner
from exo.shared.types.chunks import TokenChunk
from exo.shared.types.events import (
ChunkGenerated,
@@ -180,7 +180,7 @@ def _run(tasks: Iterable[Task]):
task_receiver.close = nothin
task_receiver.join = nothin
with unittest.mock.patch(
"exo.worker.runner.runner.mx.distributed.all_gather",
"exo.worker.runner.llm_inference.runner.mx.distributed.all_gather",
make_nothin(mx.array([1])),
):
mlx_runner.main(
@@ -224,7 +224,10 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerConnected()),
TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Running),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoading()),
RunnerStatusUpdated(
runner_id=RUNNER_1_ID,
runner_status=RunnerLoading(layers_loaded=0, total_layers=32),
),
TaskAcknowledged(task_id=LOAD_TASK_ID),
TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Complete),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoaded()),
@@ -4,7 +4,7 @@ from exo.shared.types.worker.runner_response import (
GenerationResponse,
ToolCallResponse,
)
from exo.worker.runner.runner import parse_gpt_oss
from exo.worker.runner.llm_inference.runner import parse_gpt_oss
# Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer.
# These are stable since they come from the model's vocabulary.
@@ -4,8 +4,8 @@ from collections.abc import Generator
from typing import Any
from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse
from exo.worker.runner.runner import parse_tool_calls
from exo.worker.runner.tool_parsers import make_mlx_parser
from exo.worker.runner.llm_inference.runner import parse_tool_calls
from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
def _make_responses(
+7 -5
View File
@@ -17,6 +17,13 @@ git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
exit 1
}
hosts=("$@")
for host; do
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
done
wait
cleanup() {
for host in "${hosts[@]}"; do
ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
@@ -26,11 +33,6 @@ cleanup() {
}
trap 'cleanup' EXIT INT TERM
for host; do
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
done
wait
for host; do
ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
"EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
+125
View File
@@ -0,0 +1,125 @@
#!/usr/bin/env bash
# Test that models added via API get trust_remote_code=false
# Run this against a running exo instance.
# Usage: ./test_trust_remote_code_attack.sh [host:port]
set -uo pipefail
HOST="${1:-localhost:52415}"
MODEL_ID="KevTheHermit/security-testing"
CUSTOM_CARDS_DIR="$HOME/.exo/custom_model_cards"
CARD_FILE="$CUSTOM_CARDS_DIR/KevTheHermit--security-testing.toml"
echo "=== Test: trust_remote_code attack via API ==="
echo "Target: $HOST"
echo ""
# Clean up RCE proof from previous runs
rm -f /tmp/exo-rce-proof.txt
# Step 0: Clean up any stale card from previous runs
if [ -f "$CARD_FILE" ]; then
echo "[0] Removing stale card from previous run ..."
curl -s -X DELETE \
"http://$HOST/models/custom/$(python3 -c 'import urllib.parse; print(urllib.parse.quote("'"$MODEL_ID"'", safe=""))')" >/dev/null
rm -f "$CARD_FILE"
echo " Done"
echo ""
fi
# Step 1: Add the malicious model via API
echo "[1] Adding model via POST /models/add ..."
ADD_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "http://$HOST/models/add" \
-H "Content-Type: application/json" \
-d "{\"model_id\":\"$MODEL_ID\"}")
HTTP_CODE=$(echo "$ADD_RESPONSE" | tail -1)
BODY=$(echo "$ADD_RESPONSE" | sed '$d')
echo " HTTP $HTTP_CODE"
if [ "$HTTP_CODE" -ge 400 ]; then
echo " Model add failed (HTTP $HTTP_CODE) — that's fine if model doesn't exist on HF."
echo " Response: $BODY"
echo ""
echo "RESULT: Model was rejected at add time. Attack blocked."
exit 0
fi
# Step 2: Verify the saved TOML has trust_remote_code = false
echo ""
echo "[2] Checking saved model card TOML ..."
if [ ! -f "$CARD_FILE" ]; then
echo " FAIL: Card file not found at $CARD_FILE"
exit 1
fi
if grep -q 'trust_remote_code = false' "$CARD_FILE"; then
echo " SAFE: trust_remote_code = false (fix is active)"
else
echo " VULNERABLE: trust_remote_code is not false — remote code WILL be trusted"
fi
echo " Contents:"
cat "$CARD_FILE"
# Step 3: Place the instance
echo ""
echo "[3] Attempting POST /place_instance ..."
PLACE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "http://$HOST/place_instance" \
-H "Content-Type: application/json" \
-d "{\"model_id\":\"$MODEL_ID\"}")
PLACE_CODE=$(echo "$PLACE_RESPONSE" | tail -1)
PLACE_BODY=$(echo "$PLACE_RESPONSE" | sed '$d')
echo " HTTP $PLACE_CODE"
echo " Response: $PLACE_BODY"
# Step 3b: Send a chat completion to actually trigger tokenizer loading
echo ""
echo "[3b] Sending chat completion to trigger tokenizer load ..."
CHAT_RESPONSE=$(curl -s -w "\n%{http_code}" --max-time 30 -X POST "http://$HOST/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"max_tokens\":1}")
CHAT_CODE=$(echo "$CHAT_RESPONSE" | tail -1)
CHAT_BODY=$(echo "$CHAT_RESPONSE" | sed '$d')
echo " HTTP $CHAT_CODE"
echo " Response: $CHAT_BODY"
echo ""
echo "[3c] Checking for RCE proof ..."
sleep 5
if [ -f /tmp/exo-rce-proof.txt ]; then
echo " VULNERABLE: Remote code executed!"
echo " Contents:"
cat /tmp/exo-rce-proof.txt
else
echo " SAFE: /tmp/exo-rce-proof.txt does not exist — remote code was NOT executed"
fi
# Step 4: Clean up — delete instance and custom model
echo ""
echo "[4] Cleaning up ..."
# Find and delete any instance for this model
INSTANCE_ID=$(curl -s "http://$HOST/state" | python3 -c "
import sys, json
state = json.load(sys.stdin)
for iid, wrapper in state.get('instances', {}).items():
for tag, inst in wrapper.items():
sa = inst.get('shardAssignments', {})
if sa.get('modelId', '') == '$MODEL_ID':
print(iid)
sys.exit(0)
" 2>/dev/null || true)
if [ -n "$INSTANCE_ID" ]; then
echo " Deleting instance $INSTANCE_ID ..."
curl -s -X DELETE "http://$HOST/instance/$INSTANCE_ID" >/dev/null
echo " Done"
else
echo " No instance found to delete"
fi
echo " Deleting custom model card ..."
curl -s -X DELETE \
"http://$HOST/models/custom/$(python3 -c 'import urllib.parse; print(urllib.parse.quote("'"$MODEL_ID"'", safe=""))')" >/dev/null
echo " Done"
echo ""
echo "=== DONE ==="
Generated
+8 -6
View File
@@ -363,7 +363,7 @@ wheels = [
[[package]]
name = "exo"
version = "0.3.0"
version = "0.3.68"
source = { editable = "." }
dependencies = [
{ name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -378,7 +378,7 @@ dependencies = [
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.7.dev20260218+14841977", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -451,6 +451,7 @@ dependencies = [
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
@@ -461,6 +462,7 @@ requires-dist = [
{ name = "huggingface-hub", specifier = ">=0.33.4" },
{ name = "jinja2", specifier = ">=3.1.0" },
{ name = "loguru", specifier = ">=0.7.3" },
{ name = "protobuf", specifier = ">=5.29.0" },
{ name = "tiktoken", specifier = ">=0.12.0" },
{ name = "transformers", specifier = ">=5.0.0" },
]
@@ -1023,7 +1025,7 @@ dependencies = [
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.7.dev20260218+14841977", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1070,8 +1072,8 @@ cuda13 = [
[[package]]
name = "mlx"
version = "0.30.7.dev20260218+14841977"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }
version = "0.30.7.dev20260225+257d5692"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
resolution-markers = [
"sys_platform == 'darwin'",
]
@@ -1106,7 +1108,7 @@ version = "0.30.7"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.7.dev20260218+14841977", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },