Compare commits

...

75 Commits

Author SHA1 Message Date
Evan 9f184de322 clean 2026-03-04 18:50:32 +00:00
Miguel Miranda Dias 8485805042 fix(worker): emit error chunks when a runner dies mid-command (#1645)
Closes #1586

## Summary
- track in-flight tasks in `RunnerSupervisor` (not only unacknowledged
pending tasks)
- when `_check_runner()` detects a crashed runner, emit
`ChunkGenerated(ErrorChunk)` for each in-flight command task
(`TextGeneration`, `ImageGeneration`, `ImageEdits`)
- keep existing `RunnerStatusUpdated(RunnerFailed)` emission so
planner/state still transition correctly
- add a unit test for supervisor crash path to ensure an error chunk is
emitted before failed runner status

## Why
`#1586` reports streams that can hang forever when runners crash during
warmup/loading. This keeps failure signaling at the runner-supervisor
layer, matching maintainer guidance in the issue thread.

## Validation
- attempted: `uv run pytest
src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py`
- blocked locally by environment disk exhaustion while uv tried to
materialize heavy CUDA wheels (`No space left on device` during
`nvidia-cudnn-cu13` extraction)

I kept the change scoped and added a targeted unit test for the failure
path.

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-04 17:15:58 +00:00
ciaranbor 4de8f801c7 #Add reasoning parms to chat completion and responses APIs (#1654)
## Motivation

Adds reasoning_effort parameter support to both the Chat Completions and
Responses APIs, aligning with the OpenAI spec and enabling thinking
control for gpt-oss models

## Changes

- Added ReasoningEffort literal type ("none" | "minimal" | "low" |
"medium" | "high" | "xhigh") and a resolve_reasoning_params() helper
that cross-derives reasoning_effort ↔ enable_thinking when only one is
provided
- Added reasoning_effort field to ChatCompletionRequest and reasoning
(with Reasoning model) + enable_thinking to ResponsesRequest
- Both adapters now call resolve_reasoning_params() before building
TextGenerationTaskParams
- reasoning_effort is passed through to the MLX chat template as a
template variable

## Why It Works

resolve_reasoning_params is a pure function that normalises the two
overlapping knobs (reasoning_effort and enable_thinking) into a
consistent pair, so downstream code always has both values regardless of
which the caller supplied.

## Test Plan

### Automated Testing

Added test_resolve_reasoning_params.py with 10 test cases covering:
both-None, both-set passthrough, enable_thinking → effort derivation,
and effort → enable_thinking derivation for every ReasoningEffort
variant.
2026-03-04 14:27:49 +00:00
Owleksiy 5777bf3c39 fix: coerce tool-call argument types from tool schema (#1651)
Apply schema-aware coercion to parsed tool-call arguments so
Hermes-style toolcalls can still return typed JSON (e.g. integer ids).

 - pass request tools into parse_tool_calls
 - coerce parsed argument values by function parameters schema
 - add unit tests for coercion and unknown-tool passthrough

## Motivation

Models that use Hermes-based toolcall syntax (Qwen3.5) can't reliably
call tools with non-string parameters
Example tool:
```json
    {
      "type": "function",
      "function": {
        "name": "process",
        "description": "Manage background processes",
        "parameters": {
          "type": "object",
          "properties": {
            "action": {
              "type": "string",
              "enum": ["spawn", "output", "kill", "list"]
            },
            "id": {
              "type": "integer",
              "description": "Process id"
            },
            "command": {
              "type": "string",
              "description": "Command to run for spawn"
            }
          },
          "required": ["action"],
          "additionalProperties": false
        }
      }
    }
```
Model transcript:
```
<tool_call>
<function=process>
<parameter=action>
output
</parameter>
<parameter=id>
0
</parameter>
</function>
</tool_call>
```
And the API returns:

`{"id":"a8f11689-d840-4ca5-ab1d-ead3678a11a9","name":"process","arguments":"{\"action\":
\"output\", \"id\": \"0\"}"}}`

Tool definition declared `id` as `integer`, the model output is
type-agnostic, and the translation layer treats everything as a string.

The same Qwen3.5-27B on OpenRouter and GPT-4.1-mini on openai obey the
function signature and emit correct call:
```
{"name":"process","arguments":"{\"action\": \"output\", \"id\": 0}"}}
```

Steps to reproduce:
```
❯ curl -sS -v http://localhost:52415/v1/chat/completions \
  -H 'Content-Type: application/json' \  -d @- <<'JSON'
{
  "model": "mlx-community/Qwen3.5-27B-4bit",
  "stream": false,
  "temperature": 0,
  "messages": [
    {
      "role": "user",
      "content": "Call the process tool with action=output and id=0. Do not explain anything. Just make the tool call."
    }
  ],
  "tools": [
    {
      "type": "function",
      "function": {
        "name": "process",
        "description": "Manage background processes",
        "parameters": {
          "type": "object",
          "properties": {
            "action": {
              "type": "string",
              "enum": ["spawn", "output", "kill", "list"]
            },
            "id": {
              "type": "integer",
              "description": "Process id"
            },
            "command": {
              "type": "string",
              "description": "Command to run for spawn"
            }
          },
          "required": ["action"],
          "additionalProperties": false
        }
      }
    }
  ],
  "tool_choice": {
    "type": "function",
    "function": {
      "name": "process"
    }
  }
}
JSON
```
Look for type of `id` function call

## Changes

Function call parameters are now converted to the types that the
function declaration has

## Why It Works

We now explicitly convert types where we know it (and skip if we don't)

## Test Plan

### Manual Testing
Create Qwen3.5-<ANY> instance, send the curl command above. Check that
'id' is now serialized as a number

### Automated Testing
Unit tests to cover basic type conversions

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-04 12:22:53 +00:00
Evan Quiney 886192f1e6 ignore closed resource errors when trying to cancel a task (#1652)
fixes an occasional crash during the shutdown of a failed instance.
2026-03-03 16:48:43 +00:00
Evan Quiney d914acd64e check if we have a task before we delete it (#1634)
caused a crash we should instead be logging
2026-03-03 15:32:12 +00:00
rltakashige 37296c8249 Refactor runner for implementing batching (#1632)
## Motivation

Batching will require us to send tasks concurrently and queue them up.
Our current infrastructure cannot handle that all. This PR gets us
closer to this by allowing multiple tasks to be sent in parallel and
then queuing up tasks.

## Changes

Change Plan logic
Make runner main into a class
Add a "BatchGenerator" to which tasks can be submitted (although tasks
are handled sequentially) and sent back through an MpSender.
Refactor runner to accept tasks during generation
Keep the generator threading
Separate the runner into several files for better readability

## Test Plan

### Manual Testing
Tested manually, needs a lot more automated testing. Cancellation still
works on a single device. Needs checking on multiple devices.

### Automated Testing

---------

Co-authored-by: Evan Quiney <evanev7@gmail.com>
2026-03-03 14:38:55 +00:00
Daiz 28817d3ee3 Add support for Qwen3.5 (#1644)
## Motivation

Qwen3.5 MoE models (e.g., `Qwen3.5-397B-A17B-6bit`) are now supported by
`mlx-lm` via `qwen3_5_moe` model type, but exo lacks tensor parallel
sharding support for this architecture. This prevents running large
Qwen3.5 models across multiple nodes.

Qwen3.5 uses a GatedDeltaNet hybrid attention mechanism similar to
Qwen3-Next, but with a different projection layout — separate
`in_proj_qkv`, `in_proj_z`, `in_proj_b`, `in_proj_a` instead of
Qwen3-Next's combined `in_proj_qkvz` and `in_proj_ba`. This requires
architecture-aware sharding logic.

## Changes (evan summary)

- enable qwen3_5 dense + moe tensor parallelism from config
- defensively skip evalling _cache.keys if it doesn't exist
- ignore kwargs in qwen35 pipeline masking and ensure pipeline segments match global model parameters for mask creation
- add sharding for qwen3_5 moe linear attention
- added another 6 million model cards

## Why It Works

Qwen3.5's GatedDeltaNet has an `in_proj_qkv` linear layer with three
concatenated sections: `[q(key_dim), k(key_dim), v(value_dim)]`. A naive
contiguous split (`segments=1`) would slice across section boundaries,
corrupting q/k/v values and producing garbled output.

By passing `segments=[key_dim, key_dim + key_dim]` to `shard_linear()`,
each section is split independently before distributing across devices.
This ensures every rank receives correctly aligned q, k, and v
components.

The remaining separate projections (`in_proj_z`, `in_proj_b`,
`in_proj_a`) and the MoE layers follow the same `all_to_sharded` /
`sharded_to_all` pattern already used for Qwen3-Next.

Some pipeline splits didn't include an ssm layer or a linear layer resulting in a subset of the model acting like it shouldn't create the appropriate masks for the next layer - we patch the model to manually create such masks.

## Test Plan

tensor sharded 2,3,4 models & pipeline sharded 2,3,4 with simple eval.

---------

Co-authored-by: hw <hw@hwStudio1.local>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
Co-authored-by: Evan <evanev7@gmail.com>
2026-03-03 14:31:57 +00:00
Alex Cheema 0e1b9501c3 fix: mini topology sidebar navigates home on click (#1616)
## Motivation

Clicking on the mini topology view in the right sidebar during chat was
selecting/filtering nodes instead of navigating back to the home view.
The entire mini topology section is wrapped in a button that calls
`handleGoHome()`, but individual node clicks were intercepting the event
via `stopPropagation()`.

## Changes

- Removed `onNodeClick={togglePreviewNodeFilter}` from the mini
sidebar's `TopologyGraph` component
- Added `pointer-events-none` to the topology graph container so all
clicks pass through to the parent button

## Why It Works

`TopologyGraph` only registers click/hover handlers when `onNodeClick`
is provided. Removing it disables node interaction entirely in the mini
view. The `pointer-events-none` class ensures no SVG element can
intercept clicks — they all bubble up to the parent `<button
onclick={handleGoHome}>`, which resets the chat state and returns to the
main topology view.

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Start exo and open the dashboard
- Start a chat to trigger the mini topology sidebar
- Click on the mini topology view → should navigate back to home
- Confirm nodes in the mini topology are not selectable or hoverable

### Automated Testing
- Dashboard builds successfully (`npm run build`)
2026-03-03 10:46:52 +00:00
Mustafa Alp Yılmaz f0d4ccbeb3 feat: add POST /v1/cancel/{command_id} endpoint (#1579)
## Summary

- Adds explicit `POST /v1/cancel/{command_id}` REST endpoint to cancel
in-flight text and image generation commands by ID
- Previously, cancellation only worked via HTTP disconnect (client
closes SSE connection, triggering `anyio.get_cancelled_exc_class()`).
Non-streaming clients and external consumers had no way to cleanly
cancel active generation
- The endpoint looks up the command in active generation queues, sends
`TaskCancelled` to notify workers, and closes the sender stream. Returns
404 (OpenAI error format) if the command is not found or already
completed

## Changes

| File | Change |
|------|--------|
| `src/exo/shared/types/api.py` | Add `CancelCommandResponse` model |
| `src/exo/master/api.py` | Import, route registration, `cancel_command`
handler |
| `src/exo/master/tests/test_cancel_command.py` | 3 test cases (404,
text cancel, image cancel) |
| `docs/api.md` | Document new endpoint + update summary table |

## Design Decisions

- Uses `self._send()` (not raw `command_sender.send()`) — respects API
pause state during elections
- Uses `raise HTTPException` — feeds into exo's centralized OpenAI-style
error handler
- Returns typed `CancelCommandResponse` — consistent with
`CreateInstanceResponse` / `DeleteInstanceResponse` patterns
- `sender.close()` is idempotent so concurrent cancel requests for the
same command are harmless

## Test Plan

- [x] `test_cancel_nonexistent_command_returns_404` — verifies 404 with
OpenAI error format
- [x] `test_cancel_active_text_generation` — verifies 200,
`sender.close()` called, `TaskCancelled` sent with correct
`cancelled_command_id`
- [x] `test_cancel_active_image_generation` — same verification for
image queue
- [x] `basedpyright` — 0 new errors
- [x] `ruff check` — all checks passed
- [x] `pytest` — 3/3 new tests pass, no regressions

---------

Co-authored-by: Evan <evanev7@gmail.com>
2026-03-03 10:38:52 +00:00
wysie 858dc808df fix: replace Master event_sender after EventRouter recreation (#1630) (#1637)
## Motivation
Fix for https://github.com/exo-explore/exo/issues/1630.
When loading a model that spans 2 nodes as the first action after
startup, Exo crashes with anyio.BrokenResourceError in
Master._command_processor. The receive end of the Master's event_sender
channel is already closed (open_receive_channels=0 ) when the Master
tries to send an InstanceCreated event.

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

## Changes
In _elect_loop, when is_new_master is True and the EventRouter is
replaced, the Master's event_sender is now swapped in-place with a fresh
sender from the new EventRouter. This mirrors the existing pattern where
the Worker and DownloadCoordinator are already recreated with new
channels from the new router.
<!-- Describe what you changed in detail -->

## Why It Works
When is_new_master fires (which happens on every startup when a second
node joins), the old EventRouter is shut down and a new one is created.
Shutting down the old router cancels its _ingest coroutine, which closes
the receive end of every channel created by event_router.sender(). The
Worker and DownloadCoordinator are already recreated with senders from
the new router, but the Master was not — its event_sender still pointed
to the dead channel from the old router.
This patch replaces self.master.event_sender with a new sender from the
new EventRouter, so events flow through a live channel. We swap in-place
rather than recreating the Master to preserve its DiskEventLog, internal
state, and running tasks.
The reason loading a single-node model first appeared to work around the
issue is a timing race: if the model loads fast enough, the Master sends
its events before the election fires and kills the old channel. With a
2-node model, placement and RDMA setup take longer, so the election
completes first and the channel is already dead by the time the Master
tries to send.
<!-- 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) -->
2 x Mac Studio M4 Max 128GB, RDMA
<!-- What you did: -->
Loaded a model that requires 2 nodes and it works properly now.
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->

Co-authored-by: Wysie <wysie@users.noreply.github.com>
2026-03-02 10:05:47 +00:00
ciaranbor 635118ef24 Support trace deletion in dashboard (#1628)
## Motivation

Enable using the dashboard to delete traces

## Changes

- Backend (src/exo/master/api.py): Added POST /v1/traces/delete endpoint
that deletes trace files by task ID, with path traversal protection
- API types (src/exo/shared/types/api.py): Added DeleteTracesRequest and
DeleteTracesResponse models
  - Dashboard store (app.svelte.ts): Added deleteTraces() method
- Traces page (+page.svelte): Added multi-select UI (click to select,
select all/deselect all) with a bulk delete button and confirmation
dialog

## Why It Works

- Uses existing _get_trace_path helper (now hardened against path
traversal) to resolve and delete trace files
- Dashboard selection state is reactive via Svelte 5 runes; deleted
traces refresh the list automatically

## Test Plan

### Manual Testing

- Select individual traces, select all, delete with confirmation dialog
  - Verify traces are removed from disk and list refreshes
2026-02-27 17:29:57 +00:00
Jake Hillion dc0bb5e13b fmt: add taplo TOML formatter to treefmt configuration 2026-02-27 17:19:48 +00:00
rltakashige 152a27ea5d Fix pipeline mismatched send after 1587 (#1629)
## Motivation

Tests caught a bug. It was a real bug.
2026-02-26 16:48:34 +00:00
rltakashige db36bd5ac6 Add custom prefill for pipeline (#1587)
## Motivation

Since we need to do distributed communications between prefill step
sizes, the out-of-the-box stream_generate that we currently use prevents
pipeline parallel models from doing overlapped computation. While this
was technically a regression, this communication is necessary for
cancellation, and we will need various distributed communications in the
future (e.g. for coordinating batching).

500 lines are for one testing file, so the diffs aren't as bad as they
look!

## Changes

Added a special prefill function for pipeline parallel models
Edited the model to handle 
Added a test to verify this new prefill and the original prefill produce
identical results
Improved type stubs to remove some type: ignores 

## Why It Works
<img width="768" height="1246" alt="image"
src="https://github.com/user-attachments/assets/8986ff17-ac23-4a02-9bd7-e6253a0ca799"
/>

## Test Plan

### Manual Testing
Needs more testing, but seems good so far.

### Automated Testing
Passes CI, considerable speedup seen in benchmarks (up to 1.98x) on
prefill speed.

Before:
<img width="3280" height="1238" alt="image"
src="https://github.com/user-attachments/assets/9abc1cbc-ecdb-4e48-a675-2c4cb04a32a0"
/>


After:
<img width="3344" height="1236" alt="image"
src="https://github.com/user-attachments/assets/e03c7987-41b4-4950-9ac3-2840e774ce30"
/>
2026-02-26 16:00:38 +00:00
Evan Quiney 639243aa09 event router (#1572)
replace the nack & resend logic in the worker/download coordinator with
a dedicated subsystem in front of the topic router. this centralizes
that logic (and the concept of system ids) to reduce replication in the
codebase. each system reading or writing events now gets a clean stream
of events in and can trust written events will be retried until
acknowledged.
2026-02-26 14:17:02 +00:00
Evan Quiney db73c4fd5d move messaging into rust (#1549)
the main body of the rust refactor. fixes the tokio panic on shutdown.
simplifies the networking module significantly. doesn't touch lp2p
behaviour
2026-02-26 13:58:22 +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
Alex Cheema e32b649d2f fix: enable psutil fallback for memory monitoring when macmon is missing (#1478)
## Summary
- On macOS, memory monitoring relied exclusively on `macmon` — the
psutil fallback was explicitly disabled (`memory_poll_rate = None`)
- When `macmon` is not installed (e.g., mac-mini-2 through mac-mini-4 in
our cluster), **no memory data was reported**, causing nodes to show 0GB
memory in the cluster state
- This blocked the scheduler from placing shards on those nodes since it
had no memory data to work with
- Fix: when `macmon` is not found on Darwin, fall back to psutil-based
memory polling (`memory_poll_rate = 1`)

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

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

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

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

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

---------

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

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

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

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

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

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

## Test Plan

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

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

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

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

## Test Plan

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


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

---------

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

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

## Test Plan

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

Tested that image models scale much better.

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

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

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

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

## Test Plan

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

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

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

## Changes

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

## Test Plan

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

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

---------

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

## Changes

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

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

## Test Plan

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

tested with Llama-3.3-70B, prefill progress events still show up in the
dashboard as usual
2026-02-19 18:23:28 +00:00
Evan Quiney 4c4c6ce99f simplify rust ident module
this is partly dead code, partly narrowing the rust-python boundary in
prep for future rewrites. no testing as this is all type safe
refactoring.
2026-02-19 17:19:31 +00:00
173 changed files with 16489 additions and 6659 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.
+5 -3
View File
@@ -73,9 +73,11 @@ class GenerationResponse:
finish_reason: Optional[str] = ...
def maybe_quantize_kv_cache(
prompt_cache, quantized_kv_start, kv_group_size, kv_bits
): # -> None:
...
prompt_cache: Any,
quantized_kv_start: int | None,
kv_group_size: int | None,
kv_bits: int | None,
) -> None: ...
def generate_step(
prompt: mx.array,
model: nn.Module,
+19 -30
View File
@@ -16,7 +16,7 @@ class Cache(Protocol):
self, keys: mx.array, values: mx.array
) -> tuple[mx.array, mx.array]: ...
@property
def state(self) -> tuple[mx.array, mx.array]: ...
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v) -> None: ...
@@ -92,13 +92,14 @@ class _BaseCache(Cache):
values: mx.array
offset: int
@property
def state(self) -> tuple[mx.array, mx.array]: ...
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v) -> None: ...
@property
def meta_state(self) -> Literal[""]: ...
@meta_state.setter
def meta_state(self, v) -> None: ...
def trim(self, n: int) -> int: ...
def is_trimmable(self) -> Literal[False]: ...
@classmethod
def from_state(cls, state, meta_state) -> Self: ...
@@ -114,15 +115,13 @@ class ConcatenateKVCache(_BaseCache):
def update_and_fetch(self, keys, values): # -> tuple[Any | array, Any | array]:
...
@property
def state(self): # -> tuple[Any | array | None, Any | array | None]:
...
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v): # -> None:
...
def is_trimmable(self): # -> Literal[True]:
...
def trim(self, n): # -> int:
...
def trim(self, n: int) -> int: ...
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
...
@@ -132,10 +131,7 @@ class QuantizedKVCache(_BaseCache):
def update_and_fetch(self, keys, values): # -> Any:
...
@property
def state(
self,
): # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
...
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v): # -> None:
...
@@ -147,8 +143,7 @@ class QuantizedKVCache(_BaseCache):
...
def is_trimmable(self): # -> Literal[True]:
...
def trim(self, n): # -> int:
...
def trim(self, n: int) -> int: ...
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
...
@@ -160,18 +155,18 @@ class KVCache(_BaseCache):
@property
def state(
self,
) -> tuple[array, array]: ...
) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v) -> None: ...
def is_trimmable(self): # -> Literal[True]:
...
def trim(self, n): # -> int:
...
def trim(self, n: int) -> int: ...
def to_quantized(
self, group_size: int = ..., bits: int = ...
) -> QuantizedKVCache: ...
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
...
def make_mask(
self, *args: Any, **kwargs: Any
) -> mx.array | Literal["causal"] | None: ...
class RotatingKVCache(_BaseCache):
step = ...
@@ -183,8 +178,7 @@ class RotatingKVCache(_BaseCache):
@property
def state(
self,
): # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
...
) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v): # -> None:
...
@@ -196,8 +190,7 @@ class RotatingKVCache(_BaseCache):
...
def is_trimmable(self): # -> bool:
...
def trim(self, n): # -> int:
...
def trim(self, n: int) -> int: ...
def to_quantized(
self, group_size: int = ..., bits: int = ...
) -> QuantizedKVCache: ...
@@ -212,8 +205,7 @@ class ArraysCache(_BaseCache):
...
def __getitem__(self, idx): ...
@property
def state(self): # -> list[Any | array] | list[array]:
...
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@state.setter
def state(self, v): # -> None:
...
@@ -227,8 +219,7 @@ class ArraysCache(_BaseCache):
In-place extend this cache with the other cache.
"""
def make_mask(self, N: int): # -> array | None:
...
def make_mask(self, N: int) -> mx.array | None: ...
class MambaCache(ArraysCache):
def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -239,8 +230,7 @@ class ChunkedKVCache(KVCache):
...
def update_and_fetch(self, keys, values): # -> tuple[array, array]:
...
def trim(self, n): # -> int:
...
def trim(self, n: int) -> int: ...
@property
def meta_state(self): # -> tuple[str, ...]:
...
@@ -253,10 +243,9 @@ class CacheList(_BaseCache):
def __getitem__(self, idx): ...
def is_trimmable(self): # -> bool:
...
def trim(self, n): ...
def trim(self, n: int) -> int: ...
@property
def state(self): # -> list[Any]:
...
def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
@state.setter
def state(self, v): # -> None:
...
+153
View File
@@ -0,0 +1,153 @@
from dataclasses import dataclass
from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .cache import ArraysCache, KVCache
from .qwen3_next import (
Qwen3NextAttention as Attention,
Qwen3NextMLP as MLP,
Qwen3NextRMSNormGated as RMSNormGated,
Qwen3NextSparseMoeBlock,
)
SparseMoeBlock = Qwen3NextSparseMoeBlock
from .switch_layers import SwitchGLU
@dataclass
class TextModelArgs:
model_type: str
hidden_size: int
intermediate_size: int
num_hidden_layers: int
num_attention_heads: int
rms_norm_eps: float
vocab_size: int
num_key_value_heads: int
max_position_embeddings: int
linear_num_value_heads: int
linear_num_key_heads: int
linear_key_head_dim: int
linear_value_head_dim: int
linear_conv_kernel_dim: int
tie_word_embeddings: bool
attention_bias: bool
head_dim: Optional[int]
full_attention_interval: int
num_experts: int
num_experts_per_tok: int
decoder_sparse_step: int
shared_expert_intermediate_size: int
moe_intermediate_size: int
norm_topk_prob: bool
rope_parameters: Optional[dict[str, Any]]
partial_rotary_factor: float
rope_theta: float
rope_scaling: Optional[dict[str, Any]]
@classmethod
def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
def __post_init__(self) -> None: ...
class GatedDeltaNet(nn.Module):
hidden_size: int
num_v_heads: int
num_k_heads: int
head_k_dim: int
head_v_dim: int
key_dim: int
value_dim: int
conv_kernel_size: int
conv_dim: int
conv1d: nn.Conv1d
in_proj_qkv: nn.Linear
in_proj_z: nn.Linear
in_proj_b: nn.Linear
in_proj_a: nn.Linear
dt_bias: mx.array
A_log: mx.array
norm: RMSNormGated
out_proj: nn.Linear
def __init__(self, config: TextModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class DecoderLayer(nn.Module):
is_linear: bool
linear_attn: GatedDeltaNet
self_attn: Attention
input_layernorm: nn.RMSNorm
post_attention_layernorm: nn.RMSNorm
mlp: MLP | SparseMoeBlock
def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class Qwen3_5TextModel(nn.Module):
embed_tokens: nn.Embedding
layers: list[DecoderLayer]
norm: nn.RMSNorm
ssm_idx: int
fa_idx: int
def __init__(self, args: TextModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
input_embeddings: Optional[mx.array] = None,
) -> mx.array: ...
class TextModel(nn.Module):
args: TextModelArgs
model_type: str
model: Qwen3_5TextModel
lm_head: nn.Linear
def __init__(self, args: TextModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
input_embeddings: Optional[mx.array] = None,
) -> mx.array: ...
@property
def layers(self) -> list[DecoderLayer]: ...
def make_cache(self) -> list[ArraysCache | KVCache]: ...
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
@dataclass
class ModelArgs:
model_type: str
text_config: dict[str, Any]
@classmethod
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
class Model(nn.Module):
args: ModelArgs
model_type: str
language_model: TextModel
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
input_embeddings: Optional[mx.array] = None,
) -> mx.array: ...
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
@property
def layers(self) -> list[DecoderLayer]: ...
def make_cache(self) -> list[ArraysCache | KVCache]: ...
@@ -0,0 +1,19 @@
from dataclasses import dataclass
from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .cache import ArraysCache, KVCache
from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
@dataclass
class ModelArgs:
model_type: str
text_config: dict[str, Any]
@classmethod
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
class Model(Qwen3_5Model):
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
@@ -7,6 +7,15 @@ import mlx.nn as nn
from .switch_layers import SwitchGLU
class Qwen3NextRMSNormGated(nn.Module):
eps: float
weight: mx.array
def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
def __call__(
self, hidden_states: mx.array, gate: mx.array | None = None
) -> mx.array: ...
class Qwen3NextMLP(nn.Module):
gate_proj: nn.Linear
down_proj: nn.Linear
Generated
+2258 -1831
View File
File diff suppressed because it is too large Load Diff
+1 -2
View File
@@ -3,7 +3,6 @@ resolver = "3"
members = [
"rust/networking",
"rust/exo_pyo3_bindings",
"rust/util",
]
[workspace.package]
@@ -24,7 +23,6 @@ opt-level = 3
[workspace.dependencies]
## Crate members as common dependencies
networking = { path = "rust/networking" }
util = { path = "rust/util" }
# Macro dependecies
extend = "1.2"
@@ -34,6 +32,7 @@ delegate = "0.13"
keccak-const = "0.2"
# Async dependencies
async-stream = "0.3"
tokio = "1.46"
futures-lite = "2.6.1"
futures-timer = "3.0"
+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
}
}
+1 -3
View File
@@ -2,6 +2,4 @@
#
# Lists the suite files to include. Each file defines benchmarks
# with shared constraints, topology, and default args.
include = [
"single-m3-ultra.toml",
]
include = ["single-m3-ultra.toml"]
+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")
+7 -6
View File
@@ -4,12 +4,13 @@ version = "0.1.0"
description = "Benchmarking tool for exo distributed inference"
requires-python = ">=3.13"
dependencies = [
"httpx>=0.27.0",
"loguru>=0.7.3",
"transformers>=5.0.0",
"huggingface-hub>=0.33.4",
"tiktoken>=0.12.0",
"jinja2>=3.1.0",
"httpx>=0.27.0",
"loguru>=0.7.3",
"transformers>=5.0.0",
"huggingface-hub>=0.33.4",
"tiktoken>=0.12.0",
"jinja2>=3.1.0",
"protobuf>=5.29.0",
]
[build-system]
+4 -4
View File
@@ -2,10 +2,10 @@
#
# Shared constraints applied to ALL benchmarks in this file.
constraints = [
"All(MacOsBuild(=25D125))",
"Hosts(=1)",
"All(Chip(m3_ultra))",
"All(GpuCores(=80))",
"All(MacOsBuild(=25D125))",
"Hosts(=1)",
"All(Chip(m3_ultra))",
"All(GpuCores(=80))",
]
[topology]
+65
View File
@@ -202,6 +202,23 @@
filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5));
}
/* Onboarding step 2: connection line between devices */
.onboarding-connection-line {
stroke: oklch(0.85 0.18 85 / 0.5);
stroke-width: 1.5px;
stroke-dasharray: 6, 6;
animation: flowAnimation 1s linear infinite;
filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4));
}
/* Onboarding step 4: red connection line for disconnect */
.onboarding-connection-line-red {
stroke: rgba(220, 38, 38, 0.7);
stroke-width: 1.5px;
stroke-dasharray: 6, 6;
filter: drop-shadow(0 0 2px rgba(220, 38, 38, 0.3));
}
.graph-link-active {
stroke: oklch(0.85 0.18 85 / 0.8);
stroke-width: 2px;
@@ -320,3 +337,51 @@ input:focus, textarea:focus {
transform: translate(400px, 400px);
}
}
/* Onboarding smooth fade-in */
@keyframes onb-fade-in {
from { opacity: 0; transform: translateY(8px); }
to { opacity: 1; transform: none; }
}
/* Opacity-only fade-in — no transform, safe for containers with fixed-position children */
@keyframes onb-fade-opacity {
from { opacity: 0; }
to { opacity: 1; }
}
/* Respect reduced motion preference */
@media (prefers-reduced-motion: reduce) {
.shooting-star,
.shooting-star::before {
animation: none !important;
opacity: 0 !important;
}
.graph-link {
animation: none;
}
.status-pulse {
animation: none;
}
.cursor-blink {
animation: none;
}
.onboarding-connection-line {
animation: none;
}
.onboarding-connection-line-red {
animation: none;
}
[style*="onb-fade-in"],
[style*="onb-fade-opacity"] {
animation: none !important;
opacity: 1 !important;
}
*,
*::before,
*::after {
transition-duration: 0.01ms !important;
animation-duration: 0.01ms !important;
animation-iteration-count: 1 !important;
}
}
+43 -174
View File
@@ -7,8 +7,6 @@
editingImage,
clearEditingImage,
selectedChatModel,
setSelectedChatModel,
instances,
ttftMs,
tps,
totalTokens,
@@ -29,6 +27,19 @@
showModelSelector?: boolean;
modelTasks?: Record<string, string[]>;
modelCapabilities?: Record<string, string[]>;
onSend?: () => void;
onAutoSend?: (
content: string,
files?: {
id: string;
name: string;
type: string;
textContent?: string;
preview?: string;
}[],
) => void;
onOpenModelPicker?: () => void;
modelDisplayOverride?: string;
}
let {
@@ -39,6 +50,10 @@
showModelSelector = false,
modelTasks = {},
modelCapabilities = {},
onSend,
onAutoSend,
onOpenModelPicker,
modelDisplayOverride,
}: Props = $props();
let message = $state("");
@@ -49,27 +64,12 @@
const thinkingEnabled = $derived(thinkingEnabledStore());
let loading = $derived(isLoading());
const currentModel = $derived(selectedChatModel());
const instanceData = $derived(instances());
const currentTtft = $derived(ttftMs());
const currentTps = $derived(tps());
const currentTokens = $derived(totalTokens());
const currentEditingImage = $derived(editingImage());
const isEditMode = $derived(currentEditingImage !== null);
// Custom dropdown state
let isModelDropdownOpen = $state(false);
let dropdownButtonRef: HTMLButtonElement | undefined = $state();
let dropdownPosition = $derived(() => {
if (!dropdownButtonRef || !isModelDropdownOpen)
return { top: 0, left: 0, width: 0 };
const rect = dropdownButtonRef.getBoundingClientRect();
return {
top: rect.top,
left: rect.left,
width: rect.width,
};
});
// Accept all supported file types
const acceptString = getAcceptString(["image", "text", "pdf"]);
@@ -122,73 +122,14 @@
uploadedFiles.length > 0),
);
// Extract available models from running instances
const availableModels = $derived(() => {
const models: Array<{ id: string; label: string; isImageModel: boolean }> =
[];
for (const [, instance] of Object.entries(instanceData)) {
const modelId = getInstanceModelId(instance);
if (
modelId &&
modelId !== "Unknown" &&
!models.some((m) => m.id === modelId)
) {
models.push({
id: modelId,
label: modelId.split("/").pop() || modelId,
isImageModel: modelSupportsImageGeneration(modelId),
});
}
}
return models;
});
// Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
let previousModelIds: Set<string> = new Set();
// Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
$effect(() => {
const models = availableModels();
const currentModelIds = new Set(models.map((m) => m.id));
if (models.length > 0) {
// Find newly added models (in current but not in previous)
const newModels = models.filter((m) => !previousModelIds.has(m.id));
// If no model selected, select the first available
if (!currentModel) {
setSelectedChatModel(models[0].id);
}
// If current model is stale (no longer has a running instance), reset to first available
else if (!models.some((m) => m.id === currentModel)) {
setSelectedChatModel(models[0].id);
}
// If a new model was just added, select it
else if (newModels.length > 0 && previousModelIds.size > 0) {
setSelectedChatModel(newModels[0].id);
}
} else {
// No instances running - clear the selected model
if (currentModel) {
setSelectedChatModel("");
}
}
// Update previous model IDs for next comparison
previousModelIds = currentModelIds;
});
function getInstanceModelId(instanceWrapped: unknown): string {
if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
const keys = Object.keys(instanceWrapped as Record<string, unknown>);
if (keys.length === 1) {
const instance = (instanceWrapped as Record<string, unknown>)[
keys[0]
] as { shardAssignments?: { modelId?: string } };
return instance?.shardAssignments?.modelId || "";
}
return "";
}
// Short label for the currently selected model
const currentModelLabel = $derived(
currentModel
? currentModel.split("/").pop() || currentModel
: modelDisplayOverride
? modelDisplayOverride.split("/").pop() || modelDisplayOverride
: "",
);
async function handleFiles(files: File[]) {
if (files.length === 0) return;
@@ -275,6 +216,15 @@
uploadedFiles = [];
resetTextareaHeight();
// When onAutoSend is provided, the parent controls all send logic
// (including launching non-running models before sending)
if (onAutoSend) {
onAutoSend(content, files);
onSend?.();
setTimeout(() => textareaRef?.focus(), 10);
return;
}
// Use image editing if in edit mode
if (isEditMode && currentEditingImage && content) {
editImage(content, currentEditingImage.imageDataUrl);
@@ -306,6 +256,8 @@
);
}
onSend?.();
// Refocus the textarea after sending
setTimeout(() => textareaRef?.focus(), 10);
}
@@ -416,7 +368,7 @@
{/if}
<!-- Model selector (when enabled) -->
{#if showModelSelector && availableModels().length > 0}
{#if showModelSelector}
<div
class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
>
@@ -425,33 +377,22 @@
class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
>MODEL:</span
>
<!-- Custom dropdown -->
<!-- Model button — opens the full model picker -->
<div class="relative flex-1 max-w-xs">
<button
bind:this={dropdownButtonRef}
type="button"
onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
? 'border-exo-yellow/70'
: ''}"
onclick={() => onOpenModelPicker?.()}
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
>
{#if availableModels().find((m) => m.id === currentModel)}
<span class="text-exo-yellow truncate"
>{availableModels().find((m) => m.id === currentModel)
?.label}</span
>
{:else if availableModels().length > 0}
<span class="text-exo-yellow truncate"
>{availableModels()[0].label}</span
{#if currentModelLabel}
<span class="text-exo-yellow truncate">{currentModelLabel}</span
>
{:else}
<span class="text-exo-light-gray/50">— SELECT MODEL —</span>
{/if}
</button>
<div
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
? 'rotate-180'
: ''}"
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
>
<svg
class="w-3 h-3 text-exo-yellow/60"
@@ -468,78 +409,6 @@
</svg>
</div>
</div>
{#if isModelDropdownOpen}
<!-- Backdrop to close dropdown -->
<button
type="button"
class="fixed inset-0 z-[9998] cursor-default"
onclick={() => (isModelDropdownOpen = false)}
aria-label="Close dropdown"
></button>
<!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
<div
class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
style="bottom: calc(100vh - {dropdownPosition()
.top}px + 4px); left: {dropdownPosition()
.left}px; width: {dropdownPosition().width}px;"
>
<div class="py-1">
{#each availableModels() as model}
<button
type="button"
onclick={() => {
setSelectedChatModel(model.id);
isModelDropdownOpen = false;
}}
class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
model.id
? 'bg-transparent text-exo-yellow'
: 'text-exo-light-gray hover:text-exo-yellow'}"
>
{#if currentModel === model.id}
<svg
class="w-3 h-3 flex-shrink-0"
fill="currentColor"
viewBox="0 0 20 20"
>
<path
fill-rule="evenodd"
d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
clip-rule="evenodd"
/>
</svg>
{:else}
<span class="w-3"></span>
{/if}
{#if model.isImageModel}
<svg
class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
aria-label="Image generation model"
>
<rect
x="3"
y="3"
width="18"
height="18"
rx="2"
ry="2"
/>
<circle cx="8.5" cy="8.5" r="1.5" />
<polyline points="21 15 16 10 5 21" />
</svg>
{/if}
<span class="truncate flex-1">{model.label}</span>
</button>
{/each}
</div>
</div>
{/if}
</div>
<!-- Thinking toggle -->
{#if modelSupportsThinking()}
@@ -807,8 +807,8 @@
>
AWAITING INPUT
</p>
<p class="text-sm sm:text-xs text-exo-light-gray tracking-wider mt-1">
ENTER A QUERY TO BEGIN
<p class="text-xs text-white/30 tracking-wider mt-1.5 font-mono">
Type a message below &middot; Shift+Enter for newline
</p>
</div>
{/if}
@@ -823,6 +823,7 @@
onclick={scrollToBottom}
class="sticky bottom-4 left-1/2 -translate-x-1/2 w-10 h-10 rounded-full bg-exo-dark-gray/90 border border-exo-medium-gray/50 flex items-center justify-center text-exo-light-gray hover:text-exo-yellow hover:border-exo-yellow/50 transition-all shadow-lg cursor-pointer z-10"
title="Scroll to bottom"
aria-label="Scroll to bottom of messages"
>
<svg
class="w-5 h-5"
@@ -0,0 +1,401 @@
<script lang="ts" module>
export interface ChatModelInfo {
id: string;
name: string;
base_model: string;
storage_size_megabytes: number;
capabilities: string[];
family: string;
quantization: string;
}
// Auto mode tier list (for when user just starts typing)
export const AUTO_TIERS: string[][] = [
// Tier 1 (frontier)
["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
// Tier 2 (excellent)
[
"Kimi K2",
"Qwen3 235B",
"MiniMax M2.5",
"Step 3.5 Flash",
"Qwen3 Next 80B",
],
// Tier 3 (great)
[
"GLM 4.7",
"MiniMax M2.1",
"Qwen3 Coder 480B",
"GLM 4.5 Air",
"Llama 3.3 70B",
],
// Tier 4 (good)
["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
// Tier 5 (small/fast)
[
"Llama 3.1 8B",
"GPT-OSS 20B",
"Llama 3.2 3B",
"Qwen3 0.6B",
"Llama 3.2 1B",
],
];
/** Return the tier index (0 = best) for a base_model name. */
export function getAutoTierIndex(baseModel: string): number {
for (let i = 0; i < AUTO_TIERS.length; i++) {
if (AUTO_TIERS[i].includes(baseModel)) return i;
}
return AUTO_TIERS.length; // not in any tier → lowest priority
}
/** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
export function pickAutoModel(
modelList: ChatModelInfo[],
memoryGB: number,
): ChatModelInfo | null {
for (const tier of AUTO_TIERS) {
const candidates: ChatModelInfo[] = [];
for (const baseModel of tier) {
const variants = modelList
.filter(
(m) =>
m.base_model === baseModel &&
(m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
(m.storage_size_megabytes || 0) > 0,
)
.sort(
(a, b) =>
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
);
if (variants[0]) candidates.push(variants[0]);
}
if (candidates.length > 0) {
candidates.sort(
(a, b) =>
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
);
return candidates[0];
}
}
return null;
}
</script>
<script lang="ts">
interface CategoryRecommendation {
category: "coding" | "writing" | "agentic" | "biggest";
label: string;
model: ChatModelInfo | null;
tooltip: string;
}
interface Props {
models: ChatModelInfo[];
clusterLabel: string;
totalMemoryGB: number;
onSelect: (modelId: string, category: string) => void;
onAddModel: () => void;
class?: string;
}
let {
models,
clusterLabel,
totalMemoryGB,
onSelect,
onAddModel,
class: className = "",
}: Props = $props();
// --- Hardcoded Rankings ---
const CODING_RANKING = [
"Qwen3 Coder Next",
"Qwen3 Coder 480B",
"Qwen3 30B",
"GPT-OSS 20B",
"Llama 3.1 8B",
"Llama 3.2 3B",
"Qwen3 0.6B",
];
const WRITING_RANKING = [
"Kimi K2.5",
"Kimi K2",
"Qwen3 Next 80B",
"Llama 3.3 70B",
"MiniMax M2.5",
"GLM 4.5 Air",
"GLM 4.7 Flash",
"GPT-OSS 20B",
"Llama 3.1 8B",
"Llama 3.2 3B",
"Qwen3 0.6B",
];
const AGENTIC_RANKING = [
"DeepSeek V3.1",
"GLM-5",
"Qwen3 235B",
"Step 3.5 Flash",
"GLM 4.7",
"MiniMax M2.1",
"GPT-OSS 120B",
"Llama 3.3 70B",
"Llama 3.1 70B",
"GLM 4.7 Flash",
"GPT-OSS 20B",
"Qwen3 30B",
"Llama 3.1 8B",
"Llama 3.2 3B",
"Qwen3 0.6B",
];
function getModelSizeGB(m: ChatModelInfo): number {
return (m.storage_size_megabytes || 0) / 1024;
}
function fitsInMemory(m: ChatModelInfo): boolean {
return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
}
/** For a given base_model name, find the biggest quant variant that fits in memory. */
function pickBestVariant(baseModel: string): ChatModelInfo | null {
const variants = models
.filter((m) => m.base_model === baseModel && fitsInMemory(m))
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
return variants[0] ?? null;
}
/** Walk a ranked list of base_model names, return the first that has a fitting variant. */
function pickFromRanking(ranking: string[]): ChatModelInfo | null {
for (const baseModel of ranking) {
const pick = pickBestVariant(baseModel);
if (pick) return pick;
}
return null;
}
/** Pick the single biggest model that fits. */
function pickBiggest(): ChatModelInfo | null {
const fitting = models
.filter((m) => fitsInMemory(m))
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
return fitting[0] ?? null;
}
const recommendations = $derived.by((): CategoryRecommendation[] => {
return [
{
category: "coding",
label: "Best for Coding",
model: pickFromRanking(CODING_RANKING),
tooltip:
"Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
},
{
category: "writing",
label: "Best for Writing",
model: pickFromRanking(WRITING_RANKING),
tooltip: "Ranked by creative writing quality and instruction following",
},
{
category: "agentic",
label: "Best Agentic",
model: pickFromRanking(AGENTIC_RANKING),
tooltip: "Ranked by reasoning, planning, and tool-use capability",
},
{
category: "biggest",
label: "Biggest",
model: pickBiggest(),
tooltip: "Largest model that fits in your available memory",
},
];
});
function formatSize(mb: number): string {
const gb = mb / 1024;
if (gb >= 100) return `${Math.round(gb)} GB`;
return `${gb.toFixed(1)} GB`;
}
// Category icons (SVG paths)
const categoryIcons: Record<string, string> = {
coding:
"M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
writing:
"M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
agentic:
"M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
biggest:
"M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
};
let hoveredTooltip = $state<string | null>(null);
let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
function showTooltip(category: string, e: MouseEvent | FocusEvent) {
hoveredTooltip = category;
const target = e.currentTarget as HTMLElement;
const rect = target.getBoundingClientRect();
tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
}
function hideTooltip() {
hoveredTooltip = null;
tooltipAnchor = null;
}
</script>
<div class="flex flex-col items-center justify-center gap-6 {className}">
<!-- Header -->
<div class="text-center">
<p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
Recommended for your
</p>
<p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
</div>
<!-- Category Cards Grid -->
<div class="grid grid-cols-2 gap-3 w-full max-w-md">
{#each recommendations as rec}
{#if rec.model}
<button
type="button"
onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
>
<!-- Category icon + label -->
<div class="flex items-center gap-2 w-full">
<svg
class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="1.5"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={categoryIcons[rec.category]}
/>
</svg>
<span
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
>
{rec.label}
</span>
<!-- Info tooltip -->
<div class="ml-auto flex-shrink-0">
<span
role="button"
tabindex="-1"
class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
onmouseleave={() => hideTooltip()}
onclick={(e: MouseEvent) => {
e.stopPropagation();
if (hoveredTooltip === rec.category) {
hideTooltip();
} else {
showTooltip(rec.category, e);
}
}}
>
<svg
class="w-3.5 h-3.5"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<circle cx="12" cy="12" r="10" />
<path d="M12 16v-4m0-4h.01" />
</svg>
</span>
</div>
</div>
<!-- Model name + size -->
<div class="w-full">
<p class="text-sm text-white font-mono truncate">
{rec.model.base_model}
</p>
<p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
{formatSize(rec.model.storage_size_megabytes)}
{#if rec.model.quantization}
<span class="text-exo-light-gray/40"
>&middot; {rec.model.quantization}</span
>
{/if}
</p>
</div>
</button>
{:else}
<!-- No model fits for this category -->
<div
class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
>
<div class="flex items-center gap-2">
<svg
class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="1.5"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={categoryIcons[rec.category]}
/>
</svg>
<span
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
>{rec.label}</span
>
</div>
<p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
</div>
{/if}
{/each}
</div>
<!-- Add Model Button -->
<button
type="button"
onclick={onAddModel}
class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
>
<svg
class="w-3.5 h-3.5"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
</svg>
Add Model
</button>
<!-- Auto hint -->
<p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
Or just start typing &mdash; we'll pick the best model automatically
</p>
</div>
<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
{#if hoveredTooltip && tooltipAnchor}
{@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
{#if rec}
<div
class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
8}px; transform: translate(-50%, -100%);"
>
{rec.tooltip}
</div>
{/if}
{/if}
+29 -29
View File
@@ -2,7 +2,6 @@
import {
conversations,
activeConversationId,
createConversation,
loadConversation,
deleteConversation,
deleteAllConversations,
@@ -17,9 +16,15 @@
interface Props {
class?: string;
onNewChat?: () => void;
onSelectConversation?: () => void;
}
let { class: className = "" }: Props = $props();
let {
class: className = "",
onNewChat,
onSelectConversation,
}: Props = $props();
const conversationList = $derived(conversations());
const activeId = $derived(activeConversationId());
@@ -42,10 +47,11 @@
);
function handleNewChat() {
createConversation();
onNewChat?.();
}
function handleSelectConversation(id: string) {
onSelectConversation?.();
loadConversation(id);
}
@@ -185,11 +191,7 @@
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (
instanceTag === "MlxIbvInstance" ||
instanceTag === "MlxJacclInstance"
)
instanceType = "MLX RDMA";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
let sharding: string | null = null;
const inst = instance as {
@@ -307,7 +309,7 @@
<div class="py-2">
<div class="px-4 py-2">
<span
class="text-sm text-white/70 font-mono tracking-wider uppercase"
class="text-xs text-exo-light-gray font-mono tracking-wider uppercase"
>
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"}
</span>
@@ -376,39 +378,37 @@
onkeydown={(e) =>
e.key === "Enter" &&
handleSelectConversation(conversation.id)}
class="group w-full flex items-center justify-between p-2 rounded mb-1 transition-all text-left cursor-pointer
class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer
{activeId === conversation.id
? 'bg-transparent border border-exo-yellow/30'
: 'hover:border-exo-yellow/20 border border-transparent'}"
? 'bg-exo-yellow/5 border border-exo-yellow/30'
: 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}"
>
<div class="flex-1 min-w-0 pr-2">
<div
class="text-sm truncate {activeId === conversation.id
class="text-sm font-medium truncate {activeId ===
conversation.id
? 'text-exo-yellow'
: 'text-white/90'}"
: 'text-white'}"
>
{conversation.name}
</div>
<div class="text-sm text-white/50 mt-0.5">
<div class="text-xs text-white/60 mt-0.5">
{formatDate(conversation.updatedAt)}
</div>
<div class="text-sm text-white/70 truncate">
<div class="text-xs text-exo-light-gray truncate">
{info.modelLabel}
</div>
<div class="text-xs text-white/60 font-mono">
Strategy: <span class="text-white/80"
>{info.strategyLabel}</span
>
</div>
{#if stats}
<div class="text-xs text-white/60 font-mono mt-1">
{#if stats.ttftMs}<span class="text-white/40">TTFT</span>
{stats.ttftMs.toFixed(
0,
)}ms{/if}{#if stats.ttftMs && stats.tps}<span
class="text-white/30 mx-1.5"></span
>{/if}{#if stats.tps}{stats.tps.toFixed(1)}
<span class="text-white/40">tok/s</span>{/if}
<div class="text-xs text-white/70 font-mono mt-1">
{#if stats.ttftMs}<span class="text-white/50">TTFT</span>
<span class="text-exo-yellow/80"
>{stats.ttftMs.toFixed(0)}ms</span
>{/if}{#if stats.ttftMs && stats.tps}<span
class="text-white/30 mx-1.5">·</span
>{/if}{#if stats.tps}<span class="text-exo-yellow/80"
>{stats.tps.toFixed(1)}</span
>
<span class="text-white/50">tok/s</span>{/if}
</div>
{/if}
</div>
@@ -0,0 +1,20 @@
<script lang="ts">
import { isConnected } from "$lib/stores/app.svelte";
import { slide } from "svelte/transition";
const connected = $derived(isConnected());
</script>
{#if !connected}
<div
transition:slide={{ duration: 200 }}
class="relative z-50 flex items-center justify-center gap-2 px-4 py-2 bg-red-950/80 border-b border-red-500/30"
role="alert"
aria-live="assertive"
>
<div class="w-2 h-2 bg-red-500 rounded-full animate-pulse"></div>
<span class="text-xs font-mono text-red-300 tracking-wider uppercase">
Connection lost &mdash; Reconnecting to backend&hellip;
</span>
</div>
{/if}
@@ -0,0 +1,277 @@
<script lang="ts">
/**
* DeviceIcon — renders a device icon as an SVG <g> element.
* Uses the exact same proportional math as TopologyGraph.svelte
* so that devices look identical in both the topology view and
* the onboarding animation.
*
* Must be placed inside an <svg> element.
*/
interface Props {
/** "macbook pro" | "mac studio" | "mac mini" etc. */
deviceType: string;
/** Center X coordinate in SVG space */
cx: number;
/** Center Y coordinate in SVG space */
cy: number;
/** Base sizing factor (equivalent to TopologyGraph's nodeRadius) */
size?: number;
/** RAM usage 0100 */
ramPercent?: number;
/** Unique id suffix for clip-path ids */
uid?: string;
}
let {
deviceType,
cx,
cy,
size = 60,
ramPercent = 60,
uid = "dev",
}: Props = $props();
// Apple logo path — same constant used by TopologyGraph
const APPLE_LOGO_PATH =
"M788.1 340.9c-5.8 4.5-108.2 62.2-108.2 190.5 0 148.4 130.3 200.9 134.2 202.2-.6 3.2-20.7 71.9-68.7 141.9-42.8 61.6-87.5 123.1-155.5 123.1s-85.5-39.5-164-39.5c-76.5 0-103.7 40.8-165.9 40.8s-105.6-57-155.5-127C46.7 790.7 0 663 0 541.8c0-194.4 126.4-297.5 250.8-297.5 66.1 0 121.2 43.4 162.7 43.4 39.5 0 101.1-46 176.3-46 28.5 0 130.9 2.6 198.3 99.2zm-234-181.5c31.1-36.9 53.1-88.1 53.1-139.3 0-7.1-.6-14.3-1.9-20.1-50.6 1.9-110.8 33.7-147.1 75.8-28.5 32.4-55.1 83.6-55.1 135.5 0 7.8 1.3 15.6 1.9 18.1 3.2.6 8.4 1.3 13.6 1.3 45.4 0 102.5-30.4 135.5-71.3z";
const LOGO_NATIVE_WIDTH = 814;
const LOGO_NATIVE_HEIGHT = 1000;
const wireColor = "rgba(179,179,179,0.8)";
const strokeWidth = 1.5;
const modelLower = $derived(deviceType.toLowerCase());
// ── Mac Studio dimensions (same ratios as TopologyGraph) ──
const studioW = $derived(size * 1.25);
const studioH = $derived(size * 0.85);
const studioX = $derived(cx - studioW / 2);
const studioY = $derived(cy - studioH / 2);
const studioCorner = 4;
const studioTopH = $derived(studioH * 0.15);
// Studio front panel details
const studioSlotH = $derived(studioH * 0.14);
const studioVSlotW = $derived(studioW * 0.05);
const studioVSlotY = $derived(
studioY + studioTopH + (studioH - studioTopH) * 0.6,
);
const studioVSlot1X = $derived(studioX + studioW * 0.18);
const studioVSlot2X = $derived(studioX + studioW * 0.28);
const studioHSlotW = $derived(studioW * 0.2);
const studioHSlotX = $derived(studioX + studioW * 0.5 - studioHSlotW / 2);
// Studio memory fill
const studioMemTotalH = $derived(studioH - studioTopH);
const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
// ── MacBook dimensions (same ratios as TopologyGraph) ──
const mbW = $derived((size * 1.6 * 0.85) / 1.15);
const mbH = $derived(size * 0.85);
const mbX = $derived(cx - mbW / 2);
const mbY = $derived(cy - mbH / 2);
const mbScreenH = $derived(mbH * 0.7);
const mbBaseH = $derived(mbH * 0.3);
const mbScreenW = $derived(mbW * 0.85);
const mbScreenX = $derived(cx - mbScreenW / 2);
const mbBezel = 3;
// MacBook memory fill
const mbMemTotalH = $derived(mbScreenH - mbBezel * 2);
const mbMemH = $derived((ramPercent / 100) * mbMemTotalH);
// Apple logo sizing
const mbLogoTargetH = $derived(mbScreenH * 0.22);
const mbLogoScale = $derived(mbLogoTargetH / LOGO_NATIVE_HEIGHT);
const mbLogoX = $derived(cx - (LOGO_NATIVE_WIDTH * mbLogoScale) / 2);
const mbLogoY = $derived(
mbY + mbScreenH / 2 - (LOGO_NATIVE_HEIGHT * mbLogoScale) / 2,
);
// MacBook base (trapezoidal)
const mbBaseY = $derived(mbY + mbScreenH);
const mbBaseTopW = $derived(mbScreenW);
const mbBaseBottomW = $derived(mbW);
const mbBaseTopX = $derived(cx - mbBaseTopW / 2);
const mbBaseBottomX = $derived(cx - mbBaseBottomW / 2);
// Keyboard
const mbKbX = $derived(mbBaseTopX + 6);
const mbKbY = $derived(mbBaseY + 3);
const mbKbW = $derived(mbBaseTopW - 12);
const mbKbH = $derived(mbBaseH * 0.55);
// Trackpad
const mbTpW = $derived(mbBaseTopW * 0.4);
const mbTpX = $derived(cx - mbTpW / 2);
const mbTpY = $derived(mbBaseY + mbKbH + 5);
const mbTpH = $derived(mbBaseH * 0.3);
// Clip IDs
const screenClipId = $derived(`di-screen-${uid}`);
const studioClipId = $derived(`di-studio-${uid}`);
</script>
{#if modelLower === "mac studio" || modelLower === "mac mini"}
<!-- Mac Studio / Mac Mini -->
<defs>
<clipPath id={studioClipId}>
<rect
x={studioX}
y={studioY + studioTopH}
width={studioW}
height={studioH - studioTopH}
rx={studioCorner - 1}
/>
</clipPath>
</defs>
<!-- Main body -->
<rect
x={studioX}
y={studioY}
width={studioW}
height={studioH}
rx={studioCorner}
fill="#1a1a1a"
stroke={wireColor}
stroke-width={strokeWidth}
/>
<!-- Memory fill -->
{#if ramPercent > 0}
<rect
x={studioX}
y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
width={studioW}
height={studioMemH}
fill="rgba(255,215,0,0.75)"
clip-path="url(#{studioClipId})"
/>
{/if}
<!-- Top surface divider -->
<line
x1={studioX}
y1={studioY + studioTopH}
x2={studioX + studioW}
y2={studioY + studioTopH}
stroke="rgba(179,179,179,0.3)"
stroke-width="0.5"
/>
<!-- Front panel: vertical slots -->
<rect
x={studioVSlot1X - studioVSlotW / 2}
y={studioVSlotY}
width={studioVSlotW}
height={studioSlotH}
fill="rgba(0,0,0,0.35)"
rx="1.5"
/>
<rect
x={studioVSlot2X - studioVSlotW / 2}
y={studioVSlotY}
width={studioVSlotW}
height={studioSlotH}
fill="rgba(0,0,0,0.35)"
rx="1.5"
/>
<!-- Horizontal slot (SD card) -->
<rect
x={studioHSlotX}
y={studioVSlotY}
width={studioHSlotW}
height={studioSlotH * 0.6}
fill="rgba(0,0,0,0.35)"
rx="1"
/>
{:else}
<!-- MacBook Pro -->
<defs>
<clipPath id={screenClipId}>
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel}
width={mbScreenW - mbBezel * 2}
height={mbScreenH - mbBezel * 2}
rx="2"
/>
</clipPath>
</defs>
<!-- Screen outer frame -->
<rect
x={mbScreenX}
y={mbY}
width={mbScreenW}
height={mbScreenH}
rx="3"
fill="#1a1a1a"
stroke={wireColor}
stroke-width={strokeWidth}
/>
<!-- Screen inner (dark) -->
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel}
width={mbScreenW - mbBezel * 2}
height={mbScreenH - mbBezel * 2}
rx="2"
fill="#0a0a12"
/>
<!-- Memory fill on screen -->
{#if ramPercent > 0}
<rect
x={mbScreenX + mbBezel}
y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
width={mbScreenW - mbBezel * 2}
height={mbMemH}
fill="rgba(255,215,0,0.85)"
clip-path="url(#{screenClipId})"
/>
{/if}
<!-- Apple logo -->
<path
d={APPLE_LOGO_PATH}
transform="translate({mbLogoX}, {mbLogoY}) scale({mbLogoScale})"
fill="#FFFFFF"
opacity="0.9"
/>
<!-- Keyboard base (trapezoidal) -->
<path
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
fill="#2c2c2c"
stroke={wireColor}
stroke-width="1"
/>
<!-- Keyboard area -->
<rect
x={mbKbX}
y={mbKbY}
width={mbKbW}
height={mbKbH}
fill="rgba(0,0,0,0.2)"
rx="2"
/>
<!-- Trackpad -->
<rect
x={mbTpX}
y={mbTpY}
width={mbTpW}
height={mbTpH}
fill="rgba(255,255,255,0.08)"
rx="2"
/>
{/if}
+65 -20
View File
@@ -6,6 +6,10 @@
export let showSidebarToggle = false;
export let sidebarVisible = true;
export let onToggleSidebar: (() => void) | null = null;
export let downloadProgress: {
count: number;
percentage: number;
} | null = null;
function handleHome(): void {
if (onHome) {
@@ -33,13 +37,17 @@
<div class="absolute left-6 top-1/2 -translate-y-1/2">
<button
onclick={handleToggleSidebar}
class="p-2 rounded border border-exo-medium-gray/40 hover:border-exo-yellow/50 transition-colors cursor-pointer"
class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer"
title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
aria-label={sidebarVisible
? "Hide conversation sidebar"
: "Show conversation sidebar"}
aria-pressed={sidebarVisible}
>
<svg
class="w-5 h-5 {sidebarVisible
? 'text-exo-yellow'
: 'text-exo-medium-gray'}"
: 'text-exo-light-gray'}"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
@@ -75,18 +83,19 @@
<img
src="/exo-logo.png"
alt="EXO"
class="h-18 drop-shadow-[0_0_20px_rgba(255,215,0,0.5)]"
class="h-18 drop-shadow-[0_0_4px_rgba(255,215,0,0.3)]"
/>
</button>
<!-- Right: Home + Downloads -->
<div
<nav
class="absolute right-6 top-1/2 -translate-y-1/2 flex items-center gap-4"
aria-label="Main navigation"
>
{#if showHome}
<button
onclick={handleHome}
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
title="Back to topology view"
>
<svg
@@ -107,23 +116,59 @@
{/if}
<a
href="/#/downloads"
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
title="View downloads overview"
>
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M12 3v12" />
<path d="M7 12l5 5 5-5" />
<path d="M5 21h14" />
</svg>
{#if downloadProgress}
<!-- Compact download progress indicator -->
<div class="relative w-4 h-4 flex-shrink-0">
<svg class="w-4 h-4 -rotate-90" viewBox="0 0 20 20">
<circle
cx="10"
cy="10"
r="8"
fill="none"
stroke="currentColor"
stroke-width="2"
opacity="0.2"
/>
<circle
cx="10"
cy="10"
r="8"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-dasharray={2 * Math.PI * 8}
stroke-dashoffset={2 *
Math.PI *
8 *
(1 - downloadProgress.percentage / 100)}
class="text-blue-400 transition-all duration-300"
/>
</svg>
<div
class="absolute inset-0 flex items-center justify-center text-[6px] font-mono text-blue-400"
>
{downloadProgress.count}
</div>
</div>
{:else}
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M12 3v12" />
<path d="M7 12l5 5 5-5" />
<path d="M5 21h14" />
</svg>
{/if}
Downloads
</a>
</div>
</nav>
</header>
+29 -3
View File
@@ -21,7 +21,7 @@
} | null;
nodes?: Record<string, NodeInfo>;
sharding?: "Pipeline" | "Tensor";
runtime?: "MlxRing" | "MlxIbv" | "MlxJaccl";
runtime?: "MlxRing" | "MlxJaccl";
onLaunch?: () => void;
tags?: string[];
apiPreview?: PlacementPreview | null;
@@ -348,7 +348,7 @@
// Debug mode state
const isDebugMode = $derived(debugMode());
const topology = $derived(topologyData());
const isRdma = $derived(runtime === "MlxIbv" || runtime === "MlxJaccl");
const isRdma = $derived(runtime === "MlxJaccl");
// Get interface name for an IP from node data
function getInterfaceForIp(nodeId: string, ip?: string): string | null {
@@ -567,20 +567,46 @@
<div class="flex items-center gap-1.5 mb-2">
<span
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
title={sharding === "Pipeline"
? "Pipeline: splits model into sequential stages across devices. Lower network overhead."
: "Tensor: splits each layer across devices. Best with high-bandwidth connections (Thunderbolt)."}
>
{sharding}
</span>
<span
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
title={runtime === "MlxRing"
? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
>
{runtime === "MlxRing"
? "MLX Ring"
: runtime === "MlxIbv" || runtime === "MlxJaccl"
: runtime === "MlxJaccl"
? "MLX RDMA"
: runtime}
</span>
</div>
<!-- Download Status -->
{#if isDownloading && progress}
<div class="mb-2 space-y-1">
<div class="flex items-center justify-between text-xs font-mono">
<span class="text-blue-400 tracking-wider uppercase">Downloading</span
>
<span class="text-white/60"
>{percentage.toFixed(1)}% &middot; {formatSpeed(progress.speed)}
&middot; {formatEta(progress.etaMs)}</span
>
</div>
<div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
<div
class="h-full bg-blue-500/70 transition-all duration-300"
style="width: {percentage}%"
></div>
</div>
</div>
{/if}
<!-- Mini Topology Preview -->
{#if placementPreview().nodes.length > 0}
{@const preview = placementPreview()}
@@ -6,6 +6,7 @@
capabilities: string[];
sizeRange: { min: number; max: number } | null;
downloadedOnly: boolean;
readyOnly: boolean;
}
type ModelFilterPopoverProps = {
@@ -190,34 +191,58 @@
</div>
</div>
<!-- Downloaded only -->
<!-- Availability filters -->
<div>
<h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
<div class="flex flex-wrap gap-1.5">
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
<span class="ml-1">Downloaded</span>
</button>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
<span class="ml-1">Downloaded</span>
</button>
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, readyOnly: !filters.readyOnly })}
>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<circle cx="12" cy="12" r="10" />
<path d="m9 12 2 2 4-4" />
</svg>
<span class="ml-1">Ready</span>
</button>
</div>
</div>
<!-- Size range -->
@@ -41,6 +41,7 @@
onShowInfo: (group: ModelGroup) => void;
downloadStatusMap?: Map<string, DownloadAvailability>;
launchedAt?: number;
instanceStatuses?: Record<string, { status: string; statusClass: string }>;
};
let {
@@ -56,6 +57,7 @@
onShowInfo,
downloadStatusMap,
launchedAt,
instanceStatuses = {},
}: ModelPickerGroupProps = $props();
// Group-level download status: show if any variant is downloaded
@@ -92,6 +94,12 @@
const anyVariantFits = $derived(
group.variants.some((v) => canModelFit(v.id)),
);
// Check if any variant has an active instance (ready, loading, downloading)
const anyVariantHasInstance = $derived(
instanceStatuses
? group.variants.some((v) => instanceStatuses[v.id] != null)
: false,
);
const groupFitStatus = $derived.by((): ModelFitStatus => {
let hasClusterCapacityOnly = false;
for (const variant of group.variants) {
@@ -122,16 +130,35 @@
!group.hasMultipleVariants &&
group.variants.some((v) => v.id === selectedModelId),
);
// Group-level instance status: show the "best" status across all variants
const groupInstanceStatus = $derived.by(() => {
if (!instanceStatuses) return null;
const readyStatuses = ["READY", "LOADED", "RUNNING"];
const loadingStatuses = ["LOADING", "WARMING UP"];
let bestStatus: { status: string; statusClass: string } | null = null;
for (const variant of group.variants) {
const s = instanceStatuses[variant.id];
if (!s) continue;
if (readyStatuses.includes(s.status)) return s; // Ready is best
if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
bestStatus = s;
}
}
return bestStatus;
});
</script>
<div
class="border-b border-white/5 last:border-b-0 {!anyVariantFits
class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
!anyVariantHasInstance
? 'opacity-50'
: ''}"
>
<!-- Main row -->
<div
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
anyVariantHasInstance
? 'hover:bg-white/5 cursor-pointer'
: 'cursor-not-allowed'} {isMainSelected
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
@@ -141,7 +168,7 @@
onToggleExpand();
} else {
const modelId = group.variants[0]?.id;
if (modelId && canModelFit(modelId)) {
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
onSelectModel(modelId);
}
}
@@ -155,7 +182,7 @@
onToggleExpand();
} else {
const modelId = group.variants[0]?.id;
if (modelId && canModelFit(modelId)) {
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
onSelectModel(modelId);
}
}
@@ -346,6 +373,47 @@
</span>
{/if}
<!-- Instance status badge -->
{#if groupInstanceStatus}
{#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
<span class="flex-shrink-0" title="Running">
<svg
class="w-3 h-3 text-green-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{:else if groupInstanceStatus.status === "DOWNLOADING"}
<span class="flex-shrink-0 animate-pulse" title="Downloading">
<svg
class="w-3.5 h-3.5 text-blue-400"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
<polyline points="7 10 12 15 17 10" />
<line x1="12" y1="15" x2="12" y2="3" />
</svg>
</span>
{:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
<span class="flex-shrink-0 animate-pulse" title="Loading">
<svg
class="w-3 h-3 text-yellow-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{/if}
{/if}
<!-- Check mark if selected (single-variant) -->
{#if isMainSelected}
<svg
@@ -420,20 +488,30 @@
{#each group.variants as variant}
{@const fitStatus = getModelFitStatus(variant.id)}
{@const modelCanFit = canModelFit(variant.id)}
{@const variantHasInstance = instanceStatuses[variant.id] != null}
{@const isSelected = selectedModelId === variant.id}
<button
type="button"
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
<div
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
!variantHasInstance
? 'opacity-50 cursor-not-allowed'
: 'cursor-pointer'} {isSelected
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
: 'border-l-2 border-transparent'}"
disabled={!modelCanFit}
role="button"
tabindex="0"
onclick={() => {
if (modelCanFit) {
if (modelCanFit || variantHasInstance) {
onSelectModel(variant.id);
}
}}
onkeydown={(e) => {
if (e.key === "Enter" || e.key === " ") {
e.preventDefault();
if (modelCanFit) {
onSelectModel(variant.id);
}
}
}}
>
<!-- Quantization badge -->
<span
@@ -478,6 +556,48 @@
{/if}
{/if}
<!-- Instance status badge -->
{#if instanceStatuses[variant.id]}
{@const instStatus = instanceStatuses[variant.id]}
{#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
<span class="flex-shrink-0" title="Running">
<svg
class="w-3 h-3 text-green-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{:else if instStatus.status === "DOWNLOADING"}
<span class="flex-shrink-0 animate-pulse" title="Downloading">
<svg
class="w-3.5 h-3.5 text-blue-400"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
<polyline points="7 10 12 15 17 10" />
<line x1="12" y1="15" x2="12" y2="3" />
</svg>
</span>
{:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
<span class="flex-shrink-0 animate-pulse" title="Loading">
<svg
class="w-3 h-3 text-yellow-400"
viewBox="0 0 12 12"
fill="currentColor"
>
<circle cx="6" cy="6" r="5" />
</svg>
</span>
{/if}
{/if}
<!-- Check mark if selected -->
{#if isSelected}
<svg
@@ -490,7 +610,36 @@
/>
</svg>
{/if}
</button>
<!-- Info button -->
<button
type="button"
class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
onclick={(e) => {
e.stopPropagation();
onShowInfo({
id: variant.id,
name: variant.name || variant.id,
capabilities: group.capabilities,
family: group.family,
variants: [variant],
smallestVariant: variant,
hasMultipleVariants: false,
});
}}
title="View variant details"
>
<svg
class="w-4 h-4 text-white/30 hover:text-white/50"
viewBox="0 0 24 24"
fill="currentColor"
>
<path
d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
/>
</svg>
</button>
</div>
{/each}
</div>
{/if}
@@ -36,6 +36,7 @@
capabilities: string[];
sizeRange: { min: number; max: number } | null;
downloadedOnly: boolean;
readyOnly: boolean;
}
interface HuggingFaceModel {
@@ -49,6 +50,11 @@
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
export type InstanceStatus = {
status: string;
statusClass: string;
};
type ModelPickerModalProps = {
isOpen: boolean;
models: ModelInfo[];
@@ -75,6 +81,7 @@
macmon_info?: { memory?: { ram_total?: number } };
}
>;
instanceStatuses?: Record<string, InstanceStatus>;
};
let {
@@ -96,6 +103,7 @@
usedMemoryGB,
downloadsData,
topologyNodes,
instanceStatuses = {},
}: ModelPickerModalProps = $props();
// Local state
@@ -107,6 +115,7 @@
capabilities: [],
sizeRange: null,
downloadedOnly: false,
readyOnly: false,
});
let infoGroup = $state<ModelGroup | null>(null);
@@ -440,6 +449,16 @@
);
}
// Filter to ready/running models only
if (filters.readyOnly) {
result = result.filter((g) =>
g.variants.some((v) => {
const s = instanceStatuses[v.id];
return s && s.statusClass === "ready";
}),
);
}
// Sort: fits-now first, then fits-cluster-capacity, then too-large
result.sort((a, b) => {
const getGroupFitRank = (group: ModelGroup): number => {
@@ -512,6 +531,18 @@
);
});
// Split filtered groups into recommended (fits_now) and others for visual separation
const recommendedGroups = $derived(
filteredGroups.filter((g) =>
g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
),
);
const otherGroups = $derived(
filteredGroups.filter(
(g) => !g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
),
);
function toggleGroupExpanded(groupId: string) {
const next = new Set(expandedGroups);
if (next.has(groupId)) {
@@ -538,13 +569,19 @@
}
function clearFilters() {
filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
filters = {
capabilities: [],
sizeRange: null,
downloadedOnly: false,
readyOnly: false,
};
}
const hasActiveFilters = $derived(
filters.capabilities.length > 0 ||
filters.sizeRange !== null ||
filters.downloadedOnly,
filters.downloadedOnly ||
filters.readyOnly,
);
</script>
@@ -813,6 +850,7 @@
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
{instanceStatuses}
/>
{/each}
{/if}
@@ -840,7 +878,34 @@
{/if}
</div>
{:else}
{#each filteredGroups as group}
<!-- Recommended for your cluster -->
{#if recommendedGroups.length > 0 && otherGroups.length > 0 && !searchQuery.trim()}
<div
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-green-950/60 border-b border-green-500/20 backdrop-blur-sm"
>
<svg
class="w-3.5 h-3.5 text-green-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<span
class="text-xs font-mono text-green-400 tracking-wider uppercase"
>Recommended for your cluster</span
>
<span class="text-xs font-mono text-green-400/50"
>— fits in available memory</span
>
</div>
{/if}
{#each recommendedGroups as group}
<ModelPickerGroup
{group}
isExpanded={expandedGroups.has(group.id)}
@@ -853,6 +918,34 @@
{onToggleFavorite}
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
{instanceStatuses}
/>
{/each}
<!-- Other models -->
{#if otherGroups.length > 0 && recommendedGroups.length > 0 && !searchQuery.trim()}
<div
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-exo-dark-gray/80 border-y border-exo-medium-gray/20 backdrop-blur-sm"
>
<span
class="text-xs font-mono text-white/40 tracking-wider uppercase"
>Other models</span
>
</div>
{/if}
{#each otherGroups as group}
<ModelPickerGroup
{group}
isExpanded={expandedGroups.has(group.id)}
isFavorite={favorites.has(group.id)}
{selectedModelId}
{canModelFit}
{getModelFitStatus}
onToggleExpand={() => toggleGroupExpanded(group.id)}
onSelectModel={handleSelect}
{onToggleFavorite}
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
{instanceStatuses}
/>
{/each}
{/if}
@@ -875,6 +968,11 @@
>Downloaded</span
>
{/if}
{#if filters.readyOnly}
<span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
>Ready</span
>
{/if}
{#if filters.sizeRange}
<span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
{filters.sizeRange.min}GB - {filters.sizeRange.max}GB
@@ -14,6 +14,21 @@
: 0,
);
const etaText = $derived.by(() => {
if (progress.processed <= 0 || progress.total <= 0) return null;
const elapsedMs = performance.now() - progress.startedAt;
if (elapsedMs < 200) return null; // need a minimum sample window
const tokensPerMs = progress.processed / elapsedMs;
const remainingTokens = progress.total - progress.processed;
const remainingMs = remainingTokens / tokensPerMs;
const remainingSec = Math.ceil(remainingMs / 1000);
if (remainingSec <= 0) return null;
if (remainingSec < 60) return `~${remainingSec}s remaining`;
const mins = Math.floor(remainingSec / 60);
const secs = remainingSec % 60;
return `~${mins}m ${secs}s remaining`;
});
function formatTokenCount(count: number | undefined): string {
if (count == null) return "0";
if (count >= 1000) {
@@ -40,8 +55,11 @@
style="width: {percentage}%"
></div>
</div>
<div class="text-right text-xs text-exo-light-gray/70 mt-0.5 font-mono">
{percentage}%
<div
class="flex items-center justify-between text-xs text-exo-light-gray/70 mt-0.5 font-mono"
>
<span>{etaText ?? ""}</span>
<span>{percentage}%</span>
</div>
</div>
@@ -0,0 +1,117 @@
<script lang="ts">
import { toasts, dismissToast, type Toast } from "$lib/stores/toast.svelte";
import { fly, fade } from "svelte/transition";
import { flip } from "svelte/animate";
const items = $derived(toasts());
const typeStyles: Record<
Toast["type"],
{ border: string; icon: string; iconColor: string }
> = {
success: {
border: "border-l-green-500",
icon: "M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z",
iconColor: "text-green-400",
},
error: {
border: "border-l-red-500",
icon: "M12 9v3.75m9-.75a9 9 0 11-18 0 9 9 0 0118 0zm-9 3.75h.008v.008H12v-.008z",
iconColor: "text-red-400",
},
warning: {
border: "border-l-yellow-500",
icon: "M12 9v3.75m-9.303 3.376c-.866 1.5.217 3.374 1.948 3.374h14.71c1.73 0 2.813-1.874 1.948-3.374L13.949 3.378c-.866-1.5-3.032-1.5-3.898 0L2.697 16.126z",
iconColor: "text-yellow-400",
},
info: {
border: "border-l-blue-500",
icon: "M11.25 11.25l.041-.02a.75.75 0 011.063.852l-.708 2.836a.75.75 0 001.063.853l.041-.021M21 12a9 9 0 11-18 0 9 9 0 0118 0zm-9-3.75h.008v.008H12V8.25z",
iconColor: "text-blue-400",
},
};
</script>
{#if items.length > 0}
<div
class="fixed bottom-6 right-6 z-[9999] flex flex-col gap-2 pointer-events-none"
role="log"
aria-live="polite"
aria-label="Notifications"
>
{#each items as toast (toast.id)}
{@const style = typeStyles[toast.type]}
<div
class="pointer-events-auto max-w-sm w-80 bg-exo-dark-gray/95 backdrop-blur-sm border border-exo-medium-gray/60 border-l-[3px] {style.border} rounded shadow-lg shadow-black/40"
in:fly={{ x: 80, duration: 250 }}
out:fade={{ duration: 150 }}
animate:flip={{ duration: 200 }}
role="alert"
>
<div class="flex items-start gap-3 px-4 py-3">
<!-- Icon -->
<svg
class="w-5 h-5 flex-shrink-0 mt-0.5 {style.iconColor}"
fill="none"
viewBox="0 0 24 24"
stroke-width="1.5"
stroke="currentColor"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={style.icon}
/>
</svg>
<!-- Message -->
<p class="flex-1 text-sm text-white/90 font-mono leading-snug">
{toast.message}
</p>
<!-- Dismiss button -->
<button
onclick={() => dismissToast(toast.id)}
class="flex-shrink-0 p-0.5 text-white/40 hover:text-white/80 transition-colors cursor-pointer"
aria-label="Dismiss notification"
>
<svg
class="w-4 h-4"
fill="none"
viewBox="0 0 24 24"
stroke-width="2"
stroke="currentColor"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M6 18L18 6M6 6l12 12"
/>
</svg>
</button>
</div>
<!-- Auto-dismiss progress bar -->
{#if toast.duration > 0}
<div class="h-0.5 bg-white/5 rounded-b overflow-hidden">
<div
class="h-full {style.border.replace('border-l-', 'bg-')}/60"
style="animation: shrink {toast.duration}ms linear forwards"
></div>
</div>
{/if}
</div>
{/each}
</div>
{/if}
<style>
@keyframes shrink {
from {
width: 100%;
}
to {
width: 0%;
}
}
</style>
+1
View File
@@ -12,3 +12,4 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
export { default as ChatModelSelector } from "./ChatModelSelector.svelte";
+120 -31
View File
@@ -168,7 +168,7 @@ export interface ModelDownloadStatus {
export interface PlacementPreview {
model_id: string;
sharding: "Pipeline" | "Tensor";
instance_meta: "MlxRing" | "MlxIbv" | "MlxJaccl";
instance_meta: "MlxRing" | "MlxJaccl";
instance: unknown | null;
memory_delta_by_node: Record<string, number> | null;
error: string | null;
@@ -219,7 +219,6 @@ interface RawStateResponse {
string,
{
MlxRingInstance?: Instance;
MlxIbvInstance?: Instance;
MlxJacclInstance?: Instance;
}
>;
@@ -250,6 +249,11 @@ interface RawStateResponse {
>;
// Thunderbolt bridge cycles (nodes with bridge enabled forming loops)
thunderboltBridgeCycles?: string[][];
// Disk usage per node
nodeDisk?: Record<
string,
{ total: { inBytes: number }; available: { inBytes: number } }
>;
}
export interface MessageAttachment {
@@ -276,6 +280,8 @@ export interface TokenData {
export interface PrefillProgress {
processed: number;
total: number;
/** Timestamp (performance.now()) when prefill started. */
startedAt: number;
}
export interface Message {
@@ -583,6 +589,12 @@ class AppStore {
// Image editing state
editingImage = $state<EditingImage | null>(null);
/** True when the backend is reachable. */
isConnected = $state<boolean>(true);
/** Number of consecutive fetch failures. */
private consecutiveFailures = 0;
private static readonly CONNECTION_LOST_THRESHOLD = 3;
private fetchInterval: ReturnType<typeof setInterval> | null = null;
private previewsInterval: ReturnType<typeof setInterval> | null = null;
private lastConversationPersistTs = 0;
@@ -835,6 +847,10 @@ class AppStore {
this.thinkingEnabled = conversation.enableThinking ?? true;
this.refreshConversationModelFromInstances();
// Sync global selection to the loaded conversation's model so reactive
// effects in +page.svelte can determine the correct chat launch state.
this.selectedChatModel = conversation.modelId || "";
return true;
}
@@ -905,11 +921,7 @@ class AppStore {
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (
instanceTag === "MlxIbvInstance" ||
instanceTag === "MlxJacclInstance"
)
instanceType = "MLX RDMA";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
let sharding: string | null = null;
const inst = instance as {
@@ -1288,7 +1300,19 @@ class AppStore {
// Thunderbolt bridge status per node
this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
this.lastUpdate = Date.now();
// Connection recovered
if (!this.isConnected) {
this.isConnected = true;
}
this.consecutiveFailures = 0;
} catch (error) {
this.consecutiveFailures++;
if (
this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD &&
this.isConnected
) {
this.isConnected = false;
}
console.error("Error fetching state:", error);
}
}
@@ -1652,11 +1676,12 @@ class AppStore {
if (!reader) throw new Error("No response body");
let fullContent = prefixText;
let streamedThinking = "";
const collectedTokens: TokenData[] = [...tokensToKeep];
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string };
delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -1677,6 +1702,7 @@ class AppStore {
(parsed) => {
const choice = parsed.choices?.[0];
const delta = choice?.delta?.content;
const thinkingDelta = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -1695,7 +1721,11 @@ class AppStore {
}
}
if (delta) {
if (thinkingDelta) {
streamedThinking += thinkingDelta;
}
if (delta || thinkingDelta) {
if (firstTokenTime === null) {
firstTokenTime = performance.now();
this.ttftMs = firstTokenTime - requestStartTime;
@@ -1709,9 +1739,14 @@ class AppStore {
this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000;
}
fullContent += delta;
const { displayContent, thinkingContent } =
if (delta) {
fullContent += delta;
}
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(fullContent);
const combinedThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
if (this.activeConversationId === targetConversationId) {
this.currentResponse = displayContent;
@@ -1723,7 +1758,7 @@ class AppStore {
messageId,
(m) => {
m.content = displayContent;
m.thinking = thinkingContent || undefined;
m.thinking = combinedThinking || undefined;
m.tokens = [...collectedTokens];
},
);
@@ -1735,11 +1770,14 @@ class AppStore {
// Final update
if (this.conversationExists(targetConversationId)) {
const { displayContent, thinkingContent } =
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(fullContent);
const finalThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
this.updateConversationMessage(targetConversationId, messageId, (m) => {
m.content = displayContent;
m.thinking = thinkingContent || undefined;
m.thinking = finalThinking || undefined;
m.tokens = [...collectedTokens];
if (this.ttftMs !== null) m.ttftMs = this.ttftMs;
if (this.tps !== null) m.tps = this.tps;
@@ -1815,7 +1853,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content =
"Error: No model available. Please launch an instance first.";
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.";
},
);
this.syncActiveMessagesIfNeeded(targetConversationId);
@@ -1847,11 +1885,12 @@ class AppStore {
}
let streamedContent = "";
let streamedThinking = "";
const collectedTokens: TokenData[] = [];
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string };
delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -1872,6 +1911,7 @@ class AppStore {
(parsed) => {
const choice = parsed.choices?.[0];
const delta = choice?.delta?.content;
const thinkingDelta = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -1890,10 +1930,19 @@ class AppStore {
}
}
if (delta) {
streamedContent += delta;
const { displayContent, thinkingContent } =
if (thinkingDelta) {
streamedThinking += thinkingDelta;
}
if (delta || thinkingDelta) {
if (delta) {
streamedContent += delta;
}
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const combinedThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
// Only update currentResponse if target conversation is active
if (this.activeConversationId === targetConversationId) {
@@ -1906,7 +1955,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = combinedThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -1918,14 +1967,17 @@ class AppStore {
// Final cleanup of the message (if conversation still exists)
if (this.conversationExists(targetConversationId)) {
const { displayContent, thinkingContent } =
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const finalThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
this.updateConversationMessage(
targetConversationId,
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = finalThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -2272,7 +2324,7 @@ class AppStore {
const modelToUse = this.getModelForRequest();
if (!modelToUse) {
throw new Error(
"No model selected and no running instances available. Please launch an instance first.",
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.",
);
}
@@ -2317,10 +2369,11 @@ class AppStore {
}
let streamedContent = "";
let streamedThinking = "";
interface ChatCompletionChunk {
choices?: Array<{
delta?: { content?: string };
delta?: { content?: string; reasoning_content?: string };
logprobs?: {
content?: Array<{
token: string;
@@ -2348,6 +2401,7 @@ class AppStore {
const choice = parsed.choices?.[0];
const tokenContent = choice?.delta?.content;
const thinkingContent = choice?.delta?.reasoning_content;
// Collect logprobs data
const logprobsContent = choice?.logprobs?.content;
@@ -2366,7 +2420,11 @@ class AppStore {
}
}
if (tokenContent) {
if (thinkingContent) {
streamedThinking += thinkingContent;
}
if (tokenContent || thinkingContent) {
// Track first token for TTFT
if (firstTokenTime === null) {
firstTokenTime = performance.now();
@@ -2383,11 +2441,16 @@ class AppStore {
this.tps = (tokenCount / elapsed) * 1000;
}
streamedContent += tokenContent;
if (tokenContent) {
streamedContent += tokenContent;
}
// Strip thinking tags for display and extract thinking content
const { displayContent, thinkingContent } =
// Use stripThinkingTags as fallback for any <think> tags still in content
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const combinedThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
// Only update currentResponse if target conversation is active
if (this.activeConversationId === targetConversationId) {
@@ -2400,7 +2463,7 @@ class AppStore {
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = combinedThinking || undefined;
msg.tokens = [...collectedTokens];
},
);
@@ -2420,6 +2483,7 @@ class AppStore {
this.prefillProgress = {
processed: inner.processed_tokens,
total: inner.total_tokens,
startedAt: this.prefillProgress?.startedAt ?? performance.now(),
};
},
},
@@ -2436,14 +2500,17 @@ class AppStore {
// Final cleanup of the message (if conversation still exists)
if (this.conversationExists(targetConversationId)) {
const { displayContent, thinkingContent } =
const { displayContent, thinkingContent: tagThinking } =
this.stripThinkingTags(streamedContent);
const finalThinking = [streamedThinking, tagThinking]
.filter(Boolean)
.join("\n\n");
this.updateConversationMessage(
targetConversationId,
assistantMessage.id,
(msg) => {
msg.content = displayContent;
msg.thinking = thinkingContent || undefined;
msg.thinking = finalThinking || undefined;
msg.tokens = [...collectedTokens];
// Store performance metrics on the message
if (this.ttftMs !== null) {
@@ -3091,6 +3158,23 @@ class AppStore {
return (await response.json()) as TraceStatsResponse;
}
/**
* Delete traces by task IDs
*/
async deleteTraces(
taskIds: string[],
): Promise<{ deleted: string[]; notFound: string[] }> {
const response = await fetch("/v1/traces/delete", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ taskIds }),
});
if (!response.ok) {
throw new Error(`Failed to delete traces: ${response.status}`);
}
return await response.json();
}
/**
* Get the URL for the raw trace file (for Perfetto)
*/
@@ -3200,6 +3284,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;
@@ -3231,3 +3318,5 @@ export const fetchTraceStats = (taskId: string) =>
appStore.fetchTraceStats(taskId);
export const getTraceRawUrl = (taskId: string) =>
appStore.getTraceRawUrl(taskId);
export const deleteTraces = (taskIds: string[]) =>
appStore.deleteTraces(taskIds);
+87
View File
@@ -0,0 +1,87 @@
/**
* Toast notification store - Global notification system for the EXO dashboard.
*
* Usage:
* import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte";
* addToast({ type: "success", message: "Model launched" });
* addToast({ type: "error", message: "Connection lost", persistent: true });
*/
type ToastType = "success" | "error" | "warning" | "info";
export interface Toast {
id: string;
type: ToastType;
message: string;
/** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */
duration: number;
createdAt: number;
}
interface ToastInput {
type: ToastType;
message: string;
/** If true, toast stays until manually dismissed. Default: false. */
persistent?: boolean;
/** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */
duration?: number;
}
const DEFAULT_DURATIONS: Record<ToastType, number> = {
success: 4000,
info: 4000,
warning: 6000,
error: 6000,
};
let toastList = $state<Toast[]>([]);
const timers = new Map<string, ReturnType<typeof setTimeout>>();
function generateId(): string {
return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
}
export function addToast(input: ToastInput): string {
const id = generateId();
const duration = input.persistent
? 0
: (input.duration ?? DEFAULT_DURATIONS[input.type]);
const toast: Toast = {
id,
type: input.type,
message: input.message,
duration,
createdAt: Date.now(),
};
toastList = [...toastList, toast];
if (duration > 0) {
const timer = setTimeout(() => dismissToast(id), duration);
timers.set(id, timer);
}
return id;
}
export function dismissToast(id: string): void {
const timer = timers.get(id);
if (timer) {
clearTimeout(timer);
timers.delete(id);
}
toastList = toastList.filter((t) => t.id !== id);
}
/** Dismiss all toasts matching a message (useful for dedup). */
export function dismissByMessage(message: string): void {
const matching = toastList.filter((t) => t.message === message);
for (const t of matching) {
dismissToast(t.id);
}
}
export function toasts(): Toast[] {
return toastList;
}
+4
View File
@@ -1,5 +1,7 @@
<script lang="ts">
import "../app.css";
import ToastContainer from "$lib/components/ToastContainer.svelte";
import ConnectionBanner from "$lib/components/ConnectionBanner.svelte";
let { children } = $props();
</script>
@@ -10,5 +12,7 @@
</svelte:head>
<div class="min-h-screen bg-background text-foreground">
<ConnectionBanner />
{@render children?.()}
<ToastContainer />
</div>
File diff suppressed because it is too large Load Diff
+156 -42
View File
@@ -29,7 +29,12 @@
etaMs: number;
modelDirectory?: string;
}
| { kind: "pending"; modelDirectory?: string }
| {
kind: "pending";
downloaded: number;
total: number;
modelDirectory?: string;
}
| { kind: "failed"; modelDirectory?: string }
| { kind: "not_present" };
@@ -74,7 +79,6 @@
if (typeof value === "number") return value;
if (value && typeof value === "object") {
const v = value as Record<string, unknown>;
if (typeof v.in_bytes === "number") return v.in_bytes;
if (typeof v.inBytes === "number") return v.inBytes;
}
return 0;
@@ -231,23 +235,14 @@
undefined;
let cell: CellStatus;
if (tag === "DownloadCompleted") {
const totalBytes = getBytes(
payload.total_bytes ?? payload.totalBytes,
);
const totalBytes = getBytes(payload.total);
cell = { kind: "completed", totalBytes, modelDirectory };
} else if (tag === "DownloadOngoing") {
const rawProgress =
payload.download_progress ?? payload.downloadProgress ?? {};
const prog = rawProgress as Record<string, unknown>;
const totalBytes = getBytes(
prog.total_bytes ??
prog.totalBytes ??
payload.total_bytes ??
payload.totalBytes,
);
const downloadedBytes = getBytes(
prog.downloaded_bytes ?? prog.downloadedBytes,
);
const totalBytes = getBytes(prog.total ?? payload.total);
const downloadedBytes = getBytes(prog.downloaded);
const speed = (prog.speed as number) ?? 0;
const etaMs =
(prog.eta_ms as number) ?? (prog.etaMs as number) ?? 0;
@@ -265,7 +260,20 @@
} else if (tag === "DownloadFailed") {
cell = { kind: "failed", modelDirectory };
} else {
cell = { kind: "pending", modelDirectory };
const downloaded = getBytes(
payload.downloaded ??
payload.downloaded_bytes ??
payload.downloadedBytes,
);
const total = getBytes(
payload.total ?? payload.total_bytes ?? payload.totalBytes,
);
cell = {
kind: "pending",
downloaded,
total,
modelDirectory,
};
}
const existing = row.cells[nodeId];
@@ -275,14 +283,51 @@
}
}
function rowSortKey(row: ModelRow): number {
// in progress (4) -> completed (3) -> paused (2) -> not started (1) -> not present (0)
let best = 0;
for (const cell of Object.values(row.cells)) {
let score = 0;
if (cell.kind === "downloading") score = 4;
else if (cell.kind === "completed") score = 3;
else if (cell.kind === "pending" && cell.downloaded > 0)
score = 2; // paused
else if (cell.kind === "pending" || cell.kind === "failed") score = 1; // not started
if (score > best) best = score;
}
return best;
}
function totalCompletedBytes(row: ModelRow): number {
let total = 0;
for (const cell of Object.values(row.cells)) {
if (cell.kind === "completed") total += cell.totalBytes;
}
return total;
}
const rows = Array.from(rowMap.values()).sort((a, b) => {
const aCompleted = Object.values(a.cells).filter(
(c) => c.kind === "completed",
).length;
const bCompleted = Object.values(b.cells).filter(
(c) => c.kind === "completed",
).length;
if (aCompleted !== bCompleted) return bCompleted - aCompleted;
const aPriority = rowSortKey(a);
const bPriority = rowSortKey(b);
if (aPriority !== bPriority) return bPriority - aPriority;
// Within completed or paused, sort by biggest size first
if (aPriority === 3 && bPriority === 3) {
const sizeDiff = totalCompletedBytes(b) - totalCompletedBytes(a);
if (sizeDiff !== 0) return sizeDiff;
}
if (aPriority === 2 && bPriority === 2) {
const aSize = Math.max(
...Object.values(a.cells).map((c) =>
c.kind === "pending" ? c.total : 0,
),
);
const bSize = Math.max(
...Object.values(b.cells).map((c) =>
c.kind === "pending" ? c.total : 0,
),
);
if (aSize !== bSize) return bSize - aSize;
}
return a.modelId.localeCompare(b.modelId);
});
@@ -367,7 +412,7 @@
<div>{col.label}</div>
{#if col.diskAvailable != null}
<div
class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
>
{formatBytes(col.diskAvailable)} free
</div>
@@ -391,7 +436,7 @@
</div>
{#if row.prettyName}
<div
class="text-[10px] text-exo-light-gray/60"
class="text-[10px] text-white/60"
title={row.modelId}
>
{row.modelId}
@@ -405,7 +450,7 @@
title="View model details"
>
<svg
class="w-4 h-4 text-white/30 hover:text-white/60"
class="w-4 h-4 text-white/60 hover:text-white/80"
viewBox="0 0 24 24"
fill="currentColor"
>
@@ -424,11 +469,11 @@
<td class="px-4 py-3 text-center align-middle">
{#if cell.kind === "completed"}
<div
class="flex flex-col items-center gap-0.5"
class="flex flex-col items-center gap-1"
title="Completed ({formatBytes(cell.totalBytes)})"
>
<svg
class="w-5 h-5 text-green-400"
class="w-7 h-7 text-green-400"
viewBox="0 0 20 20"
fill="currentColor"
>
@@ -438,18 +483,18 @@
clip-rule="evenodd"
></path>
</svg>
<span class="text-[10px] text-exo-light-gray/70"
<span class="text-xs text-white/70"
>{formatBytes(cell.totalBytes)}</span
>
<button
type="button"
class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
onclick={() =>
deleteDownload(col.nodeId, row.modelId)}
title="Delete from this node"
>
<svg
class="w-3.5 h-3.5"
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
@@ -472,11 +517,11 @@
cell.speed,
)} - ETA {formatEta(cell.etaMs)}"
>
<span class="text-exo-yellow text-xs font-medium"
<span class="text-exo-yellow text-sm font-medium"
>{clampPercent(cell.percentage).toFixed(1)}%</span
>
<div
class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
>
<div
class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
@@ -485,24 +530,93 @@
).toFixed(1)}%"
></div>
</div>
<span class="text-[9px] text-exo-light-gray/60"
<span class="text-[10px] text-white/70"
>{formatSpeed(cell.speed)}</span
>
</div>
{:else if cell.kind === "pending"}
<div
class="flex flex-col items-center gap-0.5"
title="Download pending"
class="flex flex-col items-center gap-1"
title={cell.downloaded > 0
? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
: "Download pending"}
>
<span class="text-exo-light-gray/50 text-sm">...</span>
{#if cell.downloaded > 0 && cell.total > 0}
<span class="text-white/70 text-xs"
>{formatBytes(cell.downloaded)} / {formatBytes(
cell.total,
)}</span
>
<div
class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
>
<div
class="h-full bg-exo-light-gray/40 rounded-full"
style="width: {(
(cell.downloaded / cell.total) *
100
).toFixed(1)}%"
></div>
</div>
{#if row.shardMetadata}
<button
type="button"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Resume download on this node"
>
<svg
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
stroke-width="2"
>
<path
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
stroke-linecap="round"
stroke-linejoin="round"
></path>
</svg>
</button>
{:else}
<span class="text-white/50 text-[10px]">paused</span
>
{/if}
{:else if row.shardMetadata}
<button
type="button"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Start download on this node"
>
<svg
class="w-6 h-6"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
stroke-width="2"
>
<path
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
stroke-linecap="round"
stroke-linejoin="round"
></path>
</svg>
</button>
{:else}
<span class="text-white/40 text-sm">...</span>
{/if}
</div>
{:else if cell.kind === "failed"}
<div
class="flex flex-col items-center gap-0.5"
class="flex flex-col items-center gap-1"
title="Download failed"
>
<svg
class="w-5 h-5 text-red-400"
class="w-7 h-7 text-red-400"
viewBox="0 0 20 20"
fill="currentColor"
>
@@ -515,13 +629,13 @@
{#if row.shardMetadata}
<button
type="button"
class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Retry download on this node"
>
<svg
class="w-3.5 h-3.5"
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
@@ -547,13 +661,13 @@
{#if row.shardMetadata}
<button
type="button"
class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
onclick={() =>
startDownload(col.nodeId, row.shardMetadata!)}
title="Download to this node"
>
<svg
class="w-3.5 h-3.5"
class="w-5 h-5"
viewBox="0 0 20 20"
fill="none"
stroke="currentColor"
+90 -3
View File
@@ -3,6 +3,7 @@
import {
listTraces,
getTraceRawUrl,
deleteTraces,
type TraceListItem,
} from "$lib/stores/app.svelte";
import HeaderNav from "$lib/components/HeaderNav.svelte";
@@ -10,6 +11,51 @@
let traces = $state<TraceListItem[]>([]);
let loading = $state(true);
let error = $state<string | null>(null);
let selectedIds = $state<Set<string>>(new Set());
let deleting = $state(false);
let allSelected = $derived(
traces.length > 0 && selectedIds.size === traces.length,
);
function toggleSelect(taskId: string) {
const next = new Set(selectedIds);
if (next.has(taskId)) {
next.delete(taskId);
} else {
next.add(taskId);
}
selectedIds = next;
}
function toggleSelectAll() {
if (allSelected) {
selectedIds = new Set();
} else {
selectedIds = new Set(traces.map((t) => t.taskId));
}
}
async function handleDelete() {
if (selectedIds.size === 0) return;
const count = selectedIds.size;
if (
!confirm(
`Delete ${count} trace${count === 1 ? "" : "s"}? This cannot be undone.`,
)
)
return;
deleting = true;
try {
await deleteTraces([...selectedIds]);
selectedIds = new Set();
await refresh();
} catch (e) {
error = e instanceof Error ? e.message : "Failed to delete traces";
} finally {
deleting = false;
}
}
function formatBytes(bytes: number): string {
if (!bytes || bytes <= 0) return "0B";
@@ -109,6 +155,16 @@
</h1>
</div>
<div class="flex items-center gap-3">
{#if selectedIds.size > 0}
<button
type="button"
class="text-xs font-mono text-red-400 hover:text-red-300 transition-colors uppercase border border-red-500/40 px-2 py-1 rounded"
onclick={handleDelete}
disabled={deleting}
>
{deleting ? "Deleting..." : `Delete (${selectedIds.size})`}
</button>
{/if}
<button
type="button"
class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
@@ -143,14 +199,41 @@
</div>
{:else}
<div class="space-y-3">
<div class="flex items-center gap-2 px-1">
<button
type="button"
class="text-xs font-mono uppercase transition-colors {allSelected
? 'text-exo-yellow'
: 'text-exo-light-gray hover:text-exo-yellow'}"
onclick={toggleSelectAll}
>
{allSelected ? "Deselect all" : "Select all"}
</button>
</div>
{#each traces as trace}
{@const isSelected = selectedIds.has(trace.taskId)}
<!-- svelte-ignore a11y_no_static_element_interactions -->
<div
class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
role="button"
tabindex="0"
class="w-full text-left rounded border-l-2 border-r border-t border-b transition-all p-4 flex items-center justify-between gap-4 cursor-pointer {isSelected
? 'bg-exo-yellow/10 border-l-exo-yellow border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30'
: 'bg-exo-black/30 border-l-transparent border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30 hover:bg-white/[0.03]'}"
onclick={() => toggleSelect(trace.taskId)}
onkeydown={(e) => {
if (e.key === "Enter" || e.key === " ") {
e.preventDefault();
toggleSelect(trace.taskId);
}
}}
>
<div class="min-w-0 flex-1">
<a
href="#/traces/{trace.taskId}"
class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
class="text-sm font-mono transition-colors truncate block {isSelected
? 'text-exo-yellow'
: 'text-white hover:text-exo-yellow'}"
onclick={(e) => e.stopPropagation()}
>
{trace.taskId}
</a>
@@ -160,7 +243,11 @@
)}
</div>
</div>
<div class="flex items-center gap-2 shrink-0">
<!-- svelte-ignore a11y_click_events_have_key_events -->
<div
class="flex items-center gap-2 shrink-0"
onclick={(e) => e.stopPropagation()}
>
<a
href="#/traces/{trace.taskId}"
class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
+22
View File
@@ -183,6 +183,27 @@ Same schema as `/v1/chat/completions`.
**Response:**
Chat completion plus benchmarking metrics.
### Cancel Command
**POST** `/v1/cancel/{command_id}`
Cancels an active generation command (text or image). Notifies workers and closes the stream.
**Path parameters:**
* `command_id`: string, ID of the command to cancel
**Response (example):**
```json
{
"message": "Command cancelled.",
"command_id": "cmd-abc-123"
}
```
Returns 404 if the command is not found or already completed.
## 5. Image Generation & Editing
### Image Generation
@@ -266,6 +287,7 @@ GET /v1/models
POST /v1/chat/completions
POST /bench/chat/completions
POST /v1/cancel/{command_id}
POST /v1/images/generations
POST /bench/images/generations
+1
View File
@@ -108,6 +108,7 @@
package = pkgsSwift.swiftPackages.swift-format;
};
shfmt.enable = true;
taplo.enable = true;
};
};
+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}")
+52 -52
View File
@@ -1,35 +1,35 @@
[project]
name = "exo"
version = "0.3.0"
version = "0.3.68"
description = "Exo"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"aiofiles>=24.1.0",
"aiohttp>=3.12.14",
"types-aiofiles>=24.1.0.20250708",
"pydantic>=2.11.7",
"fastapi>=0.116.1",
"filelock>=3.18.0",
"rustworkx>=0.17.1",
"huggingface-hub>=0.33.4",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo_pyo3_bindings", # rust bindings
"anyio==4.11.0",
"mlx; sys_platform == 'darwin'",
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
"mlx-lm==0.30.7",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
"openai-harmony>=0.0.8",
"httpx>=0.28.1",
"tomlkit>=0.14.0",
"pillow>=11.0,<12.0", # compatibility with mflux
"mflux==0.15.5",
"python-multipart>=0.0.21",
"msgspec>=0.19.0",
"zstandard>=0.23.0",
"aiofiles>=24.1.0",
"aiohttp>=3.12.14",
"types-aiofiles>=24.1.0.20250708",
"pydantic>=2.11.7",
"fastapi>=0.116.1",
"filelock>=3.18.0",
"rustworkx>=0.17.1",
"huggingface-hub>=0.33.4",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo_pyo3_bindings", # rust bindings
"anyio==4.11.0",
"mlx; sys_platform == 'darwin'",
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
"mlx-lm",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
"openai-harmony>=0.0.8",
"httpx>=0.28.1",
"tomlkit>=0.14.0",
"pillow>=11.0,<12.0", # compatibility with mflux
"mflux==0.15.5",
"python-multipart>=0.0.21",
"msgspec>=0.19.0",
"zstandard>=0.23.0",
]
[project.scripts]
@@ -38,12 +38,12 @@ exo = "exo.main:main"
# dependencies only required for development
[dependency-groups]
dev = [
"basedpyright>=1.29.0",
"pyinstaller>=6.17.0",
"pytest>=8.4.0",
"pytest-asyncio>=1.0.0",
"pytest-env",
"ruff>=0.11.13",
"basedpyright>=1.29.0",
"pyinstaller>=6.17.0",
"pytest>=8.4.0",
"pytest-asyncio>=1.0.0",
"pytest-env",
"ruff>=0.11.13",
]
# mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
@@ -57,15 +57,12 @@ dev = [
###
[tool.uv.workspace]
members = [
"rust/exo_pyo3_bindings",
"bench",
]
members = ["rust/exo_pyo3_bindings", "bench"]
[tool.uv.sources]
exo_pyo3_bindings = { workspace = true }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", rev = "834fac934c4e04de9b3d723e2b9287a2c60cfd4a" }
# Uncomment to use local mlx/mlx-lm development versions:
# mlx = { path = "/Users/Shared/mlx", editable=true }
# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
@@ -95,7 +92,15 @@ reportUnnecessaryTypeIgnoreComment = "error"
pythonVersion = "3.13"
pythonPlatform = "Darwin"
exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github"]
exclude = [
"**/.venv",
"**/venv",
"**/__pycache__",
"**/exo_scripts",
"**/.direnv",
"**/rust",
"**/.github",
]
stubPath = ".mlx_typings"
[[tool.basedpyright.executionEnvironments]]
@@ -109,17 +114,18 @@ root = "src"
[tool.uv]
required-version = ">=0.8.6"
prerelease = "allow"
environments = [
"sys_platform == 'darwin'",
"sys_platform == 'linux'",
]
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
###
# ruff configuration
###
[tool.ruff]
extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
extend-exclude = [
"shared/protobufs/**",
"*mlx_typings/**",
"rust/exo_pyo3_bindings/**",
]
[tool.ruff.lint]
extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
@@ -127,13 +133,7 @@ extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
[tool.pytest.ini_options]
pythonpath = "."
asyncio_mode = "auto"
markers = [
"slow: marks tests as slow (deselected by default)"
]
env = [
"EXO_TESTS=1"
]
markers = ["slow: marks tests as slow (deselected by default)"]
env = ["EXO_TESTS=1"]
addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
filterwarnings = [
"ignore:builtin type Swig:DeprecationWarning",
]
filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]
@@ -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
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
n_layers = 48
hidden_size = 3072
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 69593314272
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
n_layers = 48
hidden_size = 3072
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "6bit"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 100120675296
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
n_layers = 48
hidden_size = 3072
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 130648036320
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
n_layers = 48
hidden_size = 3072
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "bf16"
base_model = "Qwen3.5 122B A10B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 245125640160
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-27B-4bit"
n_layers = 64
hidden_size = 5120
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 27B"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 16054266848
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-27B-8bit"
n_layers = 64
hidden_size = 5120
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 27B"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 29500943328
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
n_layers = 24
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 2B"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 2662787264
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
n_layers = 40
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 20391405152
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
n_layers = 40
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 35B A3B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 37721130592
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
n_layers = 60
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 397B A17B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 223860768352
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
n_layers = 60
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "6bit"
base_model = "Qwen3.5 397B A17B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 322946674272
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
n_layers = 60
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 397B A17B"
capabilities = ["text", "thinking", "thinking_toggle"]
[storage_size]
in_bytes = 422032580192
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-9B-4bit"
n_layers = 32
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "4bit"
base_model = "Qwen3.5 9B"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 5950062560
@@ -0,0 +1,12 @@
model_id = "mlx-community/Qwen3.5-9B-8bit"
n_layers = 32
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "qwen"
quantization = "8bit"
base_model = "Qwen3.5 9B"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 10426433504
+15 -20
View File
@@ -5,7 +5,6 @@ edition = { workspace = true }
publish = false
[lib]
doctest = false
path = "src/lib.rs"
name = "exo_pyo3_bindings"
@@ -26,37 +25,33 @@ networking = { workspace = true }
# interop
pyo3 = { version = "0.27.2", features = [
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
# "nightly", # enables better-supported GIL integration
"experimental-async", # async support in #[pyfunction] & #[pymethods]
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
# "nightly", # enables better-supported GIL integration
"experimental-async", # async support in #[pyfunction] & #[pymethods]
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
# integrations with other libraries
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
# "ordered-float", "rust_decimal", "smallvec",
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
# integrations with other libraries
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
# "ordered-float", "rust_decimal", "smallvec",
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
] }
pyo3-stub-gen = { version = "0.17.2" }
pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime", "testing"] }
pyo3-async-runtimes = { version = "0.27.0", features = [
"attributes",
"tokio-runtime",
"testing",
] }
pyo3-log = "0.13.2"
# macro dependencies
extend = { workspace = true }
delegate = { workspace = true }
# async runtime
tokio = { workspace = true, features = ["full", "tracing"] }
futures-lite = { workspace = true }
# utility dependencies
util = { workspace = true }
# Tracing
log = { workspace = true }
env_logger = "0.11"
# Networking
libp2p = { workspace = true, features = ["full"] }
pin-project = "1.1.10"
-1
View File
@@ -1 +0,0 @@
TODO: do something here....
+34 -161
View File
@@ -2,7 +2,6 @@
# ruff: noqa: E501, F401
import builtins
import enum
import typing
@typing.final
@@ -11,138 +10,39 @@ class AllQueuesFullError(builtins.Exception):
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class ConnectionUpdate:
@property
def update_type(self) -> ConnectionUpdateType:
r"""
Whether this is a connection or disconnection event
"""
@property
def peer_id(self) -> PeerId:
r"""
Identity of the peer that we have connected to or disconnected from.
"""
@property
def remote_ipv4(self) -> builtins.str:
r"""
Remote connection's IPv4 address.
"""
@property
def remote_tcp_port(self) -> builtins.int:
r"""
Remote connection's TCP port.
"""
@typing.final
class Keypair:
r"""
Identity keypair of a node.
"""
@staticmethod
def generate_ed25519() -> Keypair:
def generate() -> Keypair:
r"""
Generate a new Ed25519 keypair.
"""
@staticmethod
def generate_ecdsa() -> Keypair:
def from_bytes(bytes: bytes) -> Keypair:
r"""
Generate a new ECDSA keypair.
"""
@staticmethod
def generate_secp256k1() -> Keypair:
r"""
Generate a new Secp256k1 keypair.
"""
@staticmethod
def from_protobuf_encoding(bytes: bytes) -> Keypair:
r"""
Decode a private key from a protobuf structure and parse it as a `Keypair`.
"""
@staticmethod
def rsa_from_pkcs8(bytes: bytes) -> Keypair:
r"""
Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
format (i.e. unencrypted) as defined in [RFC5208].
[RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
"""
@staticmethod
def secp256k1_from_der(bytes: bytes) -> Keypair:
r"""
Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
structure as defined in [RFC5915].
[RFC5915]: https://tools.ietf.org/html/rfc5915
"""
@staticmethod
def ed25519_from_bytes(bytes: bytes) -> Keypair: ...
def to_protobuf_encoding(self) -> bytes:
r"""
Encode a private key as protobuf structure.
"""
def to_peer_id(self) -> PeerId:
r"""
Convert the `Keypair` into the corresponding `PeerId`.
"""
@typing.final
class Multiaddr:
r"""
Representation of a Multiaddr.
"""
@staticmethod
def empty() -> Multiaddr:
r"""
Create a new, empty multiaddress.
"""
@staticmethod
def with_capacity(n: builtins.int) -> Multiaddr:
r"""
Create a new, empty multiaddress with the given capacity.
"""
@staticmethod
def from_bytes(bytes: bytes) -> Multiaddr:
r"""
Parse a `Multiaddr` value from its byte slice representation.
"""
@staticmethod
def from_string(string: builtins.str) -> Multiaddr:
r"""
Parse a `Multiaddr` value from its string representation.
"""
def len(self) -> builtins.int:
r"""
Return the length in bytes of this multiaddress.
"""
def is_empty(self) -> builtins.bool:
r"""
Returns true if the length of this multiaddress is 0.
Construct an Ed25519 keypair from secret key bytes
"""
def to_bytes(self) -> bytes:
r"""
Return a copy of this [`Multiaddr`]'s byte representation.
Get the secret key bytes underlying the keypair
"""
def to_string(self) -> builtins.str:
def to_node_id(self) -> builtins.str:
r"""
Convert a Multiaddr to a string.
Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
"""
@typing.final
class MessageTooLargeError(builtins.Exception):
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class NetworkingHandle:
def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
async def connection_update_recv(self) -> ConnectionUpdate:
r"""
Receives the next `ConnectionUpdate` from networking.
"""
async def connection_update_recv_many(self, limit: builtins.int) -> builtins.list[ConnectionUpdate]:
r"""
Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
will sleep until a `ConnectionUpdate`s is sent.
"""
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
r"""
Subscribe to a `GossipSub` topic.
@@ -161,18 +61,7 @@ class NetworkingHandle:
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
"""
async def gossipsub_recv(self) -> tuple[builtins.str, bytes]:
r"""
Receives the next message from the `GossipSub` network.
"""
async def gossipsub_recv_many(self, limit: builtins.int) -> builtins.list[tuple[builtins.str, bytes]]:
r"""
Receives at most `limit` messages from the `GossipSub` network and returns them.
For `limit = 0`, an empty collection of messages will be returned immediately.
For `limit > 0`, if there are no messages in the channel's queue this method
will sleep until a message is sent.
"""
async def recv(self) -> PyFromSwarm: ...
@typing.final
class NoPeersSubscribedToTopicError(builtins.Exception):
@@ -180,42 +69,26 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class PeerId:
r"""
Identifier of a peer of the network.
class PyFromSwarm:
@typing.final
class Connection(PyFromSwarm):
__match_args__ = ("peer_id", "connected",)
@property
def peer_id(self) -> builtins.str: ...
@property
def connected(self) -> builtins.bool: ...
def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
"""
@staticmethod
def random() -> PeerId:
r"""
Generates a random peer ID from a cryptographically secure PRNG.
This is useful for randomly walking on a DHT, or for testing purposes.
"""
@staticmethod
def from_bytes(bytes: bytes) -> PeerId:
r"""
Parses a `PeerId` from bytes.
"""
def to_bytes(self) -> bytes:
r"""
Returns a raw bytes representation of this `PeerId`.
"""
def to_base58(self) -> builtins.str:
r"""
Returns a base-58 encoded string of this `PeerId`.
"""
def __repr__(self) -> builtins.str: ...
def __str__(self) -> builtins.str: ...
@typing.final
class ConnectionUpdateType(enum.Enum):
r"""
Connection or disconnection event discriminant type.
"""
Connected = ...
Disconnected = ...
@typing.final
class Message(PyFromSwarm):
__match_args__ = ("origin", "topic", "data",)
@property
def origin(self) -> builtins.str: ...
@property
def topic(self) -> builtins.str: ...
@property
def data(self) -> bytes: ...
def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
...
+4 -7
View File
@@ -4,21 +4,18 @@ build-backend = "maturin"
[project]
name = "exo_pyo3_bindings"
version = "0.1.0"
version = "0.2.0"
description = "Add your description here"
readme = "README.md"
authors = [
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
{ name = "Evan Quiney", email = "evanev7@gmail.com" },
]
requires-python = ">=3.13"
dependencies = []
[dependency-groups]
dev = [
"exo_pyo3_bindings",
"pytest>=8.4.0",
"pytest-asyncio>=1.0.0",
]
dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
[tool.maturin]
#purelib = true
+6 -21
View File
@@ -1,37 +1,22 @@
//! SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
//!
use pin_project::pin_project;
//! See: <https://pyo3.rs/v0.27.2/async-await.html#detaching-from-the-interpreter-across-await>
use pyo3::prelude::*;
use std::{
future::Future,
pin::Pin,
pin::{Pin, pin},
task::{Context, Poll},
};
/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
#[pin_project]
#[repr(transparent)]
pub(crate) struct AllowThreads<F>(#[pin] F);
impl<F> AllowThreads<F>
where
Self: Future,
{
pub fn new(f: F) -> Self {
Self(f)
}
}
pub struct AllowThreads<F>(pub(crate) F);
impl<F> Future for AllowThreads<F>
where
F: Future + Send,
F: Future + Unpin + Send,
F::Output: Send,
{
type Output = F::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let waker = cx.waker();
Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker))))
Python::attach(|py| py.detach(|| pin!(&mut self.0).poll(&mut Context::from_waker(waker))))
}
}
+17 -129
View File
@@ -1,8 +1,6 @@
use crate::ext::ResultExt as _;
use libp2p::PeerId;
use libp2p::identity::Keypair;
use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _};
use pyo3::types::PyBytes;
use pyo3::types::{PyBytes, PyBytesMethods as _};
use pyo3::{Bound, PyResult, Python, pyclass, pymethods};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
@@ -18,142 +16,32 @@ pub struct PyKeypair(pub Keypair);
impl PyKeypair {
/// Generate a new Ed25519 keypair.
#[staticmethod]
fn generate_ed25519() -> Self {
fn generate() -> Self {
Self(Keypair::generate_ed25519())
}
/// Generate a new ECDSA keypair.
/// Construct an Ed25519 keypair from secret key bytes
#[staticmethod]
fn generate_ecdsa() -> Self {
Self(Keypair::generate_ecdsa())
}
/// Generate a new Secp256k1 keypair.
#[staticmethod]
fn generate_secp256k1() -> Self {
Self(Keypair::generate_secp256k1())
}
/// Decode a private key from a protobuf structure and parse it as a `Keypair`.
#[staticmethod]
fn from_protobuf_encoding(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::from_protobuf_encoding(&bytes).pyerr()?))
}
/// Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo`
/// format (i.e. unencrypted) as defined in [RFC5208].
///
/// [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5
#[staticmethod]
fn rsa_from_pkcs8(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::rsa_from_pkcs8(&mut bytes).pyerr()?))
}
/// Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey`
/// structure as defined in [RFC5915].
///
/// [RFC5915]: https://tools.ietf.org/html/rfc5915
#[staticmethod]
fn secp256k1_from_der(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::secp256k1_from_der(&mut bytes).pyerr()?))
}
#[staticmethod]
fn ed25519_from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let mut bytes = Vec::from(bytes.as_bytes());
Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?))
}
/// Encode a private key as protobuf structure.
fn to_protobuf_encoding<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
let bytes = self.0.to_protobuf_encoding().pyerr()?;
/// Get the secret key bytes underlying the keypair
fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
let bytes = self
.0
.clone()
.try_into_ed25519()
.pyerr()?
.secret()
.as_ref()
.to_vec();
Ok(PyBytes::new(py, &bytes))
}
/// Convert the `Keypair` into the corresponding `PeerId`.
fn to_peer_id(&self) -> PyPeerId {
PyPeerId(self.0.public().to_peer_id())
}
// /// Hidden constructor for pickling support. TODO: figure out how to do pickling...
// #[gen_stub(skip)]
// #[new]
// fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
// Self::from_protobuf_encoding(bytes)
// }
//
// #[gen_stub(skip)]
// fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
// *self = Self::from_protobuf_encoding(state)?;
// Ok(())
// }
//
// #[gen_stub(skip)]
// fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
// self.to_protobuf_encoding(py)
// }
//
// #[gen_stub(skip)]
// pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> {
// Ok((self.to_protobuf_encoding(py)?,))
// }
}
/// Identifier of a peer of the network.
///
/// The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer
/// as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md).
#[gen_stub_pyclass]
#[pyclass(name = "PeerId", frozen)]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct PyPeerId(pub PeerId);
#[gen_stub_pymethods]
#[pymethods]
#[allow(clippy::needless_pass_by_value)]
impl PyPeerId {
/// Generates a random peer ID from a cryptographically secure PRNG.
///
/// This is useful for randomly walking on a DHT, or for testing purposes.
#[staticmethod]
fn random() -> Self {
Self(PeerId::random())
}
/// Parses a `PeerId` from bytes.
#[staticmethod]
fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
let bytes = Vec::from(bytes.as_bytes());
Ok(Self(PeerId::from_bytes(&bytes).pyerr()?))
}
/// Returns a raw bytes representation of this `PeerId`.
fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> {
let bytes = self.0.to_bytes();
PyBytes::new(py, &bytes)
}
/// Returns a base-58 encoded string of this `PeerId`.
fn to_base58(&self) -> String {
self.0.to_base58()
}
fn __repr__(&self) -> String {
format!("PeerId({})", self.to_base58())
}
fn __str__(&self) -> String {
self.to_base58()
/// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
fn to_node_id(&self) -> String {
self.0.public().to_peer_id().to_base58()
}
}
pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyKeypair>()?;
m.add_class::<PyPeerId>()?;
Ok(())
}
+1 -166
View File
@@ -1,170 +1,5 @@
//! TODO: crate documentation
//!
//! this is here as a placeholder documentation
//!
//!
mod allow_threading;
mod ident;
mod networking;
use crate::ident::ident_submodule;
use crate::networking::networking_submodule;
use pyo3::prelude::PyModule;
use pyo3::{Bound, PyResult, pyclass, pymodule};
use pyo3_stub_gen::define_stub_info_gatherer;
/// Namespace for all the constants used by this crate.
pub(crate) mod r#const {
pub const MPSC_CHANNEL_SIZE: usize = 1024;
}
/// Namespace for crate-wide extension traits/methods
pub(crate) mod ext {
use crate::allow_threading::AllowThreads;
use extend::ext;
use pyo3::exceptions::{PyConnectionError, PyRuntimeError};
use pyo3::types::PyBytes;
use pyo3::{Py, PyErr, PyResult, Python};
use tokio::runtime::Runtime;
use tokio::sync::mpsc;
use tokio::sync::mpsc::error::TryRecvError;
use tokio::task::JoinHandle;
#[ext(pub, name = ByteArrayExt)]
impl [u8] {
fn pybytes(&self) -> Py<PyBytes> {
Python::attach(|py| PyBytes::new(py, self).unbind())
}
}
#[ext(pub, name = ResultExt)]
impl<T, E> Result<T, E>
where
E: ToString,
{
fn pyerr(self) -> PyResult<T> {
self.map_err(|e| PyRuntimeError::new_err(e.to_string()))
}
}
pub trait FutureExt: Future + Sized {
/// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
fn allow_threads_py(self) -> AllowThreads<Self>
where
AllowThreads<Self>: Future,
{
AllowThreads::new(self)
}
}
impl<T: Future> FutureExt for T {}
#[ext(pub, name = PyErrExt)]
impl PyErr {
fn receiver_channel_closed() -> Self {
PyConnectionError::new_err("Receiver channel closed unexpectedly")
}
}
#[ext(pub, name = PyResultExt)]
impl<T> PyResult<T> {
fn write_unraisable(self) -> Option<T> {
Python::attach(|py| self.write_unraisable_with(py))
}
fn write_unraisable_with(self, py: Python<'_>) -> Option<T> {
match self {
Ok(v) => Some(v),
Err(e) => {
// write error back to python
e.write_unraisable(py, None);
None
}
}
}
}
#[ext(pub, name = TokioRuntimeExt)]
impl Runtime {
fn spawn_with_scope<F>(&self, py: Python<'_>, future: F) -> PyResult<JoinHandle<F::Output>>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?;
Ok(self.spawn(pyo3_async_runtimes::tokio::scope(locals, future)))
}
}
#[ext(pub, name = TokioMpscSenderExt)]
impl<T> mpsc::Sender<T> {
/// Sends a value, waiting until there is capacity.
///
/// A successful send occurs when it is determined that the other end of the
/// channel has not hung up already. An unsuccessful send would be one where
/// the corresponding receiver has already been closed.
async fn send_py(&self, value: T) -> PyResult<()> {
self.send(value)
.await
.map_err(|_| PyErr::receiver_channel_closed())
}
}
#[ext(pub, name = TokioMpscReceiverExt)]
impl<T> mpsc::Receiver<T> {
/// Receives the next value for this receiver.
async fn recv_py(&mut self) -> PyResult<T> {
self.recv().await.ok_or_else(PyErr::receiver_channel_closed)
}
/// Receives at most `limit` values for this receiver and returns them.
///
/// For `limit = 0`, an empty collection of messages will be returned immediately.
/// For `limit > 0`, if there are no messages in the channel's queue this method
/// will sleep until a message is sent.
async fn recv_many_py(&mut self, limit: usize) -> PyResult<Vec<T>> {
// get updates from receiver channel
let mut updates = Vec::with_capacity(limit);
let received = self.recv_many(&mut updates, limit).await;
// if we received zero items, then the channel was unexpectedly closed
if limit != 0 && received == 0 {
return Err(PyErr::receiver_channel_closed());
}
Ok(updates)
}
/// Tries to receive the next value for this receiver.
fn try_recv_py(&mut self) -> PyResult<Option<T>> {
match self.try_recv() {
Ok(v) => Ok(Some(v)),
Err(TryRecvError::Empty) => Ok(None),
Err(TryRecvError::Disconnected) => Err(PyErr::receiver_channel_closed()),
}
}
}
}
/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule(name = "exo_pyo3_bindings")]
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
// install logger
pyo3_log::init();
// TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
// work with maturin, where the types generate correctly, in the right folder, without
// too many importing issues...
ident_submodule(m)?;
networking_submodule(m)?;
// top-level constructs
// TODO: ...
Ok(())
}
mod allow_threading;
define_stub_info_gatherer!(stub_info);
-571
View File
@@ -1,571 +0,0 @@
#![allow(
clippy::multiple_inherent_impl,
clippy::unnecessary_wraps,
clippy::unused_self,
clippy::needless_pass_by_value
)]
use crate::r#const::MPSC_CHANNEL_SIZE;
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
use crate::ident::{PyKeypair, PyPeerId};
use crate::pyclass;
use libp2p::futures::StreamExt as _;
use libp2p::gossipsub;
use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
use libp2p::swarm::SwarmEvent;
use networking::discovery;
use networking::swarm::create_swarm;
use pyo3::prelude::{PyModule, PyModuleMethods as _};
use pyo3::types::PyBytes;
use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymethods};
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
use std::net::IpAddr;
use tokio::sync::{Mutex, mpsc, oneshot};
mod exception {
use pyo3::types::PyTuple;
use pyo3::{exceptions::PyException, prelude::*};
use pyo3_stub_gen::derive::*;
#[gen_stub_pyclass]
#[pyclass(frozen, extends=PyException, name="NoPeersSubscribedToTopicError")]
pub struct PyNoPeersSubscribedToTopicError {}
impl PyNoPeersSubscribedToTopicError {
const MSG: &'static str = "\
No peers are currently subscribed to receive messages on this topic. \
Wait for peers to subscribe or check your network connectivity.";
/// Creates a new [ `PyErr` ] of this type.
///
/// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3"
pub(crate) fn new_err() -> PyErr {
PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyNoPeersSubscribedToTopicError {
#[new]
#[pyo3(signature = (*args))]
#[allow(unused_variables)]
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
Self {}
}
fn __repr__(&self) -> String {
format!("PeerId(\"{}\")", Self::MSG)
}
fn __str__(&self) -> String {
Self::MSG.to_string()
}
}
#[gen_stub_pyclass]
#[pyclass(frozen, extends=PyException, name="AllQueuesFullError")]
pub struct PyAllQueuesFullError {}
impl PyAllQueuesFullError {
const MSG: &'static str =
"All libp2p peers are unresponsive, resend the message or reconnect.";
/// Creates a new [ `PyErr` ] of this type.
///
/// [`PyErr`] : https://docs.rs/pyo3/latest/pyo3/struct.PyErr.html "PyErr in pyo3"
pub(crate) fn new_err() -> PyErr {
PyErr::new::<Self, _>(()) // TODO: check if this needs to be replaced???
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyAllQueuesFullError {
#[new]
#[pyo3(signature = (*args))]
#[allow(unused_variables)]
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
Self {}
}
fn __repr__(&self) -> String {
format!("PeerId(\"{}\")", Self::MSG)
}
fn __str__(&self) -> String {
Self::MSG.to_string()
}
}
}
/// Connection or disconnection event discriminant type.
#[gen_stub_pyclass_enum]
#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
#[derive(Debug, Clone, PartialEq)]
enum PyConnectionUpdateType {
Connected = 0,
Disconnected,
}
#[gen_stub_pyclass]
#[pyclass(frozen, name = "ConnectionUpdate")]
#[derive(Debug, Clone)]
struct PyConnectionUpdate {
/// Whether this is a connection or disconnection event
#[pyo3(get)]
update_type: PyConnectionUpdateType,
/// Identity of the peer that we have connected to or disconnected from.
#[pyo3(get)]
peer_id: PyPeerId,
/// Remote connection's IPv4 address.
#[pyo3(get)]
remote_ipv4: String,
/// Remote connection's TCP port.
#[pyo3(get)]
remote_tcp_port: u16,
}
enum ToTask {
GossipsubSubscribe {
topic: String,
result_tx: oneshot::Sender<PyResult<bool>>,
},
GossipsubUnsubscribe {
topic: String,
result_tx: oneshot::Sender<bool>,
},
GossipsubPublish {
topic: String,
data: Vec<u8>,
result_tx: oneshot::Sender<PyResult<MessageId>>,
},
}
#[allow(clippy::enum_glob_use)]
async fn networking_task(
mut swarm: networking::swarm::Swarm,
mut to_task_rx: mpsc::Receiver<ToTask>,
connection_update_tx: mpsc::Sender<PyConnectionUpdate>,
gossipsub_message_tx: mpsc::Sender<(String, Vec<u8>)>,
) {
use SwarmEvent::*;
use ToTask::*;
use networking::swarm::BehaviourEvent::*;
log::info!("RUST: networking task started");
loop {
tokio::select! {
message = to_task_rx.recv() => {
// handle closed channel
let Some(message) = message else {
log::info!("RUST: channel closed");
break;
};
// dispatch incoming messages
match message {
GossipsubSubscribe { topic, result_tx } => {
// try to subscribe
let result = swarm.behaviour_mut()
.gossipsub.subscribe(&IdentTopic::new(topic));
// send response oneshot
if let Err(e) = result_tx.send(result.pyerr()) {
log::error!("RUST: could not subscribe to gossipsub topic since channel already closed: {e:?}");
continue;
}
}
GossipsubUnsubscribe { topic, result_tx } => {
// try to unsubscribe from the topic
let result = swarm.behaviour_mut()
.gossipsub.unsubscribe(&IdentTopic::new(topic));
// send response oneshot (or exit if connection closed)
if let Err(e) = result_tx.send(result) {
log::error!("RUST: could not unsubscribe from gossipsub topic since channel already closed: {e:?}");
continue;
}
}
GossipsubPublish { topic, data, result_tx } => {
// try to publish the data -> catch NoPeersSubscribedToTopic error & convert to correct exception
let result = swarm.behaviour_mut().gossipsub.publish(
IdentTopic::new(topic), data);
let pyresult: PyResult<MessageId> = if let Err(PublishError::NoPeersSubscribedToTopic) = result {
Err(exception::PyNoPeersSubscribedToTopicError::new_err())
} else if let Err(PublishError::AllQueuesFull(_)) = result {
Err(exception::PyAllQueuesFullError::new_err())
} else {
result.pyerr()
};
// send response oneshot (or exit if connection closed)
if let Err(e) = result_tx.send(pyresult) {
log::error!("RUST: could not publish gossipsub message since channel already closed: {e:?}");
continue;
}
}
}
}
// architectural solution to this problem:
// create keep_alive behavior who's job it is to dial peers discovered by mDNS (and drop when expired)
// -> it will emmit TRUE connected/disconnected events consumable elsewhere
//
// gossipsub will feed off-of dial attempts created by networking, and that will bootstrap its' peers list
// then for actual communication it will dial those peers if need-be
swarm_event = swarm.select_next_some() => {
match swarm_event {
Behaviour(Gossipsub(gossipsub::Event::Message {
message: Message {
topic,
data,
..
},
..
})) => {
// topic-ID is just the topic hash!!! (since we used identity hasher)
let message = (topic.into_string(), data);
// send incoming message to channel (or exit if connection closed)
if let Err(e) = gossipsub_message_tx.send(message).await {
log::error!("RUST: could not send incoming gossipsub message since channel already closed: {e}");
continue;
}
},
Behaviour(Discovery(discovery::Event::ConnectionEstablished { peer_id, remote_ip, remote_tcp_port, .. })) => {
// grab IPv4 string
let remote_ipv4 = match remote_ip {
IpAddr::V4(ip) => ip.to_string(),
IpAddr::V6(ip) => {
log::warn!("RUST: ignoring connection to IPv6 address: {ip}");
continue;
}
};
// send connection event to channel (or exit if connection closed)
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
update_type: PyConnectionUpdateType::Connected,
peer_id: PyPeerId(peer_id),
remote_ipv4,
remote_tcp_port,
}).await {
log::error!("RUST: could not send connection update since channel already closed: {e}");
continue;
}
},
Behaviour(Discovery(discovery::Event::ConnectionClosed { peer_id, remote_ip, remote_tcp_port, .. })) => {
// grab IPv4 string
let remote_ipv4 = match remote_ip {
IpAddr::V4(ip) => ip.to_string(),
IpAddr::V6(ip) => {
log::warn!("RUST: ignoring disconnection from IPv6 address: {ip}");
continue;
}
};
// send disconnection event to channel (or exit if connection closed)
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
update_type: PyConnectionUpdateType::Disconnected,
peer_id: PyPeerId(peer_id),
remote_ipv4,
remote_tcp_port,
}).await {
log::error!("RUST: could not send connection update since channel already closed: {e}");
continue;
}
},
e => {
log::info!("RUST: other event {e:?}");
}
}
}
}
}
log::info!("RUST: networking task stopped");
}
#[gen_stub_pyclass]
#[pyclass(name = "NetworkingHandle")]
#[derive(Debug)]
struct PyNetworkingHandle {
// channels
to_task_tx: Option<mpsc::Sender<ToTask>>,
connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
}
impl Drop for PyNetworkingHandle {
fn drop(&mut self) {
// TODO: may or may not need to await a "kill-signal" oneshot channel message,
// to ensure that the networking task is done BEFORE exiting the clear function...
// but this may require GIL?? and it may not be safe to call GIL here??
self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
}
}
#[allow(clippy::expect_used)]
impl PyNetworkingHandle {
fn new(
to_task_tx: mpsc::Sender<ToTask>,
connection_update_rx: mpsc::Receiver<PyConnectionUpdate>,
gossipsub_message_rx: mpsc::Receiver<(String, Vec<u8>)>,
) -> Self {
Self {
to_task_tx: Some(to_task_tx),
connection_update_rx: Mutex::new(connection_update_rx),
gossipsub_message_rx: Mutex::new(gossipsub_message_rx),
}
}
const fn to_task_tx(&self) -> &mpsc::Sender<ToTask> {
self.to_task_tx
.as_ref()
.expect("The sender should only be None after de-initialization.")
}
}
#[gen_stub_pymethods]
#[pymethods]
impl PyNetworkingHandle {
// NOTE: `async fn`s here that use `.await` will wrap the future in `.allow_threads_py()`
// immediately beforehand to release the interpreter.
// SEE: https://pyo3.rs/v0.26.0/async-await.html#detaching-from-the-interpreter-across-await
// ---- Lifecycle management methods ----
#[new]
fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
use pyo3_async_runtimes::tokio::get_runtime;
// create communication channels
let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
let (connection_update_tx, connection_update_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
let (gossipsub_message_tx, gossipsub_message_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
// get identity
let identity = identity.borrow().0.clone();
// create networking swarm (within tokio context!! or it crashes)
let swarm = get_runtime()
.block_on(async { create_swarm(identity) })
.pyerr()?;
// spawn tokio task running the networking logic
get_runtime().spawn(async move {
networking_task(
swarm,
to_task_rx,
connection_update_tx,
gossipsub_message_tx,
)
.await;
});
Ok(Self::new(
to_task_tx,
connection_update_rx,
gossipsub_message_rx,
))
}
#[gen_stub(skip)]
const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
Ok(()) // This is needed purely so `__clear__` can work
}
#[gen_stub(skip)]
fn __clear__(&mut self) {
// TODO: may or may not need to await a "kill-signal" oneshot channel message,
// to ensure that the networking task is done BEFORE exiting the clear function...
// but this may require GIL?? and it may not be safe to call GIL here??
self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
}
// ---- Connection update receiver methods ----
/// Receives the next `ConnectionUpdate` from networking.
async fn connection_update_recv(&self) -> PyResult<PyConnectionUpdate> {
self.connection_update_rx
.lock()
.allow_threads_py() // allow-threads-aware async call
.await
.recv_py()
.allow_threads_py() // allow-threads-aware async call
.await
}
/// Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
///
/// For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
/// For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
/// will sleep until a `ConnectionUpdate`s is sent.
async fn connection_update_recv_many(&self, limit: usize) -> PyResult<Vec<PyConnectionUpdate>> {
self.connection_update_rx
.lock()
.allow_threads_py() // allow-threads-aware async call
.await
.recv_many_py(limit)
.allow_threads_py() // allow-threads-aware async call
.await
}
// TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
// so its too dangerous to expose just yet. figure out a better semantics for handling this,
// so things don't randomly block
// /// Tries to receive the next `ConnectionUpdate` from networking.
// fn connection_update_try_recv(&self) -> PyResult<Option<PyConnectionUpdate>> {
// self.connection_update_rx.blocking_lock().try_recv_py()
// }
//
// /// Checks if the `ConnectionUpdate` channel is empty.
// fn connection_update_is_empty(&self) -> bool {
// self.connection_update_rx.blocking_lock().is_empty()
// }
//
// /// Returns the number of `ConnectionUpdate`s in the channel.
// fn connection_update_len(&self) -> usize {
// self.connection_update_rx.blocking_lock().len()
// }
// ---- Gossipsub management methods ----
/// Subscribe to a `GossipSub` topic.
///
/// Returns `True` if the subscription worked. Returns `False` if we were already subscribed.
async fn gossipsub_subscribe(&self, topic: String) -> PyResult<bool> {
let (tx, rx) = oneshot::channel();
// send off request to subscribe
self.to_task_tx()
.send_py(ToTask::GossipsubSubscribe {
topic,
result_tx: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & return any errors
rx.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())?
}
/// Unsubscribes from a `GossipSub` topic.
///
/// Returns `True` if we were subscribed to this topic. Returns `False` if we were not subscribed.
async fn gossipsub_unsubscribe(&self, topic: String) -> PyResult<bool> {
let (tx, rx) = oneshot::channel();
// send off request to unsubscribe
self.to_task_tx()
.send_py(ToTask::GossipsubUnsubscribe {
topic,
result_tx: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & convert any errors
rx.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())
}
/// Publishes a message with multiple topics to the `GossipSub` network.
///
/// If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
async fn gossipsub_publish(&self, topic: String, data: Py<PyBytes>) -> PyResult<()> {
let (tx, rx) = oneshot::channel();
// send off request to subscribe
let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
self.to_task_tx()
.send_py(ToTask::GossipsubPublish {
topic,
data,
result_tx: tx,
})
.allow_threads_py() // allow-threads-aware async call
.await?;
// wait for response & return any errors => ignore messageID for now!!!
let _ = rx
.allow_threads_py() // allow-threads-aware async call
.await
.map_err(|_| PyErr::receiver_channel_closed())??;
Ok(())
}
// ---- Gossipsub message receiver methods ----
/// Receives the next message from the `GossipSub` network.
async fn gossipsub_recv(&self) -> PyResult<(String, Py<PyBytes>)> {
self.gossipsub_message_rx
.lock()
.allow_threads_py() // allow-threads-aware async call
.await
.recv_py()
.allow_threads_py() // allow-threads-aware async call
.await
.map(|(t, d)| (t, d.pybytes()))
}
/// Receives at most `limit` messages from the `GossipSub` network and returns them.
///
/// For `limit = 0`, an empty collection of messages will be returned immediately.
/// For `limit > 0`, if there are no messages in the channel's queue this method
/// will sleep until a message is sent.
async fn gossipsub_recv_many(&self, limit: usize) -> PyResult<Vec<(String, Py<PyBytes>)>> {
Ok(self
.gossipsub_message_rx
.lock()
.allow_threads_py() // allow-threads-aware async call
.await
.recv_many_py(limit)
.allow_threads_py() // allow-threads-aware async call
.await?
.into_iter()
.map(|(t, d)| (t, d.pybytes()))
.collect())
}
// TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
// so its too dangerous to expose just yet. figure out a better semantics for handling this,
// so things don't randomly block
// /// Tries to receive the next message from the `GossipSub` network.
// fn gossipsub_try_recv(&self) -> PyResult<Option<(String, Py<PyBytes>)>> {
// Ok(self
// .gossipsub_message_rx
// .blocking_lock()
// .try_recv_py()?
// .map(|(t, d)| (t, d.pybytes())))
// }
//
// /// Checks if the `GossipSub` message channel is empty.
// fn gossipsub_is_empty(&self) -> bool {
// self.gossipsub_message_rx.blocking_lock().is_empty()
// }
//
// /// Returns the number of `GossipSub` messages in the channel.
// fn gossipsub_len(&self) -> usize {
// self.gossipsub_message_rx.blocking_lock().len()
// }
}
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
m.add_class::<exception::PyAllQueuesFullError>()?;
m.add_class::<PyConnectionUpdateType>()?;
m.add_class::<PyConnectionUpdate>()?;
m.add_class::<PyConnectionUpdateType>()?;
m.add_class::<PyNetworkingHandle>()?;
Ok(())
}
-54
View File
@@ -1,54 +0,0 @@
#[cfg(test)]
mod tests {
use core::mem::drop;
use core::option::Option::Some;
use core::time::Duration;
use tokio;
use tokio::sync::mpsc;
#[tokio::test]
async fn test_drop_channel() {
struct Ping;
let (tx, mut rx) = mpsc::channel::<Ping>(10);
let _ = tokio::spawn(async move {
println!("TASK: entered");
loop {
tokio::select! {
result = rx.recv() => {
match result {
Some(_) => {
println!("TASK: pinged");
}
None => {
println!("TASK: closing channel");
break;
}
}
}
_ = tokio::time::sleep(Duration::from_secs_f32(0.1)) => {
println!("TASK: heartbeat");
}
}
}
println!("TASK: exited");
});
let tx2 = tx.clone();
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
tx.send(Ping).await.expect("Should not fail");
drop(tx);
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
tx2.send(Ping).await.expect("Should not fail");
drop(tx2);
tokio::time::sleep(Duration::from_secs_f32(0.11)).await;
}
}
+15 -13
View File
@@ -1,16 +1,20 @@
import asyncio
import pytest
from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
from exo_pyo3_bindings import (
Keypair,
NetworkingHandle,
NoPeersSubscribedToTopicError,
PyFromSwarm,
)
@pytest.mark.asyncio
async def test_sleep_on_multiple_items() -> None:
print("PYTHON: starting handle")
h = NetworkingHandle(Keypair.generate_ed25519())
h = NetworkingHandle(Keypair.generate())
ct = asyncio.create_task(_await_cons(h))
mt = asyncio.create_task(_await_msg(h))
rt = asyncio.create_task(_await_recv(h))
# sleep for 4 ticks
for i in range(4):
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
print("caught it", e)
async def _await_cons(h: NetworkingHandle):
async def _await_recv(h: NetworkingHandle):
while True:
c = await h.connection_update_recv()
print(f"PYTHON: connection update: {c}")
async def _await_msg(h: NetworkingHandle):
while True:
m = await h.gossipsub_recv()
print(f"PYTHON: message: {m}")
event = await h.recv()
match event:
case PyFromSwarm.Connection() as c:
print(f"PYTHON: connection update: {c}")
case PyFromSwarm.Message() as m:
print(f"PYTHON: message: {m}")
+5 -9
View File
@@ -5,7 +5,6 @@ edition = { workspace = true }
publish = false
[lib]
doctest = false
name = "networking"
path = "src/lib.rs"
@@ -13,25 +12,22 @@ path = "src/lib.rs"
workspace = true
[dependencies]
# datastructures
either = { workspace = true }
# macro dependencies
extend = { workspace = true }
delegate = { workspace = true }
# async
tokio = { workspace = true, features = ["full"] }
futures-lite = { workspace = true }
futures-timer = { workspace = true }
# utility dependencies
util = { workspace = true }
tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
keccak-const = { workspace = true }
pin-project = "1.1.10"
# tracing/logging
log = { workspace = true }
iroh = "0.96.1"
iroh-gossip = "0.96.0"
iroh-blobs = "0.98.0"
iroh-docs = "0.96.0"
# networking
libp2p = { workspace = true, features = ["full"] }
-65
View File
@@ -1,7 +1,3 @@
use futures_lite::StreamExt;
use libp2p::{gossipsub, identity, swarm::SwarmEvent};
use networking::{discovery, swarm};
use tokio::{io, io::AsyncBufReadExt as _, select};
use tracing_subscriber::EnvFilter;
use tracing_subscriber::filter::LevelFilter;
@@ -10,65 +6,4 @@ async fn main() {
let _ = tracing_subscriber::fmt()
.with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
.try_init();
// Configure swarm
let mut swarm =
swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
// Create a Gossipsub topic & subscribe
let topic = gossipsub::IdentTopic::new("test-net");
swarm
.behaviour_mut()
.gossipsub
.subscribe(&topic)
.expect("Subscribing to topic failed");
// Read full lines from stdin
let mut stdin = io::BufReader::new(io::stdin()).lines();
println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
// Kick it off
loop {
select! {
// on gossipsub outgoing
Ok(Some(line)) = stdin.next_line() => {
if let Err(e) = swarm
.behaviour_mut().gossipsub
.publish(topic.clone(), line.as_bytes()) {
println!("Publish error: {e:?}");
}
}
event = swarm.next() => match event {
// on gossipsub incoming
Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
propagation_source: peer_id,
message_id: id,
message,
}))) => println!(
"\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
String::from_utf8_lossy(&message.data),
),
// on discovery
Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
discovery::Event::ConnectionEstablished {
peer_id, connection_id, remote_ip, remote_tcp_port
} => {
println!("\n\nConnected to: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
}
discovery::Event::ConnectionClosed {
peer_id, connection_id, remote_ip, remote_tcp_port
} => {
eprintln!("\n\nDisconnected from: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
}
}
// ignore outgoing errors: those are normal
e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
// otherwise log any other event
e => { log::info!("Other event {e:?}"); }
}
}
}
}
-44
View File
@@ -1,44 +0,0 @@
https://github.com/ml-explore/mlx/commit/3fe98bacc7640d857acf3539f1d21b47a32e5609
^raw sockets distributed -> `<net/ndrv.h>` -> https://newosxbook.com/code/xnu-3247.1.106/bsd/net/ndrv.h.auto.html
--> header file for a networking component found in the macOS kernel (XNU) that defines structures for network device driver registration, specifically the ndrv_demux_desc and ndrv_protocol_desc structures used for demultiplexing protocol data at the network interface level. It specifies how to describe protocol data, such as an Ethernet type or a SNAP header, and how to associate these descriptions with a specific protocol family to receive matching packets.
--> Used to bind an NDRV socket so that packets that match given protocol demux descriptions can be received.
--> An NDRV socket is a special kind of socket in the Darwin/macOS operating system's XNU kernel, used for low-level network packet manipulation and binding to specific protocols for packet processing. It allows user-space applications or drivers to directly write Layer 2 (L2) network packets or interact with the network stack at a lower level, often by binding to protocol descriptors like the ndrv_protocol_desc. This type of socket is used for functions such as capturing and injecting packets, especially in network infrastructure software like routers or for kernel-level network monitoring and security tools.
--> also called PF_NDRV sockets --> https://newosxbook.com/bonus/vol1ch16.html
----> they are conceptually similar to https://scapy.disruptivelabs.in/networking/socket-interface PF_RAW or PF_PACKET
https://stackoverflow.com/questions/17169298/af-packet-on-osx
^AF_PACKET duplicates the packets as soon as it receives them from the physical layer (for incoming packets) or just before sending them out to the physical layer (for outgoing packets). -> this is on Linux only
^it doesn't exist on OS X so you can use /dev/bpfX (Berkeley Packet Filter) for sniffing
https://www.unix.com/man_page/mojave/4/ip/
^OS X manpages for IP
https://developer.apple.com/documentation/kernel/implementing_drivers_system_extensions_and_kexts
^driver kit, system extensions & kexts for macOS
----
To set up a Linux system to use a Thunderbolt connection as a network device, connect the two computers with a Thunderbolt cable, load the thunderbolt-net kernel module (usually automatic but modprobe is an option for manual loading), and then the operating system will create virtual Ethernet interfaces (e.g., thunderbolt0) for networking. You can then use standard tools like ifconfig or your desktop environment's network manager to configure these new interfaces for a link-local network.
--> https://gist.github.com/geosp/80fbd39e617b7d1d9421683df4ea224a
----> here is a guide on how to set up thunderbolt-ethernet on linux
----> I may be able to steal the thunderbolt-net code ideas to implement a kernel module for MacOS
https://chatgpt.com/s/t_68af8e41a8548191993281a014f846a7
^GPT discussion about making socket interface
https://chatgpt.com/s/t_68afb798a85c8191973c02a0fa7a48a3 --> link-local address,,??
https://chatgpt.com/s/t_68afb02987e08191b2b0044d3667ece2
^GPT discussion about accessing TB on MacOS low level interactions
--------------------------------
https://www.intel.com/content/www/us/en/support/articles/000098893/software.html
^Thunderbolt Share & Thunderbolt Networking Mode => intel's equivalent of thunderbolt bridge
---------------------------------
https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/
-->fake ethernet devices on MacOS -> omg??? we can detect thunderbolt bridge, then bind to it, then re-expose it as fake ethernet??
-->ps: https://chatgpt.com/s/t_68afb2b25fb881919526763fb5d7359c, AF/PF_NDRV are one and the same!!!
-->https://github.com/zerotier/ZeroTierOne/blob/dev/osdep/MacEthernetTapAgent.c
-382
View File
@@ -1,382 +0,0 @@
use crate::ext::MultiaddrExt;
use delegate::delegate;
use either::Either;
use futures_lite::FutureExt;
use futures_timer::Delay;
use libp2p::core::transport::PortUse;
use libp2p::core::{ConnectedPoint, Endpoint};
use libp2p::swarm::behaviour::ConnectionEstablished;
use libp2p::swarm::dial_opts::DialOpts;
use libp2p::swarm::{
CloseConnection, ConnectionClosed, ConnectionDenied, ConnectionHandler,
ConnectionHandlerSelect, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
THandlerOutEvent, ToSwarm, dummy,
};
use libp2p::{Multiaddr, PeerId, identity, mdns};
use std::collections::{BTreeSet, HashMap};
use std::convert::Infallible;
use std::io;
use std::net::IpAddr;
use std::task::{Context, Poll};
use std::time::Duration;
use util::wakerdeque::WakerDeque;
const RETRY_CONNECT_INTERVAL: Duration = Duration::from_secs(5);
mod managed {
use libp2p::swarm::NetworkBehaviour;
use libp2p::{identity, mdns, ping};
use std::io;
use std::time::Duration;
const MDNS_RECORD_TTL: Duration = Duration::from_secs(2_500);
const MDNS_QUERY_INTERVAL: Duration = Duration::from_secs(1_500);
const PING_TIMEOUT: Duration = Duration::from_millis(2_500);
const PING_INTERVAL: Duration = Duration::from_millis(2_500);
#[derive(NetworkBehaviour)]
pub struct Behaviour {
mdns: mdns::tokio::Behaviour,
ping: ping::Behaviour,
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
Ok(Self {
mdns: mdns_behaviour(keypair)?,
ping: ping_behaviour(),
})
}
}
fn mdns_behaviour(keypair: &identity::Keypair) -> io::Result<mdns::tokio::Behaviour> {
use mdns::{Config, tokio};
// mDNS config => enable IPv6
let mdns_config = Config {
ttl: MDNS_RECORD_TTL,
query_interval: MDNS_QUERY_INTERVAL,
// enable_ipv6: true, // TODO: for some reason, TCP+mDNS don't work well with ipv6?? figure out how to make work
..Default::default()
};
let mdns_behaviour = tokio::Behaviour::new(mdns_config, keypair.public().to_peer_id());
Ok(mdns_behaviour?)
}
fn ping_behaviour() -> ping::Behaviour {
ping::Behaviour::new(
ping::Config::new()
.with_timeout(PING_TIMEOUT)
.with_interval(PING_INTERVAL),
)
}
}
/// Events for when a listening connection is truly established and truly closed.
#[derive(Debug, Clone)]
pub enum Event {
ConnectionEstablished {
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
},
ConnectionClosed {
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
},
}
/// Discovery behavior that wraps mDNS to produce truly discovered durable peer-connections.
///
/// The behaviour operates as such:
/// 1) All true (listening) connections/disconnections are tracked, emitting corresponding events
/// to the swarm.
/// 1) mDNS discovered/expired peers are tracked; discovered but not connected peers are dialed
/// immediately, and expired but connected peers are disconnected from immediately.
/// 2) Every fixed interval: discovered but not connected peers are dialed, and expired but
/// connected peers are disconnected from.
pub struct Behaviour {
// state-tracking for managed behaviors & mDNS-discovered peers
managed: managed::Behaviour,
mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
retry_delay: Delay, // retry interval
// pending events to emmit => waker-backed Deque to control polling
pending_events: WakerDeque<ToSwarm<Event, Infallible>>,
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
Ok(Self {
managed: managed::Behaviour::new(keypair)?,
mdns_discovered: HashMap::new(),
retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
pending_events: WakerDeque::new(),
})
}
fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
self.pending_events.push_back(ToSwarm::Dial {
opts: DialOpts::peer_id(peer_id).addresses(vec![addr]).build(),
})
}
fn close_connection(&mut self, peer_id: PeerId, connection: ConnectionId) {
// push front to make this IMMEDIATE
self.pending_events.push_front(ToSwarm::CloseConnection {
peer_id,
connection: CloseConnection::One(connection),
})
}
fn handle_mdns_discovered(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
for (p, ma) in peers {
self.dial(p, ma.clone()); // always connect
// get peer's multi-addresses or insert if missing
let Some(mas) = self.mdns_discovered.get_mut(&p) else {
self.mdns_discovered.insert(p, BTreeSet::from([ma]));
continue;
};
// multiaddress should never already be present - else something has gone wrong
let is_new_addr = mas.insert(ma);
assert!(is_new_addr, "cannot discover a discovered peer");
}
}
fn handle_mdns_expired(&mut self, peers: Vec<(PeerId, Multiaddr)>) {
for (p, ma) in peers {
// at this point, we *must* have the peer
let mas = self
.mdns_discovered
.get_mut(&p)
.expect("nonexistent peer cannot expire");
// at this point, we *must* have the multiaddress
let was_present = mas.remove(&ma);
assert!(was_present, "nonexistent multiaddress cannot expire");
// if empty, remove the peer-id entirely
if mas.is_empty() {
self.mdns_discovered.remove(&p);
}
}
}
fn on_connection_established(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
) {
// send out connected event
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::ConnectionEstablished {
peer_id,
connection_id,
remote_ip,
remote_tcp_port,
}));
}
fn on_connection_closed(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
remote_ip: IpAddr,
remote_tcp_port: u16,
) {
// send out disconnected event
self.pending_events
.push_back(ToSwarm::GenerateEvent(Event::ConnectionClosed {
peer_id,
connection_id,
remote_ip,
remote_tcp_port,
}));
}
}
impl NetworkBehaviour for Behaviour {
type ConnectionHandler =
ConnectionHandlerSelect<dummy::ConnectionHandler, THandler<managed::Behaviour>>;
type ToSwarm = Event;
// simply delegate to underlying mDNS behaviour
delegate! {
to self.managed {
fn handle_pending_inbound_connection(&mut self, connection_id: ConnectionId, local_addr: &Multiaddr, remote_addr: &Multiaddr) -> Result<(), ConnectionDenied>;
fn handle_pending_outbound_connection(&mut self, connection_id: ConnectionId, maybe_peer: Option<PeerId>, addresses: &[Multiaddr], effective_role: Endpoint) -> Result<Vec<Multiaddr>, ConnectionDenied>;
}
}
fn handle_established_inbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
local_addr: &Multiaddr,
remote_addr: &Multiaddr,
) -> Result<THandler<Self>, ConnectionDenied> {
Ok(ConnectionHandler::select(
dummy::ConnectionHandler,
self.managed.handle_established_inbound_connection(
connection_id,
peer,
local_addr,
remote_addr,
)?,
))
}
#[allow(clippy::needless_question_mark)]
fn handle_established_outbound_connection(
&mut self,
connection_id: ConnectionId,
peer: PeerId,
addr: &Multiaddr,
role_override: Endpoint,
port_use: PortUse,
) -> Result<THandler<Self>, ConnectionDenied> {
Ok(ConnectionHandler::select(
dummy::ConnectionHandler,
self.managed.handle_established_outbound_connection(
connection_id,
peer,
addr,
role_override,
port_use,
)?,
))
}
fn on_connection_handler_event(
&mut self,
peer_id: PeerId,
connection_id: ConnectionId,
event: THandlerOutEvent<Self>,
) {
match event {
Either::Left(ev) => libp2p::core::util::unreachable(ev),
Either::Right(ev) => {
self.managed
.on_connection_handler_event(peer_id, connection_id, ev)
}
}
}
// hook into these methods to drive behavior
fn on_swarm_event(&mut self, event: FromSwarm) {
self.managed.on_swarm_event(event); // let mDNS handle swarm events
// handle swarm events to update internal state:
match event {
FromSwarm::ConnectionEstablished(ConnectionEstablished {
peer_id,
connection_id,
endpoint,
..
}) => {
let remote_address = match endpoint {
ConnectedPoint::Dialer { address, .. } => address,
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
};
if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
// handle connection established event which is filtered correctly
self.on_connection_established(peer_id, connection_id, ip, port)
}
}
FromSwarm::ConnectionClosed(ConnectionClosed {
peer_id,
connection_id,
endpoint,
..
}) => {
let remote_address = match endpoint {
ConnectedPoint::Dialer { address, .. } => address,
ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
};
if let Some((ip, port)) = remote_address.try_to_tcp_addr() {
// handle connection closed event which is filtered correctly
self.on_connection_closed(peer_id, connection_id, ip, port)
}
}
// since we are running TCP/IP transport layer, we are assuming that
// no address changes can occur, hence encountering one is a fatal error
FromSwarm::AddressChange(a) => {
unreachable!("unhandlable: address change encountered: {:?}", a)
}
_ => {}
}
}
fn poll(&mut self, cx: &mut Context) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
// delegate to managed behaviors for any behaviors they need to perform
match self.managed.poll(cx) {
Poll::Ready(ToSwarm::GenerateEvent(e)) => {
match e {
// handle discovered and expired events from mDNS
managed::BehaviourEvent::Mdns(e) => match e.clone() {
mdns::Event::Discovered(peers) => {
self.handle_mdns_discovered(peers);
}
mdns::Event::Expired(peers) => {
self.handle_mdns_expired(peers);
}
},
// handle ping events => if error then disconnect
managed::BehaviourEvent::Ping(e) => {
if let Err(_) = e.result {
self.close_connection(e.peer, e.connection.clone())
}
}
}
// since we just consumed an event, we should immediately wake just in case
// there are more events to come where that came from
cx.waker().wake_by_ref();
}
// forward any other mDNS event to the swarm or its connection handler(s)
Poll::Ready(e) => {
return Poll::Ready(
e.map_out(|_| unreachable!("events returning to swarm already handled"))
.map_in(Either::Right),
);
}
Poll::Pending => {}
}
// retry connecting to all mDNS peers periodically (fails safely if already connected)
if self.retry_delay.poll(cx).is_ready() {
for (p, mas) in self.mdns_discovered.clone() {
for ma in mas {
self.dial(p, ma)
}
}
self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
}
// send out any pending events from our own service
if let Some(e) = self.pending_events.pop_front(cx) {
return Poll::Ready(e.map_in(Either::Left));
}
// wait for pending events
Poll::Pending
}
}
-44
View File
@@ -1,44 +0,0 @@
//! TODO: crate documentation
//!
//! this is here as a placeholder documentation
//!
//!
pub mod discovery;
pub mod swarm;
/// Namespace for all the type/trait aliases used by this crate.
pub(crate) mod alias {
use std::error::Error;
pub type AnyError = Box<dyn Error + Send + Sync + 'static>;
pub type AnyResult<T> = Result<T, AnyError>;
}
/// Namespace for crate-wide extension traits/methods
pub(crate) mod ext {
use extend::ext;
use libp2p::Multiaddr;
use libp2p::multiaddr::Protocol;
use std::net::IpAddr;
#[ext(pub, name = MultiaddrExt)]
impl Multiaddr {
/// If the multiaddress corresponds to a TCP address, extracts it
fn try_to_tcp_addr(&self) -> Option<(IpAddr, u16)> {
let mut ps = self.into_iter();
let ip = if let Some(p) = ps.next() {
match p {
Protocol::Ip4(ip) => IpAddr::V4(ip),
Protocol::Ip6(ip) => IpAddr::V6(ip),
_ => return None,
}
} else {
return None;
};
let Some(Protocol::Tcp(port)) = ps.next() else {
return None;
};
Some((ip, port))
}
}
}
-143
View File
@@ -1,143 +0,0 @@
use crate::alias;
use crate::swarm::transport::tcp_transport;
pub use behaviour::{Behaviour, BehaviourEvent};
use libp2p::{SwarmBuilder, identity};
pub type Swarm = libp2p::Swarm<Behaviour>;
/// The current version of the network: this prevents devices running different versions of the
/// software from interacting with each other.
///
/// TODO: right now this is a hardcoded constant; figure out what the versioning semantics should
/// even be, and how to inject the right version into this config/initialization. E.g. should
/// this be passed in as a parameter? What about rapidly changing versions in debug builds?
/// this is all VERY very hard to figure out and needs to be mulled over as a team.
pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
/// Create and configure a swarm which listens to all ports on OS
pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
let mut swarm = SwarmBuilder::with_existing_identity(keypair)
.with_tokio()
.with_other_transport(tcp_transport)?
.with_behaviour(Behaviour::new)?
.build();
// Listen on all interfaces and whatever port the OS assigns
swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
Ok(swarm)
}
mod transport {
use crate::alias;
use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR};
use futures_lite::{AsyncRead, AsyncWrite};
use keccak_const::Sha3_256;
use libp2p::core::muxing;
use libp2p::core::transport::Boxed;
use libp2p::pnet::{PnetError, PnetOutput};
use libp2p::{PeerId, Transport, identity, noise, pnet, yamux};
use std::{env, sync::LazyLock};
/// Key used for networking's private network; parametrized on the [`NETWORK_VERSION`].
/// See [`pnet_upgrade`] for more.
static PNET_PRESHARED_KEY: LazyLock<[u8; 32]> = LazyLock::new(|| {
let builder = Sha3_256::new().update(b"exo_discovery_network");
if let Ok(var) = env::var(OVERRIDE_VERSION_ENV_VAR) {
let bytes = var.into_bytes();
builder.update(&bytes)
} else {
builder.update(NETWORK_VERSION)
}
.finalize()
});
/// Make the Swarm run on a private network, as to not clash with public libp2p nodes and
/// also different-versioned instances of this same network.
/// This is implemented as an additional "upgrade" ontop of existing [`libp2p::Transport`] layers.
async fn pnet_upgrade<TSocket>(
socket: TSocket,
_: impl Sized,
) -> Result<PnetOutput<TSocket>, PnetError>
where
TSocket: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
use pnet::{PnetConfig, PreSharedKey};
PnetConfig::new(PreSharedKey::new(*PNET_PRESHARED_KEY))
.handshake(socket)
.await
}
/// TCP/IP transport layer configuration.
pub fn tcp_transport(
keypair: &identity::Keypair,
) -> alias::AnyResult<Boxed<(PeerId, muxing::StreamMuxerBox)>> {
use libp2p::{
core::upgrade::Version,
tcp::{Config, tokio},
};
// `TCP_NODELAY` enabled => avoid latency
let tcp_config = Config::default().nodelay(true);
// V1 + lazy flushing => 0-RTT negotiation
let upgrade_version = Version::V1Lazy;
// Noise is faster than TLS + we don't care much for security
let noise_config = noise::Config::new(keypair)?;
// Use default Yamux config for multiplexing
let yamux_config = yamux::Config::default();
// Create new Tokio-driven TCP/IP transport layer
let base_transport = tokio::Transport::new(tcp_config)
.and_then(pnet_upgrade)
.upgrade(upgrade_version)
.authenticate(noise_config)
.multiplex(yamux_config);
// Return boxed transport (to flatten complex type)
Ok(base_transport.boxed())
}
}
mod behaviour {
use crate::{alias, discovery};
use libp2p::swarm::NetworkBehaviour;
use libp2p::{gossipsub, identity};
/// Behavior of the Swarm which composes all desired behaviors:
/// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
#[derive(NetworkBehaviour)]
pub struct Behaviour {
pub discovery: discovery::Behaviour,
pub gossipsub: gossipsub::Behaviour,
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
Ok(Self {
discovery: discovery::Behaviour::new(keypair)?,
gossipsub: gossipsub_behaviour(keypair),
})
}
}
fn gossipsub_behaviour(keypair: &identity::Keypair) -> gossipsub::Behaviour {
use gossipsub::{ConfigBuilder, MessageAuthenticity, ValidationMode};
// build a gossipsub network behaviour
// => signed message authenticity + strict validation mode means the message-ID is
// automatically provided by gossipsub w/out needing to provide custom message-ID function
gossipsub::Behaviour::new(
MessageAuthenticity::Signed(keypair.clone()),
ConfigBuilder::default()
.max_transmit_size(1024 * 1024)
.validation_mode(ValidationMode::Strict)
.build()
.expect("the configuration should always be valid"),
)
.expect("creating gossipsub behavior should always work")
}
}
-7
View File
@@ -1,7 +0,0 @@
// maybe this will hold test in the future...??
#[cfg(test)]
mod tests {
#[test]
fn does_nothing() {}
}
-15
View File
@@ -1,15 +0,0 @@
[package]
name = "util"
version = { workspace = true }
edition = { workspace = true }
publish = false
[lib]
doctest = false
name = "util"
path = "src/lib.rs"
[lints]
workspace = true
[dependencies]
-1
View File
@@ -1 +0,0 @@
pub mod wakerdeque;
-55
View File
@@ -1,55 +0,0 @@
use std::collections::VecDeque;
use std::fmt::{Debug, Formatter};
use std::task::{Context, Waker};
/// A wrapper around [`VecDeque`] which wakes (if it can) on any `push_*` methods,
/// and updates the internally stored waker by consuming [`Context`] on any `pop_*` methods.
pub struct WakerDeque<T> {
waker: Option<Waker>,
deque: VecDeque<T>,
}
impl<T: Debug> Debug for WakerDeque<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
self.deque.fmt(f)
}
}
impl<T> WakerDeque<T> {
pub fn new() -> Self {
Self {
waker: None,
deque: VecDeque::new(),
}
}
fn update(&mut self, cx: &mut Context<'_>) {
self.waker = Some(cx.waker().clone());
}
fn wake(&mut self) {
let Some(ref mut w) = self.waker else { return };
w.wake_by_ref();
self.waker = None;
}
pub fn pop_front(&mut self, cx: &mut Context<'_>) -> Option<T> {
self.update(cx);
self.deque.pop_front()
}
pub fn pop_back(&mut self, cx: &mut Context<'_>) -> Option<T> {
self.update(cx);
self.deque.pop_back()
}
pub fn push_front(&mut self, value: T) {
self.wake();
self.deque.push_front(value);
}
pub fn push_back(&mut self, value: T) {
self.wake();
self.deque.push_back(value);
}
}
+84 -75
View File
@@ -1,31 +1,28 @@
import asyncio
import socket
from dataclasses import dataclass, field
from typing import Iterator
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
from exo.shared.types.events import (
Event,
ForwarderEvent,
NodeDownloadProgress,
)
from exo.shared.types.worker.downloads import (
@@ -35,36 +32,29 @@ from exo.shared.types.worker.downloads import (
DownloadPending,
DownloadProgress,
)
from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.channels import Receiver, Sender, channel
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
from exo.utils.channels import Receiver, Sender
from exo.utils.task_group import TaskGroup
@dataclass
class DownloadCoordinator:
node_id: NodeId
session_id: SessionId
shard_downloader: ShardDownloader
download_command_receiver: Receiver[ForwarderDownloadCommand]
local_event_sender: Sender[ForwarderEvent]
event_index_counter: Iterator[int]
event_sender: Sender[Event]
offline: bool = False
# Local state
download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
# 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:
@@ -80,7 +70,7 @@ class DownloadCoordinator:
completed = DownloadCompleted(
shard_metadata=callback_shard,
node_id=self.node_id,
total_bytes=progress.total_bytes,
total=progress.total,
model_directory=self._model_dir(model_id),
)
self.download_status[model_id] = completed
@@ -113,45 +103,16 @@ 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._emit_existing_download_progress)
finally:
for task in self.active_downloads.values():
task.cancel()
def shutdown(self) -> None:
self._tg.cancel_scope.cancel()
self._tg.cancel_tasks()
async def _command_processor(self) -> None:
with self.download_command_receiver as commands:
@@ -185,6 +146,25 @@ class DownloadCoordinator:
)
return
# Check EXO_MODELS_PATH for pre-downloaded models
found_path = resolve_model_in_path(model_id)
if found_path is not None:
logger.info(
f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}"
)
completed = DownloadCompleted(
shard_metadata=shard,
node_id=self.node_id,
total=shard.model_card.storage_size,
model_directory=str(found_path),
read_only=True,
)
self.download_status[model_id] = completed
await self.event_sender.send(
NodeDownloadProgress(download_progress=completed)
)
return
# Emit pending status
progress = DownloadPending(
shard_metadata=shard,
@@ -203,7 +183,7 @@ class DownloadCoordinator:
completed = DownloadCompleted(
shard_metadata=shard,
node_id=self.node_id,
total_bytes=initial_progress.total_bytes,
total=initial_progress.total,
model_directory=self._model_dir(model_id),
)
self.download_status[model_id] = completed
@@ -269,6 +249,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")
@@ -297,21 +286,6 @@ class DownloadCoordinator:
)
del self.download_status[model_id]
async def _forward_events(self) -> None:
with self.event_receiver as events:
async for event in events:
idx = next(self.event_index_counter)
fe = ForwarderEvent(
origin_idx=idx,
origin=self.node_id,
session=self.session_id,
event=event,
)
logger.debug(
f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
)
await self.local_event_sender.send(fe)
async def _emit_existing_download_progress(self) -> None:
try:
while True:
@@ -332,19 +306,21 @@ class DownloadCoordinator:
status: DownloadProgress = DownloadCompleted(
node_id=self.node_id,
shard_metadata=progress.shard,
total_bytes=progress.total_bytes,
total=progress.total,
model_directory=self._model_dir(
progress.shard.model_card.model_id
),
)
elif progress.status in ["in_progress", "not_started"]:
if progress.downloaded_bytes_this_session.in_bytes == 0:
if progress.downloaded_this_session.in_bytes == 0:
status = DownloadPending(
node_id=self.node_id,
shard_metadata=progress.shard,
model_directory=self._model_dir(
progress.shard.model_card.model_id
),
downloaded=progress.downloaded,
total=progress.total,
)
else:
status = DownloadOngoing(
@@ -364,6 +340,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."
)
+116 -77
View File
@@ -20,7 +20,6 @@ from huggingface_hub import (
)
from loguru import logger
from pydantic import (
DirectoryPath,
TypeAdapter,
)
@@ -31,7 +30,7 @@ from exo.download.huggingface_utils import (
get_hf_endpoint,
get_hf_token,
)
from exo.shared.constants import EXO_MODELS_DIR
from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
from exo.shared.models.model_cards import ModelTask
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
@@ -80,9 +79,9 @@ def map_repo_file_download_progress_to_download_progress_data(
repo_file_download_progress: RepoFileDownloadProgress,
) -> DownloadProgressData:
return DownloadProgressData(
downloaded_bytes=repo_file_download_progress.downloaded,
downloaded_bytes_this_session=repo_file_download_progress.downloaded_this_session,
total_bytes=repo_file_download_progress.total,
downloaded=repo_file_download_progress.downloaded,
downloaded_this_session=repo_file_download_progress.downloaded_this_session,
total=repo_file_download_progress.total,
completed_files=1 if repo_file_download_progress.status == "complete" else 0,
total_files=1,
speed=repo_file_download_progress.speed,
@@ -95,9 +94,9 @@ def map_repo_download_progress_to_download_progress_data(
repo_download_progress: RepoDownloadProgress,
) -> DownloadProgressData:
return DownloadProgressData(
total_bytes=repo_download_progress.total_bytes,
downloaded_bytes=repo_download_progress.downloaded_bytes,
downloaded_bytes_this_session=repo_download_progress.downloaded_bytes_this_session,
total=repo_download_progress.total,
downloaded=repo_download_progress.downloaded,
downloaded_this_session=repo_download_progress.downloaded_this_session,
completed_files=repo_download_progress.completed_files,
total_files=repo_download_progress.total_files,
speed=repo_download_progress.overall_speed,
@@ -111,7 +110,27 @@ def map_repo_download_progress_to_download_progress_data(
)
def build_model_path(model_id: ModelId) -> DirectoryPath:
def resolve_model_in_path(model_id: ModelId) -> Path | None:
"""Search EXO_MODELS_PATH directories for a pre-existing model.
Checks each directory for the normalized name (org--model). A candidate
is only returned if ``is_model_directory_complete`` confirms all weight
files are present.
"""
if EXO_MODELS_PATH is None:
return None
normalized = model_id.normalize()
for search_dir in EXO_MODELS_PATH:
candidate = search_dir / normalized
if candidate.is_dir() and is_model_directory_complete(candidate):
return candidate
return None
def build_model_path(model_id: ModelId) -> Path:
found = resolve_model_in_path(model_id)
if found is not None:
return found
return EXO_MODELS_DIR / model_id.normalize()
@@ -142,7 +161,7 @@ async def delete_model(model_id: ModelId) -> bool:
async def seed_models(seed_dir: str | Path):
"""Move model in resources folder of app to .cache/huggingface/hub"""
"""Move models from resources folder to EXO_MODELS_DIR."""
source_dir = Path(seed_dir)
dest_dir = await ensure_models_dir()
for path in source_dir.iterdir():
@@ -158,6 +177,72 @@ async def seed_models(seed_dir: str | Path):
logger.error(traceback.format_exc())
def _scan_model_directory(
model_dir: Path, recursive: bool = False
) -> list[FileListEntry] | None:
"""Scan a local model directory and build a file list.
Requires at least one ``*.safetensors.index.json``. Every weight file
referenced by the index that is missing on disk gets ``size=None``.
"""
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
if not index_files:
return None
entries_by_path: dict[str, FileListEntry] = {}
if recursive:
for dirpath, _, filenames in os.walk(model_dir):
for filename in filenames:
if filename.endswith(".partial"):
continue
full_path = Path(dirpath) / filename
rel_path = str(full_path.relative_to(model_dir))
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=full_path.stat().st_size,
)
else:
for item in model_dir.iterdir():
if item.is_file() and not item.name.endswith(".partial"):
entries_by_path[item.name] = FileListEntry(
type="file",
path=item.name,
size=item.stat().st_size,
)
# Add expected weight files from index that haven't been downloaded yet
for index_file in index_files:
try:
index_data = ModelSafetensorsIndex.model_validate_json(
index_file.read_text()
)
relative_dir = index_file.parent.relative_to(model_dir)
for filename in set(index_data.weight_map.values()):
rel_path = (
str(relative_dir / filename)
if relative_dir != Path(".")
else filename
)
if rel_path not in entries_by_path:
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=None,
)
except Exception:
continue
return list(entries_by_path.values())
def is_model_directory_complete(model_dir: Path) -> bool:
"""Check if a model directory contains all required weight files."""
file_list = _scan_model_directory(model_dir, recursive=True)
return file_list is not None and all(f.size is not None for f in file_list)
async def _build_file_list_from_local_directory(
model_id: ModelId,
recursive: bool = False,
@@ -172,59 +257,7 @@ async def _build_file_list_from_local_directory(
if not await aios.path.exists(model_dir):
return None
def _scan() -> list[FileListEntry] | None:
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
if not index_files:
return None
entries_by_path: dict[str, FileListEntry] = {}
if recursive:
for dirpath, _, filenames in os.walk(model_dir):
for filename in filenames:
if filename.endswith(".partial"):
continue
full_path = Path(dirpath) / filename
rel_path = str(full_path.relative_to(model_dir))
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=full_path.stat().st_size,
)
else:
for item in model_dir.iterdir():
if item.is_file() and not item.name.endswith(".partial"):
entries_by_path[item.name] = FileListEntry(
type="file",
path=item.name,
size=item.stat().st_size,
)
# Add expected weight files from index that haven't been downloaded yet
for index_file in index_files:
try:
index_data = ModelSafetensorsIndex.model_validate_json(
index_file.read_text()
)
relative_dir = index_file.parent.relative_to(model_dir)
for filename in set(index_data.weight_map.values()):
rel_path = (
str(relative_dir / filename)
if relative_dir != Path(".")
else filename
)
if rel_path not in entries_by_path:
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=None,
)
except Exception:
continue
return list(entries_by_path.values())
file_list = await asyncio.to_thread(_scan)
file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive)
if not file_list:
return None
return file_list
@@ -281,9 +314,13 @@ async def fetch_file_list_with_cache(
_fetched_file_lists_this_session.add(cache_key)
return file_list
except Exception as e:
logger.opt(exception=e).warning(
"Ran into exception when fetching file list from HF."
)
if await aios.path.exists(cache_file):
logger.warning(
f"No internet and no cached file list for {model_id} - using local file list"
f"No cached file list for {model_id} - using local file list"
)
async with aiofiles.open(cache_file, "r") as f:
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
@@ -578,19 +615,20 @@ def calculate_repo_progress(
file_progress: dict[str, RepoFileDownloadProgress],
all_start_time: float,
) -> RepoDownloadProgress:
all_total_bytes = sum((p.total.in_bytes for p in file_progress.values()), 0)
all_downloaded_bytes = sum(
(p.downloaded.in_bytes for p in file_progress.values()), 0
all_total = sum((p.total for p in file_progress.values()), Memory.from_bytes(0))
all_downloaded = sum(
(p.downloaded for p in file_progress.values()), Memory.from_bytes(0)
)
all_downloaded_bytes_this_session = sum(
(p.downloaded_this_session.in_bytes for p in file_progress.values()), 0
all_downloaded_this_session = sum(
(p.downloaded_this_session for p in file_progress.values()),
Memory.from_bytes(0),
)
elapsed_time = time.time() - all_start_time
all_speed = (
all_downloaded_bytes_this_session / elapsed_time if elapsed_time > 0 else 0
all_downloaded_this_session.in_bytes / elapsed_time if elapsed_time > 0 else 0
)
all_eta = (
timedelta(seconds=(all_total_bytes - all_downloaded_bytes) / all_speed)
timedelta(seconds=(all_total - all_downloaded).in_bytes / all_speed)
if all_speed > 0
else timedelta(seconds=0)
)
@@ -609,11 +647,9 @@ def calculate_repo_progress(
[p for p in file_progress.values() if p.downloaded == p.total]
),
total_files=len(file_progress),
downloaded_bytes=Memory.from_bytes(all_downloaded_bytes),
downloaded_bytes_this_session=Memory.from_bytes(
all_downloaded_bytes_this_session
),
total_bytes=Memory.from_bytes(all_total_bytes),
downloaded=all_downloaded,
downloaded_this_session=all_downloaded_this_session,
total=all_total,
overall_speed=all_speed,
overall_eta=all_eta,
status=status,
@@ -791,6 +827,7 @@ async def download_shard(
for file in filtered_file_list:
downloaded_bytes = await get_downloaded_size(target_dir / file.path)
final_file_exists = await aios.path.exists(target_dir / file.path)
file_progress[file.path] = RepoFileDownloadProgress(
repo_id=shard.model_card.model_id,
repo_revision=revision,
@@ -800,7 +837,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
+3 -8
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
@@ -107,9 +102,9 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
),
completed_files=0,
total_files=0,
downloaded_bytes=Memory.from_bytes(0),
downloaded_bytes_this_session=Memory.from_bytes(0),
total_bytes=Memory.from_bytes(0),
downloaded=Memory.from_bytes(0),
downloaded_this_session=Memory.from_bytes(0),
total=Memory.from_bytes(0),
overall_speed=0,
overall_eta=timedelta(seconds=0),
status="complete",
+53 -42
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
@@ -17,6 +15,7 @@ from exo.download.coordinator import DownloadCoordinator
from exo.download.impl_shard_downloader import exo_shard_downloader
from exo.master.api import API # TODO: should API be in master?
from exo.master.main import Master
from exo.routing.event_router import EventRouter
from exo.routing.router import Router, get_node_id_keypair
from exo.shared.constants import EXO_LOG
from exo.shared.election import Election, ElectionResult
@@ -24,12 +23,14 @@ 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
@dataclass
class Node:
router: Router
event_router: EventRouter
download_coordinator: DownloadCoordinator | None
worker: Worker | None
election: Election # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
@@ -38,14 +39,13 @@ class Node:
api: API | None
node_id: NodeId
event_index_counter: Iterator[int]
offline: bool
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
@classmethod
async def create(cls, args: "Args") -> "Self":
async def create(cls, args: "Args") -> Self:
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_peer_id().to_base58())
node_id = NodeId(keypair.to_node_id())
session_id = SessionId(master_node_id=node_id, election_clock=0)
router = Router.create(keypair)
await router.register_topic(topics.GLOBAL_EVENTS)
@@ -54,21 +54,22 @@ class Node:
await router.register_topic(topics.ELECTION_MESSAGES)
await router.register_topic(topics.CONNECTION_MESSAGES)
await router.register_topic(topics.DOWNLOAD_COMMANDS)
event_router = EventRouter(
session_id,
command_sender=router.sender(topics.COMMANDS),
external_outbound=router.sender(topics.LOCAL_EVENTS),
external_inbound=router.receiver(topics.GLOBAL_EVENTS),
)
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),
event_sender=event_router.sender(),
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,
)
else:
@@ -77,9 +78,8 @@ class Node:
if args.spawn_api:
api = API(
node_id,
session_id,
port=args.api_port,
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
event_receiver=event_router.receiver(),
command_sender=router.sender(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
election_receiver=router.receiver(topics.ELECTION_MESSAGES),
@@ -90,12 +90,10 @@ class Node:
if not args.no_worker:
worker = Worker(
node_id,
session_id,
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
local_event_sender=router.sender(topics.LOCAL_EVENTS),
event_receiver=event_router.receiver(),
event_sender=event_router.sender(),
command_sender=router.sender(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
event_index_counter=event_index_counter,
)
else:
worker = None
@@ -104,6 +102,7 @@ class Node:
master = Master(
node_id,
session_id,
event_sender=event_router.sender(),
global_event_sender=router.sender(topics.GLOBAL_EVENTS),
local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
command_receiver=router.receiver(topics.COMMANDS),
@@ -126,6 +125,7 @@ class Node:
return cls(
router,
event_router,
download_coordinator,
worker,
election,
@@ -133,7 +133,6 @@ class Node:
master,
api,
node_id,
event_index_counter,
args.offline,
)
@@ -142,6 +141,7 @@ class Node:
signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
tg.start_soon(self.router.run)
tg.start_soon(self.event_router.run)
tg.start_soon(self.election.run)
if self.download_coordinator:
tg.start_soon(self.download_coordinator.run)
@@ -155,11 +155,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:
@@ -176,6 +176,17 @@ class Node:
# - Shutdown and re-create the worker
# - Shut down and re-create the API
if result.is_new_master:
await anyio.sleep(0)
self.event_router.shutdown()
self.event_router = EventRouter(
result.session_id,
self.router.sender(topics.COMMANDS),
self.router.receiver(topics.GLOBAL_EVENTS),
self.router.sender(topics.LOCAL_EVENTS),
)
self._tg.start_soon(self.event_router.run)
if (
result.session_id.master_node_id == self.node_id
and self.master is not None
@@ -189,6 +200,7 @@ class Node:
self.master = Master(
self.node_id,
result.session_id,
event_sender=self.event_router.sender(),
global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
command_receiver=self.router.receiver(topics.COMMANDS),
@@ -211,20 +223,15 @@ class Node:
f"Node {result.session_id.master_node_id} elected master"
)
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),
event_sender=self.event_router.sender(),
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,
)
self._tg.start_soon(self.download_coordinator.run)
@@ -233,20 +240,16 @@ class Node:
# TODO: add profiling etc to resource monitor
self.worker = Worker(
self.node_id,
result.session_id,
global_event_receiver=self.router.receiver(
topics.GLOBAL_EVENTS
),
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
event_receiver=self.event_router.receiver(),
event_sender=self.event_router.sender(),
command_sender=self.router.sender(topics.COMMANDS),
download_command_sender=self.router.sender(
topics.DOWNLOAD_COMMANDS
),
event_index_counter=self.event_index_counter,
)
self._tg.start_soon(self.worker.run)
if self.api:
self.api.reset(result.session_id, result.won_clock)
self.api.reset(result.won_clock, self.event_router.receiver())
else:
if self.api:
self.api.unpause(result.won_clock)
@@ -258,7 +261,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 +279,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 +299,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 +350,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()
+29 -5
View File
@@ -26,7 +26,11 @@ from exo.shared.types.chunks import (
ToolCallChunk,
)
from exo.shared.types.common import CommandId
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.text_generation import (
InputMessage,
TextGenerationTaskParams,
resolve_reasoning_params,
)
def chat_request_to_text_generation(
@@ -59,7 +63,11 @@ def chat_request_to_text_generation(
chat_template_messages.append({"role": "system", "content": content})
else:
# Skip messages with no meaningful content
if msg.content is None and msg.thinking is None and msg.tool_calls is None:
if (
msg.content is None
and msg.reasoning_content is None
and msg.tool_calls is None
):
continue
if msg.role in ("user", "assistant", "developer"):
@@ -71,6 +79,10 @@ def chat_request_to_text_generation(
dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
chat_template_messages.append(dumped)
resolved_effort, resolved_thinking = resolve_reasoning_params(
request.reasoning_effort, request.enable_thinking
)
return TextGenerationTaskParams(
model=request.model,
input=input_messages
@@ -85,7 +97,8 @@ def chat_request_to_text_generation(
seed=request.seed,
stream=request.stream,
tools=request.tools,
enable_thinking=request.enable_thinking,
reasoning_effort=resolved_effort,
enable_thinking=resolved_thinking,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
@@ -111,6 +124,11 @@ def chunk_to_response(
]
)
if chunk.is_thinking:
delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text)
else:
delta = ChatCompletionMessage(role="assistant", content=chunk.text)
return ChatCompletionResponse(
id=command_id,
created=int(time.time()),
@@ -118,7 +136,7 @@ def chunk_to_response(
choices=[
StreamingChoiceResponse(
index=0,
delta=ChatCompletionMessage(role="assistant", content=chunk.text),
delta=delta,
logprobs=logprobs,
finish_reason=chunk.finish_reason,
)
@@ -208,6 +226,7 @@ async def collect_chat_response(
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ChatCompletionResponse."""
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_calls: list[ToolCall] = []
logprobs_content: list[LogprobsContentItem] = []
model: str | None = None
@@ -228,7 +247,10 @@ async def collect_chat_response(
if model is None:
model = chunk.model
last_usage = chunk.usage or last_usage
text_parts.append(chunk.text)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.logprob is not None:
logprobs_content.append(
LogprobsContentItem(
@@ -258,6 +280,7 @@ async def collect_chat_response(
raise ValueError(error_message)
combined_text = "".join(text_parts)
combined_thinking = "".join(thinking_parts) if thinking_parts else None
assert model is not None
yield ChatCompletionResponse(
@@ -270,6 +293,7 @@ async def collect_chat_response(
message=ChatCompletionMessage(
role="assistant",
content=combined_text,
reasoning_content=combined_thinking,
tool_calls=tool_calls if tool_calls else None,
),
logprobs=Logprobs(content=logprobs_content)
+112 -21
View File
@@ -1,6 +1,7 @@
"""Claude Messages API adapter for converting requests/responses."""
import json
import re
from collections.abc import AsyncGenerator
from typing import Any
@@ -28,6 +29,8 @@ from exo.shared.types.claude_api import (
ClaudeStopReason,
ClaudeTextBlock,
ClaudeTextDelta,
ClaudeThinkingBlock,
ClaudeThinkingDelta,
ClaudeToolResultBlock,
ClaudeToolUseBlock,
ClaudeUsage,
@@ -61,6 +64,22 @@ def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str:
return "".join(sub_block.text for sub_block in block.content)
# Matches "x-anthropic-billing-header: ...;" (with optional trailing newline)
# or similar telemetry headers that change every request and break KV prefix caching.
_VOLATILE_HEADER_RE = re.compile(r"^x-anthropic-[^\n]*;\n?", re.MULTILINE)
def _strip_volatile_headers(text: str) -> str:
"""Remove Anthropic billing/telemetry headers from system prompt text.
Claude Code prepends headers like 'x-anthropic-billing-header: cc_version=...;
cc_entrypoint=...; cch=...;' that contain per-request content hashes. These
change every request and break KV prefix caching (the prefix diverges at ~20
tokens instead of matching thousands of conversation tokens).
"""
return _VOLATILE_HEADER_RE.sub("", text)
def claude_request_to_text_generation(
request: ClaudeMessagesRequest,
) -> TextGenerationTaskParams:
@@ -73,6 +92,8 @@ def claude_request_to_text_generation(
instructions = request.system
else:
instructions = "".join(block.text for block in request.system)
instructions = _strip_volatile_headers(instructions)
chat_template_messages.append({"role": "system", "content": instructions})
# Convert messages to input
@@ -85,12 +106,15 @@ def claude_request_to_text_generation(
# Process structured content blocks
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_calls: list[dict[str, Any]] = []
tool_results: list[ClaudeToolResultBlock] = []
for block in msg.content:
if isinstance(block, ClaudeTextBlock):
text_parts.append(block.text)
elif isinstance(block, ClaudeThinkingBlock):
thinking_parts.append(block.thinking)
elif isinstance(block, ClaudeToolUseBlock):
tool_calls.append(
{
@@ -106,6 +130,7 @@ def claude_request_to_text_generation(
tool_results.append(block)
content = "".join(text_parts)
reasoning_content = "".join(thinking_parts) if thinking_parts else None
# Build InputMessage from text content
if msg.role in ("user", "assistant"):
@@ -113,9 +138,14 @@ def claude_request_to_text_generation(
# Build chat_template_messages preserving tool structure
if tool_calls:
chat_template_messages.append(
{"role": "assistant", "content": content, "tool_calls": tool_calls}
)
chat_msg: dict[str, Any] = {
"role": "assistant",
"content": content,
"tool_calls": tool_calls,
}
if reasoning_content:
chat_msg["reasoning_content"] = reasoning_content
chat_template_messages.append(chat_msg)
elif tool_results:
for tr in tool_results:
chat_template_messages.append(
@@ -126,7 +156,10 @@ def claude_request_to_text_generation(
}
)
else:
chat_template_messages.append({"role": msg.role, "content": content})
chat_msg = {"role": msg.role, "content": content}
if reasoning_content:
chat_msg["reasoning_content"] = reasoning_content
chat_template_messages.append(chat_msg)
# Convert Claude tool definitions to OpenAI-style function tools
tools: list[dict[str, Any]] | None = None
@@ -143,6 +176,10 @@ def claude_request_to_text_generation(
for tool in request.tools
]
enable_thinking: bool | None = None
if request.thinking is not None:
enable_thinking = request.thinking.type in ("enabled", "adaptive")
return TextGenerationTaskParams(
model=request.model,
input=input_messages
@@ -156,6 +193,7 @@ def claude_request_to_text_generation(
stop=request.stop_sequences,
stream=request.stream,
tools=tools,
enable_thinking=enable_thinking,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
@@ -173,6 +211,7 @@ async def collect_claude_response(
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ClaudeMessagesResponse."""
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_use_blocks: list[ClaudeToolUseBlock] = []
stop_reason: ClaudeStopReason | None = None
last_usage: Usage | None = None
@@ -200,7 +239,10 @@ async def collect_claude_response(
stop_reason = "tool_use"
continue
text_parts.append(chunk.text)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -209,9 +251,12 @@ async def collect_claude_response(
raise ValueError(error_message)
combined_text = "".join(text_parts)
combined_thinking = "".join(thinking_parts)
# Build content blocks
content: list[ClaudeContentBlock] = []
if combined_thinking:
content.append(ClaudeThinkingBlock(thinking=combined_thinking))
if combined_text:
content.append(ClaudeTextBlock(text=combined_text))
content.extend(tool_use_blocks)
@@ -256,16 +301,16 @@ async def generate_claude_stream(
start_event = ClaudeMessageStartEvent(message=initial_message)
yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n"
# content_block_start for text block at index 0
block_start = ClaudeContentBlockStartEvent(
index=0, content_block=ClaudeTextBlock(text="")
)
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
output_tokens = 0
stop_reason: ClaudeStopReason | None = None
last_usage: Usage | None = None
next_block_index = 1 # text block is 0, tool blocks start at 1
next_block_index = 0
# Track whether we've started thinking/text blocks
thinking_block_started = False
thinking_block_index = -1
text_block_started = False
text_block_index = -1
async for chunk in chunk_stream:
if isinstance(chunk, PrefillProgressChunk):
@@ -310,12 +355,45 @@ async def generate_claude_stream(
output_tokens += 1 # Count each chunk as one token
# content_block_delta
delta_event = ClaudeContentBlockDeltaEvent(
index=0,
delta=ClaudeTextDelta(text=chunk.text),
)
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
if chunk.is_thinking:
# Start thinking block on first thinking token
if not thinking_block_started:
thinking_block_started = True
thinking_block_index = next_block_index
next_block_index += 1
block_start = ClaudeContentBlockStartEvent(
index=thinking_block_index,
content_block=ClaudeThinkingBlock(thinking=""),
)
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
delta_event = ClaudeContentBlockDeltaEvent(
index=thinking_block_index,
delta=ClaudeThinkingDelta(thinking=chunk.text),
)
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
else:
# Close thinking block when transitioning to text
if thinking_block_started and text_block_index == -1:
block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
# Start text block on first text token
if not text_block_started:
text_block_started = True
text_block_index = next_block_index
next_block_index += 1
block_start = ClaudeContentBlockStartEvent(
index=text_block_index,
content_block=ClaudeTextBlock(text=""),
)
yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n"
delta_event = ClaudeContentBlockDeltaEvent(
index=text_block_index,
delta=ClaudeTextDelta(text=chunk.text),
)
yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n"
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -324,9 +402,22 @@ async def generate_claude_stream(
if last_usage is not None:
output_tokens = last_usage.completion_tokens
# content_block_stop for text block
block_stop = ClaudeContentBlockStopEvent(index=0)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
# Close any open blocks
if thinking_block_started and text_block_index == -1:
block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
if text_block_started:
block_stop = ClaudeContentBlockStopEvent(index=text_block_index)
yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n"
if not thinking_block_started and not text_block_started:
empty_start = ClaudeContentBlockStartEvent(
index=0, content_block=ClaudeTextBlock(text="")
)
yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n"
empty_stop = ClaudeContentBlockStopEvent(index=0)
yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n"
# message_delta
message_delta = ClaudeMessageDeltaEvent(
+456
View File
@@ -0,0 +1,456 @@
from __future__ import annotations
import json
from collections.abc import AsyncGenerator
from typing import Any
from exo.shared.types.chunks import (
ErrorChunk,
PrefillProgressChunk,
TokenChunk,
ToolCallChunk,
)
from exo.shared.types.common import CommandId
from exo.shared.types.ollama_api import (
OllamaChatRequest,
OllamaChatResponse,
OllamaDoneReason,
OllamaGenerateRequest,
OllamaGenerateResponse,
OllamaMessage,
OllamaToolCall,
OllamaToolFunction,
)
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
def _map_done_reason(
finish_reason: str | None,
) -> OllamaDoneReason | None:
if finish_reason is None:
return None
if finish_reason == "stop":
return "stop"
if finish_reason == "length":
return "length"
if finish_reason in ("tool_calls", "function_call"):
return "tool_call"
if finish_reason == "error":
return "error"
return "stop"
def _try_parse_json(value: str) -> dict[str, Any] | str:
try:
return json.loads(value) # type: ignore
except json.JSONDecodeError:
return value
def _build_tool_calls(chunk: ToolCallChunk) -> list[OllamaToolCall]:
tool_calls: list[OllamaToolCall] = []
for index, tool in enumerate(chunk.tool_calls):
# tool.arguments is always str; try to parse as JSON dict for Ollama format
arguments: dict[str, Any] | str = _try_parse_json(tool.arguments)
tool_calls.append(
OllamaToolCall(
id=tool.id,
type="function",
function=OllamaToolFunction(
name=tool.name, arguments=arguments, index=index
),
)
)
return tool_calls
def _get_usage(
chunk: TokenChunk | ToolCallChunk,
) -> tuple[int | None, int | None]:
"""Extract (prompt_eval_count, eval_count) from a chunk."""
if chunk.usage is not None:
return (chunk.usage.prompt_tokens, chunk.usage.completion_tokens)
if chunk.stats is not None:
return (chunk.stats.prompt_tokens, chunk.stats.generation_tokens)
return (None, None)
def ollama_request_to_text_generation(
request: OllamaChatRequest,
) -> TextGenerationTaskParams:
"""Convert Ollama chat request to exo's internal text generation format."""
instructions: str | None = None
input_messages: list[InputMessage] = []
chat_template_messages: list[dict[str, Any]] = []
tool_message_index = 0
for msg in request.messages:
content = msg.content or ""
if msg.role == "system":
if instructions is None:
instructions = content
else:
instructions = f"{instructions}\n{content}"
chat_template_messages.append({"role": "system", "content": content})
continue
if msg.role in ("user", "assistant") and (
msg.content is not None or msg.thinking is not None or msg.tool_calls
):
input_messages.append(InputMessage(role=msg.role, content=content))
dumped: dict[str, Any] = {"role": msg.role, "content": content}
if msg.thinking is not None:
dumped["thinking"] = msg.thinking
if msg.tool_calls is not None:
tool_calls_list: list[dict[str, Any]] = []
for tc in msg.tool_calls:
function: dict[str, Any] = {
"name": tc.function.name,
"arguments": (
json.dumps(tc.function.arguments)
if isinstance(tc.function.arguments, dict)
else tc.function.arguments
),
}
if tc.function.index is not None:
function["index"] = tc.function.index
tool_call: dict[str, Any] = {"function": function}
if tc.id is not None:
tool_call["id"] = tc.id
if tc.type is not None:
tool_call["type"] = tc.type
tool_calls_list.append(tool_call)
dumped["tool_calls"] = tool_calls_list
if msg.name is not None:
dumped["name"] = msg.name
if msg.role == "tool":
tool_message_index += 1
tool_call_id = msg.tool_name or msg.name or f"tool_{tool_message_index}"
dumped["tool_call_id"] = tool_call_id
if msg.tool_name is not None:
dumped["tool_name"] = msg.tool_name
chat_template_messages.append(dumped)
options = request.options
return TextGenerationTaskParams(
model=request.model,
input=input_messages
if input_messages
else [InputMessage(role="user", content="")],
instructions=instructions,
max_output_tokens=options.num_predict if options else None,
temperature=options.temperature if options else None,
top_p=options.top_p if options else None,
top_k=options.top_k if options else None,
stop=options.stop if options else None,
seed=options.seed if options else None,
stream=request.stream,
tools=request.tools,
enable_thinking=request.think,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
)
async def generate_ollama_chat_stream(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str, None]:
"""Generate streaming responses in Ollama format (newline-delimited JSON)."""
thinking_parts: list[str] = []
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
error_response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant", content=chunk.error_message
),
done=True,
done_reason="error",
)
yield f"{error_response.model_dump_json(exclude_none=True)}\n"
return
case ToolCallChunk():
prompt_eval, eval_count = _get_usage(chunk)
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant",
content="",
tool_calls=_build_tool_calls(chunk),
thinking="".join(thinking_parts) if thinking_parts else None,
),
done=True,
done_reason="tool_call",
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
return
case TokenChunk():
done = chunk.finish_reason is not None
if chunk.is_thinking:
thinking_parts.append(chunk.text)
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant", content="", thinking=chunk.text
),
done=False,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
elif done:
prompt_eval, eval_count = _get_usage(chunk)
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(
role="assistant",
content=chunk.text,
),
done=True,
done_reason=_map_done_reason(chunk.finish_reason),
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
else:
response = OllamaChatResponse(
model=str(chunk.model),
message=OllamaMessage(role="assistant", content=chunk.text),
done=False,
)
yield f"{response.model_dump_json(exclude_none=True)}\n"
if done:
return
async def collect_ollama_chat_response(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str]:
"""Collect streaming chunks into a single non-streaming Ollama response.
Returns an AsyncGenerator[str] (single yield) for consistency with FastAPI
StreamingResponse cancellation handling.
"""
text_parts: list[str] = []
thinking_parts: list[str] = []
tool_calls: list[OllamaToolCall] = []
model: str | None = None
finish_reason: str | None = None
prompt_eval_count: int | None = None
eval_count: int | None = None
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
raise ValueError(chunk.error_message or "Internal server error")
case TokenChunk():
if model is None:
model = str(chunk.model)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.finish_reason is not None:
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
case ToolCallChunk():
if model is None:
model = str(chunk.model)
tool_calls.extend(_build_tool_calls(chunk))
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
combined_text = "".join(text_parts)
combined_thinking = "".join(thinking_parts) if thinking_parts else None
assert model is not None
yield OllamaChatResponse(
model=model,
message=OllamaMessage(
role="assistant",
content=combined_text,
thinking=combined_thinking,
tool_calls=tool_calls if tool_calls else None,
),
done=True,
done_reason=_map_done_reason(finish_reason),
prompt_eval_count=prompt_eval_count,
eval_count=eval_count,
).model_dump_json(exclude_none=True)
return
# ── /api/generate ──
def ollama_generate_request_to_text_generation(
request: OllamaGenerateRequest,
) -> TextGenerationTaskParams:
"""Convert Ollama generate request to exo's internal text generation format."""
chat_template_messages: list[dict[str, Any]] = []
if request.system:
chat_template_messages.append({"role": "system", "content": request.system})
chat_template_messages.append({"role": "user", "content": request.prompt})
options = request.options
return TextGenerationTaskParams(
model=request.model,
input=[InputMessage(role="user", content=request.prompt)],
instructions=request.system,
max_output_tokens=options.num_predict if options else None,
temperature=options.temperature if options else None,
top_p=options.top_p if options else None,
top_k=options.top_k if options else None,
stop=options.stop if options else None,
seed=options.seed if options else None,
stream=request.stream,
enable_thinking=request.think,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
)
async def generate_ollama_generate_stream(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str, None]:
"""Generate streaming responses for /api/generate in Ollama NDJSON format."""
thinking_parts: list[str] = []
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
resp = OllamaGenerateResponse(
model=str(chunk.model),
response="",
done=True,
done_reason="error",
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
return
case ToolCallChunk():
# generate endpoint doesn't support tools; emit as done
prompt_eval, eval_count = _get_usage(chunk)
resp = OllamaGenerateResponse(
model=str(chunk.model),
response="",
done=True,
done_reason="stop",
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
return
case TokenChunk():
done = chunk.finish_reason is not None
if chunk.is_thinking:
thinking_parts.append(chunk.text)
resp = OllamaGenerateResponse(
model=str(chunk.model),
response="",
thinking=chunk.text,
done=False,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
elif done:
prompt_eval, eval_count = _get_usage(chunk)
resp = OllamaGenerateResponse(
model=str(chunk.model),
response=chunk.text,
done=True,
done_reason=_map_done_reason(chunk.finish_reason),
prompt_eval_count=prompt_eval,
eval_count=eval_count,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
else:
resp = OllamaGenerateResponse(
model=str(chunk.model),
response=chunk.text,
done=False,
)
yield f"{resp.model_dump_json(exclude_none=True)}\n"
if done:
return
async def collect_ollama_generate_response(
_command_id: CommandId,
chunk_stream: AsyncGenerator[
ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None
],
) -> AsyncGenerator[str]:
"""Collect chunks into a single non-streaming /api/generate response."""
text_parts: list[str] = []
thinking_parts: list[str] = []
model: str | None = None
finish_reason: str | None = None
prompt_eval_count: int | None = None
eval_count: int | None = None
async for chunk in chunk_stream:
match chunk:
case PrefillProgressChunk():
continue
case ErrorChunk():
raise ValueError(chunk.error_message or "Internal server error")
case TokenChunk():
if model is None:
model = str(chunk.model)
if chunk.is_thinking:
thinking_parts.append(chunk.text)
else:
text_parts.append(chunk.text)
if chunk.finish_reason is not None:
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
case ToolCallChunk():
if model is None:
model = str(chunk.model)
finish_reason = chunk.finish_reason
prompt_eval_count, eval_count = _get_usage(chunk)
assert model is not None
yield OllamaGenerateResponse(
model=model,
response="".join(text_parts),
thinking="".join(thinking_parts) if thinking_parts else None,
done=True,
done_reason=_map_done_reason(finish_reason),
prompt_eval_count=prompt_eval_count,
eval_count=eval_count,
).model_dump_json(exclude_none=True)
return
+219 -31
View File
@@ -29,6 +29,12 @@ from exo.shared.types.openai_responses import (
ResponseOutputItemAddedEvent,
ResponseOutputItemDoneEvent,
ResponseOutputText,
ResponseReasoningItem,
ResponseReasoningSummaryPartAddedEvent,
ResponseReasoningSummaryPartDoneEvent,
ResponseReasoningSummaryText,
ResponseReasoningSummaryTextDeltaEvent,
ResponseReasoningSummaryTextDoneEvent,
ResponsesRequest,
ResponsesResponse,
ResponsesStreamEvent,
@@ -36,7 +42,11 @@ from exo.shared.types.openai_responses import (
ResponseTextDoneEvent,
ResponseUsage,
)
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
from exo.shared.types.text_generation import (
InputMessage,
TextGenerationTaskParams,
resolve_reasoning_params,
)
def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -113,6 +123,11 @@ def responses_request_to_text_generation(
)
built_chat_template = chat_template_messages if chat_template_messages else None
effort_from_reasoning = request.reasoning.effort if request.reasoning else None
resolved_effort, resolved_thinking = resolve_reasoning_params(
effort_from_reasoning, request.enable_thinking
)
return TextGenerationTaskParams(
model=request.model,
input=input_value,
@@ -126,6 +141,8 @@ def responses_request_to_text_generation(
stop=request.stop,
seed=request.seed,
chat_template_messages=built_chat_template or request.chat_template_messages,
reasoning_effort=resolved_effort,
enable_thinking=resolved_thinking,
)
@@ -141,7 +158,9 @@ async def collect_responses_response(
"""Collect all token chunks and return a single ResponsesResponse."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
reasoning_id = f"rs_{command_id}"
accumulated_text = ""
thinking_parts: list[str] = []
function_call_items: list[ResponseFunctionCallItem] = []
last_usage: Usage | None = None
error_message: str | None = None
@@ -168,6 +187,10 @@ async def collect_responses_response(
)
continue
if chunk.is_thinking:
thinking_parts.append(chunk.text)
continue
accumulated_text += chunk.text
if error_message is not None:
@@ -182,13 +205,21 @@ async def collect_responses_response(
total_tokens=last_usage.total_tokens,
)
output: list[ResponseItem] = [
output: list[ResponseItem] = []
if thinking_parts:
output.append(
ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))],
)
)
output.append(
ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text=accumulated_text)],
status="completed",
)
]
)
output.extend(function_call_items)
yield ResponsesResponse(
@@ -212,6 +243,7 @@ async def generate_responses_stream(
"""Generate OpenAI Responses API streaming events from TokenChunks."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
reasoning_id = f"rs_{command_id}"
seq = count(1)
# response.created
@@ -233,32 +265,17 @@ async def generate_responses_stream(
)
yield _format_sse(in_progress_event)
# response.output_item.added
initial_item = ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text="")],
status="in_progress",
)
item_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq), output_index=0, item=initial_item
)
yield _format_sse(item_added)
# response.content_part.added
initial_part = ResponseOutputText(text="")
part_added = ResponseContentPartAddedEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
content_index=0,
part=initial_part,
)
yield _format_sse(part_added)
accumulated_text = ""
accumulated_thinking = ""
function_call_items: list[ResponseFunctionCallItem] = []
last_usage: Usage | None = None
next_output_index = 1 # message item is at 0
next_output_index = 0
# Track dynamic block creation
reasoning_started = False
reasoning_output_index = -1
message_started = False
message_output_index = -1
async for chunk in chunk_stream:
if isinstance(chunk, PrefillProgressChunk):
@@ -327,23 +344,184 @@ async def generate_responses_stream(
next_output_index += 1
continue
if chunk.is_thinking:
# Start reasoning block on first thinking token
if not reasoning_started:
reasoning_started = True
reasoning_output_index = next_output_index
next_output_index += 1
# response.output_item.added for reasoning
reasoning_item = ResponseReasoningItem(
id=reasoning_id,
summary=[],
status="in_progress",
)
rs_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq),
output_index=reasoning_output_index,
item=reasoning_item,
)
yield _format_sse(rs_added)
# response.reasoning_summary_part.added
part_added = ResponseReasoningSummaryPartAddedEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
part=ResponseReasoningSummaryText(text=""),
)
yield _format_sse(part_added)
accumulated_thinking += chunk.text
# response.reasoning_summary_text.delta
rs_delta = ResponseReasoningSummaryTextDeltaEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
delta=chunk.text,
)
yield _format_sse(rs_delta)
continue
# Close reasoning block when transitioning to text
if reasoning_started and not message_started:
# response.reasoning_summary_text.done
rs_text_done = ResponseReasoningSummaryTextDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
text=accumulated_thinking,
)
yield _format_sse(rs_text_done)
# response.reasoning_summary_part.done
rs_part_done = ResponseReasoningSummaryPartDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
part=ResponseReasoningSummaryText(text=accumulated_thinking),
)
yield _format_sse(rs_part_done)
# response.output_item.done for reasoning
rs_item_done = ResponseOutputItemDoneEvent(
sequence_number=next(seq),
output_index=reasoning_output_index,
item=ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
),
)
yield _format_sse(rs_item_done)
# Start message block on first text token
if not message_started:
message_started = True
message_output_index = next_output_index
next_output_index += 1
initial_item = ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text="")],
status="in_progress",
)
item_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq),
output_index=message_output_index,
item=initial_item,
)
yield _format_sse(item_added)
initial_part = ResponseOutputText(text="")
part_added = ResponseContentPartAddedEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=message_output_index,
content_index=0,
part=initial_part,
)
yield _format_sse(part_added)
accumulated_text += chunk.text
# response.output_text.delta
delta_event = ResponseTextDeltaEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
output_index=message_output_index,
content_index=0,
delta=chunk.text,
)
yield _format_sse(delta_event)
# Close reasoning block if it was never followed by text
if reasoning_started and not message_started:
rs_text_done = ResponseReasoningSummaryTextDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
text=accumulated_thinking,
)
yield _format_sse(rs_text_done)
rs_part_done = ResponseReasoningSummaryPartDoneEvent(
sequence_number=next(seq),
item_id=reasoning_id,
output_index=reasoning_output_index,
summary_index=0,
part=ResponseReasoningSummaryText(text=accumulated_thinking),
)
yield _format_sse(rs_part_done)
rs_item_done = ResponseOutputItemDoneEvent(
sequence_number=next(seq),
output_index=reasoning_output_index,
item=ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
),
)
yield _format_sse(rs_item_done)
# If no message block was started, create one now (empty text)
if not message_started:
message_output_index = next_output_index
next_output_index += 1
initial_item = ResponseMessageItem(
id=item_id,
content=[ResponseOutputText(text="")],
status="in_progress",
)
item_added = ResponseOutputItemAddedEvent(
sequence_number=next(seq),
output_index=message_output_index,
item=initial_item,
)
yield _format_sse(item_added)
initial_part = ResponseOutputText(text="")
part_added_evt = ResponseContentPartAddedEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=message_output_index,
content_index=0,
part=initial_part,
)
yield _format_sse(part_added_evt)
# response.output_text.done
text_done = ResponseTextDoneEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
output_index=message_output_index,
content_index=0,
text=accumulated_text,
)
@@ -354,7 +532,7 @@ async def generate_responses_stream(
part_done = ResponseContentPartDoneEvent(
sequence_number=next(seq),
item_id=item_id,
output_index=0,
output_index=message_output_index,
content_index=0,
part=final_part,
)
@@ -367,7 +545,9 @@ async def generate_responses_stream(
status="completed",
)
item_done = ResponseOutputItemDoneEvent(
sequence_number=next(seq), output_index=0, item=final_message_item
sequence_number=next(seq),
output_index=message_output_index,
item=final_message_item,
)
yield _format_sse(item_done)
@@ -381,7 +561,15 @@ async def generate_responses_stream(
)
# response.completed
output: list[ResponseItem] = [final_message_item]
output: list[ResponseItem] = []
if reasoning_started:
output.append(
ResponseReasoningItem(
id=reasoning_id,
summary=[ResponseReasoningSummaryText(text=accumulated_thinking)],
)
)
output.append(final_message_item)
output.extend(function_call_items)
final_response = ResponsesResponse(
id=response_id,
+291 -72
View File
@@ -11,8 +11,7 @@ from typing import Annotated, Literal, cast
from uuid import uuid4
import anyio
from anyio import BrokenResourceError, create_task_group
from anyio.abc import TaskGroup
from anyio import BrokenResourceError
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -32,6 +31,14 @@ from exo.master.adapters.claude import (
collect_claude_response,
generate_claude_stream,
)
from exo.master.adapters.ollama import (
collect_ollama_chat_response,
collect_ollama_generate_response,
generate_ollama_chat_stream,
generate_ollama_generate_stream,
ollama_generate_request_to_text_generation,
ollama_request_to_text_generation,
)
from exo.master.adapters.responses import (
collect_responses_response,
generate_responses_stream,
@@ -43,6 +50,7 @@ from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
from exo.shared.constants import (
DASHBOARD_DIR,
EXO_CACHE_HOME,
EXO_EVENT_LOG_DIR,
EXO_IMAGE_CACHE_DIR,
EXO_MAX_CHUNK_SIZE,
@@ -65,6 +73,7 @@ from exo.shared.types.api import (
BenchChatCompletionResponse,
BenchImageGenerationResponse,
BenchImageGenerationTaskParams,
CancelCommandResponse,
ChatCompletionChoice,
ChatCompletionMessage,
ChatCompletionRequest,
@@ -73,6 +82,8 @@ from exo.shared.types.api import (
CreateInstanceResponse,
DeleteDownloadResponse,
DeleteInstanceResponse,
DeleteTracesRequest,
DeleteTracesResponse,
ErrorInfo,
ErrorResponse,
FinishReason,
@@ -132,16 +143,27 @@ 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, SystemId
from exo.shared.types.events import (
ChunkGenerated,
Event,
ForwarderEvent,
IndexedEvent,
PrefillProgress,
TracesMerged,
)
from exo.shared.types.memory import Memory
from exo.shared.types.ollama_api import (
OllamaChatRequest,
OllamaChatResponse,
OllamaGenerateRequest,
OllamaGenerateResponse,
OllamaModelDetails,
OllamaModelTag,
OllamaPsModel,
OllamaPsResponse,
OllamaShowRequest,
OllamaShowResponse,
OllamaTagsResponse,
)
from exo.shared.types.openai_responses import (
ResponsesRequest,
ResponsesResponse,
@@ -152,9 +174,10 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
from exo.utils.banner import print_startup_banner
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.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:
@@ -174,11 +197,9 @@ class API:
def __init__(
self,
node_id: NodeId,
session_id: SessionId,
*,
port: int,
# Ideally this would be a MasterForwarderEvent but type system says no :(
global_event_receiver: Receiver[ForwarderEvent],
event_receiver: Receiver[IndexedEvent],
command_sender: Sender[ForwarderCommand],
download_command_sender: Sender[ForwarderDownloadCommand],
# This lets us pause the API if an election is running
@@ -186,13 +207,12 @@ 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
self.event_receiver = event_receiver
self.election_receiver = election_receiver
self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
self.node_id: NodeId = node_id
self.session_id: SessionId = session_id
self.last_completed_election: int = 0
self.port = port
@@ -230,18 +250,20 @@ 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):
def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
logger.info("Resetting API State")
self._event_log.close()
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
self.state = State()
self.session_id = new_session_id
self.event_buffer = OrderedBuffer[Event]()
self._system_id = SystemId()
self._text_generation_queues = {}
self._image_generation_queues = {}
self.unpause(result_clock)
self.event_receiver.close()
self.event_receiver = event_receiver
self._tg.start_soon(self._apply_state)
def unpause(self, result_clock: int):
logger.info("Unpausing API")
@@ -301,14 +323,33 @@ class API:
self.app.get("/images/{image_id}")(self.get_image)
self.app.post("/v1/messages", response_model=None)(self.claude_messages)
self.app.post("/v1/responses", response_model=None)(self.openai_responses)
self.app.post("/v1/cancel/{command_id}")(self.cancel_command)
# Ollama API
self.app.head("/ollama/")(self.ollama_version)
self.app.head("/ollama/api/version")(self.ollama_version)
self.app.post("/ollama/api/chat", response_model=None)(self.ollama_chat)
self.app.post("/ollama/api/api/chat", response_model=None)(self.ollama_chat)
self.app.post("/ollama/api/v1/chat", response_model=None)(self.ollama_chat)
self.app.post("/ollama/api/generate", response_model=None)(self.ollama_generate)
self.app.get("/ollama/api/tags")(self.ollama_tags)
self.app.get("/ollama/api/api/tags")(self.ollama_tags)
self.app.get("/ollama/api/v1/tags")(self.ollama_tags)
self.app.post("/ollama/api/show")(self.ollama_show)
self.app.get("/ollama/api/ps")(self.ollama_ps)
self.app.get("/ollama/api/version")(self.ollama_version)
self.app.get("/state")(lambda: self.state)
self.app.get("/events")(self.stream_events)
self.app.post("/download/start")(self.start_download)
self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
self.app.get("/v1/traces")(self.list_traces)
self.app.post("/v1/traces/delete")(self.delete_traces)
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(
@@ -528,6 +569,25 @@ class API:
instance_id=instance_id,
)
async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
"""Cancel an active command by closing its stream and notifying workers."""
sender = self._text_generation_queues.get(
command_id
) or self._image_generation_queues.get(command_id)
if sender is None:
raise HTTPException(
status_code=404,
detail="Command not found or already completed",
)
await self._send(TaskCancelled(cancelled_command_id=command_id))
sender.close()
return CancelCommandResponse(
message="Command cancelled.",
command_id=command_id,
)
async def _token_chunk_stream(
self, command_id: CommandId
) -> AsyncGenerator[
@@ -554,7 +614,7 @@ class API:
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
@@ -902,7 +962,7 @@ class API:
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
@@ -988,7 +1048,7 @@ class API:
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
ForwarderCommand(origin=self._system_id, command=command)
)
raise
finally:
@@ -1294,6 +1354,163 @@ class API:
media_type="application/json",
)
async def _ollama_root(self) -> JSONResponse:
"""Respond to HEAD / from Ollama CLI connectivity checks."""
return JSONResponse(content="Ollama is running")
async def ollama_chat(
self, request: Request
) -> OllamaChatResponse | StreamingResponse:
"""Ollama Chat API — accepts JSON regardless of Content-Type."""
body = await request.body()
payload = OllamaChatRequest.model_validate_json(body)
task_params = ollama_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
command = TextGeneration(task_params=task_params)
await self._send(command)
if payload.stream:
return StreamingResponse(
generate_ollama_chat_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/x-ndjson",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
else:
return StreamingResponse(
collect_ollama_chat_response(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
async def ollama_generate(
self, request: Request
) -> OllamaGenerateResponse | StreamingResponse:
"""Ollama Generate API — accepts JSON regardless of Content-Type."""
body = await request.body()
payload = OllamaGenerateRequest.model_validate_json(body)
task_params = ollama_generate_request_to_text_generation(payload)
resolved_model = await self._resolve_and_validate_text_model(
ModelId(task_params.model)
)
task_params = task_params.model_copy(update={"model": resolved_model})
command = TextGeneration(task_params=task_params)
await self._send(command)
if payload.stream:
return StreamingResponse(
generate_ollama_generate_stream(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/x-ndjson",
headers={
"Cache-Control": "no-cache",
"Connection": "close",
"X-Accel-Buffering": "no",
},
)
else:
return StreamingResponse(
collect_ollama_generate_response(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
async def ollama_tags(self) -> OllamaTagsResponse:
"""Returns list of models in Ollama tags format. We return the downloaded ones only."""
def none_if_empty(value: str) -> str | None:
return value or None
downloaded_model_ids: set[str] = set()
for node_downloads in self.state.downloads.values():
for dl in node_downloads:
if isinstance(dl, DownloadCompleted):
downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
cards = [
c for c in await get_model_cards() if c.model_id in downloaded_model_ids
]
now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
return OllamaTagsResponse(
models=[
OllamaModelTag(
name=str(card.model_id),
model=str(card.model_id),
modified_at=now,
size=card.storage_size.in_bytes,
digest="sha256:000000000000",
details=OllamaModelDetails(
family=none_if_empty(card.family),
quantization_level=none_if_empty(card.quantization),
),
)
for card in cards
]
)
async def ollama_show(self, request: Request) -> OllamaShowResponse:
"""Returns model information in Ollama show format."""
body = await request.body()
payload = OllamaShowRequest.model_validate_json(body)
model_name = payload.name or payload.model
if not model_name:
raise HTTPException(status_code=400, detail="name or model is required")
try:
card = await ModelCard.load(ModelId(model_name))
except Exception as exc:
raise HTTPException(
status_code=404, detail=f"Model not found: {model_name}"
) from exc
return OllamaShowResponse(
modelfile=f"FROM {card.model_id}",
template="{{ .Prompt }}",
details=OllamaModelDetails(
family=card.family or None,
quantization_level=card.quantization or None,
),
)
async def ollama_ps(self) -> OllamaPsResponse:
"""Returns list of running models (active instances)."""
models: list[OllamaPsModel] = []
seen: set[str] = set()
for instance in self.state.instances.values():
model_id = str(instance.shard_assignments.model_id)
if model_id in seen:
continue
seen.add(model_id)
models.append(
OllamaPsModel(
name=model_id,
model=model_id,
size=0,
)
)
return OllamaPsResponse(models=models)
async def ollama_version(self) -> dict[str, str]:
"""Returns version information for Ollama API compatibility."""
return {"version": "exo v1.0"}
def _calculate_total_available_memory(self) -> Memory:
"""Calculate total available memory across all nodes in bytes."""
total_available = Memory()
@@ -1323,7 +1540,7 @@ class API:
name=card.model_id.short(),
description="",
tags=[],
storage_size_megabytes=int(card.storage_size.in_mb),
storage_size_megabytes=card.storage_size.in_mb,
supports_tensor=card.supports_tensor,
tasks=[task.value for task in card.tasks],
is_custom=is_custom_card(card.model_id),
@@ -1394,8 +1611,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)
@@ -1410,7 +1626,7 @@ class API:
finally:
self._event_log.close()
self.command_sender.close()
self.global_event_receiver.close()
self.event_receiver.close()
async def run_api(self, ev: anyio.Event):
cfg = Config()
@@ -1427,52 +1643,31 @@ class API:
)
async def _apply_state(self):
with self.global_event_receiver as events:
async for f_event in events:
if f_event.origin != self.session_id.master_node_id:
continue
self.event_buffer.ingest(f_event.origin_idx, f_event.event)
for idx, event in self.event_buffer.drain_indexed():
self._event_log.append(event)
self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
with self.event_receiver as events:
async for i_event in events:
self._event_log.append(i_event.event)
self.state = apply(self.state, i_event)
event = i_event.event
if isinstance(event, ChunkGenerated):
if queue := self._image_generation_queues.get(
event.command_id, None
):
assert isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except BrokenResourceError:
self._image_generation_queues.pop(
event.command_id, None
)
if queue := self._text_generation_queues.get(
event.command_id, None
):
assert not isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except BrokenResourceError:
self._text_generation_queues.pop(event.command_id, None)
elif isinstance(event, PrefillProgress):
if queue := self._text_generation_queues.get(
event.command_id, None
):
try:
await queue.send(
PrefillProgressChunk(
model=event.model,
processed_tokens=event.processed_tokens,
total_tokens=event.total_tokens,
)
)
except BrokenResourceError:
self._text_generation_queues.pop(event.command_id, None)
if isinstance(event, TracesMerged):
self._save_merged_trace(event)
if isinstance(event, ChunkGenerated):
if queue := self._image_generation_queues.get(
event.command_id, None
):
assert isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except BrokenResourceError:
self._image_generation_queues.pop(event.command_id, None)
if queue := self._text_generation_queues.get(
event.command_id, None
):
assert not isinstance(event.chunk, ImageChunk)
try:
await queue.send(event.chunk)
except BrokenResourceError:
self._text_generation_queues.pop(event.command_id, None)
if isinstance(event, TracesMerged):
self._save_merged_trace(event)
def _save_merged_trace(self, event: TracesMerged) -> None:
traces = [
@@ -1508,12 +1703,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(
@@ -1536,8 +1731,12 @@ class API:
await self._send_download(command)
return DeleteDownloadResponse(command_id=command.command_id)
def _get_trace_path(self, task_id: str) -> Path:
return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
@staticmethod
def _get_trace_path(task_id: str) -> Path:
trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
return trace_path
async def list_traces(self) -> TraceListResponse:
traces: list[TraceListItem] = []
@@ -1635,3 +1834,23 @@ class API:
media_type="application/json",
filename=f"trace_{task_id}.json",
)
async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse:
deleted: list[str] = []
not_found: list[str] = []
for task_id in request.task_ids:
trace_path = self._get_trace_path(task_id)
if trace_path.exists():
trace_path.unlink()
deleted.append(task_id)
else:
not_found.append(task_id)
return DeleteTracesResponse(deleted=deleted, not_found=not_found)
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})
+31 -47
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,
@@ -60,8 +60,9 @@ from exo.shared.types.tasks import (
TextGeneration as TextGenerationTask,
)
from exo.shared.types.worker.instances import InstanceId
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.channels import Receiver, Sender
from exo.utils.event_buffer import MultiSourceBuffer
from exo.utils.task_group import TaskGroup
class Master:
@@ -71,26 +72,23 @@ class Master:
session_id: SessionId,
*,
command_receiver: Receiver[ForwarderCommand],
local_event_receiver: Receiver[ForwarderEvent],
global_event_sender: Sender[ForwarderEvent],
event_sender: Sender[Event],
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.node_id = node_id
self.session_id = session_id
self.state = State()
self._tg: TaskGroup = TaskGroup()
self.command_task_mapping: dict[CommandId, TaskId] = {}
self.command_receiver = command_receiver
self.local_event_receiver = local_event_receiver
self.global_event_sender = global_event_sender
self.download_command_sender = download_command_sender
send, recv = channel[Event]()
self.event_sender: Sender[Event] = send
self._loopback_event_receiver: Receiver[Event] = recv
self._loopback_event_sender: Sender[ForwarderEvent] = (
local_event_receiver.clone_sender()
)
self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
self.event_sender = event_sender
self._system_id = SystemId()
self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
self._expected_ranks: dict[TaskId, set[int]] = {}
@@ -102,19 +100,16 @@ class Master:
async with self._tg as tg:
tg.start_soon(self._event_processor)
tg.start_soon(self._command_processor)
tg.start_soon(self._loopback_processor)
tg.start_soon(self._plan)
finally:
self._event_log.close()
self.global_event_sender.close()
self.local_event_receiver.close()
self.command_receiver.close()
self._loopback_event_sender.close()
self._loopback_event_receiver.close()
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 +283,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)
@@ -333,17 +328,22 @@ class Master:
task_id=task_id,
)
)
case TaskFinished():
generated_events.append(
TaskDeleted(
task_id=self.command_task_mapping[
command.finished_command_id
]
else:
logger.warning(
f"Nonexistent command {command.cancelled_command_id} cancelled"
)
)
self.command_task_mapping.pop(
command.finished_command_id, None
)
case TaskFinished():
if (
task_id := self.command_task_mapping.pop(
command.finished_command_id, None
)
) is not None:
generated_events.append(TaskDeleted(task_id=task_id))
else:
logger.warning(
f"Finished command {command.finished_command_id} finished"
)
case RequestEventLog():
# We should just be able to send everything, since other buffers will ignore old messages
# rate limit to 1000 at a time
@@ -407,27 +407,11 @@ class Master:
self._event_log.append(event)
await self._send_event(indexed)
async def _loopback_processor(self) -> None:
# this would ideally not be necessary.
# this is WAY less hacky than how I was working around this before
local_index = 0
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}"),
origin_idx=local_index,
session=self.session_id,
event=event,
)
)
local_index += 1
# This function is re-entrant, take care!
async def _send_event(self, event: IndexedEvent):
# Convenience method since this line is ugly
await self.global_event_sender.send(
ForwarderEvent(
GlobalForwarderEvent(
origin=self.node_id,
origin_idx=event.idx,
session=self.session_id,
+26 -6
View File
@@ -112,7 +112,11 @@ def place_instance(
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
]
if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
if command.instance_meta == InstanceMeta.MlxJaccl:
if not smallest_rdma_cycles:
raise ValueError(
"Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
)
smallest_cycles = smallest_rdma_cycles
cycles_with_leaf_nodes: list[Cycle] = [
@@ -129,6 +133,11 @@ def place_instance(
),
)
# Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
if len(selected_cycle) == 1:
command.instance_meta = InstanceMeta.MlxRing
command.sharding = Sharding.Pipeline
shard_assignments = get_shard_assignments(
command.model_card, selected_cycle, command.sharding, node_memory
)
@@ -138,18 +147,29 @@ def place_instance(
instance_id = InstanceId()
target_instances = dict(deepcopy(current_instances))
if len(selected_cycle) == 1:
command.instance_meta = InstanceMeta.MlxRing
# TODO: Single node instances
match command.instance_meta:
case InstanceMeta.MlxJaccl:
# TODO(evan): shard assignments should contain information about ranks, this is ugly
def get_device_rank(node_id: NodeId) -> int:
runner_id = shard_assignments.node_to_runner[node_id]
shard_metadata = shard_assignments.runner_to_shard.get(runner_id)
assert shard_metadata is not None
return shard_metadata.device_rank
zero_node_ids = [
node_id
for node_id in selected_cycle.node_ids
if get_device_rank(node_id) == 0
]
assert len(zero_node_ids) == 1
coordinator_node_id = zero_node_ids[0]
mlx_jaccl_devices = get_mlx_jaccl_devices_matrix(
[node_id for node_id in selected_cycle],
cycle_digraph,
)
mlx_jaccl_coordinators = get_mlx_jaccl_coordinators(
coordinator=selected_cycle.node_ids[0],
coordinator=coordinator_node_id,
coordinator_port=random_ephemeral_port(),
cycle_digraph=cycle_digraph,
node_network=node_network,
+32 -16
View File
@@ -102,22 +102,21 @@ def _allocate_and_validate_layers(
layer_allocations = allocate_layers_proportionally(
total_layers=model_card.n_layers,
memory_fractions=[
node_memory[node_id].ram_available.in_bytes / total_memory.in_bytes
for node_id in node_ids
node_memory[node_id].ram_available / total_memory for node_id in node_ids
],
)
total_storage_bytes = model_card.storage_size.in_bytes
total_storage = model_card.storage_size
total_layers = model_card.n_layers
for i, node_id in enumerate(node_ids):
node_layers = layer_allocations[i]
required_memory = (total_storage_bytes * node_layers) // total_layers
available_memory = node_memory[node_id].ram_available.in_bytes
required_memory = (total_storage * node_layers) // total_layers
available_memory = node_memory[node_id].ram_available
if required_memory > available_memory:
raise ValueError(
f"Node {i} ({node_id}) has insufficient memory: "
f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
f"but only has {available_memory / (1024**3):.2f} GB available"
f"requires {required_memory.in_gb:.2f} GB for {node_layers} layers, "
f"but only has {available_memory.in_gb:.2f} GB available"
)
return layer_allocations
@@ -342,6 +341,7 @@ def _find_ip_prioritised(
other_node_id: NodeId,
cycle_digraph: Topology,
node_network: Mapping[NodeId, NodeNetworkInfo],
ring: bool,
) -> str | None:
"""Find an IP address between nodes with prioritization.
@@ -354,13 +354,27 @@ def _find_ip_prioritised(
ip_to_type = {
iface.ip_address: iface.interface_type for iface in other_network.interfaces
}
priority = {
"ethernet": 0,
"wifi": 1,
"unknown": 2,
"maybe_ethernet": 3,
"thunderbolt": 4,
}
# Ring should prioritise fastest connection. As a best-effort, we prioritise TB.
# TODO: Profile and get actual connection speeds.
if ring:
priority = {
"thunderbolt": 0,
"maybe_ethernet": 1,
"ethernet": 2,
"wifi": 3,
"unknown": 4,
}
# RDMA prefers ethernet coordinator
else:
priority = {
"ethernet": 0,
"wifi": 1,
"unknown": 2,
"maybe_ethernet": 3,
"thunderbolt": 4,
}
return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
@@ -400,7 +414,7 @@ def get_mlx_ring_hosts_by_node(
continue
connection_ip = _find_ip_prioritised(
node_id, other_node_id, cycle_digraph, node_network
node_id, other_node_id, cycle_digraph, node_network, ring=True
)
if connection_ip is None:
raise ValueError(
@@ -431,7 +445,9 @@ def get_mlx_jaccl_coordinators(
if n == coordinator:
return "0.0.0.0"
ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network)
ip = _find_ip_prioritised(
n, coordinator, cycle_digraph, node_network, ring=False
)
if ip is not None:
return ip

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