Compare commits

...

91 Commits

Author SHA1 Message Date
Alex Cheema 9cacef39f0 feat: add --bootstrap-peer flag to bypass mDNS for peer discovery
macOS TCC blocks mDNS multicast from SSH sessions, preventing
SSH-spawned exo processes from discovering peers. This adds a
--bootstrap-peer IP:PORT flag that directly dials a known peer,
bypassing mDNS. The flag is additive — mDNS still runs alongside.

Changes across Rust networking, PyO3 bindings, and Python CLI:
- discovery.rs: store bootstrap addresses, dial on startup and retry
- swarm.rs: thread bootstrap_peers through to Behaviour
- networking.rs: accept optional bootstrap_peers in PyO3 constructor
- router.py: pass bootstrap_peers to NetworkingHandle
- main.py: parse --bootstrap-peer CLI arg into multiaddr format

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-16 10:24:00 -08:00
Alex Cheema 3df896dec1 fix: cancel active tasks on meta-instance cascade delete
Previously, DeleteMetaInstance cascade-deleted backing instances without
cancelling their active tasks, leaving orphaned task references. Now emits
TaskStatusUpdated(Cancelled) for Pending/Running tasks before InstanceDeleted.

Also adds lifecycle logging for meta-instance operations, a GET /meta_instances
endpoint, and 2 regression tests.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-16 09:51:55 -08:00
Alex Cheema 29c3489f3e Merge remote-tracking branch 'origin/main' into alexcheema/mlx-distributed-transfer
Resolve merge conflicts between meta-instance feature and task cancellation:
- commands.py: Include both CreateMetaInstance/DeleteMetaInstance and TaskCancelled
- plan.py: Add _cancel_tasks from main, keep all_runners param in _create_runner
- runner_supervisor.py: Merge cancel_sender/cancelled fields with pipe relay fields
- reconcile.py: Pass empty tasks dict to get_transition_events (new param from main)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-16 05:38:47 -08:00
Evan Quiney 2759e92334 api cancellation (#1276)
closing the http request to the api now
- sends a cancellation from the api
- writes that canellation in the master
- worker plans off the cancellation
- runner observes that cancellation after every generation step (+1
communication per token)
- cancellation happens synchronously to prevent gpu locks

closes #61

---------

Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-16 05:30:07 -08:00
Alex Cheema 19f606a379 fix: dashboard derived reactivity bug and tautological check in meta-instance UI
1. DERIVED REACTIVITY BUG: `unifiedDisplayItems` used `$derived(fn)` which
   made the derived value the function itself instead of its result. Svelte
   never tracked reactive dependencies in the function body, so the instance
   list didn't update when metaInstances or instances changed. Fixed by using
   `$derived.by(fn)` and removing the `()` call-sites in the template.

2. TAUTOLOGICAL CHECK: In `getMetaInstancePlacingStatus`, the `lastError ? ...
   : null` guard inside the `failures > 0` branch was always true because
   `lastFailureError` and `consecutiveFailures` are always set together in
   `apply_instance_retrying` and `apply_instance_deleted`. Removed the dead
   `: null` branch.

Also fixes pyright errors in test file by using proper pytest.MonkeyPatch type.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-16 05:24:49 -08:00
Alex Cheema c74b46dca0 fix: add timeout and error handling for ModelCard.load in MetaInstanceReconciler
ModelCard.load() does async I/O inside the 1-second reconcile loop. A slow
or failing load blocked all reconciliation (health checks, node timeouts,
other meta-instances). Adds a 10-second timeout, per-meta-instance error
handling with MetaInstancePlacementFailed events, and documents the
intentional early return in apply_instance_retrying.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-16 05:11:29 -08:00
Jake Hillion 131fb141a6 bench: add --danger-delete-downloads flag with planning phase
exo bench previously relied on the worker's plan loop to download
models, which could fail silently or run into disk space issues during
benchmarking. This made it difficult to diagnose download failures.

Added a planning phase that runs before benchmarking to explicitly
handle downloads. It checks available disk space on each node via the
/state endpoint and starts downloads via POST /download/start. When
the --danger-delete-downloads flag is set and there's insufficient
space, it deletes existing models from smallest to largest until
there's room for the benchmark model.

Test plan:
- CI

```
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --pp 128,2048,4096 --tg 128 --stdout --settle-timeout 10 --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-02-16 12:12:11.807 | INFO     | __main__:main:710 - pp/tg mode: combinations (product) - 3 pairs
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
2026-02-16 12:12:13.455 | DEBUG    | __main__:main:725 - [exo-bench] loaded tokenizer: mlx-community/gpt-oss-120b-MXFP4-Q8 for prompt sizer
2026-02-16 12:12:13.473 | DEBUG    | __main__:main:761 - exo-bench model: short_id=gpt-oss-120b-MXFP4-Q8 full_id=mlx-community/gpt-oss-120b-MXFP4-Q8
2026-02-16 12:12:13.473 | INFO     | __main__:main:762 - placements: 1
2026-02-16 12:12:13.474 | INFO     | __main__:main:764 -   - Pipeline / MlxRing / nodes=1
2026-02-16 12:12:13.474 | INFO     | __main__:main:771 - Planning phase: checking downloads...
Traceback (most recent call last):
  File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 885, in <module>
    raise SystemExit(main())
                     ~~~~^^
  File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 772, in main
    run_planning_phase(
    ~~~~~~~~~~~~~~~~~~^
        client,
        ^^^^^^^
    ...<4 lines>...
        settle_deadline,
        ^^^^^^^^^^^^^^^^
    )
    ^
  File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 367, in run_planning_phase
    raise RuntimeError(
    ...<2 lines>...
    )
RuntimeError: Insufficient disk on 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj: need 65GB, have 55GB. Use --danger-delete-downloads to free space.
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --pp 128,2048,4096 --tg 128 --stdout --settle-timeout 10 --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8 --danger-delete-downloads
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-02-16 12:12:19.626 | INFO     | __main__:main:710 - pp/tg mode: combinations (product) - 3 pairs
2026-02-16 12:12:21.262 | DEBUG    | __main__:main:725 - [exo-bench] loaded tokenizer: mlx-community/gpt-oss-120b-MXFP4-Q8 for prompt sizer
2026-02-16 12:12:21.280 | DEBUG    | __main__:main:761 - exo-bench model: short_id=gpt-oss-120b-MXFP4-Q8 full_id=mlx-community/gpt-oss-120b-MXFP4-Q8
2026-02-16 12:12:21.280 | INFO     | __main__:main:762 - placements: 1
2026-02-16 12:12:21.280 | INFO     | __main__:main:764 -   - Pipeline / MlxRing / nodes=1
2026-02-16 12:12:21.280 | INFO     | __main__:main:771 - Planning phase: checking downloads...
2026-02-16 12:12:21.336 | INFO     | __main__:run_planning_phase:386 - Deleting mlx-community/Qwen3-0.6B-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (335MB)
2026-02-16 12:12:21.350 | INFO     | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-1B-Instruct-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (679MB)
2026-02-16 12:12:21.363 | INFO     | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-3B-Instruct-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (1740MB)
2026-02-16 12:12:21.373 | INFO     | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-3B-Instruct-8bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (3264MB)
2026-02-16 12:12:21.384 | INFO     | __main__:run_planning_phase:386 - Deleting mlx-community/GLM-4.7-Flash-8bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (30366MB)
2026-02-16 12:12:21.413 | INFO     | __main__:run_planning_phase:407 - Started download on 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj
```

It's not pretty but it works!
2026-02-16 13:06:38 +00:00
Evan Quiney 2d8bfc2e3c fix: PlaceInstanceParams broken field validator
our field validator for PlaceInstance was wrong - we can just rely on default behaviour here anyway!
2026-02-16 03:58:43 -08:00
ciaranbor 042999f728 Ciaran/message deletion (#1409)
## Motivation

When a user deletes a message during an active streamed generation, it
can cause unexpected behavior. The delete confirmation text was also
misleading — it said "all responses after it" only for user messages,
which didn't accurately describe the behavior (all messages after the
deleted one are removed, regardless of role)

## Changes

- Prevent deletion during streaming: Disabled the delete button and
blocked handleDeleteClick when loading is true, with a visual indication
(dimmed button, cursor-not-allowed, tooltip change)
- Clarified delete confirmation text: Replaced role-specific wording
with a simpler, accurate message:
    - Last message: "Delete this message?"
- Any other message: "Delete this message and all messages after it?"

## Why It Works

Guarding on the loading state at both the click handler and the button's
disabled attribute ensures no deletion can be triggered while a response
is being streamed

## Test Plan

### Manual Testing

- Verify the delete button is visually disabled and non-clickable while
a response is streaming
- Verify the tooltip shows "Cannot delete while generating" during
streaming
  - Verify the last message shows "Delete this message?" confirmation
- Verify non-last messages show "Delete this message and all messages
after it?" confirmation
  - Verify deletion works normally when not streaming
2026-02-16 11:46:41 +00:00
ciaranbor b61dc2eb35 Prevent image editing without image input (#1410)
## Motivation

Models that only support image editing (ImageToImage but not
TextToImage) would silently attempt text-to-image generation when a user
submitted a text prompt without an attached image

## Changes

- Added an early return guard in handleSubmit() that prevents submission
when the selected model only supports image editing and no image is
attached (isEditOnlyWithoutImage)
- Fixed the text-to-image generation branch to use the more specific
modelSupportsTextToImage() check instead of the broad isImageModel(),
ensuring only models with TextToImage capability trigger generation from
text alone
- The existing isEditOnlyWithoutImage derived state (which was already
used for UI hints like placeholder text and button disabling) now also
blocks the actual submit path

## Why It Works

The text-to-image fallback now correctly checks
modelSupportsTextToImage() directly, so edit-only models no longer fall
through to the generation path


## Test Plan

### Manual Testing

- Select an edit-only image model (e.g., one with only ImageToImage
capability)
- Verify the send button is disabled and placeholder reads "Attach an
image to edit..." when no image is attached
  - Attach an image and verify the form becomes submittable
- Select a text-to-image model and verify text-only prompts still
trigger generation normally
  - Ensure pressing `enter` doesn't bypass check
2026-02-16 11:39:59 +00:00
rltakashige 36a7115b6f Pass usage and generation stats through all adapters correctly (#1461)
## Motivation

Exo is not returning usage stats correctly at the moment.

## Changes

- Correctly pass usage stats instead of generation stats.
- Pass usage stats within tool calls.

## Test Plan

### Manual Testing
Needs manual testing.

### Automated Testing
Passes CI.
2026-02-16 11:20:04 +00:00
Jake Hillion 0b7d88b43b python: add hermetic basedpyright typecheck to nix flake check
The existing CI typecheck job used `uv run basedpyright` which depends
on a non-hermetic uv sync step. This replaces it with a fully hermetic
typecheck as a Nix flake check using the uv2nix virtual environment.

Added a `typecheckVenv` with dev dependencies, a `linuxOverlay` to
ignore native shared library deps (NVIDIA, torch, triton, mlx) that
aren't needed at type-check time, and `passthru` preservation plus
`.pyi` stub copying on the `exo-pyo3-bindings` overlay so basedpyright
can resolve the Rust bindings types. Also guarded the `mlx` Nix build
override to macOS only since it requires Metal. Removed the old
non-hermetic `typecheck` CI job since `nix flake check` now covers it.

The hermetic check ensures type checking uses exactly the locked
dependency versions and catches type errors without requiring a
working uv/pip environment.

Test plan:
- CI (`nix flake check` runs on x86_64-linux, aarch64-linux, aarch64-darwin)
- Verified `nix build ".#checks.x86_64-linux.typecheck"` passes with 0 errors
2026-02-16 11:09:23 +00:00
Alex Cheema 7028f99d7e fix: validate placement before sending command to prevent silent failures (BUG-001c)
The place_instance API endpoint used fire-and-forget: it sent the command
and returned HTTP 200 immediately. On a fresh cluster start, the master's
state often lacks topology/memory data, so placement raises ValueError
which was silently caught and logged. The caller never learned it failed.

Two fixes:
- API: validate placement locally before sending, return HTTP 400 on
  failure instead of silently accepting an unprocessable command
- Master: emit MetaInstancePlacementFailed on immediate placement error
  in CreateMetaInstance handler so the error surfaces in state right away

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 16:13:38 -08:00
Alex Cheema 68d9debcb6 fix: place_instance now uses all available nodes to prevent OOM (BUG-001d)
The placement algorithm previously selected the smallest viable cycle,
causing large models to be distributed across too few nodes and running
out of memory. Changed get_smallest_cycles to get_largest_cycles so that
all healthy nodes are utilized, spreading layers more evenly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 15:13:03 -08:00
Alex Cheema 9bf5f86b9e Revert "feat: add task cancellation for client disconnect handling (BUG-001)"
This reverts commit 2a75672a4a.
2026-02-15 13:48:01 -08:00
Alex Cheema 2a75672a4a feat: add task cancellation for client disconnect handling (BUG-001)
- Add TaskCancelled command and Cancelled task status
- Detect API client disconnects in master/api.py
- Handle TaskCancelled in master state machine
- Add _cancel_tasks to worker for graceful task cleanup
- Add cancel_receiver to runner for inference abort
- Add mx_any helper in MLX utils for cancellable operations
- Guard instance lookup in worker to prevent KeyError
- Update tests for cancellation flow

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-15 13:43:15 -08:00
Alex Cheema 8a6344fba3 Fix model discovery reporting DownloadPending for fully-downloaded models
On startup, _emit_existing_download_progress() used
downloaded_bytes_this_session to decide between DownloadPending and
DownloadOngoing. Since downloaded_bytes_this_session is always 0 on
startup (it tracks the current session only), fully-downloaded models
were incorrectly reported as DownloadPending.

Now checks actual disk state: if downloaded_bytes >= total_bytes, emit
DownloadCompleted regardless of session bytes. This fixes the UI showing
models as pending when they're already available.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-14 09:39:06 -08:00
Alex Cheema 316a7a344e fix: use force=True for multiprocessing set_start_method
Prevents RuntimeError when the context has already been set,
e.g. when Terminal.app reuses a tab or the process restarts.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-14 07:02:32 -08:00
Alex Cheema 16b58d6946 fix: eliminate command/reconciler interleaving race in meta-instance
Two race conditions existed in the meta-instance lifecycle:

1. CreateMetaInstance buffered MetaInstanceCreated but didn't apply it
   before awaiting ModelCard.load(). The reconciler could interleave
   during the await, leading to duplicate placements.

   Fix: apply MetaInstanceCreated eagerly via _apply_and_broadcast,
   then re-check satisfaction after the await so placement uses fresh
   state and skips if the reconciler already handled it.

2. delete_meta_instance (API handler) sent DeleteMetaInstance then
   read self.state.instances for cascade deletion. State was stale,
   so backing instances created between the send and the read were
   missed — permanently orphaning them.

   Fix: move cascade delete into the command processor's
   DeleteMetaInstance handler, where InstanceDeleted events are
   generated atomically with MetaInstanceDeleted.

Reproduced on 4-node Mac Mini cluster: 28K anomalies in stress test
including 21 permanently orphaned instances. After fix, the cascade
delete and placement are race-free.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-13 12:50:57 -08:00
Alex Cheema 39741907c7 test: add 25 edge-case tests for MetaInstance lifecycle
Cover retry logic, error handling, backward compatibility,
concurrent scenarios, placement error tracking, and serialization.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-13 12:17:41 -08:00
Alex Cheema e5c73c564c Merge remote-tracking branch 'origin/main' into alexcheema/meta-instance 2026-02-13 10:10:38 -08:00
Alex Cheema 1c3cc699d4 fix: add missing getModelFitStatus prop to Recent tab (#1470)
## Summary
- Clicking the **Recent** tab in the Model Picker crashed with
`TypeError: e.getModelFitStatus is not a function`
- The `ModelPickerGroup` component in the Recent tab section was missing
the `{getModelFitStatus}` prop, while all other tabs (e.g., the main
model list) passed it correctly
- Added the missing `{getModelFitStatus}` prop so the Recent tab renders
without errors, matching the behavior of the other tabs

## Test plan
- [ ] Open the dashboard and click **SELECT MODEL**
- [ ] Switch to the **Recent** tab — verify it renders without crashing
- [ ] Confirm model fit status indicators display correctly on recent
models
- [ ] Verify the other tabs (All, Favorites) still work as before

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

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-13 16:37:32 +00:00
rltakashige 5a28642790 Add support for Step 3.5 flash! (#1460)
## Motivation

Working version of #1366 

## Changes

Add Step 3.5 Flash

## Test Plan

### Manual Testing
Works!

### Automated Testing
Running two processes tensor/pipeline sharded gives same logits as
single process.
2026-02-13 12:10:18 +00:00
Alex Cheema 77799a170a Fix JACCL SideChannel bytes serialization for JSON round-trip
TaggedModel's wrap validator converts JSON→Python validation context,
which breaks strict-mode bytes deserialization from JSON strings.
Use Base64Bytes type to encode/decode bytes as base64 strings in JSON.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-12 14:23:44 -08:00
Alex Cheema 5dddd9cd2b Use named pipes (FIFOs) for JACCL SideChannel relay
Anonymous pipes from os.pipe() don't survive multiprocessing.Process
spawn on macOS (default since Python 3.8). The FD numbers are passed
but the actual file descriptors don't exist in the child process,
causing EBADF errors.

Switch to named pipes (FIFOs) which the child opens by path in the
spawned process, getting valid FDs for the C++ SideChannel.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-12 12:06:37 -08:00
Alex Cheema 6950f94109 dashboard: show macOS version in debug mode (#1454)
## Motivation

When debugging cluster issues, it's useful to see which macOS version
each node is running — especially since version mismatches can cause
compatibility problems. The OS version data is already collected by the
identity gatherer but wasn't shown in the topology graph.

## Changes

- Added OS version label (e.g. "macOS 15.2") to each node in the
topology graph when debug mode is enabled
- Renders below the existing TB and RDMA debug labels using the same
styling conventions
- Sources data from the existing `nodeIdentities` store (no backend
changes needed)

## Why It Works

The `nodeIdentities` store already contains `osVersion` for each node.
We simply read it in the `TopologyGraph` component and append a text
label in the debug section, following the exact same pattern as the TB
and RDMA labels.

## Test Plan

### Manual Testing
<!-- Hardware: MacBook Pro -->
- Enable debug mode in the dashboard
- Verify OS version label appears below TB/RDMA labels on each node
- Verify label disappears when debug mode is disabled

### Automated Testing
- Dashboard build passes (`npm run build`)

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
2026-02-12 17:56:55 +00:00
Alex Cheema d0c44273db feat: add enable_thinking toggle for thinking-capable models (#1457)
## Motivation

Fixes #1456. Models like DeepSeek V3.2, Qwen3, and GLM-4.7 always run in
thinking mode because their chat templates auto-inject `<think>`. Users
need a way to disable thinking for models that support both modes.

## Changes

**API**: Added `enable_thinking: bool | None` to `ChatCompletionRequest`
and `TextGenerationTaskParams`. Passed through the adapter to
`tokenizer.apply_chat_template()` as a kwarg (only when explicitly set,
so models without the template variable are unaffected).

**Dashboard**: Added a thinking toggle button in the chat input area.
Visible only when the selected model has both "text" and "thinking"
capabilities.

## Why It Works

Most thinking model chat templates (DeepSeek, Qwen3, GLM) accept
`enable_thinking` as a Jinja template variable. Passing
`enable_thinking=False` prevents the template from injecting `<think>`,
matching the vLLM convention.

## Test Plan

### Manual Testing
- `curl` with `"enable_thinking": false` against a thinking model — no
`<think>` in output
- Dashboard toggle visible for thinking models, hidden for text-only
models

### Automated Testing
- basedpyright: 0 errors
- ruff: clean
- pytest: 188 passed
- dashboard build: success

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-12 17:35:24 +00:00
Alex Cheema 0e08c2bfd3 Add pipe-based JACCL SideChannel relay via exo control plane
Replace fragile TCP SideChannel with anonymous pipes relayed through
exo's event-sourced control plane. RunnerSupervisor creates pipe pairs
for MlxJaccl instances, relays all_gather rounds via JacclSideChannelData/
JacclSideChannelGathered events through the master, eliminating errno=57
crashes from Thunderbolt RDMA driver instability.

Also includes dashboard RDMA warning improvements and instance retry fixes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-12 08:55:40 -08:00
Jake Hillion cc33213842 bench: add --settle-timeout for cluster startup retry (#1449)
exo_bench.py fails if started too soon after a cluster starts because
the topology hasn't populated yet, resulting in no valid placements.

Extracted the preview-fetch-and-filter logic into a
`fetch_and_filter_placements` helper and added a retry loop with
exponential backoff (1s initial, 2x multiplier, 60s cap). The new
`--settle-timeout` flag controls how long to retry (default 0 = try
once, preserving existing behaviour). Each retry logs a warning
explaining the cluster may still be settling.

Test plan:
- Tested on several freshly started clusters. This used to fail a lot,
  now it succeeds.
2026-02-12 16:38:09 +00:00
Alex Cheema f9ffdaef5f Preserve last_failure_error across instance recreation, fix RDMA banner wording
- apply_instance_created no longer clears last_failure_error so the
  error context persists while the new instance starts up
- Dashboard retryError shows the error without (N/3) prefix when
  consecutiveFailures is 0 (instance was recreated)
- Jaccl warning tooltip now says "experimental RDMA driver in macOS"

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 16:48:30 -08:00
Alex Cheema 8c2416c9ea chore: remove temporary screenshot files
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 16:40:16 -08:00
Alex Cheema e5007f619a temp: add jaccl warning screenshots for PR comment
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 16:38:53 -08:00
Alex Cheema a627f67253 dashboard: show warning banner for [jaccl] RDMA driver errors
Detect errors containing "[jaccl]" in MetaInstance failure errors and
display a red dismissible alert banner. The tooltip explains this is a
macOS RDMA driver issue and that the affected machine needs to be
restarted. Alert re-appears if a new error arrives after dismissal.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 16:38:42 -08:00
Alex Cheema f189222bfc Merge remote-tracking branch 'origin/main' into alexcheema/meta-instance
# Conflicts:
#	dashboard/src/lib/stores/app.svelte.ts
#	dashboard/src/routes/+page.svelte
2026-02-11 15:59:50 -08:00
Alex Cheema 62e8110e97 fix: prevent DownloadModel TaskCreated event flood (#1452)
## Motivation

When a model download fails repeatedly (e.g. `ContentLengthError` on a
large model like `zai-org/GLM-5`), the download coordinator accumulates
duplicate progress callbacks — one per retry cycle. Each callback
independently throttles at 1 event/sec, so after N retries, every
download progress tick generates N events instead of 1. After an hour of
failures (~60 retry cycles), this produces ~60 `NodeDownloadProgress`
events/sec, overwhelming the master, delaying heartbeats, and causing
the node to time itself out.

### The callback accumulation cycle
1. `_start_download_task()` calls
`shard_downloader.on_progress(callback)` which **appends** to a list
2. Download fails → `DownloadFailed` status set, but old callback stays
in the list
3. 60s later: `_emit_existing_download_progress()` scans disk → resets
status to `DownloadPending`
4. Worker sends new `StartDownload` → coordinator accepts (guard didn't
check `DownloadFailed`)
5. `_start_download_task()` appends **another** callback
6. Each callback has its own throttle → N callbacks = N events per
progress tick

## Changes

### Commit 1: `src/exo/worker/main.py`
Move the `DownloadModel` backoff check **before** `TaskCreated` emission
in `plan_step()`. Previously `TaskCreated` was emitted unconditionally
every 0.1s even when backoff blocked the download command.

### Commit 2: `src/exo/download/coordinator.py`
1. **Register progress callback once** in `__post_init__` instead of
per-download in `_start_download_task()`. Uses a per-model throttle dict
instead of per-callback closure variables.
2. **Add `DownloadFailed` to the `_start_download()` guard** so
redundant `_start_download_task()` calls don't happen. Retries still
work because `_emit_existing_download_progress` resets `DownloadFailed`
→ `DownloadPending` by scanning disk every 60s.

## Why It Works

The root cause was callbacks accumulating in
`ResumableShardDownloader.on_progress_callbacks` (a list that only
appends, never clears). By registering one callback per coordinator
lifetime and guarding against re-entry on `DownloadFailed`, we ensure
exactly one progress event per model per progress tick regardless of how
many retry cycles have occurred.

## Test Plan

### Manual Testing
- Verified the download retry flow: failed download → 60s scan resets
status → new `StartDownload` accepted → download retries with single
callback

### Automated Testing
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — passes
- `uv run pytest` — 188 passed

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 23:50:43 +00:00
Alex Cheema 98773437f3 Make info gatherer monitors resilient with retry loops and timeouts (#1448)
## Motivation

Info gatherer monitors could silently stop posting events, causing stale
node state after rejoins. The macmon monitor was especially fragile — it
had no retry loop, so a crash or parse error would kill it permanently.
Worse, the unhandled exception would propagate to the TaskGroup and take
down *all* sibling monitors. Additionally, none of the monitors had
timeouts on their subprocess calls, so a hung `system_profiler` or
`networksetup` could stall a monitor indefinitely.

## Changes

- Wrap `_monitor_macmon` in a `while True` retry loop with `except
Exception`, matching the pattern used by all other monitors
- Add `fail_after` timeouts to all monitor loop bodies:
- 10s for lightweight commands (`_monitor_misc`, `_watch_system_info`,
`_gather_iface_map` init)
- 30s for heavier commands (`_monitor_system_profiler_thunderbolt_data`,
`_monitor_thunderbolt_bridge_status`)
- Remove unused `CalledProcessError` and `cast` imports

## Why It Works

All monitors now follow the same resilient pattern: `while True` → `try`
with `fail_after` → `except Exception` (logs warning) → `sleep`. If a
subprocess hangs, the timeout fires and `TimeoutError` is caught by the
existing `except Exception` handler. If macmon crashes, it restarts
after the interval instead of dying permanently. No single monitor
failure can cascade to kill the others.

## Test Plan

### Manual Testing
<!-- Hardware: macOS with macmon installed -->
<!-- What you did: -->
- Run exo, kill macmon process (`kill $(pgrep macmon)`), verify it
restarts and metrics resume
- Verify all monitors continue posting events after simulated hangs

### Automated Testing
- All 188 existing tests pass
- basedpyright: 0 errors
- ruff: all checks passed

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 23:07:36 +00:00
Alex Cheema ad6d35d68a Retry runners within the same Instance instead of recreating
When runners fail for a MetaInstance-backed Instance, retry up to 3
times by restarting runners within the same Instance rather than
deleting and recreating it each time. After 3 failures, delete the
Instance so MetaInstanceReconciler can create a fresh one.

- Add InstanceRetrying event that removes runners from state (signaling
  workers to restart) and increments consecutive_failures on MetaInstance
- InstanceHealthReconciler emits InstanceRetrying when under retry limit,
  InstanceDeleted when exhausted or no MetaInstance
- Worker _kill_runner detects retry signal (runner deleted from state +
  terminal supervisor) and cleans up for _create_runner to recreate
- Worker _create_runner guards against oscillation by blocking creation
  while any peer runner has explicit terminal status
- InstanceCreated resets consecutive_failures for fresh starts

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 14:21:11 -08:00
Jake Hillion a8acb3cafb dashboard: show available disk space on downloads page
The downloads page previously only showed the approximate space used by
downloaded models (summed from completed download sizes), but did not
show how much disk space was actually available. This made it difficult
to know if a download would succeed before pressing the button.

Added disk space tracking to the InfoGatherer that polls the models
directory partition every 30 seconds. The DiskUsage type captures total
and available space, which flows through the event system to State and
is exposed via the /state API. The dashboard now displays "X on disk /
Y available" for each node in the downloads view.

Test plan:
- CI
2026-02-11 21:57:28 +00:00
Alex Cheema a0721dbe57 feat: warn when cluster nodes have mismatched macOS versions (#1436)
## Motivation

When nodes in an exo cluster run different macOS versions, inference can
produce incompatible results or fail silently. Users currently have no
way to know this from the dashboard.

## Changes

- Added `get_os_version()` to `system_info.py` that returns the macOS
version (e.g. `"15.3"`) or platform name for non-Mac nodes
- Added `os_version` field to `NodeIdentity` and
`StaticNodeInformation`, gathered once at startup
- Propagated `os_version` through the event sourcing pipeline
(`apply.py`)
- Exposed `nodeIdentities` from the dashboard store with `osVersion`
- Added a derived `macosVersionMismatch` check in `+page.svelte` that
triggers when 2+ macOS nodes report different versions
- Rendered a yellow "INCOMPATIBLE macOS VERSIONS" warning badge
(matching the existing Thunderbolt Bridge cycle warning style) with a
hover tooltip listing each node's name and version, in all three
topology view sizes (large, medium, compact)

## Why It Works

The OS version is a static property gathered once at node startup via
`platform.mac_ver()`. It flows through the existing
`StaticNodeInformation` → `NodeGatheredInfo` event → `NodeIdentity`
state pipeline, so no new event types or state fields beyond
`os_version` on `NodeIdentity` are needed. The dashboard derives the
mismatch by comparing `osVersion` across all nodes whose version looks
like a macOS version string (starts with a digit).

## Test Plan

### Manual Testing
Hardware: 4x Mac Studio M2 Ultra 512GB (s18, s17 (2), james, mike),
connected via Thunderbolt
- s18 and s17 (2) on macOS 26.2, james and mike on macOS 26.3
- Verified the "INCOMPATIBLE macOS VERSIONS" warning badge appears in
the topology view
- Verified the hover tooltip lists all four nodes with their respective
versions
- Screenshots attached in comment below

### Automated Testing
- basedpyright: 0 errors
- ruff check: all checks passed
- nix fmt: no formatting changes needed
- Dashboard builds successfully

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 21:18:59 +00:00
Alex Cheema c236d62caf Remove timestamp-based retry cooldown
Remove last_failure_at field and RETRY_COOLDOWN_SECONDS logic.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:59:39 -08:00
Alex Cheema a8069e8a30 Consolidate failure state onto MetaInstance, add 5s retry cooldown
Move placement_error, consecutive_failures, last_failure_error, and
last_failure_at directly onto the MetaInstance model instead of keeping
them as separate State mappings (meta_instance_errors, InstanceFailureInfo,
meta_instance_failure_info). Adds a 5-second cooldown between retry attempts
to prevent rapid instance churn when runners fail instantly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:55:47 -08:00
Alex Cheema 84ce555d55 Show retry attempt count with error message, e.g. (2/3)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:43:20 -08:00
Alex Cheema 50e2bcf93e fix: RDMA debug labels, TB5 info box, and rdma_ctl status detection (#1437)
## Motivation

Several RDMA/Thunderbolt UX issues in the dashboard and macOS app:

1. **Debug mode showed "? ?" for RDMA connections** — the topology view
only extracted IPs from socket connections, not RDMA interface names
2. **No way to detect if RDMA is actually enabled** — the system only
knew about TB5 hardware and RDMA topology edges, not whether `rdma_ctl`
was enabled on each node
3. **False "RDMA AVAILABLE" info box** — showed on Mac Minis with idle
TB5 ports even when RDMA was already enabled, and on single nodes with
TB5
4. **macOS app only showed local RDMA status** — ran `rdma_ctl` locally
with no visibility into other nodes in the cluster

## Changes

### Dashboard: Fix RDMA debug labels (`0abc90c4`)
- Added `sourceRdmaIface` and `sinkRdmaIface` to `TopologyEdge`
interface
- Updated `TopologyGraph.svelte` and `ModelCard.svelte` to show `RDMA
en2 → en3` instead of `? ?`

### Dashboard: TB5 RDMA info box (`a3795552`, `8ce8e173`)
- Added dismissible info box when 2+ nodes have TB5 hardware but RDMA is
disabled
- Includes setup instructions (Recovery mode → `rdma_ctl enable` →
reboot, TB5 cables, macOS version match)
- Requires 2+ exo nodes with TB5 to avoid false positives from
single-node setups

### Backend: `rdma_ctl status` detection (`ae07239b`)
- Added `RdmaCtlStatus` event to `info_gatherer.py` — runs `rdma_ctl
status` with 5s timeout, `shutil.which` guard, and `OSError` handling
(polls every 10s on macOS)
- Added `NodeRdmaCtlStatus` model to `profiling.py` and `node_rdma_ctl`
field to `State`
- Handle in `apply.py` (event apply + node timeout cleanup)
- Exposed `nodeRdmaCtl` in dashboard store (`app.svelte.ts`)
- Info box detection now uses actual RDMA status instead of TB5 link
speeds

### Dashboard: Per-node RDMA debug labels (`ae07239b`)
- Debug mode shows `RDMA:ON` (green) or `RDMA:OFF` (dim) per node in
topology view, below the TB bridge label

### macOS app: Cluster-wide RDMA status from `/state` (`a1455b61`,
`d0d77b63`)
- Added `NodeRdmaCtlStatus` to `ClusterState.swift` — decoded from
`/state` endpoint
- Replaced local-only `rdma_ctl status` check with cluster-wide
`nodeRdmaCtl` from state
- Debug section shows per-node RDMA enabled/disabled for all nodes in
the cluster
- Still shows local `ibv_devices` and `ibv_devinfo` details (device
names, active ports) for richer local debugging

## Files changed

| Area | File | Change |
|------|------|--------|
| Backend | `src/exo/utils/info_gatherer/info_gatherer.py` |
`RdmaCtlStatus` event, monitor task |
| Backend | `src/exo/shared/types/profiling.py` | `NodeRdmaCtlStatus`
model |
| Backend | `src/exo/shared/types/state.py` | `node_rdma_ctl` field |
| Backend | `src/exo/shared/apply.py` | Event handler + timeout cleanup
|
| Dashboard | `dashboard/src/lib/stores/app.svelte.ts` | `nodeRdmaCtl` +
`nodeThunderbolt` in store |
| Dashboard | `dashboard/src/routes/+page.svelte` | Info box with RDMA
detection + instructions |
| Dashboard | `dashboard/src/lib/components/TopologyGraph.svelte` | RDMA
debug labels per node + fix "? ?" |
| Dashboard | `dashboard/src/lib/components/ModelCard.svelte` | RDMA
interface display fix |
| App | `app/EXO/EXO/Models/ClusterState.swift` | `NodeRdmaCtlStatus`
struct + decode |
| App | `app/EXO/EXO/ContentView.swift` | Cluster-wide RDMA view + local
device details |
| App | `app/EXO/EXO/Services/NetworkStatusService.swift` | Remove local
`rdma_ctl`, keep `ibv_*` |

## Test Plan

- [x] `uv run basedpyright` — 0 errors
- [x] `uv run ruff check` — pass
- [x] `nix fmt` — clean
- [x] `cd dashboard && npm run build` — success
- [x] `uv run pytest` — 188 passed
- [x] Xcode build — compiles (only pre-existing `dist/exo` resource
error)
- [x] Deployed to Mac Minis — `nodeRdmaCtl` shows `enabled: true`, no
false info box
- [x] Deployed to James cluster — RDMA debug labels show correctly

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

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 20:43:11 +00:00
Alex Cheema b78ea438bc Include node friendly names in runner error messages
Each error in the combined message is now prefixed with the node's friendly
name (e.g. "MacBook Pro: OOM; Mac Studio: connection reset") so the root
cause node is easily identifiable.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:41:10 -08:00
Alex Cheema 1960b16f9f Remove permanent retry blocking, allow continuous retry batches
The dashboard % 3 logic already handles displaying retry progress in batches
(RETRYING 1/3, 2/3, 3/3, then PLACING with error, repeat). No need to
permanently block placement after 3 failures.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:35:03 -08:00
Alex Cheema 7bed91c9c2 feat: add Recent tab to model picker (#1440)
## Motivation

When frequently switching between models, it's tedious to search through
the full model list to find ones you've used before. A "Recent" tab
provides quick access to previously launched models.

## Changes

- **New store** (`dashboard/src/lib/stores/recents.svelte.ts`):
`RecentsStore` class persisting recently launched model IDs with
timestamps to localStorage (key: `exo-recent-models`). Caps at 20
entries, deduplicates on re-launch (moves to top).
- **FamilySidebar**: Added "Recent" tab between Favorites and Hub,
conditionally shown when there are recent models.
- **FamilyLogos**: Added clock/history icon for the recents tab.
- **ModelPickerModal**: Added `recentModelIds`/`hasRecents` props.
Derives single-variant `ModelGroup[]` from recent IDs and renders them
using the same `ModelPickerGroup` component as all other tabs —
consistent styling, memory grey-out, favorites, info button, download
indicators.
- **+page.svelte**: Calls `recordRecentLaunch(modelId)` after successful
instance launch. Passes reactive recent state to the modal.

## Why It Works

Follows the exact same pattern as the existing Favorites feature
(localStorage persistence, conditional tab display, reactive Svelte 5
`$state`/`$derived`). Recent models are wrapped as single-variant
`ModelGroup` objects so they reuse `ModelPickerGroup` for identical row
rendering across all tabs.

## Test Plan

### Manual Testing
<!-- Hardware: MacBook Pro -->
- Launch a model instance → reopen model picker → "Recent" tab appears
with the launched model
- Launch a second model → it appears at top of the Recent list
- Re-launch the first model → it moves back to top
- Search within the Recent tab filters the list
- Models that don't fit in memory are greyed out (same as All tab)
- Close/reopen browser → recents persist from localStorage

### Automated Testing
- Dashboard builds successfully (`npm run build`)

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-11 12:34:08 -08:00
Alex Cheema c6838c8fd8 Show retry count in exceeded retry limit message (3/3)
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:28:17 -08:00
Alex Cheema 420d9b9e76 Collect all runner error messages instead of just the last one
When multiple runners fail, concatenate all error messages with "; " so the
real error isn't hidden by generic side-effect failures from other runners.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:27:49 -08:00
Alex Cheema 13f1e9c489 Stop infinite retries after 3 failures, show errors persistently in dashboard
MetaInstanceReconciler now checks failure count before placement — after 3
consecutive failures it emits MetaInstancePlacementFailed instead of retrying
forever. Dashboard shows "Retrying after error: <msg>" in orange throughout
the retry cycle, not just during the brief window with no backing instance.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:21:11 -08:00
Alex Cheema 451a06b3d8 Add instance retry logic with max 3 retries and failure tracking
- Extend InstanceDeleted with failure_error field for runner crash info
- Add InstanceFailureInfo model tracking consecutive failures per MetaInstance
- InstanceHealthReconciler now detects runner failures (all terminal with
  at least one RunnerFailed) in addition to connection failures
- apply_instance_deleted increments failure counter for meta-bound instances
- Dashboard shows RETRYING (N/3) status with error messages, and
  "Instance re-created due to failure" after 3 consecutive failures
- Extract and display RunnerFailed error messages in instance status

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 12:09:42 -08:00
Alex Cheema 94b55d66f4 Fix MetaInstance.node_ids frozenset failing JSON deserialization
frozenset serializes to a JSON array but cannot be deserialized back
in strict mode through the TaggedModel wrap validator (list → frozenset
coercion is rejected). Changed to list[NodeId] since the model is
already frozen/immutable.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:54:56 -08:00
Alex Cheema 2b68b931c5 Send node_ids from placement preview when launching instances
The dashboard now extracts node IDs from the selected preview's
memory_delta_by_node, ensuring the backend places on exactly the
nodes the user was shown. Also reverts incorrect RDMA min_nodes >= 2
enforcement since single-node RDMA is valid.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:49:37 -08:00
Alex Cheema 4aecaa7748 Enforce min_nodes >= 2 for RDMA (MlxJaccl) instances
RDMA requires at least 2 nodes — a single-node RDMA instance is
nonsensical. Enforce this in both the dashboard (when building the
launch request) and the backend placement (when filtering cycles).
Previously, selecting RDMA would still place on 1 node because
min_nodes defaulted to 1 and the placement silently switched to Ring.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:42:21 -08:00
Alex Cheema 25e2891c30 Ensure min_nodes >= node filter size when launching
When user selects specific nodes via the filter, min_nodes should be at
least the number of filtered nodes to prevent placement from picking a
smaller cycle.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:36:51 -08:00
Alex Cheema 16345e0ffa Send node_ids from dashboard, error on RDMA when unavailable
Dashboard was not including the user's node filter in the POST to
/meta_instance, so placement ignored which nodes the user selected.
Also, placement silently fell back to Ring when RDMA was requested but
no RDMA-connected cycles were available — now raises an error that
surfaces via MetaInstancePlacementFailed.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:26:29 -08:00
Alex Cheema 3a845f90b0 Fix use_default validator silently ignoring sharding/instance_meta
The mode="plain" validator bypassed Pydantic's string-to-enum coercion,
so JSON strings like "Tensor" and "MlxJaccl" from the dashboard failed
the isinstance check and silently fell back to Pipeline/MlxRing defaults.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:05:00 -08:00
Alex Cheema dccf2440ba Add placement error feedback and per-node loading status
Show why MetaInstance placement fails instead of stuck "PLACING", and
show per-node runner status during loading for multi-node instances.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 10:01:07 -08:00
Alex Cheema f96f3f2c0f Show MetaInstance sharding/type while PLACING, fix MlxIbv references
When a MetaInstance has no backing instance yet, derive the strategy
display from the MetaInstance's own sharding and instanceMeta fields
rather than showing "Unknown (Unknown)".

Also clean up all stale MlxIbv references across the dashboard —
the backend enum is MlxJaccl.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 09:23:44 -08:00
Alex Cheema 7d54e468d5 Extract reconciler into ProcessManager protocol, fix RDMA instance type
- Replace inline _plan() with ProcessManager loop (_reconcile), tick
  every 1s instead of 10s — safe because all PMs are idempotent
- Fix dashboard sending "MlxIbv" instead of "MlxJaccl" for RDMA
  instance type, which silently fell back to MlxRing default
- Remove all stale MlxIbv references from dashboard

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 09:19:13 -08:00
Alex Cheema 124d504f95 Extract reconciler into ProcessManager protocol
Replace inline _plan() steps with a list of ProcessManagers, each
implementing async reconcile(State) -> Sequence[Event]. Tick every
1s instead of 10s — safe because all PMs are idempotent against state.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-11 09:06:41 -08:00
Hunter Bown 48caea4a23 feat: add intermediate model-fit state for cluster-capacity-only models (#1441)
## Motivation
In the model picker, non-runnable models were not clearly separated
between two different cases:

- model exceeds currently available RAM but can fit in total cluster
capacity
- model exceeds total cluster capacity

That made it harder to distinguish "not runnable right now" from "too
large for this cluster."

## Changes
- Added a tri-state fit status in dashboard model-picker flow:
  - `fits_now`
  - `fits_cluster_capacity`
  - `too_large`
- Updated dashboard logic to compute both available cluster RAM and
total cluster RAM.
- Passed fit status through picker components.
- Updated model size color mapping:
  - `fits_now` -> white/gray
  - `fits_cluster_capacity` -> orange
  - `too_large` -> red
- Updated group ordering for non-runnable models:
- orange groups (`fits_cluster_capacity`) are listed above red groups
(`too_large`).

## Why It Works
Launch safety is unchanged: selection is still gated by existing
placement feasibility (`canModelFit`), so models that cannot run now
remain disabled.

The new fit status is used for visual distinction and ordering only:
- runnable now
- fits cluster capacity but not free RAM now
- too large for cluster capacity

## Before
Non-runnable models were not clearly distinguished by temporary capacity
vs hard capacity limit.

## After
Model picker clearly separates states by both color and order:
- Runnable now (white/gray)
- Fits cluster capacity but not free RAM now (orange, disabled)
- Exceeds cluster capacity (red, disabled)

## Test Plan
### Manual testing
- Open model picker with live cluster memory telemetry.
- Verify white/gray models are selectable.
- Verify orange models are disabled and appear above red models.
- Verify red models are disabled and appear below orange models.

### Automated checks
- `npm --prefix dashboard run build` passes.
- `uv run basedpyright` passes.
- `uv run ruff check` passes.
- `npm --prefix dashboard run check` reports existing pre-change Svelte
diagnostics (same known SVG title/a11y items).
- `uv run pytest` in this local environment exits during collection due
existing `tests/start_distributed_test.py` `SystemExit` usage.
- No Python code was changed.

---------

Co-authored-by: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
2026-02-11 00:21:35 +00:00
Alex Cheema 9ab4a40989 Simplify MetaInstance binding: put meta_instance_id on Instance
The separate MetaInstanceBound event + meta_instance_backing map
introduced two bugs: stale exclusion sets in the reconciler loop and
a delete ordering race. Embedding meta_instance_id directly on
BaseInstance eliminates the binding mechanism entirely — when an
instance is created for a MetaInstance it carries the ID, when
deleted the binding is gone. No separate map, no cleanup, no races.

Also fixes delete_meta_instance to cascade-delete backing instances.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-10 16:15:29 -08:00
Alex Cheema f4329c72c2 Add explicit MetaInstance binding, slim MetaInstance to use ModelId
- Add MetaInstanceBound event and meta_instance_backing State field
  for explicit MetaInstance → Instance binding (prevents ambiguous
  linking when two MetaInstances have identical constraints)
- Replace model_card: ModelCard with model_id: ModelId on MetaInstance
  (load ModelCard on-demand at placement time)
- Add MetaInstance API endpoints (POST /meta_instance, DELETE)
- Update dashboard to use MetaInstances as primary primitive with
  unified display items merging MetaInstances and orphan instances
- Dashboard launches via MetaInstance instead of direct Instance creation

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-10 15:53:07 -08:00
Mustafa Alp Yılmaz eead50b4cd Fix setrlimit crash when hard file descriptor limit < 65535 (#1430)
## Summary

`exo` crashes on startup when the system's hard file descriptor limit is
below 65535, which occurs in macOS LaunchDaemon environments, Docker
containers, and other restricted setups.

**Root cause:** `resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft,
65535), hard))` raises `ValueError` when `hard < 65535` because the soft
limit cannot exceed the hard limit.

**Fix:**
- **`main.py`**: Clamp the target soft limit to `min(max(soft, 65535),
hard)` so it never exceeds the hard limit
- **`utils_mlx.py`**: Query current limits instead of hardcoding `(2048,
4096)`, which both crashed on restricted systems and incorrectly lowered
the hard limit when it was set higher

## Test plan

- [x] `basedpyright` passes with 0 errors
- [x] `ruff check` passes
- [x] Verify startup works on system with hard limit < 65535 (tested in
macOS LaunchDaemon with hard limit 10240)
- [x] Verify startup still works on default macOS (hard limit typically
unlimited)
2026-02-10 22:46:22 +00:00
Alex Cheema ceb76b8f6c Add MetaInstance declarative layer with connection health checking
Introduces MetaInstance as a declarative constraint ensuring an instance
matching given parameters (model, sharding, min_nodes) always exists.
The master's reconciliation loop continuously checks for unsatisfied
meta-instances and attempts placement. Connection health checking
verifies that specific IPs (MlxRing) and RDMA interfaces (MlxJaccl)
stored on instances still exist as topology edges, enabling automatic
recovery when cables are swapped or interfaces change.

Also eliminates the master's loopback event path, unifying all event
emission through _apply_and_broadcast for simpler control flow.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-10 13:43:50 -08:00
Jake Hillion 199df64cfc util: remove VecExt trait, inline at call site (#1446)
VecExt added a .map() convenience method on Vec<T> that simply called
.into_iter().map(f).collect(). This thin wrapper provided no
optimisation benefit and obscured a standard iterator pattern behind a
nightly feature gate and an extra dependency.

Replaced the single call site in exo_pyo3_bindings with the equivalent
iterator chain and removed the ext module, the extend dependency, and
the trait_alias feature gate from the util crate.

Test plan:
- CI
2026-02-10 20:45:57 +00:00
Ryuichi Leo Takashige dc7ade8052 set the mlx hash 2026-02-10 20:16:17 +00:00
Ryuichi Leo Takashige dc781497c5 update mlx to 0.30.6 2026-02-10 20:16:17 +00:00
Jake Hillion c37eb24331 util: remove dead code (#1445)
The util crate contained several unused items: NonemptyArray,
BoxedSliceExt, a blanket Sealed trait, an empty alias module, six unused
nightly feature gates, and five unused Cargo dependencies (thiserror,
once_cell, internment, derive_more, bon, recursion).

Removed all items that had no references outside their own definitions,
keeping only WakerDeque, VecExt, and the trait_alias feature gate which
are actively used by the networking and exo_pyo3_bindings crates.

Test plan:
- CI
2026-02-10 20:10:54 +00:00
Jake Hillion 8af2af6328 nix: override apple-sdk to 26.2 and enable MLX_BUILD_CPU (#1443)
The pinned nixpkgs provides apple-sdk 26.0, but building MLX requires
SDK 26.2. The upstream package reads versions.json via a relative path
at eval time, so it can't be overridden through callPackage args.

Added a thin overlay that copies the upstream apple-sdk source and
patches only metadata/versions.json to point at SDK 26.2. Also enabled
MLX_BUILD_CPU in the MLX nix build.

This avoids vendoring the entire apple-sdk package (~2200 lines) while
still getting the SDK version we need.

Test plan:
- CI
- Built and ran on two machines connected with Thunderbolt 5 - Kimi K2.5
starts in Tensor+RDMA and seems sensible.
2026-02-10 19:53:53 +00:00
rltakashige 43728b2047 Send all exo logs (#1439)
## Motivation

Log rotation adds a bunch of .zst files. Let's send them all in the bug
reports.

This PR also standardises the logs so that they all include the
timestamp.
2026-02-10 19:28:33 +00:00
rltakashige 1699fcfb9f standardise logs (#1442)
## Motivation

Standardises exo.log and event_log

## Changes

- exo.log and exo.log.zst are now in a exo_log directory.
- event_log is now timestamped and not numbered. The timestamps can be
sorted as they are YYYY-MM-DD-HH-MM

## Test Plan

### Manual Testing
Nothing crashes.
2026-02-10 19:13:25 +00:00
rltakashige 009b43c662 add log rotation for .exo/exo.log (#1438)
## Motivation

.exo.log currently contains all past history. This just makes it hard to
read and is unnecessarily expensive even on disk.

## Why It Works

Just uses loguru's rotation.

## Test Plan

### Manual Testing
exo.log is new. 
<img width="1992" height="706" alt="image"
src="https://github.com/user-attachments/assets/9b293993-1141-43e7-b58e-0ddd2d4eda2e"
/>
2026-02-10 18:55:06 +00:00
Jake Hillion 1f242e8eee gossipsub: stop silent message dropping and warn (#1434)
The 15-second publish_queue_duration caused messages in peer queues to
be silently dropped. When events are dropped, workers detect gaps in the
event index sequence and request missing events via the NACK path
(RequestEventLog), but this recovery is inefficient.

Removed the timeout configuration - gossipsub now uses its default
behavior without time-based eviction. If queue buildup is a concern,
queue size should be limited explicitly rather than dropping by timeout.

Split error handling to log AllQueuesFullError as a warning (indicates
peers are unresponsive) while keeping NoPeersSubscribedToTopicError
silent (expected during startup and network partitions).

Test plan:
- CI
2026-02-10 18:47:47 +00:00
rltakashige 64179c6fc1 Dont save to app directory (#1435)
## Motivation

App keeps losing Local Network permissions.

## Changes

Don't save stuff to the app directory anymore. Instead, save to .exo.

## Why It Works

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

## Test Plan

### Manual Testing
Before:
<img width="1512" height="106" alt="image"
src="https://github.com/user-attachments/assets/544ef57e-b626-484d-941f-2472969aa208"
/>

After:
<img width="433" height="53" alt="Screenshot 2026-02-10 at 17 43 06"
src="https://github.com/user-attachments/assets/3de2856b-cdf6-4b35-aa8f-50440686344f"
/>


### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->
2026-02-10 18:41:36 +00:00
Jake Hillion 305a3c8b70 event_log: move event log from unbounded in-memory list to disk (#1432)
The master and API event logs (list[Event]) grew unbounded in RAM for
the lifetime of the process. Events are rarely read back (only for
RequestEventLog when a new node catches up, or the dashboard /events
endpoint).

Introduced a DiskEventLog class that writes length-prefixed msgpack
records to an append-only file, using a bounded LRU cache of byte
offsets for indexed access. On close, the active file is compressed
with ZSTD and rotated into a numbered archive slot, keeping the last 5
archives (events.1.bin.zst through events.5.bin.zst). On construction,
any stale active file from a crash is rotated before opening a fresh
log. The /events API endpoint now streams the JSON array one event at a
time rather than materializing the full list in memory. Deserialization
routes msgpack through json.dumps into Pydantic's validate_json() to
get correct JSON-mode coercion (e.g. string to enum) under strict mode.

This bounds memory usage to the LRU cache (128 entries) regardless of
event volume, while still supporting efficient sequential reads from
disk when needed.

Test plan:
- CI
- New unit tests for DiskEventLog: append/read, range queries, rotation
  on close, stale file recovery, idempotent close, successive sessions,
  archive retention limit (5 max)
- Tested on a cluster with 9000 events. /events continues working.
- On disk size is 3.9MiB with ~8000 events, and the compression is very
  effective.
- Disconnected and rejoined a machine, it rejoined fine.

---------

Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
2026-02-10 17:27:32 +00:00
Alex Cheema ead19bea74 Always load image model cards into cache (#1421)
## Motivation

Follows up on #1408. Image models (FLUX, Qwen-Image, etc.) don't have a
`config.json` on HuggingFace. Previously, image model TOML cards were
only loaded into `_card_cache` when `EXO_ENABLE_IMAGE_MODELS=true`. When
the flag was off but an image model was requested (e.g., via
`get_placement_previews`), `ModelCard.load()` fell through to
`fetch_from_hf()` which tried to download `config.json` — causing
`FileNotFoundError` spam. #1408 added defensive error handling; this PR
fixes the root cause.

## Changes

**`model_cards.py`**: Always include `image_model_cards/` in
`CARD_SEARCH_PATH` so image model TOML cards are always loaded into
`_card_cache`. `ModelCard.load()` then finds them directly and never
falls through to `fetch_from_hf()`. The `EXO_ENABLE_IMAGE_MODELS` flag
now controls whether image models appear in `get_model_cards()` (the
listing) rather than whether they're loaded at all.

## Why It Works

`fetch_from_hf()` is designed for text models only (it hardcodes
`tasks=[ModelTask.TextGeneration]` and requires `config.json`). Image
models should never reach that path. By always having them in the cache,
the lookup succeeds immediately and `fetch_from_hf()` is never called.

## Test Plan

### Automated Testing
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — passes

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-10 09:11:57 -08:00
Jake Hillion 5a83e59182 dashboard: allow typing in chat input while response is generating (#1433)
The chat textarea was fully disabled during response generation,
preventing users from drafting their next message while waiting.

Removed the `disabled={loading}` attribute from the textarea element.
Submission is still blocked during generation by the early return in
`handleSubmit()` and the submit button's own disabled state.

Test plan:
- Ran on one machine. While a model was writing a really long poem, I
typed my next response. I couldn't submit it with Enter and the button
still said "Processing" greyed out. I could send the message after
generation finished.
2026-02-10 16:12:08 +00:00
Jake Hillion 5b5577bead build-app: upload DMG to S3 for non-tagged builds (#1428)
Non-tagged builds (test-app branch, manual dispatch) only uploaded the
DMG as a GitHub artifact, which requires authentication to download.

Added an early exit path that uploads the DMG with a commit hash suffix
(EXO-<sha>.dmg) for non-tagged builds, making it publicly accessible
via S3.

Test plan:
- CI
-
https://github.com/exo-explore/exo/actions/runs/21837274032/job/63011907978
  worked as intended

Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-10 15:47:49 +00:00
Evan Quiney 8314a2aa78 cleaning up the todos (#1406)
kinda closes #1400 ( a bit )
2026-02-10 12:35:29 +00:00
Alex Cheema 163cf18384 Add error handling to info gatherer monitor loops (#1422)
## Motivation

If any of the `InfoGatherer` monitor loops throw an unexpected
exception, the entire monitoring task crashes and never recovers. This
can silently stop memory, network, or Thunderbolt data collection for
the lifetime of the process.

## Changes

Wrap the body of each `while True` monitor loop in a try/except that
logs the exception as a warning and continues to the next iteration. The
sleep at the end of each loop runs regardless, providing natural backoff
before retry.

Affected methods: `_monitor_misc`,
`_monitor_system_profiler_thunderbolt_data`, `_monitor_memory_usage`,
`_watch_system_info`, `_monitor_thunderbolt_bridge_status`.

`_monitor_macmon` already had its own error handling so was left as-is.

## Why It Works

A transient error (e.g., a subprocess failing, a permission issue) in
one iteration no longer kills the loop. The warning log provides
visibility while the monitor continues collecting data on subsequent
iterations.

## Test Plan

### Automated Testing
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — passes

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-10 12:24:43 +00:00
rltakashige 2204f651c8 Yield from reachability checks (#1427)
## Motivation

check_reachable waits for all connection profile checks to be completed.
Since there are retries on failures, this can take around 20s to
resolve, preventing any instances from showing up. This feels very slow
for UX, and it slows down distributed testing.

## Changes

Made check_reachable an async generator.

## Test Plan

### Manual Testing
Works for me at least.
2026-02-10 12:18:45 +00:00
rltakashige 4abdaaf74b Address GPU timeouts (#1429)
## Motivation

For large prompts and/or slow machines, users are running into GPU
timeout errors very often.

## Changes

Only during prefill, we eval distributed operations. We don't do this
during decode to maintain decode performance.
Raise the prefill step size to 8192 because now we can (we see a speedup
here).
We also now see a 2x speedup in pipeline parallel prefill by disabling
an unnecessary all_gather during prefill.

## Why It Works

GPU timeout errors happen in the Metal backend when GPU operations take
too long without making progress.
By isolating distributed operations, we can allow them to run without
any timeouts.

## Test Plan

### Manual Testing
Doesn't GPU timeout on 100k tokens on Minimax anymore. Also tested on
Kimi.

### Automated Testing
Needs more exo bench, but I think this is a good step in the right
direction.
2026-02-10 11:53:23 +00:00
ciaranbor 2fbdb27bb1 Handle config.json not found (image models) (#1408)
## Motivation

When downloading image models, a missing config.json file triggers a
FileNotFoundError inside download_file_with_retry. This error was being
caught by the generic except Exception handler and retried 3 times
before failing. Then, the whole thing would be retried from the start

## Changes

- src/exo/download/download_utils.py: Added FileNotFoundError to the
list of immediately-raised exceptions in download_file_with_retry,
alongside HuggingFaceAuthenticationError. This prevents useless retries
when a file genuinely doesn't exist on the remote.
- src/exo/master/api.py: Wrapped ModelCard.load(model_id) in a
try/except that converts failures into an HTTPException(400) with a
descriptive error message, giving API consumers a clear error response.

## Why It Works


- FileNotFoundError is a deterministic error — the file won't appear on
retry, so re-raising immediately avoids 3 wasted download attempts with
exponential backoff.
- Catching ModelCard.load() failures and returning a 400 HTTP response
prevents unhandled exceptions from surfacing as opaque 500 errors in the
API.

## Test Plan

### Manual Testing

Verified an image model not in model cards does not cause an infinite
error loop
2026-02-07 03:34:58 +00:00
ciaranbor 3f57416dbf Add image lightbox (#1414)
## Motivation

No way to view generated or attached images at full resolution in the
dashboard

## Changes

- New ImageLightbox.svelte — fullscreen overlay with download, close
(click-outside/Escape), and transitions
- ChatMessages.svelte — all images (input attachments + generated) are
now clickable to open in lightbox; added expand button to generated
image hover overlay

## Why It Works

Single expandedImageSrc state variable drives the lightbox — set it to
show, null to hide.

## Test Plan

### Manual Testing

  - Click any image (attachment thumbnail or generated) → lightbox opens
  - Close via Escape, click-outside, or close button
  - Download button saves with correct extension
2026-02-07 01:30:03 +00:00
rltakashige 8f3681cf7e Synchronize before warmup (#1419)
## Motivation

Maybe addresses #1303 

## Changes

Add an mx barrier before warmup

## Why It Works

It might, it might not. Shouldn't break anything that's not already
broken though.

## Test Plan

### Manual Testing
The two machines I tested on were fine on GLM 4.7 Flash 8bit (the one in
exo.log in the issue). Obviously not definitive for anything, however.

<img width="594" height="878" alt="image"
src="https://github.com/user-attachments/assets/534d3ad6-16ef-4cb5-b823-43c8d4e1d3c6"
/>
2026-02-07 00:14:19 +00:00
ciaranbor 9dc4f786bd Ciaran/image model listing (#1417)
## Motivation

Image models (FLUX, Qwen Image) had no family grouping or quantization
metadata in the dashboard

## Changes

- Added family, quantization, base_model, and capabilities fields to all
18 image model TOML cards (FLUX.1 variants + Qwen Image variants)
  - Added FLUX and Qwen Image SVG logos to FamilyLogos.svelte
- Added "flux" and "qwen-image" families to the sidebar and family sort
order
- Added "Image Gen" and "Image Edit" capability filters in
ModelFilterPopover.svelte
  - Added image edit icon/badge to ModelPickerGroup.svelte
- Made the model category sidebar scrollable to accommodate the new
entries
  - Hidden scrollbars on model list panels

## Why It Works

Reuses the existing family/quantization grouping infrastructure that
LLMs already use, extending it to image models with appropriate metadata
and icons

## Test Plan

### Manual Testing

Verified image models behave like text models in the model list dialog

---------

Co-authored-by: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
2026-02-06 16:08:57 -08:00
rltakashige dcb4cabc15 Update the nix hash for mlx 0.30.5 (#1416)
## 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-06 21:27:10 +00:00
Jake Hillion d79b3a0e75 bench: make exo-bench available via nix run on all platforms (#1415)
exo-bench was gated behind isDarwin in python/parts.nix because it used
exoVenv, which pulls in MLX (Darwin-only). However, exo_bench.py is an
HTTP client that only needs loguru, transformers, huggingface-hub, and
tiktoken.

Made bench a uv workspace member with its own pyproject.toml declaring
only the minimal dependencies. Added a separate benchVenv in parts.nix
built from that workspace member, and moved exo-bench out of the
isDarwin block so it is available on all platforms.

Test plan:
- `nix run .#exo-bench -- --help` prints argparse help

---------

Co-authored-by: Ryuichi Leo Takashige <leo@exolabs.net>
Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-06 21:07:17 +00:00
Evan Quiney a2f1d48712 slow down catchup (#1407)
our event log request blasted the whole event log over libp2p, now it
just does the next 1000 messages - hopefully allowing nodes to catch up
a bit more consistently for long lived clusters

Co-authored-by: rltakashige <rl.takashige@gmail.com>
2026-02-06 20:45:27 +00:00
rltakashige 3b2f553a25 Fix kimi tool calling id (#1413)
## Motivation

Kimi produces its own tool id. It gets confused when we generate our own
id.

## Changes

Add id to tool call item and parse Kimi id properly.

## Test Plan

### Manual Testing
<img width="3198" height="522" alt="image"
src="https://github.com/user-attachments/assets/d71ec2be-7f57-49dc-a569-d304cc430f4d"
/>

Long running Kimi K2.5 cluster querying itself through OpenCode running
on the same Kimi K2.5 instance.
2026-02-06 11:33:51 -08:00
118 changed files with 6945 additions and 1314 deletions
+7 -1
View File
@@ -396,7 +396,7 @@ jobs:
path: output/EXO-${{ env.RELEASE_VERSION }}.dmg
- name: Upload to S3
if: env.SPARKLE_S3_BUCKET != '' && github.ref_type == 'tag'
if: env.SPARKLE_S3_BUCKET != ''
env:
AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
@@ -412,6 +412,12 @@ jobs:
PREFIX="${PREFIX}/"
fi
DMG_NAME="EXO-${RELEASE_VERSION}.dmg"
if [[ "${{ github.ref_type }}" != "tag" ]]; then
aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-${GITHUB_SHA}.dmg"
exit 0
fi
aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}${DMG_NAME}"
if [[ "$IS_ALPHA" != "true" ]]; then
aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"
-27
View File
@@ -8,33 +8,6 @@ on:
- main
jobs:
typecheck:
runs-on: ubuntu-latest
steps:
- name: Checkout repository
uses: actions/checkout@v4
with:
lfs: false
- uses: cachix/install-nix-action@v31
with:
nix_path: nixpkgs=channel:nixos-unstable
- uses: cachix/cachix-action@v14
name: Configure Cachix
with:
name: exo
authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
- name: Load nix develop environment
run: nix run github:nicknovitski/nix-develop/v1
- name: Sync dependencies
run: uv sync --all-packages
- name: Run type checker
run: uv run basedpyright --project pyproject.toml
nix:
name: Build and check (${{ matrix.system }})
runs-on: ${{ matrix.runner }}
+1 -1
View File
@@ -2366,7 +2366,7 @@ class custom_function:
def default_device() -> Device:
"""Get the default device."""
def default_stream(device: Device) -> Stream:
def default_stream(device: Device | DeviceType) -> Stream:
"""Get the device's default stream."""
def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
+151
View File
@@ -0,0 +1,151 @@
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs
from .switch_layers import SwitchGLU
@dataclass
class ModelArgs(BaseModelArgs):
model_type: str
hidden_size: int
num_hidden_layers: int
vocab_size: int
num_attention_heads: int
num_attention_groups: int
head_dim: int
intermediate_size: int
rms_norm_eps: float
rope_theta: float
rope_scaling: Optional[Dict[str, Any]]
max_position_embeddings: int
sliding_window: int
layer_types: Optional[List[str]]
yarn_only_types: Optional[List[str]]
partial_rotary_factors: Optional[List[float]]
attention_other_setting: Optional[Dict[str, Any]]
use_head_wise_attn_gate: bool
moe_num_experts: int
moe_top_k: int
moe_intermediate_size: int
share_expert_dim: int
moe_layers_enum: Optional[str]
moe_router_scaling_factor: float
norm_expert_weight: bool
swiglu_limits: Optional[List[float]]
swiglu_limits_shared: Optional[List[float]]
tie_word_embeddings: bool
class Step3p5MLP(nn.Module):
hidden_size: int
intermediate_size: int
gate_proj: nn.Linear
up_proj: nn.Linear
down_proj: nn.Linear
limit: Optional[float]
def __init__(
self, args: ModelArgs, intermediate_size: int, swiglu_limit: float = 0
) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class Step3p5MoEGate(nn.Module):
top_k: int
n_routed_experts: int
routed_scaling_factor: float
norm_topk_prob: bool
gate: nn.Linear
router_bias: mx.array
def __init__(self, args: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
class Step3p5MoE(nn.Module):
gate: Step3p5MoEGate
switch_mlp: SwitchGLU
share_expert: Step3p5MLP
sharding_group: Optional[mx.distributed.Group]
def __init__(self, args: ModelArgs, layer_idx: int) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class Step3p5Attention(nn.Module):
is_sliding: bool
num_heads: int
num_kv_heads: int
head_dim: int
scale: float
q_proj: nn.Linear
k_proj: nn.Linear
v_proj: nn.Linear
o_proj: nn.Linear
q_norm: nn.Module
k_norm: nn.Module
use_head_wise_attn_gate: bool
g_proj: nn.Linear
rope: nn.Module
def __init__(self, args: ModelArgs, layer_idx: int) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class Step3p5DecoderLayer(nn.Module):
self_attn: Step3p5Attention
is_sliding: bool
is_moe_layer: bool
mlp: Step3p5MLP | Step3p5MoE
input_layernorm: nn.Module
post_attention_layernorm: nn.Module
def __init__(self, args: ModelArgs, layer_idx: int) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class Step3p5Model(nn.Module):
args: ModelArgs
vocab_size: int
num_layers: int
embed_tokens: nn.Embedding
layers: list[Step3p5DecoderLayer]
norm: nn.Module
_swa_idx: Optional[int]
_full_idx: Optional[int]
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
cache: Optional[List[Any]] = None,
) -> mx.array: ...
class Model(nn.Module):
args: ModelArgs
model_type: str
model: Step3p5Model
lm_head: nn.Linear
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[List[Any]] = None,
) -> mx.array: ...
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
def shard(self, group: Optional[mx.distributed.Group] = None) -> None: ...
@property
def layers(self) -> list[Step3p5DecoderLayer]: ...
def make_cache(self) -> list[Any]: ...
@property
def cast_predicate(self) -> Any: ...
@property
def quant_predicate(self) -> Any: ...
+75
View File
@@ -119,3 +119,78 @@ From .cursorrules:
## Testing
Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests.
## Dashboard UI Testing & Screenshots
### Building and Running the Dashboard
```bash
# Build the dashboard (must be done before running exo)
cd dashboard && npm install && npm run build && cd ..
# Start exo (serves the dashboard at http://localhost:52415)
uv run exo &
sleep 8 # Wait for server to start
```
### Taking Headless Screenshots with Playwright
Use Playwright with headless Chromium for programmatic screenshots — no manual browser interaction needed.
**Setup (one-time):**
```bash
npx --yes playwright install chromium
cd /tmp && npm init -y && npm install playwright
```
**Taking screenshots:**
```javascript
// Run from /tmp where playwright is installed: cd /tmp && node -e "..."
const { chromium } = require('playwright');
(async () => {
const browser = await chromium.launch({ headless: true });
const page = await browser.newPage({ viewport: { width: 1280, height: 800 } });
await page.goto('http://localhost:52415', { waitUntil: 'networkidle' });
await page.waitForTimeout(2000);
// Inject test data into localStorage if needed (e.g., recent models)
await page.evaluate(() => {
localStorage.setItem('exo-recent-models', JSON.stringify([
{ modelId: 'mlx-community/Qwen3-30B-A3B-4bit', launchedAt: Date.now() },
]));
});
await page.reload({ waitUntil: 'networkidle' });
await page.waitForTimeout(2000);
// Interact with UI elements
await page.locator('text=SELECT MODEL').click();
await page.waitForTimeout(1000);
// Take screenshot
await page.screenshot({ path: '/tmp/screenshot.png', fullPage: false });
await browser.close();
})();
```
### Uploading Images to GitHub PRs
GitHub's API doesn't support direct image upload for PR comments. Workaround:
1. **Commit images to the branch** (temporarily):
```bash
cp /tmp/screenshot.png .
git add screenshot.png
git commit -m "temp: add screenshots for PR"
git push origin <branch>
COMMIT_SHA=$(git rev-parse HEAD)
```
2. **Post PR comment** referencing the raw image URL (uses permanent commit SHA so images survive deletion):
```bash
gh pr comment <PR_NUMBER> --body "![Screenshot](https://raw.githubusercontent.com/exo-explore/exo/${COMMIT_SHA}/screenshot.png)"
```
3. **Remove the images** from the branch:
```bash
git rm screenshot.png
git commit -m "chore: remove temporary screenshot files"
git push origin <branch>
```
The images still render in the PR comment because they reference the permanent commit SHA.
Generated
-100
View File
@@ -353,31 +353,6 @@ dependencies = [
"generic-array",
]
[[package]]
name = "bon"
version = "3.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ebeb9aaf9329dff6ceb65c689ca3db33dbf15f324909c60e4e5eef5701ce31b1"
dependencies = [
"bon-macros",
"rustversion",
]
[[package]]
name = "bon-macros"
version = "3.8.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77e9d642a7e3a318e37c2c9427b5a6a48aa1ad55dcd986f3034ab2239045a645"
dependencies = [
"darling",
"ident_case",
"prettyplease",
"proc-macro2",
"quote",
"rustversion",
"syn 2.0.111",
]
[[package]]
name = "bs58"
version = "0.5.1"
@@ -700,41 +675,6 @@ dependencies = [
"syn 2.0.111",
]
[[package]]
name = "darling"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0"
dependencies = [
"darling_core",
"darling_macro",
]
[[package]]
name = "darling_core"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4"
dependencies = [
"fnv",
"ident_case",
"proc-macro2",
"quote",
"strsim",
"syn 2.0.111",
]
[[package]]
name = "darling_macro"
version = "0.21.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81"
dependencies = [
"darling_core",
"quote",
"syn 2.0.111",
]
[[package]]
name = "data-encoding"
version = "2.9.0"
@@ -1625,12 +1565,6 @@ dependencies = [
"zerovec",
]
[[package]]
name = "ident_case"
version = "1.0.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39"
[[package]]
name = "idna"
version = "1.1.0"
@@ -1745,15 +1679,6 @@ dependencies = [
"generic-array",
]
[[package]]
name = "internment"
version = "0.8.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "636d4b0f6a39fd684effe2a73f5310df16a3fa7954c26d36833e98f44d1977a2"
dependencies = [
"hashbrown 0.15.5",
]
[[package]]
name = "inventory"
version = "0.3.21"
@@ -3307,16 +3232,6 @@ dependencies = [
"zerocopy",
]
[[package]]
name = "prettyplease"
version = "0.2.37"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b"
dependencies = [
"proc-macro2",
"syn 2.0.111",
]
[[package]]
name = "primeorder"
version = "0.13.6"
@@ -3740,12 +3655,6 @@ dependencies = [
"yasna",
]
[[package]]
name = "recursion"
version = "0.5.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9dba2197bf7b1d87b4dd460c195f4edeb45a94e82e8054f8d5f317c1f0e93ca1"
[[package]]
name = "redox_syscall"
version = "0.5.18"
@@ -4807,15 +4716,6 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "util"
version = "0.0.1"
dependencies = [
"bon",
"derive_more",
"extend",
"internment",
"once_cell",
"recursion",
"thiserror 2.0.17",
]
[[package]]
name = "uuid"
+11 -11
View File
@@ -5,21 +5,21 @@
[X] Fetching download status of all models on start
[X] Deduplication of tasks in plan_step.
[X] resolve_allow_patterns should just be wildcard now.
[] no mx_barrier in genreate.py mlx_generate at the end.
[X] no mx_barrier in genreate.py mlx_generate at the end.
[] cache assertion not needed in auto_parallel.py PipelineLastLayer.
[] GPTOSS support dropped in auto_parallel.py.
[] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
[] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
[] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
[X] GPTOSS support dropped in auto_parallel.py.
[X] sharding changed "all-to-sharded" became _all_to_sharded in auto_parallel.py.
[X] same as above with "sharded-to-all" became _sharded_to_all in auto_parallel.py.
[X] Dropped support for Ministral3Model, DeepseekV32Model, Glm4MoeModel, Qwen3NextModel, GptOssMode in auto_parallel.py.
[] Dropped prefill/decode code in auto_parallel.py and utils_mlx.py.
[X] KV_CACHE_BITS should be None to disable quantized KV cache.
[] Dropped _set_nofile_limit in utils_mlx.py.
[] We have group optional in load_mlx_items in utils_mlx.py.
[] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
[] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
[X] Dropped _set_nofile_limit in utils_mlx.py.
[X] We have group optional in load_mlx_items in utils_mlx.py.
[X] Dropped add_missing_chat_templates for GptOss in load_mlx_items in utils_mlx.py.
[X] Dropped model.make_cache in make_kv_cache in utils_mlx.py.
[X] We put cache limit back in utils_mlx.py.
[] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
[] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
[X] topology.py remove_node removes the connections after checking if node is is in self._node_id_to_rx_id_map. on beta_1 it checks after, so would remove stale connections I guess?
[X] Missing Glm 4.7 model cards (this isn't ready yet but should be picked up, probably create an issue... the blocker is transforemrs version doesn't support the tokenizer for Glm 4.7. rc-1 does but we can't upgrade as it breaks other things.)
[] try-except in _command_processor only excepts ValueError. This was silently failing leading to un-debuggable errors (we had a KeyError that was happening ). Changed this to catch Exception instead of ValueError. See exo-v2 89ae38405e0052e3c22405daf094b065878aa873 and fb99fea69b5a39017efc90c5dad0072e677455f0.
[X] In placement.py, place_instance no longer looks at model_meta.supports_tensor and check if this tensor parallel number of nodes is supported by the model's tensor dimensions.
[X] In placement.py, place_instanec, we no longer have the special case to exclude DeepSeek v3.1 pipeline parallel (it doesn't work).
-14
View File
@@ -1,28 +1,14 @@
2. Currently a lot of requests from the API are timing out, but we still process those requests internally. If an API request times out, we should cancel all corresponding tasks to that API request (why process a request with nobody listening).
3. Task cancellation. When API http request gets cancelled, it should cancel corresponding task.
4. I'd like to see profiled network latency / bandwidth.
5. I'd like to see how much bandwidth each link is using.
6. We should handle the case where one machine doesn't have the model downloaded and then other machines are waiting on it. In this case we get loads of timeout errors because the others are waiting for the one that needs to download the model.
7. Solve the problem of in continuous batching when a new prompt comes in, it will block decode of the current batch until the prefill is complete.
8. We want people to be able to copy models over to a new device without ever connecting EXO to the internet. Right now EXO require internet connection once to cache some files to check if a download is complete. Instead, we should simply check if there is a non-empty model folder locally with no .partial files. This indicates it's a fully downloaded model that can be loaded.
10. More granular control over how to deploy instances.
12. Nix is great but installing it is a pain and we have ended up in a lot of cases having PATH issues or installation issues. For example, after rebooting mike it seemed to no longer have a nix installation and needed reinstalling. It has a bunch of broken symlinks left over from nix that caused ssh to fail, making it even harder to debug. We need consistent environments (perhaps MDM) so we can guarantee nix is installed properly on each machine.
13. Memory pressure instead of memory used.
14. Show the type of each connection (TB5, Ethernet, etc.) in the UI. Refer to old exo: https://github.com/exo-explore/exo/blob/56f783b38dc6b08ce606b07a5386dc40dae00330/exo/helpers.py#L251
15. Prioritise certain connection types (or by latency). TB5 > Ethernet > WiFi. Refer to old exo: https://github.com/exo-explore/exo/blob/56f783b38dc6b08ce606b07a5386dc40dae00330/exo/helpers.py#L251
16. Dynamically switch to higher priority connection when it becomes available. Probably bring back InstanceReplacedAtomically.
17. Faster model loads by streaming model from other devices in cluster.
18. Add support for specifying the type of network connection to use in a test. Depends on 15/16.
20. Add chat completion cancellations (e.g OpenWebUI has something for cancelling an ongoing request).
23. Do we need cache_limit? We went back and forth on that a lot because we thought it might be causing issues. One problem is it sets it relative to model size. So if you have multiple models loaded in it will take the most recent model size for the cache_limit. This is problematic if you launch DeepSeek -> Llama for example.
24. further openai/lmstudio api compatibility
25. Rethink retry logic
26. Task cancellation. When API http request gets cancelled, it should cancel corresponding task.
27. Log cleanup - per-module log filters and default to DEBUG log levels
28. Validate RDMA connections with ibv_devinfo in the info gatherer
Potential refactors:
2. Topology can be simplified
Random errors we've run into:
+33 -31
View File
@@ -563,21 +563,45 @@ struct ContentView: View {
}
private var rdmaStatusView: some View {
let rdma = networkStatusService.status.rdmaStatus
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) {
Text("RDMA: \(rdmaStatusText(rdma))")
.font(.caption2)
.foregroundColor(rdmaStatusColor(rdma))
if !rdma.devices.isEmpty {
Text(" Devices: \(rdma.devices.joined(separator: ", "))")
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 !rdma.activePorts.isEmpty {
Text(" Active Ports:")
if !localPorts.isEmpty {
Text(" Local Active Ports:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(rdma.activePorts, id: \.device) { port in
ForEach(localPorts, id: \.device) { port in
Text(" \(port.device) port \(port.port): \(port.state)")
.font(.caption2)
.foregroundColor(.green)
@@ -586,28 +610,6 @@ struct ContentView: View {
}
}
private func rdmaStatusText(_ rdma: RDMAStatus) -> String {
switch rdma.rdmaCtlEnabled {
case .some(true):
return "Enabled"
case .some(false):
return "Disabled"
case nil:
return rdma.devices.isEmpty ? "Not Available" : "Available"
}
}
private func rdmaStatusColor(_ rdma: RDMAStatus) -> Color {
switch rdma.rdmaCtlEnabled {
case .some(true):
return .green
case .some(false):
return .orange
case nil:
return rdma.devices.isEmpty ? .secondary : .green
}
}
private var sendBugReportButton: some View {
VStack(alignment: .leading, spacing: 4) {
Button {
+9 -1
View File
@@ -28,6 +28,10 @@ final class ExoProcessController: ObservableObject {
}
}
static let exoDirectoryURL: URL = {
URL(fileURLWithPath: NSHomeDirectory()).appendingPathComponent(".exo")
}()
@Published private(set) var status: Status = .stopped
@Published private(set) var lastError: String?
@Published private(set) var launchCountdownSeconds: Int?
@@ -78,7 +82,11 @@ final class ExoProcessController: ObservableObject {
let child = Process()
child.executableURL = executableURL
child.currentDirectoryURL = runtimeURL
let exoHomeURL = Self.exoDirectoryURL
try? FileManager.default.createDirectory(
at: exoHomeURL, withIntermediateDirectories: true
)
child.currentDirectoryURL = exoHomeURL
child.environment = makeEnvironment(for: runtimeURL)
child.standardOutput = FileHandle.nullDevice
+10
View File
@@ -15,6 +15,7 @@ struct ClusterState: Decodable {
let nodeMemory: [String: MemoryInfo]
let nodeSystem: [String: SystemInfo]
let nodeThunderboltBridge: [String: ThunderboltBridgeStatus]
let nodeRdmaCtl: [String: NodeRdmaCtlStatus]
/// Computed property for backwards compatibility - merges granular state into NodeProfile
var nodeProfiles: [String: NodeProfile] {
@@ -65,6 +66,10 @@ struct ClusterState: Decodable {
try container.decodeIfPresent(
[String: ThunderboltBridgeStatus].self, forKey: .nodeThunderboltBridge
) ?? [:]
self.nodeRdmaCtl =
try container.decodeIfPresent(
[String: NodeRdmaCtlStatus].self, forKey: .nodeRdmaCtl
) ?? [:]
}
private enum CodingKeys: String, CodingKey {
@@ -78,6 +83,7 @@ struct ClusterState: Decodable {
case nodeMemory
case nodeSystem
case nodeThunderboltBridge
case nodeRdmaCtl
}
}
@@ -159,6 +165,10 @@ struct ThunderboltBridgeStatus: Decodable {
let serviceName: String?
}
struct NodeRdmaCtlStatus: Decodable {
let enabled: Bool
}
struct MemoryInfo: Decodable {
let ramTotal: MemoryValue?
let ramAvailable: MemoryValue?
+46 -14
View File
@@ -44,16 +44,12 @@ struct BugReportService {
let dayPrefix = Self.dayPrefixString(now)
let prefix = "reports/\(dayPrefix)/\(timestamp)/"
let logData = readLog()
let logFiles = readAllLogs()
let ifconfigText = try await captureIfconfig()
let hostName = Host.current().localizedName ?? "unknown"
let debugInfo = readDebugInfo()
async let stateResult = fetch(url: baseURL.appendingPathComponent("state"))
async let eventsResult = fetch(url: baseURL.appendingPathComponent("events"))
let stateData = try await stateResult
let eventsData = try await eventsResult
let stateData = try await fetch(url: baseURL.appendingPathComponent("state"))
// Extract cluster TB bridge status from exo state
let clusterTbBridgeStatus = extractClusterTbBridgeStatus(from: stateData)
@@ -67,12 +63,19 @@ struct BugReportService {
clusterTbBridgeStatus: clusterTbBridgeStatus
)
let uploads: [(path: String, data: Data?)] = [
("\(prefix)exo.log", logData),
let eventLogFiles = readAllEventLogs()
var uploads: [(path: String, data: Data?)] = logFiles.map { (path, data) in
("\(prefix)\(path)", data)
}
uploads.append(
contentsOf: eventLogFiles.map { (path, data) in
("\(prefix)\(path)", data as Data?)
})
uploads.append(contentsOf: [
("\(prefix)state.json", stateData),
("\(prefix)events.json", eventsData),
("\(prefix)report.json", reportJSON),
]
])
let uploadItems: [(key: String, body: Data)] = uploads.compactMap { item in
guard let body = item.data else { return nil }
@@ -149,11 +152,40 @@ struct BugReportService {
return decoded.urls
}
private func readLog() -> Data? {
let logURL = URL(fileURLWithPath: NSHomeDirectory())
private func readAllLogs() -> [(path: String, data: Data)] {
let dir = URL(fileURLWithPath: NSHomeDirectory())
.appendingPathComponent(".exo")
.appendingPathComponent("exo.log")
return try? Data(contentsOf: logURL)
.appendingPathComponent("exo_log")
var results: [(path: String, data: Data)] = []
let contents = (try? FileManager.default.contentsOfDirectory(atPath: dir.path)) ?? []
for name in contents {
if let data = try? Data(contentsOf: dir.appendingPathComponent(name)) {
results.append(("exo_log/\(name)", data))
}
}
return results
}
private func readAllEventLogs() -> [(path: String, data: Data)] {
let eventLogDir = URL(fileURLWithPath: NSHomeDirectory())
.appendingPathComponent(".exo")
.appendingPathComponent("event_log")
var results: [(path: String, data: Data)] = []
for subdir in ["master", "api"] {
let dir = eventLogDir.appendingPathComponent(subdir)
let contents =
(try? FileManager.default.contentsOfDirectory(atPath: dir.path)) ?? []
for name in contents where name.hasPrefix("events.") {
if let data = try? Data(contentsOf: dir.appendingPathComponent(name)) {
results.append(("event_log/\(subdir)/\(name)", data))
}
}
}
return results
}
private func captureIfconfig() async throws -> String {
@@ -35,28 +35,18 @@ struct NetworkStatus: Equatable {
let thunderboltBridgeState: ThunderboltState?
let bridgeInactive: Bool?
let interfaceStatuses: [InterfaceIpStatus]
let rdmaStatus: RDMAStatus
let localRdmaDevices: [String]
let localRdmaActivePorts: [RDMAPort]
static let empty = NetworkStatus(
thunderboltBridgeState: nil,
bridgeInactive: nil,
interfaceStatuses: [],
rdmaStatus: .empty
localRdmaDevices: [],
localRdmaActivePorts: []
)
}
struct RDMAStatus: Equatable {
let rdmaCtlEnabled: Bool?
let devices: [String]
let activePorts: [RDMAPort]
var isAvailable: Bool {
rdmaCtlEnabled == true || !devices.isEmpty
}
static let empty = RDMAStatus(rdmaCtlEnabled: nil, devices: [], activePorts: [])
}
struct RDMAPort: Equatable {
let device: String
let port: String
@@ -80,31 +70,11 @@ private struct NetworkStatusFetcher {
thunderboltBridgeState: readThunderboltBridgeState(),
bridgeInactive: readBridgeInactive(),
interfaceStatuses: readInterfaceStatuses(),
rdmaStatus: readRDMAStatus()
localRdmaDevices: readRDMADevices(),
localRdmaActivePorts: readRDMAActivePorts()
)
}
private func readRDMAStatus() -> RDMAStatus {
let rdmaCtlEnabled = readRDMACtlEnabled()
let devices = readRDMADevices()
let activePorts = readRDMAActivePorts()
return RDMAStatus(
rdmaCtlEnabled: rdmaCtlEnabled, devices: devices, activePorts: activePorts)
}
private func readRDMACtlEnabled() -> Bool? {
let result = runCommand(["rdma_ctl", "status"])
guard result.exitCode == 0 else { return nil }
let output = result.output.lowercased().trimmingCharacters(in: .whitespacesAndNewlines)
if output.contains("enabled") {
return true
}
if output.contains("disabled") {
return false
}
return nil
}
private func readRDMADevices() -> [String] {
let result = runCommand(["ibv_devices"])
guard result.exitCode == 0 else { return [] }
+245 -50
View File
@@ -19,6 +19,11 @@ from urllib.parse import urlencode
from loguru import logger
from transformers import AutoTokenizer
# Backoff constants for cluster settling retry
_SETTLE_INITIAL_BACKOFF_S = 1.0
_SETTLE_MAX_BACKOFF_S = 60.0
_SETTLE_BACKOFF_MULTIPLIER = 2.0
# Monkey-patch for transformers 5.x compatibility
# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
# which was moved in transformers 5.0.0rc2
@@ -283,6 +288,151 @@ def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]
raise ValueError(f"Model not found in /models: {model_arg}")
def run_planning_phase(
client: ExoClient,
full_model_id: str,
preview: dict[str, Any],
danger_delete: bool,
timeout: float,
settle_deadline: float | None,
) -> None:
"""Check disk space and ensure model is downloaded before benchmarking."""
# Get model size from /models
models = client.request_json("GET", "/models") or {}
model_bytes = 0
for m in models.get("data", []):
if m.get("hugging_face_id") == full_model_id:
model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
break
if not model_bytes:
logger.warning(
f"Could not determine size for {full_model_id}, skipping disk check"
)
return
# Get nodes from preview
inner = unwrap_instance(preview["instance"])
node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
state = client.request_json("GET", "/state")
downloads = state.get("downloads", {})
node_disk = state.get("nodeDisk", {})
for node_id in node_ids:
node_downloads = downloads.get(node_id, [])
# Check if model already downloaded on this node
already_downloaded = any(
"DownloadCompleted" in p
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
]
== full_model_id
for p in node_downloads
)
if already_downloaded:
continue
# Wait for disk info if settle_deadline is set
disk_info = node_disk.get(node_id, {})
backoff = _SETTLE_INITIAL_BACKOFF_S
while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
remaining = settle_deadline - time.monotonic()
logger.info(
f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
state = client.request_json("GET", "/state")
node_disk = state.get("nodeDisk", {})
disk_info = node_disk.get(node_id, {})
if not disk_info:
logger.warning(f"No disk info for {node_id}, skipping space check")
continue
avail = disk_info.get("available", {}).get("inBytes", 0)
if avail >= model_bytes:
continue
if not danger_delete:
raise RuntimeError(
f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
)
# Delete from smallest to largest
completed = [
(
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
],
p["DownloadCompleted"]["totalBytes"]["inBytes"],
)
for p in node_downloads
if "DownloadCompleted" in p
]
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)")
client.request_json("DELETE", f"/download/{node_id}/{del_model}")
avail += size
if avail >= model_bytes:
break
if avail < model_bytes:
raise RuntimeError(f"Could not free enough space on {node_id}")
# Start downloads (idempotent)
for node_id in node_ids:
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
shard = runner_to_shard[runner_id]
client.request_json(
"POST",
"/download/start",
body={
"targetNodeId": node_id,
"shardMetadata": shard,
},
)
logger.info(f"Started download on {node_id}")
# Wait for downloads
start = time.time()
while time.time() - start < timeout:
state = client.request_json("GET", "/state")
downloads = state.get("downloads", {})
all_done = True
for node_id in node_ids:
done = any(
"DownloadCompleted" in p
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
"modelCard"
]["modelId"]
== full_model_id
for p in downloads.get(node_id, [])
)
failed = [
p["DownloadFailed"]["errorMessage"]
for p in downloads.get(node_id, [])
if "DownloadFailed" in p
and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
"modelId"
]
== full_model_id
]
if failed:
raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
if not done:
all_done = False
if all_done:
return
time.sleep(1)
raise TimeoutError("Downloads did not complete in time")
def placement_filter(instance_meta: str, wanted: str) -> bool:
s = (instance_meta or "").lower()
if wanted == "both":
@@ -388,6 +538,66 @@ class PromptSizer:
return content, tok
def fetch_and_filter_placements(
client: ExoClient, full_model_id: str, args: argparse.Namespace
) -> list[dict[str, Any]]:
previews_resp = client.request_json(
"GET", "/instance/previews", params={"model_id": full_model_id}
)
previews = previews_resp.get("previews") or []
selected: list[dict[str, Any]] = []
for p in previews:
if p.get("error") is not None:
continue
if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
continue
if not sharding_filter(str(p.get("sharding", "")), args.sharding):
continue
instance = p.get("instance")
if not isinstance(instance, dict):
continue
n = nodes_used_in_instance(instance)
# Skip tensor ring single node as it is pointless when pipeline ring
if n == 1 and (
(args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
or (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
):
continue
if (
args.skip_pipeline_jaccl
and (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
and (
args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
)
):
continue
if (
args.skip_tensor_ring
and (
args.instance_meta == "both"
and "ring" in p.get("instance_meta", "").lower()
)
and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
):
continue
if args.min_nodes <= n <= args.max_nodes:
selected.append(p)
return selected
def main() -> int:
ap = argparse.ArgumentParser(
prog="exo-bench",
@@ -464,6 +674,17 @@ def main() -> int:
action="store_true",
help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
)
ap.add_argument(
"--settle-timeout",
type=float,
default=0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
ap.add_argument(
"--danger-delete-downloads",
action="store_true",
help="Delete existing models from smallest to largest to make room for benchmark model.",
)
args = ap.parse_args()
pp_list = parse_int_list(args.pp)
@@ -487,11 +708,6 @@ def main() -> int:
client = ExoClient(args.host, args.port, timeout_s=args.timeout)
short_id, full_model_id = resolve_model_short_id(client, args.model)
previews_resp = client.request_json(
"GET", "/instance/previews", params={"model_id": full_model_id}
)
previews = previews_resp.get("previews") or []
tokenizer = load_tokenizer_for_bench(full_model_id)
if tokenizer is None:
raise RuntimeError("[exo-bench] tokenizer load failed")
@@ -503,54 +719,23 @@ def main() -> int:
logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
raise
selected: list[dict[str, Any]] = []
for p in previews:
if p.get("error") is not None:
continue
if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
continue
if not sharding_filter(str(p.get("sharding", "")), args.sharding):
continue
settle_deadline = (
time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
)
instance = p.get("instance")
if not isinstance(instance, dict):
continue
selected = fetch_and_filter_placements(client, full_model_id, args)
n = nodes_used_in_instance(instance)
# Skip tensor ring single node as it is pointless when pipeline ring
if n == 1 and (
(args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
or (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
if not selected and settle_deadline:
backoff = _SETTLE_INITIAL_BACKOFF_S
while not selected and time.monotonic() < settle_deadline:
remaining = settle_deadline - time.monotonic()
logger.warning(
f"No valid placements yet (cluster may still be settling). "
f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
)
):
continue
if (
args.skip_pipeline_jaccl
and (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
and (
args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
)
):
continue
if (
args.skip_tensor_ring
and (
args.instance_meta == "both"
and "ring" in p.get("instance_meta", "").lower()
)
and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
):
continue
if args.min_nodes <= n <= args.max_nodes:
selected.append(p)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
selected = fetch_and_filter_placements(client, full_model_id, args)
if not selected:
logger.error("No valid placements matched your filters.")
@@ -575,6 +760,16 @@ def main() -> int:
if args.dry_run:
return 0
logger.info("Planning phase: checking downloads...")
run_planning_phase(
client,
full_model_id,
selected[0],
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
all_rows: list[dict[str, Any]] = []
for preview in selected:
+16
View File
@@ -0,0 +1,16 @@
[project]
name = "exo-bench"
version = "0.1.0"
description = "Benchmarking tool for exo distributed inference"
requires-python = ">=3.13"
dependencies = [
"loguru>=0.7.3",
"transformers>=5.0.0",
"huggingface-hub>=0.33.4",
"tiktoken>=0.12.0",
"jinja2>=3.1.0",
]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
View File
+52 -4
View File
@@ -12,6 +12,8 @@
ttftMs,
tps,
totalTokens,
thinkingEnabled as thinkingEnabledStore,
setConversationThinking,
} from "$lib/stores/app.svelte";
import ChatAttachments from "./ChatAttachments.svelte";
import ImageParamsPanel from "./ImageParamsPanel.svelte";
@@ -25,6 +27,7 @@
autofocus?: boolean;
showModelSelector?: boolean;
modelTasks?: Record<string, string[]>;
modelCapabilities?: Record<string, string[]>;
}
let {
@@ -34,6 +37,7 @@
autofocus = true,
showModelSelector = false,
modelTasks = {},
modelCapabilities = {},
}: Props = $props();
let message = $state("");
@@ -41,6 +45,7 @@
let fileInputRef: HTMLInputElement | undefined = $state();
let uploadedFiles = $state<ChatUploadedFile[]>([]);
let isDragOver = $state(false);
const thinkingEnabled = $derived(thinkingEnabledStore());
let loading = $derived(isLoading());
const currentModel = $derived(selectedChatModel());
const instanceData = $derived(instances());
@@ -95,6 +100,12 @@
);
});
const modelSupportsThinking = $derived(() => {
if (!currentModel) return false;
const caps = modelCapabilities[currentModel] || [];
return caps.includes("thinking") && caps.includes("text");
});
const isEditOnlyWithoutImage = $derived(
currentModel !== null &&
modelSupportsOnlyImageEditing(currentModel) &&
@@ -254,6 +265,7 @@
function handleSubmit() {
if ((!message.trim() && uploadedFiles.length === 0) || loading) return;
if (isEditOnlyWithoutImage) return;
const content = message.trim();
const files = [...uploadedFiles];
@@ -278,11 +290,19 @@
if (imageFile.preview) {
editImage(content, imageFile.preview);
}
} else if (isImageModel() && content) {
} else if (
currentModel &&
modelSupportsTextToImage(currentModel) &&
content
) {
// Use image generation for text-to-image models
generateImage(content);
} else {
sendMessage(content, files);
sendMessage(
content,
files,
modelSupportsThinking() ? thinkingEnabled : null,
);
}
// Refocus the textarea after sending
@@ -520,6 +540,35 @@
</div>
{/if}
</div>
<!-- Thinking toggle -->
{#if modelSupportsThinking()}
<button
type="button"
onclick={() => setConversationThinking(!thinkingEnabled)}
class="flex items-center gap-1.5 px-2 py-1 rounded text-xs font-mono tracking-wide transition-all duration-200 flex-shrink-0 cursor-pointer border {thinkingEnabled
? 'bg-exo-yellow/15 border-exo-yellow/40 text-exo-yellow'
: 'bg-exo-medium-gray/30 border-exo-medium-gray/50 text-exo-light-gray/60 hover:text-exo-light-gray'}"
title={thinkingEnabled
? "Thinking enabled — click to disable"
: "Thinking disabled — click to enable"}
>
<svg
class="w-3.5 h-3.5"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="1.5"
>
<path
d="M12 2a7 7 0 0 0-7 7c0 2.38 1.19 4.47 3 5.74V17a1 1 0 0 0 1 1h6a1 1 0 0 0 1-1v-2.26c1.81-1.27 3-3.36 3-5.74a7 7 0 0 0-7-7zM9 20h6M10 22h4"
stroke-linecap="round"
stroke-linejoin="round"
/>
</svg>
<span>{thinkingEnabled ? "THINK" : "NO THINK"}</span>
</button>
{/if}
<!-- Performance stats -->
{#if currentTtft !== null || currentTps !== null}
<div class="flex items-center gap-4 text-xs font-mono flex-shrink-0">
@@ -599,9 +648,8 @@
: isImageModel()
? "Describe the image you want to generate..."
: placeholder}
disabled={loading}
rows={1}
class="flex-1 resize-none bg-transparent text-foreground placeholder:text-exo-light-gray/60 placeholder:text-sm placeholder:tracking-[0.15em] placeholder:leading-7 focus:outline-none focus:ring-0 focus:border-none disabled:opacity-50 text-sm leading-7 font-mono"
class="flex-1 resize-none bg-transparent text-foreground placeholder:text-exo-light-gray/60 placeholder:text-sm placeholder:tracking-[0.15em] placeholder:leading-7 focus:outline-none focus:ring-0 focus:border-none text-sm leading-7 font-mono"
style="min-height: 28px; max-height: 150px;"
></textarea>
@@ -13,6 +13,7 @@
import type { MessageAttachment } from "$lib/stores/app.svelte";
import MarkdownContent from "./MarkdownContent.svelte";
import TokenHeatmap from "./TokenHeatmap.svelte";
import ImageLightbox from "./ImageLightbox.svelte";
interface Props {
class?: string;
@@ -101,6 +102,9 @@
let copiedMessageId = $state<string | null>(null);
let expandedThinkingMessageIds = $state<Set<string>>(new Set());
// Lightbox state
let expandedImageSrc = $state<string | null>(null);
// Uncertainty heatmap toggle
let heatmapMessageIds = $state<Set<string>>(new Set());
@@ -221,6 +225,7 @@
}
function handleDeleteClick(messageId: string) {
if (loading) return;
deleteConfirmId = messageId;
}
@@ -251,7 +256,7 @@
</script>
<div class="flex flex-col gap-4 sm:gap-6 {className}">
{#each messageList as message (message.id)}
{#each messageList as message, i (message.id)}
<div
class="group flex {message.role === 'user'
? 'justify-end'
@@ -313,9 +318,11 @@
<!-- Delete confirmation -->
<div class="bg-red-500/10 border border-red-500/30 rounded-lg p-3">
<p class="text-xs text-red-400 mb-3">
Delete this message{message.role === "user"
? " and all responses after it"
: ""}?
{#if i === messageList.length - 1}
Delete this message?
{:else}
Delete this message and all messages after it?
{/if}
</p>
<div class="flex gap-2 justify-end">
<button
@@ -389,10 +396,15 @@
class="flex items-center gap-2 bg-exo-dark-gray/60 border border-exo-yellow/20 rounded px-2 py-1 text-xs font-mono"
>
{#if attachment.type === "image" && attachment.preview}
<!-- svelte-ignore a11y_no_noninteractive_element_interactions, a11y_click_events_have_key_events -->
<img
src={attachment.preview}
alt={attachment.name}
class="w-12 h-12 object-cover rounded border border-exo-yellow/20"
class="w-12 h-12 object-cover rounded border border-exo-yellow/20 cursor-pointer hover:border-exo-yellow/50 transition-colors"
onclick={() => {
if (attachment.preview)
expandedImageSrc = attachment.preview;
}}
/>
{:else}
<span>{getAttachmentIcon(attachment)}</span>
@@ -466,15 +478,44 @@
<div class="mb-3">
{#each message.attachments.filter((a) => a.type === "generated-image") as attachment}
<div class="relative group/img inline-block">
<!-- svelte-ignore a11y_no_noninteractive_element_interactions, a11y_click_events_have_key_events -->
<img
src={attachment.preview}
alt=""
class="max-w-full max-h-[512px] rounded-lg border border-exo-yellow/20 shadow-lg shadow-black/20"
class="max-w-full max-h-[512px] rounded-lg border border-exo-yellow/20 shadow-lg shadow-black/20 cursor-pointer"
onclick={() => {
if (attachment.preview)
expandedImageSrc = attachment.preview;
}}
/>
<!-- Button overlay -->
<div
class="absolute top-2 right-2 flex gap-1 opacity-0 group-hover/img:opacity-100 transition-opacity"
>
<!-- Expand button -->
<button
type="button"
class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer"
onclick={() => {
if (attachment.preview)
expandedImageSrc = attachment.preview;
}}
title="Expand image"
>
<svg
class="w-4 h-4"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M4 8V4m0 0h4M4 4l5 5m11-1V4m0 0h-4m4 0l-5 5M4 16v4m0 0h4m-4 0l5-5m11 5l-5-5m5 5v-4m0 4h-4"
/>
</svg>
</button>
<!-- Edit button -->
<button
type="button"
@@ -713,8 +754,13 @@
<!-- Delete button -->
<button
onclick={() => handleDeleteClick(message.id)}
class="p-1.5 text-exo-light-gray hover:text-red-400 transition-colors rounded hover:bg-red-500/10 cursor-pointer"
title="Delete message"
disabled={loading}
class="p-1.5 transition-colors rounded {loading
? 'text-exo-light-gray/30 cursor-not-allowed'
: 'text-exo-light-gray hover:text-red-400 hover:bg-red-500/10 cursor-pointer'}"
title={loading
? "Cannot delete while generating"
: "Delete message"}
>
<svg
class="w-3.5 h-3.5"
@@ -789,3 +835,8 @@
</button>
{/if}
</div>
<ImageLightbox
src={expandedImageSrc}
onclose={() => (expandedImageSrc = null)}
/>
@@ -185,11 +185,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 {
@@ -13,6 +13,12 @@
d="M12 2l3.09 6.26L22 9.27l-5 4.87 1.18 6.88L12 17.77l-6.18 3.25L7 14.14 2 9.27l6.91-1.01L12 2z"
/>
</svg>
{:else if family === "recents"}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
d="M13 3a9 9 0 0 0-9 9H1l3.89 3.89.07.14L9 12H6c0-3.87 3.13-7 7-7s7 3.13 7 7-3.13 7-7 7c-1.93 0-3.68-.79-4.94-2.06l-1.42 1.42A8.954 8.954 0 0 0 13 21a9 9 0 0 0 0-18zm-1 5v5l4.28 2.54.72-1.21-3.5-2.08V8H12z"
/>
</svg>
{:else if family === "llama" || family === "meta"}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
@@ -58,6 +64,18 @@
d="M17.962 1.844h-4.326l-3.425 7.81H5.369V1.878H1.5V22h3.87v-8.477h6.824a3.025 3.025 0 002.743-1.75V22h3.87v-8.477a3.87 3.87 0 00-3.588-3.86v-.01h-2.125a3.94 3.94 0 002.323-2.12l2.545-5.689z"
/>
</svg>
{:else if family === "flux"}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
d="M12 2L2 19h7.5l2.5-4.5L14.5 19H22L12 2zm0 4.5L16.5 17h-3l-1.5-2.7L10.5 17h-3L12 6.5z"
/>
</svg>
{:else if family === "qwen-image"}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
d="M12.604 1.34c.393.69.784 1.382 1.174 2.075a.18.18 0 00.157.091h5.552c.174 0 .322.11.446.327l1.454 2.57c.19.337.24.478.024.837-.26.43-.513.864-.76 1.3l-.367.658c-.106.196-.223.28-.04.512l2.652 4.637c.172.301.111.494-.043.77-.437.785-.882 1.564-1.335 2.34-.159.272-.352.375-.68.37-.777-.016-1.552-.01-2.327.016a.099.099 0 00-.081.05 575.097 575.097 0 01-2.705 4.74c-.169.293-.38.363-.725.364-.997.003-2.002.004-3.017.002a.537.537 0 01-.465-.271l-1.335-2.323a.09.09 0 00-.083-.049H4.982c-.285.03-.553-.001-.805-.092l-1.603-2.77a.543.543 0 01-.002-.54l1.207-2.12a.198.198 0 000-.197 550.951 550.951 0 01-1.875-3.272l-.79-1.395c-.16-.31-.173-.496.095-.965.465-.813.927-1.625 1.387-2.436.132-.234.304-.334.584-.335a338.3 338.3 0 012.589-.001.124.124 0 00.107-.063l2.806-4.895a.488.488 0 01.422-.246c.524-.001 1.053 0 1.583-.006L11.704 1c.341-.003.724.032.9.34zm-3.432.403a.06.06 0 00-.052.03L6.254 6.788a.157.157 0 01-.135.078H3.253c-.056 0-.07.025-.041.074l5.81 10.156c.025.042.013.062-.034.063l-2.795.015a.218.218 0 00-.2.116l-1.32 2.31c-.044.078-.021.118.068.118l5.716.008c.046 0 .08.02.104.061l1.403 2.454c.046.081.092.082.139 0l5.006-8.76.783-1.382a.055.055 0 01.096 0l1.424 2.53a.122.122 0 00.107.062l2.763-.02a.04.04 0 00.035-.02.041.041 0 000-.04l-2.9-5.086a.108.108 0 010-.113l.293-.507 1.12-1.977c.024-.041.012-.062-.035-.062H9.2c-.059 0-.073-.026-.043-.077l1.434-2.505a.107.107 0 000-.114L9.225 1.774a.06.06 0 00-.053-.031zm6.29 8.02c.046 0 .058.02.034.06l-.832 1.465-2.613 4.585a.056.056 0 01-.05.029.058.058 0 01-.05-.029L8.498 9.841c-.02-.034-.01-.052.028-.054l.216-.012 6.722-.012z"
/>
</svg>
{:else if family === "huggingface"}
<svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
<path
@@ -5,15 +5,22 @@
families: string[];
selectedFamily: string | null;
hasFavorites: boolean;
hasRecents: boolean;
onSelect: (family: string | null) => void;
};
let { families, selectedFamily, hasFavorites, onSelect }: FamilySidebarProps =
$props();
let {
families,
selectedFamily,
hasFavorites,
hasRecents,
onSelect,
}: FamilySidebarProps = $props();
// Family display names
const familyNames: Record<string, string> = {
favorites: "Favorites",
recents: "Recent",
huggingface: "Hub",
llama: "Meta",
qwen: "Qwen",
@@ -22,6 +29,8 @@
glm: "GLM",
minimax: "MiniMax",
kimi: "Kimi",
flux: "FLUX",
"qwen-image": "Qwen Img",
};
function getFamilyName(family: string): string {
@@ -32,7 +41,7 @@
</script>
<div
class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[64px]"
class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[64px] overflow-y-auto scrollbar-hide"
>
<!-- All models (no filter) -->
<button
@@ -87,6 +96,31 @@
</button>
{/if}
<!-- Recent (only show if has recent models) -->
{#if hasRecents}
<button
type="button"
onclick={() => onSelect("recents")}
class="group flex flex-col items-center justify-center p-2 rounded transition-all duration-200 cursor-pointer {selectedFamily ===
'recents'
? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
: 'hover:bg-white/5 border-l-2 border-transparent'}"
title="Recently launched models"
>
<FamilyLogos
family="recents"
class={selectedFamily === "recents"
? "text-exo-yellow"
: "text-white/50 group-hover:text-white/70"}
/>
<span
class="text-[9px] font-mono mt-0.5 {selectedFamily === 'recents'
? 'text-exo-yellow'
: 'text-white/40 group-hover:text-white/60'}">Recent</span
>
</button>
{/if}
<!-- HuggingFace Hub -->
<button
type="button"
@@ -0,0 +1,96 @@
<script lang="ts">
import { fade, fly } from "svelte/transition";
import { cubicOut } from "svelte/easing";
interface Props {
src: string | null;
onclose: () => void;
}
let { src, onclose }: Props = $props();
function handleKeydown(e: KeyboardEvent) {
if (e.key === "Escape") {
onclose();
}
}
function extensionFromSrc(dataSrc: string): string {
const match = dataSrc.match(/^data:image\/(\w+)/);
if (match) return match[1] === "jpeg" ? "jpg" : match[1];
const urlMatch = dataSrc.match(/\.(\w+)(?:\?|$)/);
if (urlMatch) return urlMatch[1];
return "png";
}
function handleDownload(e: MouseEvent) {
e.stopPropagation();
if (!src) return;
const link = document.createElement("a");
link.href = src;
link.download = `image-${Date.now()}.${extensionFromSrc(src)}`;
link.click();
}
function handleClose(e: MouseEvent) {
e.stopPropagation();
onclose();
}
</script>
<svelte:window onkeydown={src ? handleKeydown : undefined} />
{#if src}
<div
class="fixed inset-0 z-50 bg-black/90 backdrop-blur-sm flex items-center justify-center"
transition:fade={{ duration: 200 }}
onclick={onclose}
role="presentation"
onintrostart={() => (document.body.style.overflow = "hidden")}
onoutroend={() => (document.body.style.overflow = "")}
>
<div class="absolute top-4 right-4 flex gap-2 z-10">
<button
type="button"
class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer transition-colors"
onclick={handleDownload}
title="Download image"
>
<svg
class="w-5 h-5"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
/>
</svg>
</button>
<button
type="button"
class="p-2 rounded-lg bg-exo-dark-gray/80 border border-exo-yellow/30 text-exo-yellow hover:bg-exo-dark-gray hover:border-exo-yellow/50 cursor-pointer transition-colors"
onclick={handleClose}
title="Close"
>
<svg class="w-5 h-5" viewBox="0 0 24 24" fill="currentColor">
<path
d="M19 6.41L17.59 5 12 10.59 6.41 5 5 6.41 10.59 12 5 17.59 6.41 19 12 13.41 17.59 19 19 17.59 13.41 12 19 6.41z"
/>
</svg>
</button>
</div>
<!-- svelte-ignore a11y_no_noninteractive_element_interactions, a11y_click_events_have_key_events -->
<img
{src}
alt=""
class="max-w-[90vw] max-h-[90vh] object-contain rounded-lg shadow-2xl"
transition:fly={{ y: 20, duration: 300, easing: cubicOut }}
onclick={(e) => e.stopPropagation()}
/>
</div>
{/if}
+13 -6
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 {
@@ -422,9 +422,16 @@
const bToACandidates: Array<{ ip: string; iface: string | null }> = [];
for (const edge of topology.edges) {
const ip = edge.sendBackIp || "?";
const iface =
edge.sendBackInterface || getInterfaceForIp(edge.source, ip);
let ip: string;
let iface: string | null;
if (edge.sourceRdmaIface || edge.sinkRdmaIface) {
ip = "RDMA";
iface = `${edge.sourceRdmaIface || "?"} \u2192 ${edge.sinkRdmaIface || "?"}`;
} else {
ip = edge.sendBackIp || "?";
iface = edge.sendBackInterface || getInterfaceForIp(edge.source, ip);
}
if (edge.source === nodeId1 && edge.target === nodeId2) {
aToBCandidates.push({ ip, iface });
@@ -568,7 +575,7 @@
>
{runtime === "MlxRing"
? "MLX Ring"
: runtime === "MlxIbv" || runtime === "MlxJaccl"
: runtime === "MlxJaccl"
? "MLX RDMA"
: runtime}
</span>
@@ -24,6 +24,8 @@
{ id: "thinking", label: "Thinking" },
{ id: "code", label: "Code" },
{ id: "vision", label: "Vision" },
{ id: "image_gen", label: "Image Gen" },
{ id: "image_edit", label: "Image Edit" },
];
// Size ranges
@@ -142,6 +144,45 @@
stroke-linejoin="round"
/><circle cx="12" cy="12" r="3" /></svg
>
{:else if cap.id === "image_gen"}
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="1.5"
><rect
x="3"
y="3"
width="18"
height="18"
rx="2"
ry="2"
stroke-linecap="round"
stroke-linejoin="round"
/><circle cx="8.5" cy="8.5" r="1.5" /><path
d="M21 15l-5-5L5 21"
stroke-linecap="round"
stroke-linejoin="round"
/></svg
>
{:else if cap.id === "image_edit"}
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="1.5"
><path
d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"
stroke-linecap="round"
stroke-linejoin="round"
/><path
d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"
stroke-linecap="round"
stroke-linejoin="round"
/></svg
>
{/if}
<span class="ml-1">{cap.label}</span>
</button>
@@ -26,6 +26,7 @@
nodeNames: string[];
nodeIds: string[];
};
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
type ModelPickerGroupProps = {
group: ModelGroup;
@@ -33,11 +34,13 @@
isFavorite: boolean;
selectedModelId: string | null;
canModelFit: (id: string) => boolean;
getModelFitStatus: (id: string) => ModelFitStatus;
onToggleExpand: () => void;
onSelectModel: (modelId: string) => void;
onToggleFavorite: (baseModelId: string) => void;
onShowInfo: (group: ModelGroup) => void;
downloadStatusMap?: Map<string, DownloadAvailability>;
launchedAt?: number;
};
let {
@@ -46,11 +49,13 @@
isFavorite,
selectedModelId,
canModelFit,
getModelFitStatus,
onToggleExpand,
onSelectModel,
onToggleFavorite,
onShowInfo,
downloadStatusMap,
launchedAt,
}: ModelPickerGroupProps = $props();
// Group-level download status: show if any variant is downloaded
@@ -72,10 +77,45 @@
return `${mb}MB`;
}
function timeAgo(ts: number): string {
const seconds = Math.floor((Date.now() - ts) / 1000);
if (seconds < 60) return "just now";
const minutes = Math.floor(seconds / 60);
if (minutes < 60) return `${minutes}m ago`;
const hours = Math.floor(minutes / 60);
if (hours < 24) return `${hours}h ago`;
const days = Math.floor(hours / 24);
return `${days}d ago`;
}
// Check if any variant can fit
const anyVariantFits = $derived(
group.variants.some((v) => canModelFit(v.id)),
);
const groupFitStatus = $derived.by((): ModelFitStatus => {
let hasClusterCapacityOnly = false;
for (const variant of group.variants) {
const fitStatus = getModelFitStatus(variant.id);
if (fitStatus === "fits_now") {
return "fits_now";
}
if (fitStatus === "fits_cluster_capacity") {
hasClusterCapacityOnly = true;
}
}
return hasClusterCapacityOnly ? "fits_cluster_capacity" : "too_large";
});
function getSizeClassForFitStatus(fitStatus: ModelFitStatus): string {
switch (fitStatus) {
case "fits_now":
return "text-white/40";
case "fits_cluster_capacity":
return "text-orange-400/80";
case "too_large":
return "text-red-400/70";
}
}
// Check if this group's model is currently selected (for single-variant groups)
const isMainSelected = $derived(
@@ -200,9 +240,42 @@
stroke-width="1.5"
title="Supports image generation"
>
<rect x="3" y="3" width="18" height="18" rx="2" ry="2" />
<rect
x="3"
y="3"
width="18"
height="18"
rx="2"
ry="2"
stroke-linecap="round"
stroke-linejoin="round"
/>
<circle cx="8.5" cy="8.5" r="1.5" />
<path d="M21 15l-5-5L5 21" />
<path
d="M21 15l-5-5L5 21"
stroke-linecap="round"
stroke-linejoin="round"
/>
</svg>
{:else if cap === "image_edit"}
<svg
class="w-3.5 h-3.5 text-white/40 flex-shrink-0"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="1.5"
title="Supports image editing"
>
<path
d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"
stroke-linecap="round"
stroke-linejoin="round"
/>
<path
d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"
stroke-linecap="round"
stroke-linejoin="round"
/>
</svg>
{/if}
{/each}
@@ -211,7 +284,14 @@
<!-- Size indicator (smallest variant) -->
{#if !group.hasMultipleVariants && group.smallestVariant?.storage_size_megabytes}
<span class="text-xs font-mono text-white/30 flex-shrink-0">
{@const singleVariantFitStatus = getModelFitStatus(
group.smallestVariant.id,
)}
<span
class="text-xs font-mono flex-shrink-0 {getSizeClassForFitStatus(
singleVariantFitStatus,
)}"
>
{formatSize(group.smallestVariant.storage_size_megabytes)}
</span>
{/if}
@@ -222,13 +302,24 @@
.map((v) => v.storage_size_megabytes || 0)
.filter((s) => s > 0)
.sort((a, b) => a - b)}
<span class="text-xs font-mono text-white/30 flex-shrink-0">
<span
class="text-xs font-mono flex-shrink-0 {getSizeClassForFitStatus(
groupFitStatus,
)}"
>
{group.variants.length} variants{#if sizes.length >= 2}{" "}({formatSize(
sizes[0],
)}-{formatSize(sizes[sizes.length - 1])}){/if}
</span>
{/if}
<!-- Time ago (for recent models) -->
{#if launchedAt}
<span class="text-xs font-mono text-white/20 flex-shrink-0">
{timeAgo(launchedAt)}
</span>
{/if}
<!-- Download availability indicator -->
{#if groupDownloadStatus && groupDownloadStatus.nodeIds.length > 0}
<span
@@ -327,6 +418,7 @@
{#if isExpanded && group.hasMultipleVariants}
<div class="bg-black/20 border-t border-white/5">
{#each group.variants as variant}
{@const fitStatus = getModelFitStatus(variant.id)}
{@const modelCanFit = canModelFit(variant.id)}
{@const isSelected = selectedModelId === variant.id}
<button
@@ -351,7 +443,11 @@
</span>
<!-- Size -->
<span class="text-xs font-mono text-white/40 flex-1">
<span
class="text-xs font-mono flex-1 {getSizeClassForFitStatus(
fitStatus,
)}"
>
{formatSize(variant.storage_size_megabytes)}
</span>
@@ -6,6 +6,7 @@
import ModelFilterPopover from "./ModelFilterPopover.svelte";
import HuggingFaceResultItem from "./HuggingFaceResultItem.svelte";
import { getNodesWithModelDownloaded } from "$lib/utils/downloads";
import { getRecentEntries } from "$lib/stores/recents.svelte";
interface ModelInfo {
id: string;
@@ -46,13 +47,18 @@
tags: string[];
}
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
type ModelPickerModalProps = {
isOpen: boolean;
models: ModelInfo[];
selectedModelId: string | null;
favorites: Set<string>;
recentModelIds?: string[];
hasRecents?: boolean;
existingModelIds: Set<string>;
canModelFit: (modelId: string) => boolean;
getModelFitStatus: (modelId: string) => ModelFitStatus;
onSelect: (modelId: string) => void;
onClose: () => void;
onToggleFavorite: (baseModelId: string) => void;
@@ -76,8 +82,11 @@
models,
selectedModelId,
favorites,
recentModelIds = [],
hasRecents: hasRecentsTab = false,
existingModelIds,
canModelFit,
getModelFitStatus,
onSelect,
onClose,
onToggleFavorite,
@@ -363,6 +372,8 @@
"deepseek",
"gpt-oss",
"llama",
"flux",
"qwen-image",
];
return Array.from(families).sort((a, b) => {
const aIdx = familyOrder.indexOf(a);
@@ -381,7 +392,11 @@
// Filter by family
if (selectedFamily === "favorites") {
result = result.filter((g) => favorites.has(g.id));
} else if (selectedFamily && selectedFamily !== "huggingface") {
} else if (
selectedFamily &&
selectedFamily !== "huggingface" &&
selectedFamily !== "recents"
) {
result = result.filter((g) => g.family === selectedFamily);
}
@@ -425,13 +440,23 @@
);
}
// Sort: models that fit first, then by size (largest first)
// Sort: fits-now first, then fits-cluster-capacity, then too-large
result.sort((a, b) => {
const aFits = a.variants.some((v) => canModelFit(v.id));
const bFits = b.variants.some((v) => canModelFit(v.id));
const getGroupFitRank = (group: ModelGroup): number => {
let hasClusterCapacityOnly = false;
for (const variant of group.variants) {
const fitStatus = getModelFitStatus(variant.id);
if (fitStatus === "fits_now") return 0;
if (fitStatus === "fits_cluster_capacity") {
hasClusterCapacityOnly = true;
}
}
return hasClusterCapacityOnly ? 1 : 2;
};
if (aFits && !bFits) return -1;
if (!aFits && bFits) return 1;
const aRank = getGroupFitRank(a);
const bRank = getGroupFitRank(b);
if (aRank !== bRank) return aRank - bRank;
return (
(b.smallestVariant.storage_size_megabytes || 0) -
@@ -445,6 +470,48 @@
// Check if any favorites exist
const hasFavorites = $derived(favorites.size > 0);
// Timestamp lookup for recent models
const recentTimestamps = $derived(
new Map(getRecentEntries().map((e) => [e.modelId, e.launchedAt])),
);
// Recent models: single-variant ModelGroups in launch order
const recentGroups = $derived.by((): ModelGroup[] => {
if (!recentModelIds || recentModelIds.length === 0) return [];
const result: ModelGroup[] = [];
for (const id of recentModelIds) {
const model = models.find((m) => m.id === id);
if (model) {
result.push({
id: model.base_model || model.id,
name: model.name || model.id,
capabilities: model.capabilities || ["text"],
family: model.family || "",
variants: [model],
smallestVariant: model,
hasMultipleVariants: false,
});
}
}
return result;
});
// Filtered recent groups (apply search query)
const filteredRecentGroups = $derived.by((): ModelGroup[] => {
if (!searchQuery.trim()) return recentGroups;
const query = searchQuery.toLowerCase().trim();
return recentGroups.filter(
(g) =>
g.name.toLowerCase().includes(query) ||
g.variants.some(
(v) =>
v.id.toLowerCase().includes(query) ||
(v.name || "").toLowerCase().includes(query) ||
(v.quantization || "").toLowerCase().includes(query),
),
);
});
function toggleGroupExpanded(groupId: string) {
const next = new Set(expandedGroups);
if (next.has(groupId)) {
@@ -602,11 +669,12 @@
families={uniqueFamilies}
{selectedFamily}
{hasFavorites}
hasRecents={hasRecentsTab}
onSelect={(family) => (selectedFamily = family)}
/>
<!-- Model list -->
<div class="flex-1 overflow-y-auto flex flex-col">
<div class="flex-1 overflow-y-auto scrollbar-hide flex flex-col">
{#if selectedFamily === "huggingface"}
<!-- HuggingFace Hub view -->
<div class="flex-1 flex flex-col min-h-0">
@@ -624,7 +692,7 @@
</div>
<!-- Results list -->
<div class="flex-1 overflow-y-auto">
<div class="flex-1 overflow-y-auto scrollbar-hide">
{#if hfIsLoadingTrending && hfTrendingModels.length === 0}
<div
class="flex items-center justify-center py-12 text-white/40"
@@ -709,6 +777,45 @@
</div>
</div>
</div>
{:else if selectedFamily === "recents"}
<!-- Recent models view -->
{#if filteredRecentGroups.length === 0}
<div
class="flex flex-col items-center justify-center h-full text-white/40 p-8"
>
<svg
class="w-12 h-12 mb-3"
viewBox="0 0 24 24"
fill="currentColor"
>
<path
d="M13 3a9 9 0 0 0-9 9H1l3.89 3.89.07.14L9 12H6c0-3.87 3.13-7 7-7s7 3.13 7 7-3.13 7-7 7c-1.93 0-3.68-.79-4.94-2.06l-1.42 1.42A8.954 8.954 0 0 0 13 21a9 9 0 0 0 0-18zm-1 5v5l4.28 2.54.72-1.21-3.5-2.08V8H12z"
/>
</svg>
<p class="font-mono text-sm">
{searchQuery
? "No matching recent models"
: "No recently launched models"}
</p>
</div>
{:else}
{#each filteredRecentGroups 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)}
launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
/>
{/each}
{/if}
{:else if filteredGroups.length === 0}
<div
class="flex flex-col items-center justify-center h-full text-white/40 p-8"
@@ -740,6 +847,7 @@
isFavorite={favorites.has(group.id)}
{selectedModelId}
{canModelFit}
{getModelFitStatus}
onToggleExpand={() => toggleGroupExpanded(group.id)}
onSelectModel={handleSelect}
{onToggleFavorite}
@@ -6,6 +6,8 @@
isTopologyMinimized,
debugMode,
nodeThunderboltBridge,
nodeRdmaCtl,
nodeIdentities,
type NodeInfo,
} from "$lib/stores/app.svelte";
@@ -31,6 +33,8 @@
const data = $derived(topologyData());
const debugEnabled = $derived(debugMode());
const tbBridgeData = $derived(nodeThunderboltBridge());
const rdmaCtlData = $derived(nodeRdmaCtl());
const identitiesData = $derived(nodeIdentities());
function getNodeLabel(nodeId: string): string {
const node = data?.nodes?.[nodeId];
@@ -333,14 +337,27 @@
if (edge.source === a) entry.aToB = true;
else entry.bToA = true;
const ip = edge.sendBackIp || "?";
const ifaceInfo = getInterfaceLabel(edge.source, ip);
let ip: string;
let ifaceLabel: string;
let missingIface: boolean;
if (edge.sourceRdmaIface || edge.sinkRdmaIface) {
ip = "RDMA";
ifaceLabel = `${edge.sourceRdmaIface || "?"} \u2192 ${edge.sinkRdmaIface || "?"}`;
missingIface = false;
} else {
ip = edge.sendBackIp || "?";
const ifaceInfo = getInterfaceLabel(edge.source, ip);
ifaceLabel = ifaceInfo.label;
missingIface = ifaceInfo.missing;
}
entry.connections.push({
from: edge.source,
to: edge.target,
ip,
ifaceLabel: ifaceInfo.label,
missingIface: ifaceInfo.missing,
ifaceLabel,
missingIface,
});
pairMap.set(key, entry);
});
@@ -1120,15 +1137,17 @@
.text(` (${ramUsagePercent.toFixed(0)}%)`);
}
// Debug mode: Show TB bridge status
// Debug mode: Show TB bridge and RDMA status
if (debugEnabled) {
let debugLabelY =
nodeInfo.y +
iconBaseHeight / 2 +
(showFullLabels ? 32 : showCompactLabels ? 26 : 22);
const debugFontSize = showFullLabels ? 9 : 7;
const debugLineHeight = showFullLabels ? 11 : 9;
const tbStatus = tbBridgeData[nodeInfo.id];
if (tbStatus) {
const tbY =
nodeInfo.y +
iconBaseHeight / 2 +
(showFullLabels ? 32 : showCompactLabels ? 26 : 22);
const tbFontSize = showFullLabels ? 9 : 7;
const tbColor = tbStatus.enabled
? "rgba(234,179,8,0.9)"
: "rgba(100,100,100,0.7)";
@@ -1136,12 +1155,46 @@
nodeG
.append("text")
.attr("x", nodeInfo.x)
.attr("y", tbY)
.attr("y", debugLabelY)
.attr("text-anchor", "middle")
.attr("fill", tbColor)
.attr("font-size", tbFontSize)
.attr("font-size", debugFontSize)
.attr("font-family", "SF Mono, Monaco, monospace")
.text(tbText);
debugLabelY += debugLineHeight;
}
const rdmaStatus = rdmaCtlData[nodeInfo.id];
if (rdmaStatus !== undefined) {
const rdmaColor = rdmaStatus.enabled
? "rgba(74,222,128,0.9)"
: "rgba(100,100,100,0.7)";
const rdmaText = rdmaStatus.enabled ? "RDMA:ON" : "RDMA:OFF";
nodeG
.append("text")
.attr("x", nodeInfo.x)
.attr("y", debugLabelY)
.attr("text-anchor", "middle")
.attr("fill", rdmaColor)
.attr("font-size", debugFontSize)
.attr("font-family", "SF Mono, Monaco, monospace")
.text(rdmaText);
debugLabelY += debugLineHeight;
}
const identity = identitiesData[nodeInfo.id];
if (identity?.osVersion) {
nodeG
.append("text")
.attr("x", nodeInfo.x)
.attr("y", debugLabelY)
.attr("text-anchor", "middle")
.attr("fill", "rgba(179,179,179,0.7)")
.attr("font-size", debugFontSize)
.attr("font-family", "SF Mono, Monaco, monospace")
.text(
`macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
);
}
}
});
+121 -10
View File
@@ -49,6 +49,7 @@ export interface NodeInfo {
};
last_macmon_update: number;
friendly_name?: string;
os_version?: string;
}
export interface TopologyEdge {
@@ -56,6 +57,8 @@ export interface TopologyEdge {
target: string;
sendBackIp?: string;
sendBackInterface?: string;
sourceRdmaIface?: string;
sinkRdmaIface?: string;
}
export interface TopologyData {
@@ -76,6 +79,8 @@ interface RawNodeIdentity {
modelId?: string;
chipId?: string;
friendlyName?: string;
osVersion?: string;
osBuildVersion?: string;
}
interface RawMemoryUsage {
@@ -163,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;
@@ -214,7 +219,6 @@ interface RawStateResponse {
string,
{
MlxRingInstance?: Instance;
MlxIbvInstance?: Instance;
MlxJacclInstance?: Instance;
}
>;
@@ -225,6 +229,19 @@ interface RawStateResponse {
nodeMemory?: Record<string, RawMemoryUsage>;
nodeSystem?: Record<string, RawSystemPerformanceProfile>;
nodeNetwork?: Record<string, RawNodeNetworkInfo>;
// Thunderbolt identifiers per node
nodeThunderbolt?: Record<
string,
{
interfaces: Array<{
rdmaInterface: string;
domainUuid: string;
linkSpeed: string;
}>;
}
>;
// RDMA ctl status per node
nodeRdmaCtl?: Record<string, { enabled: boolean }>;
// Thunderbolt bridge status per node
nodeThunderboltBridge?: Record<
string,
@@ -232,6 +249,20 @@ interface RawStateResponse {
>;
// Thunderbolt bridge cycles (nodes with bridge enabled forming loops)
thunderboltBridgeCycles?: string[][];
// MetaInstances (declarative instance constraints)
metaInstances?: Record<string, MetaInstanceData>;
}
export interface MetaInstanceData {
metaInstanceId: string;
modelId: string;
sharding: string;
instanceMeta: string;
minNodes: number;
nodeIds: string[] | null;
placementError: string | null;
consecutiveFailures: number;
lastFailureError: string | null;
}
export interface MessageAttachment {
@@ -278,6 +309,7 @@ export interface Conversation {
modelId: string | null;
sharding: string | null;
instanceType: string | null;
enableThinking: boolean | null;
}
const STORAGE_KEY = "exo-conversations";
@@ -425,6 +457,7 @@ function transformTopology(
},
last_macmon_update: Date.now() / 1000,
friendly_name: identity?.friendlyName,
os_version: identity?.osVersion,
};
}
@@ -437,6 +470,8 @@ function transformTopology(
if (!Array.isArray(edgeList)) continue;
for (const edge of edgeList) {
let sendBackIp: string | undefined;
let sourceRdmaIface: string | undefined;
let sinkRdmaIface: string | undefined;
if (edge && typeof edge === "object" && "sinkMultiaddr" in edge) {
const multiaddr = edge.sinkMultiaddr;
if (multiaddr) {
@@ -444,10 +479,23 @@ function transformTopology(
multiaddr.ip_address ||
extractIpFromMultiaddr(multiaddr.address);
}
} else if (
edge &&
typeof edge === "object" &&
"sourceRdmaIface" in edge
) {
sourceRdmaIface = edge.sourceRdmaIface;
sinkRdmaIface = edge.sinkRdmaIface;
}
if (nodes[source] && nodes[sink] && source !== sink) {
edges.push({ source, target: sink, sendBackIp });
edges.push({
source,
target: sink,
sendBackIp,
sourceRdmaIface,
sinkRdmaIface,
});
}
}
}
@@ -490,12 +538,33 @@ class AppStore {
instances = $state<Record<string, unknown>>({});
runners = $state<Record<string, unknown>>({});
downloads = $state<Record<string, unknown[]>>({});
nodeDisk = $state<
Record<
string,
{ total: { inBytes: number }; available: { inBytes: number } }
>
>({});
placementPreviews = $state<PlacementPreview[]>([]);
selectedPreviewModelId = $state<string | null>(null);
isLoadingPreviews = $state(false);
previewNodeFilter = $state<Set<string>>(new Set());
lastUpdate = $state<number | null>(null);
metaInstances = $state<Record<string, MetaInstanceData>>({});
nodeIdentities = $state<Record<string, RawNodeIdentity>>({});
thunderboltBridgeCycles = $state<string[][]>([]);
nodeThunderbolt = $state<
Record<
string,
{
interfaces: Array<{
rdmaInterface: string;
domainUuid: string;
linkSpeed: string;
}>;
}
>
>({});
nodeRdmaCtl = $state<Record<string, { enabled: boolean }>>({});
nodeThunderboltBridge = $state<
Record<
string,
@@ -551,6 +620,7 @@ class AppStore {
modelId: conversation.modelId ?? null,
sharding: conversation.sharding ?? null,
instanceType: conversation.instanceType ?? null,
enableThinking: conversation.enableThinking ?? null,
}));
}
} catch (error) {
@@ -740,6 +810,7 @@ class AppStore {
modelId: derivedModelId,
sharding: derivedSharding,
instanceType: derivedInstanceType,
enableThinking: null,
};
this.conversations.unshift(conversation);
@@ -765,6 +836,7 @@ class AppStore {
this.hasStartedChat = true;
this.isTopologyMinimized = true;
this.isSidebarOpen = true; // Auto-open sidebar when chatting
this.thinkingEnabled = conversation.enableThinking ?? true;
this.refreshConversationModelFromInstances();
return true;
@@ -837,11 +909,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 {
@@ -1206,6 +1274,17 @@ class AppStore {
if (data.downloads) {
this.downloads = data.downloads;
}
// MetaInstances
this.metaInstances = data.metaInstances ?? {};
if (data.nodeDisk) {
this.nodeDisk = data.nodeDisk;
}
// Node identities (for OS version mismatch detection)
this.nodeIdentities = data.nodeIdentities ?? {};
// Thunderbolt identifiers per node
this.nodeThunderbolt = data.nodeThunderbolt ?? {};
// RDMA ctl status per node
this.nodeRdmaCtl = data.nodeRdmaCtl ?? {};
// Thunderbolt bridge cycles
this.thunderboltBridgeCycles = data.thunderboltBridgeCycles ?? [];
// Thunderbolt bridge status per node
@@ -1869,6 +1948,11 @@ class AppStore {
}
}
/**
* Whether thinking is enabled for the current conversation
*/
thinkingEnabled = $state(true);
/**
* Selected model for chat (can be set by the UI)
*/
@@ -2047,6 +2131,7 @@ class AppStore {
textContent?: string;
preview?: string;
}[],
enableThinking?: boolean | null,
): Promise<void> {
if ((!content.trim() && (!files || files.length === 0)) || this.isLoading)
return;
@@ -2194,6 +2279,9 @@ class AppStore {
stream: true,
logprobs: true,
top_logprobs: 5,
...(enableThinking != null && {
enable_thinking: enableThinking,
}),
}),
});
@@ -2852,6 +2940,18 @@ class AppStore {
);
}
/**
* Update the thinking preference for the active conversation
*/
setConversationThinking(enabled: boolean) {
this.thinkingEnabled = enabled;
const conv = this.getActiveConversation();
if (conv) {
conv.enableThinking = enabled;
this.saveConversationsToStorage();
}
}
/**
* Start a download on a specific node
*/
@@ -2956,14 +3056,17 @@ export const tps = () => appStore.tps;
export const totalTokens = () => appStore.totalTokens;
export const topologyData = () => appStore.topologyData;
export const instances = () => appStore.instances;
export const metaInstances = () => appStore.metaInstances;
export const runners = () => appStore.runners;
export const downloads = () => appStore.downloads;
export const nodeDisk = () => appStore.nodeDisk;
export const placementPreviews = () => appStore.placementPreviews;
export const selectedPreviewModelId = () => appStore.selectedPreviewModelId;
export const isLoadingPreviews = () => appStore.isLoadingPreviews;
export const lastUpdate = () => appStore.lastUpdate;
export const isTopologyMinimized = () => appStore.isTopologyMinimized;
export const selectedChatModel = () => appStore.selectedChatModel;
export const thinkingEnabled = () => appStore.thinkingEnabled;
export const debugMode = () => appStore.getDebugMode();
export const topologyOnlyMode = () => appStore.getTopologyOnlyMode();
export const chatSidebarVisible = () => appStore.getChatSidebarVisible();
@@ -2979,7 +3082,8 @@ export const sendMessage = (
textContent?: string;
preview?: string;
}[],
) => appStore.sendMessage(content, files);
enableThinking?: boolean | null,
) => appStore.sendMessage(content, files, enableThinking);
export const generateImage = (prompt: string, modelId?: string) =>
appStore.generateImage(prompt, modelId);
export const editImage = (
@@ -3022,6 +3126,8 @@ export const deleteAllConversations = () => appStore.deleteAllConversations();
export const renameConversation = (id: string, name: string) =>
appStore.renameConversation(id, name);
export const getActiveConversation = () => appStore.getActiveConversation();
export const setConversationThinking = (enabled: boolean) =>
appStore.setConversationThinking(enabled);
// Sidebar actions
export const isSidebarOpen = () => appStore.isSidebarOpen;
@@ -3038,7 +3144,12 @@ export const setChatSidebarVisible = (visible: boolean) =>
appStore.setChatSidebarVisible(visible);
export const refreshState = () => appStore.fetchState();
// Thunderbolt bridge status
// Node identities (for OS version mismatch detection)
export const nodeIdentities = () => appStore.nodeIdentities;
// Thunderbolt & RDMA status
export const nodeThunderbolt = () => appStore.nodeThunderbolt;
export const nodeRdmaCtl = () => appStore.nodeRdmaCtl;
export const thunderboltBridgeCycles = () => appStore.thunderboltBridgeCycles;
export const nodeThunderboltBridge = () => appStore.nodeThunderboltBridge;
@@ -0,0 +1,75 @@
/**
* RecentsStore - Manages recently launched models with localStorage persistence
*/
import { browser } from "$app/environment";
const RECENTS_KEY = "exo-recent-models";
const MAX_RECENT_MODELS = 20;
interface RecentEntry {
modelId: string;
launchedAt: number;
}
class RecentsStore {
recents = $state<RecentEntry[]>([]);
constructor() {
if (browser) {
this.loadFromStorage();
}
}
private loadFromStorage() {
try {
const stored = localStorage.getItem(RECENTS_KEY);
if (stored) {
const parsed = JSON.parse(stored) as RecentEntry[];
this.recents = parsed;
}
} catch (error) {
console.error("Failed to load recent models:", error);
}
}
private saveToStorage() {
try {
localStorage.setItem(RECENTS_KEY, JSON.stringify(this.recents));
} catch (error) {
console.error("Failed to save recent models:", error);
}
}
recordLaunch(modelId: string) {
// Remove existing entry for this model (if any) to move it to top
const filtered = this.recents.filter((r) => r.modelId !== modelId);
// Prepend new entry
const next = [{ modelId, launchedAt: Date.now() }, ...filtered];
// Cap at max
this.recents = next.slice(0, MAX_RECENT_MODELS);
this.saveToStorage();
}
getRecentModelIds(): string[] {
return this.recents.map((r) => r.modelId);
}
hasAny(): boolean {
return this.recents.length > 0;
}
clearAll() {
this.recents = [];
this.saveToStorage();
}
}
export const recentsStore = new RecentsStore();
export const hasRecents = () => recentsStore.hasAny();
export const getRecentModelIds = () => recentsStore.getRecentModelIds();
export const getRecentEntries = () => recentsStore.recents;
export const recordRecentLaunch = (modelId: string) =>
recentsStore.recordLaunch(modelId);
export const clearRecents = () => recentsStore.clearAll();
File diff suppressed because it is too large Load Diff
+19 -7
View File
@@ -3,6 +3,7 @@
import {
topologyData,
downloads,
nodeDisk,
type DownloadProgress,
refreshState,
lastUpdate as lastUpdateStore,
@@ -37,10 +38,13 @@
nodeId: string;
nodeName: string;
models: ModelEntry[];
diskAvailable?: number;
diskTotal?: number;
};
const data = $derived(topologyData());
const downloadsData = $derived(downloads());
const nodeDiskData = $derived(nodeDisk());
function getNodeLabel(nodeId: string): string {
const node = data?.nodes?.[nodeId];
@@ -327,10 +331,17 @@
];
}
// Get disk info for this node
const diskInfo = nodeDiskData?.[nodeId];
const diskAvailable = diskInfo?.available?.inBytes;
const diskTotal = diskInfo?.total?.inBytes;
built.push({
nodeId,
nodeName: getNodeLabel(nodeId),
models,
diskAvailable,
diskTotal,
});
}
@@ -417,6 +428,14 @@
<div class="text-xs text-exo-light-gray font-mono truncate">
{node.nodeId}
</div>
<div class="text-xs text-exo-light-gray font-mono mt-1">
{formatBytes(
node.models
.filter((m) => m.status === "completed")
.reduce((sum, m) => sum + m.totalBytes, 0),
)} models{#if node.diskAvailable != null}
- {formatBytes(node.diskAvailable)} free{/if}
</div>
</div>
<div
class="text-xs font-mono uppercase tracking-wider whitespace-nowrap shrink-0 text-right"
@@ -429,13 +448,6 @@
/ {node.models.length} models</span
>
</div>
<div class="text-exo-light-gray normal-case tracking-normal">
{formatBytes(
node.models
.filter((m) => m.status === "completed")
.reduce((sum, m) => sum + m.totalBytes, 0),
)} on disk
</div>
</div>
</div>
+3
View File
@@ -83,6 +83,9 @@
_module.args.pkgs = import inputs.nixpkgs {
inherit system;
config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
overlays = [
(import ./nix/apple-sdk-overlay.nix)
];
};
treefmt = {
projectRootFile = "flake.nix";
+18
View File
@@ -0,0 +1,18 @@
# Overlay that builds apple-sdk with a custom versions.json (for SDK 26.2).
# The upstream nixpkgs package reads versions.json at eval time via a relative
# path, so we can't override it through callPackage args. Instead, we copy
# the upstream source and patch the one file.
final: _prev:
let
upstreamSrc = final.path + "/pkgs/by-name/ap/apple-sdk";
patchedSrc = final.runCommandLocal "apple-sdk-src-patched" { } ''
cp -r ${upstreamSrc} $out
chmod -R u+w $out
cp ${./apple-sdk/metadata/versions.json} $out/metadata/versions.json
'';
in
{
apple-sdk_26 = final.callPackage (patchedSrc + "/package.nix") {
darwinSdkMajorVersion = "26";
};
}
+26
View File
@@ -0,0 +1,26 @@
{
"14": {
"urls": [
"https://swcdn.apple.com/content/downloads/14/48/052-59890-A_I0F5YGAY0Y/p9n40hio7892gou31o1v031ng6fnm9sb3c/CLTools_macOSNMOS_SDK.pkg",
"https://web.archive.org/web/20250211001355/https://swcdn.apple.com/content/downloads/14/48/052-59890-A_I0F5YGAY0Y/p9n40hio7892gou31o1v031ng6fnm9sb3c/CLTools_macOSNMOS_SDK.pkg"
],
"version": "14.4",
"hash": "sha256-QozDiwY0Czc0g45vPD7G4v4Ra+3DujCJbSads3fJjjM="
},
"15": {
"urls": [
"https://swcdn.apple.com/content/downloads/52/01/082-41241-A_0747ZN8FHV/dectd075r63pppkkzsb75qk61s0lfee22j/CLTools_macOSNMOS_SDK.pkg",
"https://web.archive.org/web/20250530132510/https://swcdn.apple.com/content/downloads/52/01/082-41241-A_0747ZN8FHV/dectd075r63pppkkzsb75qk61s0lfee22j/CLTools_macOSNMOS_SDK.pkg"
],
"version": "15.5",
"hash": "sha256-HBiSJuw1XBUK5R/8Sj65c3rftSEvQl/O9ZZVp/g1Amo="
},
"26": {
"urls": [
"https://swcdn.apple.com/content/downloads/60/22/089-71960-A_W8BL1RUJJ6/5zkyplomhk1cm7z6xja2ktgapnhhti6wwd/CLTools_macOSNMOS_SDK.pkg",
"https://web.archive.org/web/20250915230423/https://swcdn.apple.com/content/downloads/60/22/089-71960-A_W8BL1RUJJ6/5zkyplomhk1cm7z6xja2ktgapnhhti6wwd/CLTools_macOSNMOS_SDK.pkg"
],
"version": "26.2",
"hash": "sha256-hXRlMieVv0smna5uiWRwq87IWOaPWtAjAldbi+wQXcw="
}
}
+3 -2
View File
@@ -41,7 +41,7 @@ let
mlx = stdenv.mkDerivation rec {
pname = "mlx";
version = let v = "0.30.5"; in
version = let v = "0.30.6"; 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;
@@ -50,7 +50,7 @@ let
owner = "ml-explore";
repo = "mlx";
tag = "v${version}";
hash = "sha256-OJk6jPlbaSlsUdk3ADz3tWcRzTWXRof3/q8Soe1AO6w=";
hash = "sha256-avD5EGhwgmPdXLAyQSqTO6AXk/W3ziH+f6AetjK3Sdo=";
};
patches = [
@@ -85,6 +85,7 @@ let
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
(lib.cmakeBool "MLX_BUILD_CPU" true)
(lib.cmakeBool "MLX_BUILD_METAL" true)
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
(lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}")
+6 -3
View File
@@ -17,8 +17,8 @@ dependencies = [
"loguru>=0.7.3",
"exo_pyo3_bindings", # rust bindings
"anyio==4.11.0",
"mlx==0.30.5; sys_platform == 'darwin'",
"mlx[cpu]==0.30.5; sys_platform == 'linux'",
"mlx==0.30.6; sys_platform == 'darwin'",
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
"mlx-lm==0.30.6",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
@@ -28,6 +28,8 @@ dependencies = [
"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]
@@ -57,6 +59,7 @@ dev = [
[tool.uv.workspace]
members = [
"rust/exo_pyo3_bindings",
"bench",
]
[tool.uv.sources]
@@ -129,7 +132,7 @@ markers = [
env = [
"EXO_TESTS=1"
]
addopts = "-m 'not slow'"
addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
filterwarnings = [
"ignore:builtin type Swig:DeprecationWarning",
]
+54 -5
View File
@@ -14,7 +14,9 @@
# Override overlay to inject Nix-built components
exoOverlay = final: prev: {
# Replace workspace exo_pyo3_bindings with Nix-built wheel
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
# Copy .pyi stub + py.typed marker so basedpyright can find the types.
exo-pyo3-bindings = pkgs.stdenv.mkDerivation {
pname = "exo-pyo3-bindings";
version = "0.1.0";
@@ -22,6 +24,12 @@
# Install from pre-built wheel
nativeBuildInputs = [ final.pyprojectWheelHook ];
dontStrip = true;
passthru = prev.exo-pyo3-bindings.passthru or { };
postInstall = ''
local siteDir=$out/${final.python.sitePackages}/exo_pyo3_bindings
cp ${inputs.self}/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi $siteDir/
touch $siteDir/py.typed
'';
};
};
@@ -29,17 +37,32 @@
# Overlay to provide build systems and custom packages
buildSystemsOverlay = final: prev: {
# Use our pure Nix-built MLX with Metal support
mlx = self'.packages.mlx;
# mlx-lm is a git dependency that needs setuptools
mlx-lm = prev.mlx-lm.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
} // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
# Use our pure Nix-built MLX with Metal support (macOS only)
mlx = self'.packages.mlx;
};
# Additional overlay for Linux-specific fixes (type checking env).
# Native wheels have shared lib dependencies we don't need at type-check time.
linuxOverlay = final: prev:
let
ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
in
lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
(lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
mlx = ignoreMissing prev.mlx;
torch = ignoreMissing prev.torch;
triton = ignoreMissing prev.triton;
}
);
pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
inherit python;
}).overrideScope (
@@ -48,6 +71,7 @@
overlay
exoOverlay
buildSystemsOverlay
linuxOverlay
]
);
exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default;
@@ -69,6 +93,16 @@
text = ''exec python ${path} "$@"'';
};
benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" {
exo-bench = [ ];
};
mkBenchScript = name: path: pkgs.writeShellApplication {
inherit name;
runtimeInputs = [ benchVenv ];
text = ''exec python ${path} "$@"'';
};
mkSimplePythonScript = name: path: pkgs.writeShellApplication {
inherit name;
runtimeInputs = [ pkgs.python313 ];
@@ -96,8 +130,8 @@
exo = exoPackage;
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
exo-test-env = testVenv;
exo-bench = mkPythonScript "exo-bench" (inputs.self + /bench/exo_bench.py);
} // {
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
};
@@ -108,6 +142,21 @@
${pkgs.ruff}/bin/ruff check ${inputs.self}
touch $out
'';
# Hermetic basedpyright type checking
typecheck = pkgs.runCommand "typecheck"
{
nativeBuildInputs = [
testVenv
pkgs.basedpyright
];
}
''
cd ${inputs.self}
export HOME=$TMPDIR
basedpyright --pythonpath ${testVenv}/bin/python
touch $out
'';
};
};
}
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["ImageToImage"]
family = "flux"
quantization = "4bit"
base_model = "FLUX.1 Kontext"
capabilities = ["image_edit"]
[storage_size]
in_bytes = 15475325472
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["ImageToImage"]
family = "flux"
quantization = "8bit"
base_model = "FLUX.1 Kontext"
capabilities = ["image_edit"]
[storage_size]
in_bytes = 21426029632
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["ImageToImage"]
family = "flux"
quantization = ""
base_model = "FLUX.1 Kontext"
capabilities = ["image_edit"]
[storage_size]
in_bytes = 33327437952
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = "4bit"
base_model = "FLUX.1 Krea"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 15475325472
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = "8bit"
base_model = "FLUX.1 Krea"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 21426029632
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = ""
base_model = "FLUX.1 Krea"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 33327437952
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = "4bit"
base_model = "FLUX.1 Dev"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 15475325472
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = "8bit"
base_model = "FLUX.1 Dev"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 21426029632
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = ""
base_model = "FLUX.1 Dev"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 33327437952
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = "4bit"
base_model = "FLUX.1 Schnell"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 15470210592
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = "8bit"
base_model = "FLUX.1 Schnell"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 21415799872
@@ -3,6 +3,10 @@ n_layers = 57
hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
family = "flux"
quantization = ""
base_model = "FLUX.1 Schnell"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 33306978432
@@ -4,6 +4,10 @@ hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
uses_cfg = true
family = "qwen-image"
quantization = "4bit"
base_model = "Qwen Image"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 26799533856
@@ -4,6 +4,10 @@ hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
uses_cfg = true
family = "qwen-image"
quantization = "8bit"
base_model = "Qwen Image"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 37014734400
@@ -4,6 +4,10 @@ hidden_size = 1
supports_tensor = false
tasks = ["ImageToImage"]
uses_cfg = true
family = "qwen-image"
quantization = "4bit"
base_model = "Qwen Image Edit"
capabilities = ["image_edit"]
[storage_size]
in_bytes = 26799533856
@@ -4,6 +4,10 @@ hidden_size = 1
supports_tensor = false
tasks = ["ImageToImage"]
uses_cfg = true
family = "qwen-image"
quantization = "8bit"
base_model = "Qwen Image Edit"
capabilities = ["image_edit"]
[storage_size]
in_bytes = 37014734400
@@ -4,6 +4,10 @@ hidden_size = 1
supports_tensor = false
tasks = ["ImageToImage"]
uses_cfg = true
family = "qwen-image"
quantization = ""
base_model = "Qwen Image Edit"
capabilities = ["image_edit"]
[storage_size]
in_bytes = 57445135488
@@ -4,6 +4,10 @@ hidden_size = 1
supports_tensor = false
tasks = ["TextToImage"]
uses_cfg = true
family = "qwen-image"
quantization = ""
base_model = "Qwen Image"
capabilities = ["image_gen"]
[storage_size]
in_bytes = 57445135488
@@ -0,0 +1,12 @@
model_id = "mlx-community/Step-3.5-Flash-4bit"
n_layers = 45
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "step"
quantization = "4bit"
base_model = "Step 3.5 Flash"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 114572190076
@@ -0,0 +1,12 @@
model_id = "mlx-community/Step-3.5-Flash-6bit"
n_layers = 45
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "step"
quantization = "6bit"
base_model = "Step 3.5 Flash"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 159039627774
@@ -0,0 +1,12 @@
model_id = "mlx-community/Step-3.5-Flash-8Bit"
n_layers = 45
hidden_size = 4096
supports_tensor = true
tasks = ["TextGeneration"]
family = "step"
quantization = "8bit"
base_model = "Step 3.5 Flash"
capabilities = ["text", "thinking"]
[storage_size]
in_bytes = 209082699847
+1 -1
View File
@@ -130,7 +130,7 @@ class Multiaddr:
@typing.final
class NetworkingHandle:
def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
def __new__(cls, identity: Keypair, bootstrap_peers: builtins.list[builtins.str] = ...) -> NetworkingHandle: ...
async def connection_update_recv(self) -> ConnectionUpdate:
r"""
Receives the next `ConnectionUpdate` from networking.
+18 -4
View File
@@ -22,7 +22,6 @@ use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymetho
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};
use util::ext::VecExt as _;
mod exception {
use pyo3::types::PyTuple;
@@ -343,8 +342,10 @@ impl PyNetworkingHandle {
// ---- Lifecycle management methods ----
#[new]
fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
#[pyo3(signature = (identity, bootstrap_peers=vec![]))]
fn py_new(identity: Bound<'_, PyKeypair>, bootstrap_peers: Vec<String>) -> PyResult<Self> {
use pyo3_async_runtimes::tokio::get_runtime;
use std::str::FromStr;
// create communication channels
let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
@@ -354,9 +355,20 @@ impl PyNetworkingHandle {
// get identity
let identity = identity.borrow().0.clone();
// parse bootstrap peer multiaddrs
let parsed_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
.into_iter()
.map(|s| libp2p::Multiaddr::from_str(&s))
.collect::<Result<_, _>>()
.pyerr()?;
if !parsed_peers.is_empty() {
log::info!("RUST: bootstrap peers: {:?}", parsed_peers);
}
// create networking swarm (within tokio context!! or it crashes)
let swarm = get_runtime()
.block_on(async { create_swarm(identity) })
.block_on(async { create_swarm(identity, parsed_peers) })
.pyerr()?;
// spawn tokio task running the networking logic
@@ -532,7 +544,9 @@ impl PyNetworkingHandle {
.recv_many_py(limit)
.allow_threads_py() // allow-threads-aware async call
.await?
.map(|(t, d)| (t, d.pybytes())))
.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)
+26 -4
View File
@@ -106,6 +106,9 @@ pub struct Behaviour {
managed: managed::Behaviour,
mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
/// Addresses provided via --bootstrap-peer, re-dialed on the retry loop.
bootstrap_addrs: Vec<Multiaddr>,
retry_delay: Delay, // retry interval
// pending events to emmit => waker-backed Deque to control polling
@@ -113,13 +116,21 @@ pub struct Behaviour {
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
Ok(Self {
pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
let mut behaviour = Self {
managed: managed::Behaviour::new(keypair)?,
mdns_discovered: HashMap::new(),
bootstrap_addrs: bootstrap_peers,
retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
pending_events: WakerDeque::new(),
})
};
// Immediately dial all bootstrap peers
for addr in &behaviour.bootstrap_addrs.clone() {
behaviour.dial_addr(addr.clone());
}
Ok(behaviour)
}
fn dial(&mut self, peer_id: PeerId, addr: Multiaddr) {
@@ -128,6 +139,14 @@ impl Behaviour {
})
}
/// Dial by address only — PeerId is resolved via Noise handshake on connect.
/// Used for bootstrap peers where only IP:port is known.
fn dial_addr(&mut self, addr: Multiaddr) {
self.pending_events.push_back(ToSwarm::Dial {
opts: DialOpts::unknown_peer_id().address(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 {
@@ -362,13 +381,16 @@ impl NetworkBehaviour for Behaviour {
Poll::Pending => {}
}
// retry connecting to all mDNS peers periodically (fails safely if already connected)
// retry connecting to all mDNS + bootstrap peers periodically (fails safely if already connected)
if self.retry_delay.poll_unpin(cx).is_ready() {
for (p, mas) in self.mdns_discovered.clone() {
for ma in mas {
self.dial(p, ma)
}
}
for addr in self.bootstrap_addrs.clone() {
self.dial_addr(addr)
}
self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
}
+12 -8
View File
@@ -1,7 +1,7 @@
use crate::alias;
use crate::swarm::transport::tcp_transport;
pub use behaviour::{Behaviour, BehaviourEvent};
use libp2p::{SwarmBuilder, identity};
use libp2p::{Multiaddr, SwarmBuilder, identity};
pub type Swarm = libp2p::Swarm<Behaviour>;
@@ -16,11 +16,14 @@ 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> {
pub fn create_swarm(
keypair: identity::Keypair,
bootstrap_peers: Vec<Multiaddr>,
) -> alias::AnyResult<Swarm> {
let mut swarm = SwarmBuilder::with_existing_identity(keypair)
.with_tokio()
.with_other_transport(tcp_transport)?
.with_behaviour(Behaviour::new)?
.with_behaviour(|kp| Behaviour::new(kp, bootstrap_peers))?
.build();
// Listen on all interfaces and whatever port the OS assigns
@@ -105,8 +108,7 @@ mod transport {
mod behaviour {
use crate::{alias, discovery};
use libp2p::swarm::NetworkBehaviour;
use libp2p::{gossipsub, identity};
use std::time::Duration;
use libp2p::{Multiaddr, gossipsub, identity};
/// Behavior of the Swarm which composes all desired behaviors:
/// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
@@ -117,9 +119,12 @@ mod behaviour {
}
impl Behaviour {
pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
pub fn new(
keypair: &identity::Keypair,
bootstrap_peers: Vec<Multiaddr>,
) -> alias::AnyResult<Self> {
Ok(Self {
discovery: discovery::Behaviour::new(keypair)?,
discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
gossipsub: gossipsub_behaviour(keypair),
})
}
@@ -134,7 +139,6 @@ mod behaviour {
gossipsub::Behaviour::new(
MessageAuthenticity::Signed(keypair.clone()),
ConfigBuilder::default()
.publish_queue_duration(Duration::from_secs(15))
.max_transmit_size(1024 * 1024)
.validation_mode(ValidationMode::Strict)
.build()
-10
View File
@@ -13,13 +13,3 @@ path = "src/lib.rs"
workspace = true
[dependencies]
# macro dependencies
extend = { workspace = true }
# utility dependencies
thiserror = { workspace = true }
once_cell = { workspace = true }
internment = { workspace = true }
derive_more = { workspace = true }
bon = { workspace = true }
recursion = { workspace = true }
-52
View File
@@ -1,53 +1 @@
//! TODO: crate documentation
//!
//! this is here as a placeholder documentation
//!
//!
// enable Rust-unstable features for convenience
#![feature(trait_alias)]
#![feature(stmt_expr_attributes)]
#![feature(type_alias_impl_trait)]
#![feature(specialization)]
#![feature(unboxed_closures)]
#![feature(const_trait_impl)]
#![feature(fn_traits)]
pub mod nonempty;
pub mod wakerdeque;
pub(crate) mod private {
// sealed traits support
pub trait Sealed {}
impl<T: ?Sized> Sealed for T {}
}
/// Namespace for all the type/trait aliases used by this crate.
pub(crate) mod alias {}
/// Namespace for crate-wide extension traits/methods
pub mod ext {
use extend::ext;
#[ext(pub, name = BoxedSliceExt)]
impl<T> Box<[T]> {
#[inline]
fn map<B, F>(self, f: F) -> Box<[B]>
where
F: FnMut(T) -> B,
{
self.into_iter().map(f).collect()
}
}
#[ext(pub, name = VecExt)]
impl<T> Vec<T> {
#[inline]
fn map<B, F>(self, f: F) -> Vec<B>
where
F: FnMut(T) -> B,
{
self.into_iter().map(f).collect()
}
}
}
-138
View File
@@ -1,138 +0,0 @@
use std::slice::SliceIndex;
use std::{ops, slice};
use thiserror::Error;
#[derive(Error, Debug)]
#[error("Cannot create to `NonemptyArray` because the supplied slice is empty")]
pub struct EmptySliceError;
/// A pointer to a non-empty fixed-size slice allocated on the heap.
#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
#[repr(transparent)]
pub struct NonemptyArray<T>(Box<[T]>);
#[allow(clippy::arbitrary_source_item_ordering)]
impl<T> NonemptyArray<T> {
#[inline]
pub fn singleton(value: T) -> Self {
Self(Box::new([value]))
}
#[allow(clippy::missing_errors_doc)]
#[inline]
pub fn try_from_boxed_slice<S: Into<Box<[T]>>>(
boxed_slice: S,
) -> Result<Self, EmptySliceError> {
let boxed_slice = boxed_slice.into();
if boxed_slice.is_empty() {
Err(EmptySliceError)
} else {
Ok(Self(boxed_slice))
}
}
#[must_use]
#[inline]
pub fn into_boxed_slice(self) -> Box<[T]> {
self.0
}
#[must_use]
#[inline]
pub fn to_vec(&self) -> Vec<T>
where
T: Clone,
{
self.0.to_vec()
}
#[must_use]
#[inline]
pub const fn as_slice(&self) -> &[T] {
&self.0
}
#[allow(clippy::indexing_slicing)]
#[must_use]
#[inline]
pub fn first(&self) -> &T {
&self.0[0]
}
#[allow(clippy::indexing_slicing, clippy::arithmetic_side_effects)]
#[must_use]
#[inline]
pub fn last(&self) -> &T {
&self.0[self.0.len() - 1]
}
#[must_use]
#[inline]
pub fn get<I>(&self, index: I) -> Option<&I::Output>
where
I: SliceIndex<[T]>,
{
self.0.get(index)
}
#[allow(clippy::len_without_is_empty)]
#[must_use]
#[inline]
pub const fn len(&self) -> usize {
self.0.len()
}
#[allow(clippy::iter_without_into_iter)]
#[inline]
pub fn iter(&self) -> slice::Iter<'_, T> {
self.0.iter()
}
#[allow(clippy::iter_without_into_iter)]
#[inline]
pub fn iter_mut(&mut self) -> slice::IterMut<'_, T> {
self.0.iter_mut()
}
#[inline]
#[must_use]
pub fn map<U, F: FnMut(T) -> U>(self, f: F) -> NonemptyArray<U> {
NonemptyArray(self.0.into_iter().map(f).collect())
}
}
impl<T> From<NonemptyArray<T>> for Box<[T]> {
#[inline]
fn from(value: NonemptyArray<T>) -> Self {
value.into_boxed_slice()
}
}
impl<T> ops::Index<usize> for NonemptyArray<T> {
type Output = T;
#[inline]
fn index(&self, index: usize) -> &Self::Output {
self.0.index(index)
}
}
impl<T> IntoIterator for NonemptyArray<T> {
type Item = T;
type IntoIter = std::vec::IntoIter<T>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.into_boxed_slice().into_vec().into_iter()
}
}
impl<'a, T> IntoIterator for &'a NonemptyArray<T> {
type Item = &'a T;
type IntoIter = slice::Iter<'a, T>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
+61 -44
View File
@@ -56,8 +56,49 @@ class DownloadCoordinator:
event_receiver: Receiver[Event] = field(init=False)
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
# 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]()
self.shard_downloader.on_progress(self._download_progress_callback)
async def _download_progress_callback(
self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
) -> None:
model_id = callback_shard.model_card.model_id
throttle_interval_secs = 1.0
if progress.status == "complete":
completed = DownloadCompleted(
shard_metadata=callback_shard,
node_id=self.node_id,
total_bytes=progress.total_bytes,
)
self.download_status[model_id] = completed
await self.event_sender.send(
NodeDownloadProgress(download_progress=completed)
)
if model_id in self.active_downloads:
del self.active_downloads[model_id]
self._last_progress_time.pop(model_id, None)
elif (
progress.status == "in_progress"
and current_time() - self._last_progress_time.get(model_id, 0.0)
> throttle_interval_secs
):
ongoing = DownloadOngoing(
node_id=self.node_id,
shard_metadata=callback_shard,
download_progress=map_repo_download_progress_to_download_progress_data(
progress
),
)
self.download_status[model_id] = ongoing
await self.event_sender.send(
NodeDownloadProgress(download_progress=ongoing)
)
self._last_progress_time[model_id] = current_time()
async def run(self) -> None:
logger.info("Starting DownloadCoordinator")
@@ -119,12 +160,12 @@ class DownloadCoordinator:
async def _start_download(self, shard: ShardMetadata) -> None:
model_id = shard.model_card.model_id
# Check if already downloading or complete
# Check if already downloading, complete, or recently failed
if model_id in self.download_status:
status = self.download_status[model_id]
if isinstance(status, (DownloadOngoing, DownloadCompleted)):
if isinstance(status, (DownloadOngoing, DownloadCompleted, DownloadFailed)):
logger.debug(
f"Download for {model_id} already in progress or complete, skipping"
f"Download for {model_id} already in progress, complete, or failed, skipping"
)
return
@@ -169,46 +210,6 @@ class DownloadCoordinator:
self.download_status[model_id] = status
self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
last_progress_time = 0.0
throttle_interval_secs = 1.0
async def download_progress_callback(
callback_shard: ShardMetadata, progress: RepoDownloadProgress
) -> None:
nonlocal last_progress_time
if progress.status == "complete":
completed = DownloadCompleted(
shard_metadata=callback_shard,
node_id=self.node_id,
total_bytes=progress.total_bytes,
)
self.download_status[callback_shard.model_card.model_id] = completed
await self.event_sender.send(
NodeDownloadProgress(download_progress=completed)
)
# Clean up active download tracking
if callback_shard.model_card.model_id in self.active_downloads:
del self.active_downloads[callback_shard.model_card.model_id]
elif (
progress.status == "in_progress"
and current_time() - last_progress_time > throttle_interval_secs
):
ongoing = DownloadOngoing(
node_id=self.node_id,
shard_metadata=callback_shard,
download_progress=map_repo_download_progress_to_download_progress_data(
progress
),
)
self.download_status[callback_shard.model_card.model_id] = ongoing
await self.event_sender.send(
NodeDownloadProgress(download_progress=ongoing)
)
last_progress_time = current_time()
self.shard_downloader.on_progress(download_progress_callback)
async def download_wrapper() -> None:
try:
await self.shard_downloader.ensure_shard(shard)
@@ -283,6 +284,12 @@ class DownloadCoordinator:
_,
progress,
) in self.shard_downloader.get_shard_download_status():
model_id = progress.shard.model_card.model_id
# Active downloads emit progress via the callback — don't overwrite
if model_id in self.active_downloads:
continue
if progress.status == "complete":
status: DownloadProgress = DownloadCompleted(
node_id=self.node_id,
@@ -290,7 +297,17 @@ class DownloadCoordinator:
total_bytes=progress.total_bytes,
)
elif progress.status in ["in_progress", "not_started"]:
if progress.downloaded_bytes_this_session.in_bytes == 0:
if (
progress.downloaded_bytes.in_bytes
>= progress.total_bytes.in_bytes
> 0
):
status = DownloadCompleted(
node_id=self.node_id,
shard_metadata=progress.shard,
total_bytes=progress.total_bytes,
)
elif progress.downloaded_bytes_this_session.in_bytes == 0:
status = DownloadPending(
node_id=self.node_id, shard_metadata=progress.shard
)
+2
View File
@@ -457,6 +457,8 @@ async def download_file_with_retry(
)
except HuggingFaceAuthenticationError:
raise
except FileNotFoundError:
raise
except HuggingFaceRateLimitError as e:
if attempt == n_attempts - 1:
raise e
+17 -3
View File
@@ -46,7 +46,12 @@ class Node:
keypair = get_node_id_keypair()
node_id = NodeId(keypair.to_peer_id().to_base58())
session_id = SessionId(master_node_id=node_id, election_clock=0)
router = Router.create(keypair)
bootstrap_peers: list[str] = []
if args.bootstrap_peer:
ip, port = args.bootstrap_peer.rsplit(":", 1)
bootstrap_peers.append(f"/ip4/{ip}/tcp/{port}")
logger.info(f"Bootstrap peer: {args.bootstrap_peer}")
router = Router.create(keypair, bootstrap_peers=bootstrap_peers)
await router.register_topic(topics.GLOBAL_EVENTS)
await router.register_topic(topics.LOCAL_EVENTS)
await router.register_topic(topics.COMMANDS)
@@ -251,9 +256,10 @@ class Node:
def main():
args = Args.parse()
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
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")
@@ -282,6 +288,7 @@ class Args(CamelCaseModel):
no_worker: bool = False
no_downloads: bool = False
fast_synch: bool | None = None # None = auto, True = force on, False = force off
bootstrap_peer: str | None = None
@classmethod
def parse(cls) -> Self:
@@ -342,6 +349,13 @@ class Args(CamelCaseModel):
dest="fast_synch",
help="Force MLX FAST_SYNCH off",
)
parser.add_argument(
"--bootstrap-peer",
type=str,
dest="bootstrap_peer",
default=None,
help="IP:PORT of an existing node to connect to (bypasses mDNS)",
)
args = parser.parse_args()
return cls(**vars(args)) # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically
+21 -6
View File
@@ -3,7 +3,6 @@
import time
from collections.abc import AsyncGenerator
from typing import Any
from uuid import uuid4
from exo.shared.types.api import (
ChatCompletionChoice,
@@ -18,6 +17,7 @@ from exo.shared.types.api import (
LogprobsContentItem,
StreamingChoiceResponse,
ToolCall,
Usage,
)
from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
from exo.shared.types.common import CommandId
@@ -80,6 +80,7 @@ def chat_request_to_text_generation(
seed=request.seed,
stream=request.stream,
tools=request.tools,
enable_thinking=request.enable_thinking,
chat_template_messages=chat_template_messages
if chat_template_messages
else None,
@@ -125,6 +126,8 @@ async def generate_chat_stream(
chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
) -> AsyncGenerator[str, None]:
"""Generate Chat Completions API streaming events from chunks."""
last_usage: Usage | None = None
async for chunk in chunk_stream:
if isinstance(chunk, ErrorChunk):
error_response = ErrorResponse(
@@ -138,10 +141,12 @@ async def generate_chat_stream(
yield "data: [DONE]\n\n"
return
last_usage = chunk.usage or last_usage
if isinstance(chunk, ToolCallChunk):
tool_call_deltas = [
ToolCall(
id=str(uuid4()),
id=tool.id,
index=i,
function=tool,
)
@@ -161,12 +166,15 @@ async def generate_chat_stream(
finish_reason="tool_calls",
)
],
usage=last_usage,
)
yield f"data: {tool_response.model_dump_json()}\n\n"
yield "data: [DONE]\n\n"
return
chunk_response = chunk_to_response(chunk, command_id)
if chunk.finish_reason is not None:
chunk_response = chunk_response.model_copy(update={"usage": last_usage})
yield f"data: {chunk_response.model_dump_json()}\n\n"
if chunk.finish_reason is not None:
@@ -176,7 +184,9 @@ async def generate_chat_stream(
async def collect_chat_response(
command_id: CommandId,
chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
) -> ChatCompletionResponse:
) -> AsyncGenerator[str]:
# This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
# 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] = []
tool_calls: list[ToolCall] = []
@@ -184,6 +194,7 @@ async def collect_chat_response(
model: str | None = None
finish_reason: FinishReason | None = None
error_message: str | None = None
last_usage: Usage | None = None
async for chunk in chunk_stream:
if isinstance(chunk, ErrorChunk):
@@ -193,6 +204,8 @@ async def collect_chat_response(
if model is None:
model = chunk.model
last_usage = chunk.usage or last_usage
if isinstance(chunk, TokenChunk):
text_parts.append(chunk.text)
if chunk.logprob is not None:
@@ -207,7 +220,7 @@ async def collect_chat_response(
if isinstance(chunk, ToolCallChunk):
tool_calls.extend(
ToolCall(
id=str(uuid4()),
id=tool.id,
index=i,
function=tool,
)
@@ -223,7 +236,7 @@ async def collect_chat_response(
combined_text = "".join(text_parts)
assert model is not None
return ChatCompletionResponse(
yield ChatCompletionResponse(
id=command_id,
created=int(time.time()),
model=model,
@@ -241,4 +254,6 @@ async def collect_chat_response(
finish_reason=finish_reason,
)
],
)
usage=last_usage,
).model_dump_json()
return
+21 -19
View File
@@ -3,9 +3,8 @@
import json
from collections.abc import AsyncGenerator
from typing import Any
from uuid import uuid4
from exo.shared.types.api import FinishReason
from exo.shared.types.api import FinishReason, Usage
from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
from exo.shared.types.claude_api import (
ClaudeContentBlock,
@@ -162,12 +161,14 @@ async def collect_claude_response(
command_id: CommandId,
model: str,
chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
) -> ClaudeMessagesResponse:
) -> AsyncGenerator[str]:
# This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
# 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] = []
tool_use_blocks: list[ClaudeToolUseBlock] = []
stop_reason: ClaudeStopReason | None = None
last_stats = None
last_usage: Usage | None = None
error_message: str | None = None
async for chunk in chunk_stream:
@@ -175,21 +176,21 @@ async def collect_claude_response(
error_message = chunk.error_message or "Internal server error"
break
last_usage = chunk.usage or last_usage
if isinstance(chunk, ToolCallChunk):
for tool in chunk.tool_calls:
tool_use_blocks.append(
ClaudeToolUseBlock(
id=f"toolu_{uuid4().hex[:24]}",
id=f"toolu_{tool.id}",
name=tool.name,
input=json.loads(tool.arguments), # pyright: ignore[reportAny]
)
)
last_stats = chunk.stats or last_stats
stop_reason = "tool_use"
continue
text_parts.append(chunk.text)
last_stats = chunk.stats or last_stats
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
@@ -209,11 +210,11 @@ async def collect_claude_response(
if not content:
content.append(ClaudeTextBlock(text=""))
# Use actual usage data from stats if available
input_tokens = last_stats.prompt_tokens if last_stats else 0
output_tokens = last_stats.generation_tokens if last_stats else 0
# Use actual usage data if available
input_tokens = last_usage.prompt_tokens if last_usage else 0
output_tokens = last_usage.completion_tokens if last_usage else 0
return ClaudeMessagesResponse(
yield ClaudeMessagesResponse(
id=f"msg_{command_id}",
model=model,
content=content,
@@ -222,7 +223,8 @@ async def collect_claude_response(
input_tokens=input_tokens,
output_tokens=output_tokens,
),
)
).model_dump_json()
return
async def generate_claude_stream(
@@ -250,7 +252,7 @@ async def generate_claude_stream(
output_tokens = 0
stop_reason: ClaudeStopReason | None = None
last_stats = None
last_usage: Usage | None = None
next_block_index = 1 # text block is 0, tool blocks start at 1
async for chunk in chunk_stream:
@@ -258,13 +260,14 @@ async def generate_claude_stream(
# Close text block and bail
break
last_usage = chunk.usage or last_usage
if isinstance(chunk, ToolCallChunk):
last_stats = chunk.stats or last_stats
stop_reason = "tool_use"
# Emit tool_use content blocks
for tool in chunk.tool_calls:
tool_id = f"toolu_{uuid4().hex[:24]}"
tool_id = f"toolu_{tool.id}"
tool_input_json = tool.arguments
# content_block_start for tool_use
@@ -291,7 +294,6 @@ async def generate_claude_stream(
continue
output_tokens += 1 # Count each chunk as one token
last_stats = chunk.stats or last_stats
# content_block_delta
delta_event = ClaudeContentBlockDeltaEvent(
@@ -303,9 +305,9 @@ async def generate_claude_stream(
if chunk.finish_reason is not None:
stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason)
# Use actual token count from stats if available
if last_stats is not None:
output_tokens = last_stats.generation_tokens
# Use actual token count from usage if available
if last_usage is not None:
output_tokens = last_usage.completion_tokens
# content_block_stop for text block
block_stop = ClaudeContentBlockStopEvent(index=0)
+27 -24
View File
@@ -3,8 +3,8 @@
from collections.abc import AsyncGenerator
from itertools import count
from typing import Any
from uuid import uuid4
from exo.shared.types.api import Usage
from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
from exo.shared.types.common import CommandId
from exo.shared.types.openai_responses import (
@@ -122,13 +122,15 @@ async def collect_responses_response(
command_id: CommandId,
model: str,
chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
) -> ResponsesResponse:
) -> AsyncGenerator[str]:
# This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because
# FastAPI handles the cancellation better but wouldn't auto-serialize for some reason
"""Collect all token chunks and return a single ResponsesResponse."""
response_id = f"resp_{command_id}"
item_id = f"item_{command_id}"
accumulated_text = ""
function_call_items: list[ResponseFunctionCallItem] = []
last_stats = None
last_usage: Usage | None = None
error_message: str | None = None
async for chunk in chunk_stream:
@@ -136,32 +138,32 @@ async def collect_responses_response(
error_message = chunk.error_message or "Internal server error"
break
last_usage = chunk.usage or last_usage
if isinstance(chunk, ToolCallChunk):
for tool in chunk.tool_calls:
function_call_items.append(
ResponseFunctionCallItem(
id=f"fc_{uuid4().hex[:24]}",
call_id=f"call_{uuid4().hex[:24]}",
id=f"fc_{tool.id}",
call_id=f"call_{tool.id}",
name=tool.name,
arguments=tool.arguments,
)
)
last_stats = chunk.stats or last_stats
continue
accumulated_text += chunk.text
last_stats = chunk.stats or last_stats
if error_message is not None:
raise ValueError(error_message)
# Create usage from stats if available
# Create usage from usage data if available
usage = None
if last_stats is not None:
if last_usage is not None:
usage = ResponseUsage(
input_tokens=last_stats.prompt_tokens,
output_tokens=last_stats.generation_tokens,
total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
input_tokens=last_usage.prompt_tokens,
output_tokens=last_usage.completion_tokens,
total_tokens=last_usage.total_tokens,
)
output: list[ResponseItem] = [
@@ -173,14 +175,15 @@ async def collect_responses_response(
]
output.extend(function_call_items)
return ResponsesResponse(
yield ResponsesResponse(
id=response_id,
model=model,
status="completed",
output=output,
output_text=accumulated_text,
usage=usage,
)
).model_dump_json()
return
async def generate_responses_stream(
@@ -236,18 +239,19 @@ async def generate_responses_stream(
accumulated_text = ""
function_call_items: list[ResponseFunctionCallItem] = []
last_stats = None
last_usage: Usage | None = None
next_output_index = 1 # message item is at 0
async for chunk in chunk_stream:
if isinstance(chunk, ErrorChunk):
break
last_usage = chunk.usage or last_usage
if isinstance(chunk, ToolCallChunk):
last_stats = chunk.stats or last_stats
for tool in chunk.tool_calls:
fc_id = f"fc_{uuid4().hex[:24]}"
call_id = f"call_{uuid4().hex[:24]}"
fc_id = f"fc_{tool.id}"
call_id = f"call_{tool.id}"
# response.output_item.added for function_call
fc_item = ResponseFunctionCallItem(
@@ -303,7 +307,6 @@ async def generate_responses_stream(
continue
accumulated_text += chunk.text
last_stats = chunk.stats or last_stats
# response.output_text.delta
delta_event = ResponseTextDeltaEvent(
@@ -347,13 +350,13 @@ async def generate_responses_stream(
)
yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n"
# Create usage from stats if available
# Create usage from usage data if available
usage = None
if last_stats is not None:
if last_usage is not None:
usage = ResponseUsage(
input_tokens=last_stats.prompt_tokens,
output_tokens=last_stats.generation_tokens,
total_tokens=last_stats.prompt_tokens + last_stats.generation_tokens,
input_tokens=last_usage.prompt_tokens,
output_tokens=last_usage.completion_tokens,
total_tokens=last_usage.total_tokens,
)
# response.completed
+140 -32
View File
@@ -3,7 +3,7 @@ import contextlib
import json
import random
import time
from collections.abc import AsyncGenerator, Awaitable, Callable
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator
from datetime import datetime, timezone
from http import HTTPStatus
from pathlib import Path
@@ -37,11 +37,13 @@ from exo.master.adapters.responses import (
generate_responses_stream,
responses_request_to_text_generation,
)
from exo.master.event_log import DiskEventLog
from exo.master.image_store import ImageStore
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_EVENT_LOG_DIR,
EXO_IMAGE_CACHE_DIR,
EXO_MAX_CHUNK_SIZE,
EXO_TRACING_CACHE_DIR,
@@ -69,8 +71,11 @@ from exo.shared.types.api import (
ChatCompletionResponse,
CreateInstanceParams,
CreateInstanceResponse,
CreateMetaInstanceParams,
CreateMetaInstanceResponse,
DeleteDownloadResponse,
DeleteInstanceResponse,
DeleteMetaInstanceResponse,
ErrorInfo,
ErrorResponse,
FinishReason,
@@ -113,8 +118,10 @@ from exo.shared.types.claude_api import (
from exo.shared.types.commands import (
Command,
CreateInstance,
CreateMetaInstance,
DeleteDownload,
DeleteInstance,
DeleteMetaInstance,
DownloadCommand,
ForwarderCommand,
ForwarderDownloadCommand,
@@ -123,10 +130,11 @@ from exo.shared.types.commands import (
PlaceInstance,
SendInputChunk,
StartDownload,
TaskCancelled,
TaskFinished,
TextGeneration,
)
from exo.shared.types.common import CommandId, Id, NodeId, SessionId
from exo.shared.types.common import CommandId, Id, MetaInstanceId, NodeId, SessionId
from exo.shared.types.events import (
ChunkGenerated,
Event,
@@ -135,6 +143,7 @@ from exo.shared.types.events import (
TracesMerged,
)
from exo.shared.types.memory import Memory
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.openai_responses import (
ResponsesRequest,
ResponsesResponse,
@@ -146,6 +155,8 @@ from exo.utils.banner import print_startup_banner
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
return f"image/{image_format or 'png'}"
@@ -175,7 +186,7 @@ class API:
election_receiver: Receiver[ElectionMessage],
) -> None:
self.state = State()
self._event_log: list[Event] = []
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
self.command_sender = command_sender
self.download_command_sender = download_command_sender
self.global_event_receiver = global_event_receiver
@@ -223,6 +234,8 @@ class API:
def reset(self, new_session_id: SessionId, result_clock: int):
logger.info("Resetting API State")
self._event_log.close()
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
self.state = State()
self.session_id = new_session_id
self.event_buffer = OrderedBuffer[Event]()
@@ -269,6 +282,9 @@ class API:
self.app.get("/instance/previews")(self.get_placement_previews)
self.app.get("/instance/{instance_id}")(self.get_instance)
self.app.delete("/instance/{instance_id}")(self.delete_instance)
self.app.get("/meta_instances")(self.list_meta_instances)
self.app.post("/meta_instance")(self.create_meta_instance)
self.app.delete("/meta_instance/{meta_instance_id}")(self.delete_meta_instance)
self.app.get("/models")(self.get_models)
self.app.get("/v1/models")(self.get_models)
self.app.post("/models/add")(self.add_custom_model)
@@ -289,7 +305,7 @@ class API:
self.app.post("/v1/messages", response_model=None)(self.claude_messages)
self.app.post("/v1/responses", response_model=None)(self.openai_responses)
self.app.get("/state")(lambda: self.state)
self.app.get("/events")(lambda: self._event_log)
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)
@@ -298,12 +314,27 @@ class API:
self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
async def place_instance(self, payload: PlaceInstanceParams):
model_card = await ModelCard.load(payload.model_id)
command = PlaceInstance(
model_card=await ModelCard.load(payload.model_id),
model_card=model_card,
sharding=payload.sharding,
instance_meta=payload.instance_meta,
min_nodes=payload.min_nodes,
)
# Validate placement before sending — fail fast with a clear error
# instead of silently dropping the command in the master.
try:
get_instance_placements(
command,
topology=self.state.topology,
current_instances=self.state.instances,
node_memory=self.state.node_memory,
node_network=self.state.node_network,
)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
await self._send(command)
return CreateInstanceResponse(
@@ -386,7 +417,12 @@ class API:
if len(list(self.state.topology.list_nodes())) == 0:
return PlacementPreviewResponse(previews=[])
model_card = await ModelCard.load(model_id)
try:
model_card = await ModelCard.load(model_id)
except Exception as exc:
raise HTTPException(
status_code=400, detail=f"Failed to load model card: {exc}"
) from exc
instance_combinations: list[tuple[Sharding, InstanceMeta, int]] = []
for sharding in (Sharding.Pipeline, Sharding.Tensor):
for instance_meta in (InstanceMeta.MlxRing, InstanceMeta.MlxJaccl):
@@ -510,6 +546,44 @@ class API:
instance_id=instance_id,
)
def list_meta_instances(self) -> dict[MetaInstanceId, MetaInstance]:
return dict(self.state.meta_instances)
async def create_meta_instance(
self, payload: CreateMetaInstanceParams
) -> CreateMetaInstanceResponse:
meta_instance = MetaInstance(
model_id=payload.model_id,
sharding=payload.sharding,
instance_meta=payload.instance_meta,
min_nodes=payload.min_nodes,
node_ids=payload.node_ids,
)
command = CreateMetaInstance(meta_instance=meta_instance)
await self._send(command)
return CreateMetaInstanceResponse(
message="Command received.",
command_id=command.command_id,
meta_instance_id=meta_instance.meta_instance_id,
)
async def delete_meta_instance(
self, meta_instance_id: MetaInstanceId
) -> DeleteMetaInstanceResponse:
meta = self.state.meta_instances.get(meta_instance_id)
if not meta:
raise HTTPException(status_code=404, detail="MetaInstance not found")
# Command processor handles cascade-deleting backing instances
command = DeleteMetaInstance(meta_instance_id=meta_instance_id)
await self._send(command)
return DeleteMetaInstanceResponse(
message="Command received.",
command_id=command.command_id,
meta_instance_id=meta_instance_id,
)
async def _token_chunk_stream(
self, command_id: CommandId
) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]:
@@ -529,16 +603,14 @@ class API:
break
except anyio.get_cancelled_exc_class():
# TODO: TaskCancelled
"""
self.command_sender.send_nowait(
ForwarderCommand(origin=self.node_id, command=command)
)
"""
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
)
raise
finally:
command = TaskFinished(finished_command_id=command_id)
await self._send(command)
await self._send(TaskFinished(finished_command_id=command_id))
if command_id in self._text_generation_queues:
del self._text_generation_queues[command_id]
@@ -633,11 +705,14 @@ class API:
"X-Accel-Buffering": "no",
},
)
return await collect_chat_response(
command.command_id,
self._token_chunk_stream(command.command_id),
)
else:
return StreamingResponse(
collect_chat_response(
command.command_id,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
async def bench_chat_completions(
self, payload: BenchChatCompletionRequest
@@ -653,8 +728,7 @@ class API:
command = TextGeneration(task_params=task_params)
await self._send(command)
response = await self._collect_text_generation_with_stats(command.command_id)
return response
return await self._collect_text_generation_with_stats(command.command_id)
async def _resolve_and_validate_text_model(self, model_id: ModelId) -> ModelId:
"""Validate a text model exists and return the resolved model ID.
@@ -689,6 +763,22 @@ class API:
)
return resolved_model
def stream_events(self) -> StreamingResponse:
def _generate_json_array(events: Iterator[Event]) -> Iterator[str]:
yield "["
first = True
for event in events:
if not first:
yield ","
first = False
yield event.model_dump_json()
yield "]"
return StreamingResponse(
_generate_json_array(self._event_log.read_all()),
media_type="application/json",
)
async def get_image(self, image_id: str) -> FileResponse:
stored = self._image_store.get(Id(image_id))
if stored is None:
@@ -856,6 +946,11 @@ class API:
del image_metadata[key]
except anyio.get_cancelled_exc_class():
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
)
raise
finally:
await self._send(TaskFinished(finished_command_id=command_id))
@@ -937,6 +1032,11 @@ class API:
return (images, stats if capture_stats else None)
except anyio.get_cancelled_exc_class():
command = TaskCancelled(cancelled_command_id=command_id)
with anyio.CancelScope(shield=True):
await self.command_sender.send(
ForwarderCommand(origin=self.node_id, command=command)
)
raise
finally:
await self._send(TaskFinished(finished_command_id=command_id))
@@ -1194,12 +1294,15 @@ class API:
"X-Accel-Buffering": "no",
},
)
return await collect_claude_response(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
)
else:
return StreamingResponse(
collect_claude_response(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
async def openai_responses(
self, payload: ResponsesRequest
@@ -1227,11 +1330,15 @@ class API:
},
)
return await collect_responses_response(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
)
else:
return StreamingResponse(
collect_responses_response(
command.command_id,
payload.model,
self._token_chunk_stream(command.command_id),
),
media_type="application/json",
)
def _calculate_total_available_memory(self) -> Memory:
"""Calculate total available memory across all nodes in bytes."""
@@ -1337,6 +1444,7 @@ class API:
with anyio.CancelScope(shield=True):
shutdown_ev.set()
finally:
self._event_log.close()
self.command_sender.close()
self.global_event_receiver.close()
+187
View File
@@ -0,0 +1,187 @@
import contextlib
import json
from collections import OrderedDict
from collections.abc import Iterator
from datetime import datetime, timezone
from io import BufferedRandom, BufferedReader
from pathlib import Path
import msgspec
import zstandard
from loguru import logger
from pydantic import TypeAdapter
from exo.shared.types.events import Event
_EVENT_ADAPTER: TypeAdapter[Event] = TypeAdapter(Event)
_HEADER_SIZE = 4 # uint32 big-endian
_OFFSET_CACHE_SIZE = 128
_MAX_ARCHIVES = 5
def _serialize_event(event: Event) -> bytes:
return msgspec.msgpack.encode(event.model_dump(mode="json"))
def _deserialize_event(raw: bytes) -> Event:
# Decode msgpack into a Python dict, then re-encode as JSON for Pydantic.
# Pydantic's validate_json() uses JSON-mode coercion (e.g. string -> enum)
# even under strict=True, whereas validate_python() does not. Going through
# JSON is the only way to get correct round-trip deserialization without
# disabling strict mode or adding casts everywhere.
as_json = json.dumps(msgspec.msgpack.decode(raw, type=dict))
return _EVENT_ADAPTER.validate_json(as_json)
def _unpack_header(header: bytes) -> int:
return int.from_bytes(header, byteorder="big")
def _skip_record(f: BufferedReader) -> bool:
"""Skip one length-prefixed record. Returns False on EOF."""
header = f.read(_HEADER_SIZE)
if len(header) < _HEADER_SIZE:
return False
f.seek(_unpack_header(header), 1)
return True
def _read_record(f: BufferedReader) -> Event | None:
"""Read one length-prefixed record. Returns None on EOF."""
header = f.read(_HEADER_SIZE)
if len(header) < _HEADER_SIZE:
return None
length = _unpack_header(header)
payload = f.read(length)
if len(payload) < length:
return None
return _deserialize_event(payload)
class DiskEventLog:
"""Append-only event log backed by a file on disk.
On-disk format: sequence of length-prefixed msgpack records.
Each record is [4-byte big-endian uint32 length][msgpack payload].
Uses a bounded LRU cache of event index byte offset for efficient
random access without storing an offset per event.
"""
def __init__(self, directory: Path) -> None:
self._directory = directory
self._directory.mkdir(parents=True, exist_ok=True)
self._active_path = directory / "events.bin"
self._offset_cache: OrderedDict[int, int] = OrderedDict()
self._count: int = 0
# Rotate stale active file from a previous session/crash
if self._active_path.exists():
self._rotate(self._active_path, self._directory)
self._file: BufferedRandom = open(self._active_path, "w+b") # noqa: SIM115
def _cache_offset(self, idx: int, offset: int) -> None:
self._offset_cache[idx] = offset
self._offset_cache.move_to_end(idx)
if len(self._offset_cache) > _OFFSET_CACHE_SIZE:
self._offset_cache.popitem(last=False)
def _seek_to(self, f: BufferedReader, target_idx: int) -> None:
"""Seek f to the byte offset of event target_idx, using cache or scanning forward."""
if target_idx in self._offset_cache:
self._offset_cache.move_to_end(target_idx)
f.seek(self._offset_cache[target_idx])
return
# Find the highest cached index before target_idx
scan_from_idx = 0
scan_from_offset = 0
for cached_idx in self._offset_cache:
if cached_idx < target_idx:
scan_from_idx = cached_idx
scan_from_offset = self._offset_cache[cached_idx]
# Scan forward, skipping records
f.seek(scan_from_offset)
for _ in range(scan_from_idx, target_idx):
_skip_record(f)
self._cache_offset(target_idx, f.tell())
def append(self, event: Event) -> None:
packed = _serialize_event(event)
self._file.write(len(packed).to_bytes(_HEADER_SIZE, byteorder="big"))
self._file.write(packed)
self._count += 1
def read_range(self, start: int, end: int) -> Iterator[Event]:
"""Yield events from index start (inclusive) to end (exclusive)."""
end = min(end, self._count)
if start < 0 or end < 0 or start >= end:
return
self._file.flush()
with open(self._active_path, "rb") as f:
self._seek_to(f, start)
for _ in range(end - start):
event = _read_record(f)
if event is None:
break
yield event
# Cache where we ended up so the next sequential read is a hit
if end < self._count:
self._cache_offset(end, f.tell())
def read_all(self) -> Iterator[Event]:
"""Yield all events from the log one at a time."""
if self._count == 0:
return
self._file.flush()
with open(self._active_path, "rb") as f:
for _ in range(self._count):
event = _read_record(f)
if event is None:
break
yield event
def __len__(self) -> int:
return self._count
def close(self) -> None:
"""Close the file and rotate active file to compressed archive."""
if self._file.closed:
return
self._file.close()
if self._active_path.exists() and self._count > 0:
self._rotate(self._active_path, self._directory)
elif self._active_path.exists():
self._active_path.unlink()
@staticmethod
def _rotate(source: Path, directory: Path) -> None:
"""Compress source into a timestamped archive.
Keeps at most ``_MAX_ARCHIVES`` compressed copies. Oldest beyond
the limit are deleted.
"""
try:
stamp = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S_%f")
dest = directory / f"events.{stamp}.bin.zst"
compressor = zstandard.ZstdCompressor()
with open(source, "rb") as f_in, open(dest, "wb") as f_out:
compressor.copy_stream(f_in, f_out)
source.unlink()
logger.info(f"Rotated event log: {source} -> {dest}")
# Prune oldest archives beyond the limit
archives = sorted(directory.glob("events.*.bin.zst"))
for old in archives[:-_MAX_ARCHIVES]:
old.unlink()
except Exception as e:
logger.opt(exception=e).warning(f"Failed to rotate event log {source}")
# Clean up the source even if compression fails
with contextlib.suppress(OSError):
source.unlink()
+204 -71
View File
@@ -1,9 +1,11 @@
from datetime import datetime, timedelta, timezone
from collections.abc import Sequence
from datetime import datetime, timezone
import anyio
from anyio.abc import TaskGroup
from loguru import logger
from exo.master.event_log import DiskEventLog
from exo.master.placement import (
add_instance_to_placements,
cancel_unnecessary_downloads,
@@ -11,11 +13,22 @@ from exo.master.placement import (
get_transition_events,
place_instance,
)
from exo.master.process_managers import ProcessManager
from exo.master.process_managers.instance_health import InstanceHealthReconciler
from exo.master.process_managers.meta_instance import MetaInstanceReconciler
from exo.master.process_managers.node_timeout import NodeTimeoutReconciler
from exo.master.reconcile import (
find_unsatisfied_meta_instances,
try_place_for_meta_instance,
)
from exo.shared.apply import apply
from exo.shared.constants import EXO_TRACING_ENABLED
from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
from exo.shared.models.model_cards import ModelCard
from exo.shared.types.commands import (
CreateInstance,
CreateMetaInstance,
DeleteInstance,
DeleteMetaInstance,
ForwarderCommand,
ForwarderDownloadCommand,
ImageEdits,
@@ -23,6 +36,7 @@ from exo.shared.types.commands import (
PlaceInstance,
RequestEventLog,
SendInputChunk,
TaskCancelled,
TaskFinished,
TestCommand,
TextGeneration,
@@ -34,10 +48,15 @@ from exo.shared.types.events import (
IndexedEvent,
InputChunkReceived,
InstanceDeleted,
JacclSideChannelData,
JacclSideChannelGathered,
MetaInstanceCreated,
MetaInstanceDeleted,
MetaInstancePlacementFailed,
NodeGatheredInfo,
NodeTimedOut,
TaskCreated,
TaskDeleted,
TaskStatusUpdated,
TraceEventData,
TracesCollected,
TracesMerged,
@@ -57,7 +76,8 @@ 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.shared.types.worker.runners import RunnerId
from exo.utils.channels import Receiver, Sender
from exo.utils.event_buffer import MultiSourceBuffer
@@ -81,17 +101,16 @@ class Master:
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]()
# TODO: not have this
self._event_log: list[Event] = []
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
self._expected_ranks: dict[TaskId, set[int]] = {}
self._jaccl_pending: dict[InstanceId, dict[int, dict[RunnerId, bytes]]] = {}
self._process_managers: Sequence[ProcessManager] = [
InstanceHealthReconciler(),
NodeTimeoutReconciler(),
MetaInstanceReconciler(),
]
async def run(self):
logger.info("Starting Master")
@@ -100,14 +119,12 @@ 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)
tg.start_soon(self._reconcile)
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")
@@ -278,7 +295,7 @@ class Master:
case DeleteInstance():
placement = delete_instance(command, self.state.instances)
transition_events = get_transition_events(
self.state.instances, placement
self.state.instances, placement, self.state.tasks
)
for cmd in cancel_unnecessary_downloads(
placement, self.state.downloads
@@ -289,6 +306,85 @@ class Master:
)
)
generated_events.extend(transition_events)
case CreateMetaInstance():
logger.info(
f"Creating MetaInstance for {command.meta_instance.model_id}"
f" (min_nodes={command.meta_instance.min_nodes},"
f" sharding={command.meta_instance.sharding})"
)
# Apply immediately so self.state is fresh across
# the await below and the reconciler won't race.
await self._apply_and_broadcast(
MetaInstanceCreated(meta_instance=command.meta_instance)
)
# Immediate placement attempt for responsiveness
model_card = await ModelCard.load(
command.meta_instance.model_id
)
# Re-check: reconciler may have satisfied it during the await
meta_id = command.meta_instance.meta_instance_id
still_unsatisfied = any(
m.meta_instance_id == meta_id
for m in find_unsatisfied_meta_instances(
self.state.meta_instances,
self.state.instances,
self.state.topology,
)
)
if still_unsatisfied:
result = try_place_for_meta_instance(
command.meta_instance,
model_card,
self.state.topology,
self.state.instances,
self.state.node_memory,
self.state.node_network,
)
generated_events.extend(result.events)
if result.error is not None:
generated_events.append(
MetaInstancePlacementFailed(
meta_instance_id=meta_id,
reason=result.error,
)
)
case DeleteMetaInstance():
backing_count = sum(
1
for inst in self.state.instances.values()
if inst.meta_instance_id == command.meta_instance_id
)
logger.info(
f"Deleting MetaInstance {command.meta_instance_id}"
f" (cascade-deleting {backing_count} backing instance(s))"
)
generated_events.append(
MetaInstanceDeleted(
meta_instance_id=command.meta_instance_id
)
)
# Cascade-delete backing instances atomically,
# cancelling any active tasks first.
for iid, inst in self.state.instances.items():
if inst.meta_instance_id == command.meta_instance_id:
for task in self.state.tasks.values():
if (
task.instance_id == iid
and task.task_status
in (
TaskStatus.Pending,
TaskStatus.Running,
)
):
generated_events.append(
TaskStatusUpdated(
task_status=TaskStatus.Cancelled,
task_id=task.task_id,
)
)
generated_events.append(
InstanceDeleted(instance_id=iid)
)
case PlaceInstance():
placement = place_instance(
command,
@@ -298,7 +394,7 @@ class Master:
self.state.node_network,
)
transition_events = get_transition_events(
self.state.instances, placement
self.state.instances, placement, self.state.tasks
)
generated_events.extend(transition_events)
case CreateInstance():
@@ -308,7 +404,7 @@ class Master:
self.state.instances,
)
transition_events = get_transition_events(
self.state.instances, placement
self.state.instances, placement, self.state.tasks
)
generated_events.extend(transition_events)
case SendInputChunk(chunk=chunk):
@@ -318,6 +414,18 @@ class Master:
chunk=chunk,
)
)
case TaskCancelled():
if (
task_id := self.command_task_mapping.get(
command.cancelled_command_id
)
) is not None:
generated_events.append(
TaskStatusUpdated(
task_status=TaskStatus.Cancelled,
task_id=task_id,
)
)
case TaskFinished():
generated_events.append(
TaskDeleted(
@@ -326,42 +434,45 @@ class Master:
]
)
)
if command.finished_command_id in self.command_task_mapping:
del self.command_task_mapping[
command.finished_command_id
]
self.command_task_mapping.pop(
command.finished_command_id, None
)
case RequestEventLog():
# We should just be able to send everything, since other buffers will ignore old messages
for i in range(command.since_idx, len(self._event_log)):
await self._send_event(
IndexedEvent(idx=i, event=self._event_log[i])
)
# rate limit to 1000 at a time
end = min(command.since_idx + 1000, len(self._event_log))
for i, event in enumerate(
self._event_log.read_range(command.since_idx, end),
start=command.since_idx,
):
await self._send_event(IndexedEvent(idx=i, event=event))
for event in generated_events:
await self.event_sender.send(event)
await self._apply_and_broadcast(event)
except ValueError as e:
logger.opt(exception=e).warning("Error in command processor")
# These plan loops are the cracks showing in our event sourcing architecture - more things could be commands
async def _plan(self) -> None:
async def _apply_and_broadcast(self, event: Event) -> None:
"""Apply event to state, persist to disk, and broadcast to workers.
State is updated synchronously (before any await), so callers can
rely on ``self.state`` reflecting this event immediately after the
call. Python's cooperative scheduling guarantees no interleaving
between the state read and write.
"""
logger.debug(f"Master indexing event: {str(event)[:100]}")
indexed = IndexedEvent(event=event, idx=len(self._event_log))
self.state = apply(self.state, indexed)
event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage]
self._event_log.append(event)
await self._send_event(indexed)
async def _reconcile(self) -> None:
while True:
# kill broken instances
connected_node_ids = set(self.state.topology.list_nodes())
for instance_id, instance in self.state.instances.items():
for node_id in instance.shard_assignments.node_to_runner:
if node_id not in connected_node_ids:
await self.event_sender.send(
InstanceDeleted(instance_id=instance_id)
)
break
# time out dead nodes
for node_id, time in self.state.last_seen.items():
now = datetime.now(tz=timezone.utc)
if now - time > timedelta(seconds=30):
logger.info(f"Manually removing node {node_id} due to inactivity")
await self.event_sender.send(NodeTimedOut(node_id=node_id))
await anyio.sleep(10)
for pm in self._process_managers:
events = await pm.reconcile(self.state)
for event in events:
await self._apply_and_broadcast(event)
await anyio.sleep(1)
async def _event_processor(self) -> None:
with self.local_event_receiver as local_events:
@@ -379,32 +490,15 @@ class Master:
await self._handle_traces_collected(event)
continue
logger.debug(f"Master indexing event: {str(event)[:100]}")
indexed = IndexedEvent(event=event, idx=len(self._event_log))
self.state = apply(self.state, indexed)
if isinstance(event, JacclSideChannelData):
await self._apply_and_broadcast(event)
await self._handle_jaccl_side_channel(event)
continue
event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage]
if isinstance(event, NodeGatheredInfo):
event.when = str(datetime.now(tz=timezone.utc))
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
await self._apply_and_broadcast(event)
# This function is re-entrant, take care!
async def _send_event(self, event: IndexedEvent):
@@ -436,10 +530,49 @@ class Master:
for trace_data in self._pending_traces[task_id].values():
all_trace_data.extend(trace_data)
await self.event_sender.send(
await self._apply_and_broadcast(
TracesMerged(task_id=task_id, traces=all_trace_data)
)
del self._pending_traces[task_id]
if task_id in self._expected_ranks:
del self._expected_ranks[task_id]
async def _handle_jaccl_side_channel(self, event: JacclSideChannelData) -> None:
"""Accumulate SideChannel contributions; when all runners for an instance
have submitted for the same sequence, emit JacclSideChannelGathered."""
iid = event.instance_id
seq = event.sequence
if iid not in self._jaccl_pending:
self._jaccl_pending[iid] = {}
if seq not in self._jaccl_pending[iid]:
self._jaccl_pending[iid][seq] = {}
self._jaccl_pending[iid][seq][event.runner_id] = event.data
instance = self.state.instances.get(iid)
if instance is None:
logger.warning(f"JacclSideChannelData for unknown instance {iid}")
return
expected_runners = set(instance.shard_assignments.runner_to_shard.keys())
submitted = set(self._jaccl_pending[iid][seq].keys())
logger.info(
f"JACCL side channel: instance={iid} seq={seq} "
f"submitted={len(submitted)}/{len(expected_runners)}"
)
if submitted >= expected_runners:
gathered = dict(self._jaccl_pending[iid][seq])
del self._jaccl_pending[iid][seq]
if not self._jaccl_pending[iid]:
del self._jaccl_pending[iid]
await self._apply_and_broadcast(
JacclSideChannelGathered(
instance_id=iid,
sequence=seq,
gathered_data=gathered,
)
)
+32 -9
View File
@@ -6,11 +6,11 @@ from typing import Sequence
from exo.master.placement_utils import (
Cycle,
filter_cycles_by_memory,
get_largest_cycles,
get_mlx_jaccl_coordinators,
get_mlx_jaccl_devices_matrix,
get_mlx_ring_hosts_by_node,
get_shard_assignments,
get_smallest_cycles,
)
from exo.shared.models.model_cards import ModelId
from exo.shared.topology import Topology
@@ -22,9 +22,15 @@ from exo.shared.types.commands import (
PlaceInstance,
)
from exo.shared.types.common import NodeId
from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
from exo.shared.types.events import (
Event,
InstanceCreated,
InstanceDeleted,
TaskStatusUpdated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import (
DownloadOngoing,
DownloadProgress,
@@ -100,23 +106,27 @@ def place_instance(
"Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)"
)
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
largest_cycles = get_largest_cycles(cycles_with_sufficient_memory)
smallest_rdma_cycles = [
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
largest_rdma_cycles = [
cycle for cycle in largest_cycles if topology.is_rdma_cycle(cycle)
]
if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
smallest_cycles = smallest_rdma_cycles
if command.instance_meta == InstanceMeta.MlxJaccl:
if not largest_rdma_cycles:
raise ValueError(
"Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
)
largest_cycles = largest_rdma_cycles
cycles_with_leaf_nodes: list[Cycle] = [
cycle
for cycle in smallest_cycles
for cycle in largest_cycles
if any(topology.node_is_leaf(node_id) for node_id in cycle)
]
selected_cycle = max(
cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else largest_cycles,
key=lambda cycle: sum(
(node_memory[node_id].ram_available for node_id in cycle),
start=Memory(),
@@ -186,6 +196,7 @@ def delete_instance(
def get_transition_events(
current_instances: Mapping[InstanceId, Instance],
target_instances: Mapping[InstanceId, Instance],
tasks: Mapping[TaskId, Task],
) -> Sequence[Event]:
events: list[Event] = []
@@ -201,6 +212,18 @@ def get_transition_events(
# find instances to delete
for instance_id in current_instances:
if instance_id not in target_instances:
for task in tasks.values():
if task.instance_id == instance_id and task.task_status in [
TaskStatus.Pending,
TaskStatus.Running,
]:
events.append(
TaskStatusUpdated(
task_status=TaskStatus.Cancelled,
task_id=task.task_id,
)
)
events.append(
InstanceDeleted(
instance_id=instance_id,
+3 -3
View File
@@ -37,11 +37,11 @@ def filter_cycles_by_memory(
return filtered_cycles
def get_smallest_cycles(
def get_largest_cycles(
cycles: list[Cycle],
) -> list[Cycle]:
min_nodes = min(len(cycle) for cycle in cycles)
return [cycle for cycle in cycles if len(cycle) == min_nodes]
max_nodes = max(len(cycle) for cycle in cycles)
return [cycle for cycle in cycles if len(cycle) == max_nodes]
def allocate_layers_proportionally(
@@ -0,0 +1,12 @@
from collections.abc import Sequence
from typing import Protocol, runtime_checkable
from exo.shared.types.events import Event
from exo.shared.types.state import State
@runtime_checkable
class ProcessManager(Protocol):
"""A reconciliation step that examines state and returns corrective events."""
async def reconcile(self, state: State) -> Sequence[Event]: ...
@@ -0,0 +1,62 @@
from collections.abc import Sequence
from typing import final
from loguru import logger
from exo.master.reconcile import instance_connections_healthy, instance_runners_failed
from exo.shared.types.events import Event, InstanceDeleted, InstanceRetrying
from exo.shared.types.state import State
MAX_INSTANCE_RETRIES = 3
@final
class InstanceHealthReconciler:
"""Delete instances whose network connections are broken or whose runners have all failed."""
async def reconcile(self, state: State) -> Sequence[Event]:
events: list[Event] = []
for instance_id, instance in state.instances.items():
if not instance_connections_healthy(instance, state.topology):
events.append(
InstanceDeleted(
instance_id=instance_id,
failure_error="Network connection lost",
)
)
continue
is_failed, error_message = instance_runners_failed(
instance, state.runners, state.node_identities
)
if is_failed:
# Retry within the same instance if backed by a MetaInstance
mid = instance.meta_instance_id
mi = state.meta_instances.get(mid) if mid else None
if mid and mi and mi.consecutive_failures < MAX_INSTANCE_RETRIES:
logger.info(
f"Instance {instance_id} failed (attempt"
f" {mi.consecutive_failures + 1}/{MAX_INSTANCE_RETRIES}),"
f" retrying: {error_message}"
)
events.append(
InstanceRetrying(
instance_id=instance_id,
meta_instance_id=mid,
failure_error=error_message or "Runner failed",
)
)
else:
if mid and mi:
logger.warning(
f"Instance {instance_id} exceeded retry limit"
f" ({MAX_INSTANCE_RETRIES}), deleting:"
f" {error_message}"
)
events.append(
InstanceDeleted(
instance_id=instance_id,
failure_error=error_message,
)
)
return events
@@ -0,0 +1,91 @@
from collections.abc import Sequence
from typing import final
import anyio
from loguru import logger
from exo.master.reconcile import (
find_unsatisfied_meta_instances,
try_place_for_meta_instance,
)
from exo.shared.models.model_cards import ModelCard
from exo.shared.types.events import Event, InstanceCreated, MetaInstancePlacementFailed
from exo.shared.types.state import State
from exo.shared.types.worker.instances import Instance, InstanceId
MODEL_CARD_LOAD_TIMEOUT_SECONDS = 10
@final
class MetaInstanceReconciler:
"""Place instances for unsatisfied MetaInstances."""
async def reconcile(self, state: State) -> Sequence[Event]:
all_events: list[Event] = []
# Local copy for intermediate tracking — so placement of B
# sees A's instance and doesn't double-place on same resources.
current_instances: dict[InstanceId, Instance] = dict(state.instances)
unsatisfied = find_unsatisfied_meta_instances(
state.meta_instances,
current_instances,
state.topology,
)
for meta_instance in unsatisfied:
try:
with anyio.fail_after(MODEL_CARD_LOAD_TIMEOUT_SECONDS):
model_card = await ModelCard.load(meta_instance.model_id)
except TimeoutError:
logger.warning(
f"ModelCard.load timed out for {meta_instance.model_id}, skipping this cycle"
)
continue
except Exception as exc:
logger.warning(
f"ModelCard.load failed for {meta_instance.model_id}: {exc}"
)
error = f"Failed to load model card: {exc}"
if meta_instance.placement_error != error:
all_events.append(
MetaInstancePlacementFailed(
meta_instance_id=meta_instance.meta_instance_id,
reason=error,
)
)
continue
result = try_place_for_meta_instance(
meta_instance,
model_card,
state.topology,
current_instances,
state.node_memory,
state.node_network,
)
# Update local instance map so next placement sees this one
for event in result.events:
if isinstance(event, InstanceCreated):
logger.info(
f"MetaInstance reconciler placed instance"
f" {event.instance.instance_id} for"
f" {meta_instance.model_id}"
)
current_instances[event.instance.instance_id] = event.instance
all_events.extend(result.events)
# Emit placement failure if error differs from what's already in state
if (
result.error is not None
and meta_instance.placement_error != result.error
):
logger.warning(
f"MetaInstance placement failed for"
f" {meta_instance.model_id}: {result.error}"
)
all_events.append(
MetaInstancePlacementFailed(
meta_instance_id=meta_instance.meta_instance_id,
reason=result.error,
)
)
return all_events
@@ -0,0 +1,27 @@
from collections.abc import Sequence
from datetime import datetime, timedelta, timezone
from typing import final
from loguru import logger
from exo.shared.types.events import Event, NodeTimedOut
from exo.shared.types.state import State
_DEFAULT_TIMEOUT = timedelta(seconds=30)
@final
class NodeTimeoutReconciler:
"""Time out nodes that haven't been seen recently."""
def __init__(self, timeout: timedelta = _DEFAULT_TIMEOUT) -> None:
self.timeout = timeout
async def reconcile(self, state: State) -> Sequence[Event]:
now = datetime.now(tz=timezone.utc)
events: list[Event] = []
for node_id, last_seen in state.last_seen.items():
if now - last_seen > self.timeout:
logger.info(f"Removing node {node_id} due to inactivity")
events.append(NodeTimedOut(node_id=node_id))
return events
+240
View File
@@ -0,0 +1,240 @@
from collections.abc import Mapping, Sequence
from typing import NamedTuple
from loguru import logger
from exo.master.placement import get_transition_events, place_instance
from exo.shared.models.model_cards import ModelCard
from exo.shared.topology import Topology
from exo.shared.types.commands import PlaceInstance
from exo.shared.types.common import MetaInstanceId, NodeId
from exo.shared.types.events import Event
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
from exo.shared.types.topology import RDMAConnection, SocketConnection
from exo.shared.types.worker.instances import (
BaseInstance,
Instance,
InstanceId,
MlxJacclInstance,
MlxRingInstance,
)
from exo.shared.types.worker.runners import (
RunnerFailed,
RunnerId,
RunnerShutdown,
RunnerStatus,
)
class PlacementResult(NamedTuple):
"""Result of a placement attempt: events to apply and optional error reason."""
events: Sequence[Event]
error: str | None
def _get_ring_order(instance: BaseInstance) -> list[NodeId]:
"""Reconstruct ring order from shard device_rank."""
node_ranks: list[tuple[NodeId, int]] = []
for node_id, runner_id in instance.shard_assignments.node_to_runner.items():
shard = instance.shard_assignments.runner_to_shard[runner_id]
node_ranks.append((node_id, shard.device_rank))
node_ranks.sort(key=lambda x: x[1])
return [node_id for node_id, _ in node_ranks]
def _ring_connections_healthy(instance: MlxRingInstance, topology: Topology) -> bool:
"""Check that the specific IPs used by a ring instance still exist in the topology."""
ring = _get_ring_order(instance)
n = len(ring)
for node in ring:
hosts = instance.hosts_by_node[node]
for idx in range(n):
host = hosts[idx]
if host.ip in ("0.0.0.0", "198.51.100.1"):
continue # self or placeholder
# Real connection: node → ring[idx]. Check specific IP.
connections = topology.get_all_connections_between(node, ring[idx])
if not any(
isinstance(c, SocketConnection)
and c.sink_multiaddr.ip_address == host.ip
for c in connections
):
return False
return True
def _jaccl_connections_healthy(instance: MlxJacclInstance, topology: Topology) -> bool:
"""Check that the specific RDMA interfaces used by a JACCL instance still exist."""
ring = _get_ring_order(instance)
n = len(ring)
for i in range(n):
for j in range(n):
iface = instance.jaccl_devices[i][j]
if iface is None:
continue
connections = topology.get_all_connections_between(ring[i], ring[j])
if not any(
isinstance(c, RDMAConnection) and c.source_rdma_iface == iface
for c in connections
):
return False
return True
def instance_connections_healthy(instance: Instance, topology: Topology) -> bool:
"""Check that an instance's nodes and specific connections are still in the topology."""
instance_nodes = set(instance.shard_assignments.node_to_runner.keys())
if not all(topology.contains_node(n) for n in instance_nodes):
return False
if len(instance_nodes) <= 1:
return True
match instance:
case MlxRingInstance():
return _ring_connections_healthy(instance, topology)
case MlxJacclInstance():
return _jaccl_connections_healthy(instance, topology)
def instance_runners_failed(
instance: Instance,
runners: Mapping[RunnerId, RunnerStatus],
node_identities: Mapping[NodeId, NodeIdentity],
) -> tuple[bool, str | None]:
"""Check if an instance's runners have all reached terminal failure states.
Returns ``(True, error_message)`` when ALL runners are terminal
(``RunnerFailed`` or ``RunnerShutdown``) and at least one is ``RunnerFailed``.
Returns ``(False, None)`` when runners are still active, haven't reported
yet, or all gracefully shut down (no ``RunnerFailed``).
"""
instance_runner_ids = set(instance.shard_assignments.node_to_runner.values())
if not instance_runner_ids:
return False, None
# Build reverse mapping: runner_id -> node_id
runner_to_node: dict[RunnerId, NodeId] = {
runner_id: node_id
for node_id, runner_id in instance.shard_assignments.node_to_runner.items()
}
has_any_failed = False
error_messages: list[str] = []
for runner_id in instance_runner_ids:
status = runners.get(runner_id)
if status is None:
# Runner hasn't reported yet — instance is still starting
return False, None
if isinstance(status, RunnerFailed):
has_any_failed = True
if status.error_message:
node_id = runner_to_node.get(runner_id)
name = (
node_identities[node_id].friendly_name
if node_id and node_id in node_identities
else node_id or "unknown"
)
error_messages.append(f"{name}: {status.error_message}")
elif isinstance(status, RunnerShutdown):
pass # Terminal but not a failure indicator on its own
else:
# Runner is still active (connecting, loading, running, etc.)
return False, None
if has_any_failed:
return True, "; ".join(error_messages) if error_messages else "Runner failed"
# All runners are Shutdown but none Failed — graceful shutdown, not a failure
return False, None
def instance_satisfies_meta_instance(
meta_instance: MetaInstance,
instance: Instance,
) -> bool:
"""Check if a single instance satisfies a meta-instance's constraints.
This is a pure constraint check (model, min_nodes, node_ids).
Use ``instance_connections_healthy`` separately for topology health.
"""
if instance.shard_assignments.model_id != meta_instance.model_id:
return False
instance_nodes = set(instance.shard_assignments.node_to_runner.keys())
if len(instance_nodes) < meta_instance.min_nodes:
return False
return meta_instance.node_ids is None or set(meta_instance.node_ids).issubset(
instance_nodes
)
def find_unsatisfied_meta_instances(
meta_instances: Mapping[MetaInstanceId, MetaInstance],
instances: Mapping[InstanceId, Instance],
topology: Topology,
) -> Sequence[MetaInstance]:
"""Return meta-instances that have no healthy backing instance."""
unsatisfied: list[MetaInstance] = []
for meta_id, meta_instance in meta_instances.items():
has_healthy_backing = any(
instance.meta_instance_id == meta_id
and instance_connections_healthy(instance, topology)
for instance in instances.values()
)
if not has_healthy_backing:
unsatisfied.append(meta_instance)
return unsatisfied
def try_place_for_meta_instance(
meta_instance: MetaInstance,
model_card: ModelCard,
topology: Topology,
current_instances: Mapping[InstanceId, Instance],
node_memory: Mapping[NodeId, MemoryUsage],
node_network: Mapping[NodeId, NodeNetworkInfo],
) -> PlacementResult:
"""Try to place an instance satisfying the meta-instance constraints.
Returns a :class:`PlacementResult` with events on success, or an error
reason on failure.
"""
command = PlaceInstance(
model_card=model_card,
sharding=meta_instance.sharding,
instance_meta=meta_instance.instance_meta,
min_nodes=meta_instance.min_nodes,
)
try:
target_instances = place_instance(
command,
topology,
current_instances,
node_memory,
node_network,
required_nodes=(
set(meta_instance.node_ids) if meta_instance.node_ids else None
),
)
# Tag the new instance with meta_instance_id
new_instance_ids = set(target_instances.keys()) - set(current_instances.keys())
if new_instance_ids:
new_id = next(iter(new_instance_ids))
target_instances[new_id] = target_instances[new_id].model_copy(
update={"meta_instance_id": meta_instance.meta_instance_id}
)
return PlacementResult(
events=list(get_transition_events(current_instances, target_instances, {})),
error=None,
)
except ValueError as e:
logger.debug(
f"MetaInstance placement not possible for {meta_instance.model_id}: {e}"
)
return PlacementResult(events=[], error=str(e))
+21 -5
View File
@@ -4,7 +4,11 @@ import json
from collections.abc import AsyncGenerator
from typing import Any, cast
from exo.master.adapters.claude import collect_claude_response, generate_claude_stream
from exo.master.adapters.claude import (
ClaudeMessagesResponse,
collect_claude_response,
generate_claude_stream,
)
from exo.shared.types.api import ToolCallItem
from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
from exo.shared.types.common import CommandId, ModelId
@@ -17,6 +21,18 @@ async def _chunks_to_stream(
yield chunk
async def _collect_response(
command_id: CommandId,
model: str,
chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None],
) -> ClaudeMessagesResponse:
"""Helper to consume the async generator and parse the JSON response."""
parts: list[str] = []
async for part in collect_claude_response(command_id, model, chunk_stream):
parts.append(part)
return ClaudeMessagesResponse.model_validate_json("".join(parts))
MODEL = ModelId("test-model")
COMMAND_ID = CommandId("cmd_test123")
@@ -47,7 +63,7 @@ class TestCollectClaudeResponseToolUse:
],
),
]
response = await collect_claude_response(
response = await _collect_response(
COMMAND_ID, "test-model", _chunks_to_stream(chunks)
)
@@ -77,7 +93,7 @@ class TestCollectClaudeResponseToolUse:
],
),
]
response = await collect_claude_response(
response = await _collect_response(
COMMAND_ID, "test-model", _chunks_to_stream(chunks)
)
@@ -102,7 +118,7 @@ class TestCollectClaudeResponseToolUse:
],
),
]
response = await collect_claude_response(
response = await _collect_response(
COMMAND_ID, "test-model", _chunks_to_stream(chunks)
)
@@ -116,7 +132,7 @@ class TestCollectClaudeResponseToolUse:
async def test_no_content_produces_empty_text_block(self):
chunks: list[ErrorChunk | ToolCallChunk | TokenChunk] = []
response = await collect_claude_response(
response = await _collect_response(
COMMAND_ID, "test-model", _chunks_to_stream(chunks)
)
assert len(response.content) == 1
+155
View File
@@ -0,0 +1,155 @@
from pathlib import Path
import pytest
from exo.master.event_log import DiskEventLog
from exo.shared.types.events import TestEvent
@pytest.fixture
def log_dir(tmp_path: Path) -> Path:
return tmp_path / "event_log"
def test_append_and_read_back(log_dir: Path):
log = DiskEventLog(log_dir)
events = [TestEvent() for _ in range(5)]
for e in events:
log.append(e)
assert len(log) == 5
result = list(log.read_all())
assert len(result) == 5
for original, restored in zip(events, result, strict=True):
assert original.event_id == restored.event_id
log.close()
def test_read_range(log_dir: Path):
log = DiskEventLog(log_dir)
events = [TestEvent() for _ in range(10)]
for e in events:
log.append(e)
result = list(log.read_range(3, 7))
assert len(result) == 4
for i, restored in enumerate(result):
assert events[3 + i].event_id == restored.event_id
log.close()
def test_read_range_bounds(log_dir: Path):
log = DiskEventLog(log_dir)
events = [TestEvent() for _ in range(3)]
for e in events:
log.append(e)
# Start beyond count
assert list(log.read_range(5, 10)) == []
# Negative start
assert list(log.read_range(-1, 2)) == []
# End beyond count is clamped
result = list(log.read_range(1, 100))
assert len(result) == 2
log.close()
def test_empty_log(log_dir: Path):
log = DiskEventLog(log_dir)
assert len(log) == 0
assert list(log.read_all()) == []
assert list(log.read_range(0, 10)) == []
log.close()
def _archives(log_dir: Path) -> list[Path]:
return sorted(log_dir.glob("events.*.bin.zst"))
def test_rotation_on_close(log_dir: Path):
log = DiskEventLog(log_dir)
log.append(TestEvent())
log.close()
active = log_dir / "events.bin"
assert not active.exists()
archives = _archives(log_dir)
assert len(archives) == 1
assert archives[0].stat().st_size > 0
def test_rotation_on_construction_with_stale_file(log_dir: Path):
log_dir.mkdir(parents=True, exist_ok=True)
(log_dir / "events.bin").write_bytes(b"stale data")
log = DiskEventLog(log_dir)
archives = _archives(log_dir)
assert len(archives) == 1
assert archives[0].exists()
assert len(log) == 0
log.close()
def test_empty_log_no_archive(log_dir: Path):
"""Closing an empty log should not leave an archive."""
log = DiskEventLog(log_dir)
log.close()
active = log_dir / "events.bin"
assert not active.exists()
assert _archives(log_dir) == []
def test_close_is_idempotent(log_dir: Path):
log = DiskEventLog(log_dir)
log.append(TestEvent())
log.close()
archive = _archives(log_dir)
log.close() # should not raise
assert _archives(log_dir) == archive
def test_successive_sessions(log_dir: Path):
"""Simulate two master sessions: both archives should be kept."""
log1 = DiskEventLog(log_dir)
log1.append(TestEvent())
log1.close()
first_archive = _archives(log_dir)[-1]
log2 = DiskEventLog(log_dir)
log2.append(TestEvent())
log2.append(TestEvent())
log2.close()
# Session 1 archive shifted to slot 2, session 2 in slot 1
second_archive = _archives(log_dir)[-1]
should_be_first_archive = _archives(log_dir)[-2]
assert first_archive.exists()
assert second_archive.exists()
assert first_archive != second_archive
assert should_be_first_archive == first_archive
def test_rotation_keeps_at_most_5_archives(log_dir: Path):
"""After 7 sessions, only the 5 most recent archives should remain."""
all_archives: list[Path] = []
for _ in range(7):
log = DiskEventLog(log_dir)
log.append(TestEvent())
log.close()
all_archives.append(_archives(log_dir)[-1])
for old in all_archives[:2]:
assert not old.exists()
for recent in all_archives[2:]:
assert recent.exists()
@@ -0,0 +1,778 @@
"""Edge-case and regression tests for MetaInstance lifecycle, concurrent operations, and error handling."""
import pytest
from exo.master.process_managers.instance_health import (
MAX_INSTANCE_RETRIES,
InstanceHealthReconciler,
)
from exo.master.process_managers.meta_instance import MetaInstanceReconciler
from exo.master.reconcile import (
find_unsatisfied_meta_instances,
instance_connections_healthy,
instance_runners_failed,
instance_satisfies_meta_instance,
)
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.topology import Topology
from exo.shared.types.common import Host, MetaInstanceId, NodeId
from exo.shared.types.events import (
IndexedEvent,
InstanceCreated,
InstanceDeleted,
InstanceRetrying,
MetaInstanceCreated,
MetaInstanceDeleted,
MetaInstancePlacementFailed,
TaskStatusUpdated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.profiling import NodeIdentity
from exo.shared.types.state import State
from exo.shared.types.tasks import LoadModel, TaskId, TaskStatus
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.instances import (
InstanceId,
MlxRingInstance,
)
from exo.shared.types.worker.runners import (
RunnerFailed,
RunnerId,
RunnerReady,
ShardAssignments,
)
from exo.shared.types.worker.shards import PipelineShardMetadata
# --- Helpers (copied from test_reconcile.py for independence) ---
def _model_card(model_id: str = "test-org/test-model") -> ModelCard:
return ModelCard(
model_id=ModelId(model_id),
storage_size=Memory.from_kb(1000),
n_layers=10,
hidden_size=30,
supports_tensor=True,
tasks=[ModelTask.TextGeneration],
)
def _topology(*node_ids: str, connect: bool = True) -> Topology:
t = Topology()
nodes = [NodeId(n) for n in node_ids]
for n in nodes:
t.add_node(n)
if connect and len(nodes) > 1:
for i in range(len(nodes)):
j = (i + 1) % len(nodes)
t.add_connection(
Connection(
source=nodes[i],
sink=nodes[j],
edge=SocketConnection(
sink_multiaddr=Multiaddr(
address=f"/ip4/10.0.0.{j + 1}/tcp/50000"
)
),
)
)
t.add_connection(
Connection(
source=nodes[j],
sink=nodes[i],
edge=SocketConnection(
sink_multiaddr=Multiaddr(
address=f"/ip4/10.0.0.{i + 1}/tcp/50000"
)
),
)
)
return t
def _meta_instance(
model_id: str = "test-org/test-model",
*,
min_nodes: int = 1,
node_ids: list[NodeId] | None = None,
meta_instance_id: MetaInstanceId | None = None,
consecutive_failures: int = 0,
last_failure_error: str | None = None,
placement_error: str | None = None,
) -> MetaInstance:
return MetaInstance(
meta_instance_id=meta_instance_id or MetaInstanceId(),
model_id=ModelId(model_id),
min_nodes=min_nodes,
node_ids=node_ids,
consecutive_failures=consecutive_failures,
last_failure_error=last_failure_error,
placement_error=placement_error,
)
def _instance(
model_id: str = "test-org/test-model",
node_ids: list[str] | None = None,
instance_id: InstanceId | None = None,
meta_instance_id: MetaInstanceId | None = None,
) -> tuple[InstanceId, MlxRingInstance]:
iid = instance_id or InstanceId()
nodes = node_ids or ["node-a"]
n = len(nodes)
mc = _model_card(model_id)
ephemeral_port = 50000
node_to_runner = {NodeId(nd): RunnerId() for nd in nodes}
runner_to_shard = {
runner_id: PipelineShardMetadata(
model_card=mc,
device_rank=i,
world_size=n,
start_layer=0,
end_layer=mc.n_layers,
n_layers=mc.n_layers,
)
for i, runner_id in enumerate(node_to_runner.values())
}
hosts_by_node: dict[NodeId, list[Host]] = {}
for r, node_str in enumerate(nodes):
hosts: list[Host] = []
for idx in range(n):
if idx == r:
hosts.append(Host(ip="0.0.0.0", port=ephemeral_port))
elif n > 1 and idx in ((r - 1) % n, (r + 1) % n):
hosts.append(Host(ip=f"10.0.0.{idx + 1}", port=ephemeral_port))
else:
hosts.append(Host(ip="198.51.100.1", port=0))
hosts_by_node[NodeId(node_str)] = hosts
return iid, MlxRingInstance(
instance_id=iid,
shard_assignments=ShardAssignments(
model_id=ModelId(model_id),
runner_to_shard=runner_to_shard,
node_to_runner=node_to_runner,
),
hosts_by_node=hosts_by_node,
ephemeral_port=ephemeral_port,
meta_instance_id=meta_instance_id,
)
# =============================================================================
# 1. MetaInstance lifecycle edge cases
# =============================================================================
def test_meta_instance_model_is_frozen():
"""MetaInstance should be immutable (frozen model)."""
meta = _meta_instance()
try:
meta.model_id = ModelId("something-else")
raise AssertionError("Should have raised")
except Exception:
pass # Expected — frozen model
def test_meta_instance_created_then_deleted_roundtrip():
"""Create and delete a MetaInstance through apply — state should be clean."""
state = State()
meta = _meta_instance()
state = apply(
state, IndexedEvent(idx=0, event=MetaInstanceCreated(meta_instance=meta))
)
assert meta.meta_instance_id in state.meta_instances
state = apply(
state,
IndexedEvent(
idx=1, event=MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id)
),
)
assert meta.meta_instance_id not in state.meta_instances
assert len(state.meta_instances) == 0
def test_delete_nonexistent_meta_instance_is_safe():
"""Deleting a MetaInstance that doesn't exist should not crash."""
state = State()
event = MetaInstanceDeleted(meta_instance_id=MetaInstanceId("nonexistent"))
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert len(new_state.meta_instances) == 0
def test_placement_failed_for_nonexistent_meta_instance_is_safe():
"""MetaInstancePlacementFailed for unknown ID should not crash."""
state = State()
event = MetaInstancePlacementFailed(
meta_instance_id=MetaInstanceId("nonexistent"),
reason="test",
)
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert len(new_state.meta_instances) == 0
def test_multiple_meta_instances_for_same_model():
"""Multiple MetaInstances for the same model are tracked independently."""
state = State()
meta_a = _meta_instance("test-org/model-x")
meta_b = _meta_instance("test-org/model-x")
state = apply(
state, IndexedEvent(idx=0, event=MetaInstanceCreated(meta_instance=meta_a))
)
state = apply(
state, IndexedEvent(idx=1, event=MetaInstanceCreated(meta_instance=meta_b))
)
assert len(state.meta_instances) == 2
assert meta_a.meta_instance_id in state.meta_instances
assert meta_b.meta_instance_id in state.meta_instances
# =============================================================================
# 2. Retry logic edge cases
# =============================================================================
def test_retry_counter_resets_on_successful_instance_creation():
"""When a new instance is created for a meta-instance, failures should reset."""
meta = _meta_instance(consecutive_failures=2, last_failure_error="old")
_, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(meta_instances={meta.meta_instance_id: meta})
state = apply(state, IndexedEvent(idx=0, event=InstanceCreated(instance=inst)))
mi = state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 0
# last_failure_error is preserved (for UI display)
assert mi.last_failure_error == "old"
async def test_retry_count_increments_through_full_cycle():
"""Walk through MAX_INSTANCE_RETRIES worth of retries, then verify delete."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
topology = _topology("node-a")
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
topology=topology,
)
runner_ids = list(inst.shard_assignments.node_to_runner.values())
for idx, i in enumerate(range(MAX_INSTANCE_RETRIES)):
# Simulate runners failing
state_with_runners = state.model_copy(
update={"runners": {runner_ids[0]: RunnerFailed(error_message=f"fail-{i}")}}
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state_with_runners)
assert len(events) == 1
assert isinstance(events[0], InstanceRetrying), f"iteration {i}"
state = apply(state, IndexedEvent(idx=idx, event=events[0]))
# After MAX_INSTANCE_RETRIES retries, failure counter should be at max
mi = state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == MAX_INSTANCE_RETRIES
# Next failure should result in deletion
state_with_runners = state.model_copy(
update={"runners": {runner_ids[0]: RunnerFailed(error_message="final")}}
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state_with_runners)
assert len(events) == 1
assert isinstance(events[0], InstanceDeleted)
async def test_health_reconciler_respects_exact_limit():
"""At exactly MAX_INSTANCE_RETRIES, reconciler should delete, not retry."""
meta = _meta_instance(consecutive_failures=MAX_INSTANCE_RETRIES)
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
runner_ids = list(inst.shard_assignments.node_to_runner.values())
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
runners={runner_ids[0]: RunnerFailed(error_message="OOM")},
topology=_topology("node-a"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceDeleted)
async def test_health_reconciler_at_limit_minus_one_retries():
"""At MAX_INSTANCE_RETRIES - 1, reconciler should still retry."""
meta = _meta_instance(consecutive_failures=MAX_INSTANCE_RETRIES - 1)
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
runner_ids = list(inst.shard_assignments.node_to_runner.values())
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
runners={runner_ids[0]: RunnerFailed(error_message="OOM")},
topology=_topology("node-a"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceRetrying)
# =============================================================================
# 3. Error handling edge cases
# =============================================================================
def test_runners_failed_with_empty_error_message():
"""RunnerFailed with empty error_message should still report as failed."""
_, inst = _instance(node_ids=["node-a"])
runners = {
rid: RunnerFailed(error_message="")
for rid in inst.shard_assignments.node_to_runner.values()
}
is_failed, error = instance_runners_failed(inst, runners, {})
assert is_failed is True
# Empty error message means we get the fallback
assert error == "Runner failed"
def test_runners_failed_with_none_error_message():
"""RunnerFailed with None error_message should still report as failed."""
_, inst = _instance(node_ids=["node-a"])
runners = {
rid: RunnerFailed(error_message=None)
for rid in inst.shard_assignments.node_to_runner.values()
}
is_failed, error = instance_runners_failed(inst, runners, {})
assert is_failed is True
assert error == "Runner failed"
def test_runners_failed_collects_all_error_messages():
"""With multiple failed runners, all error messages should be collected."""
_, inst = _instance(node_ids=["node-a", "node-b", "node-c"])
runner_ids = list(inst.shard_assignments.node_to_runner.values())
runners = {
runner_ids[0]: RunnerFailed(error_message="OOM on GPU 0"),
runner_ids[1]: RunnerFailed(error_message="OOM on GPU 1"),
runner_ids[2]: RunnerFailed(error_message="OOM on GPU 2"),
}
is_failed, error = instance_runners_failed(inst, runners, {})
assert is_failed is True
assert error is not None
assert "OOM on GPU 0" in error
assert "OOM on GPU 1" in error
assert "OOM on GPU 2" in error
def test_runners_failed_includes_friendly_name():
"""Error messages should include node friendly names when available."""
_, inst = _instance(node_ids=["node-a"])
node_id = NodeId("node-a")
runner_ids = list(inst.shard_assignments.node_to_runner.values())
runners = {runner_ids[0]: RunnerFailed(error_message="OOM")}
identities = {node_id: NodeIdentity(friendly_name="My Mac Studio")}
is_failed, error = instance_runners_failed(inst, runners, identities)
assert is_failed is True
assert error is not None
assert "My Mac Studio" in error
def test_instance_retrying_for_missing_instance_is_safe():
"""InstanceRetrying for an instance not in state should not crash.
NOTE: When the instance is missing, the handler returns early WITHOUT
incrementing the MetaInstance failure counter. This means stale retry
events for already-deleted instances are silently dropped. This is
acceptable since the InstanceDeleted handler already increments failures.
"""
meta = _meta_instance()
state = State(meta_instances={meta.meta_instance_id: meta})
event = InstanceRetrying(
instance_id=InstanceId("nonexistent"),
meta_instance_id=meta.meta_instance_id,
failure_error="crash",
)
new_state = apply(state, IndexedEvent(idx=0, event=event))
# Does not crash, but failure count is NOT incremented (early return)
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 0
# =============================================================================
# 4. Backward compatibility
# =============================================================================
def test_instance_without_meta_instance_id_works():
"""Instances created without meta_instance_id should still function normally."""
_, inst = _instance(node_ids=["node-a"])
assert inst.meta_instance_id is None
topology = _topology("node-a")
assert instance_connections_healthy(inst, topology) is True
def test_instance_deleted_without_meta_does_not_affect_meta_instances():
"""Deleting an instance without meta_instance_id should not affect meta_instances."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"]) # no meta_instance_id
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
event = InstanceDeleted(instance_id=iid, failure_error="crash")
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 0 # unchanged
def test_satisfies_ignores_meta_instance_id_binding():
"""instance_satisfies_meta_instance checks constraints only, not binding."""
meta = _meta_instance()
_, inst = _instance(node_ids=["node-a"]) # no meta_instance_id set
# Should match on constraints (model, min_nodes) regardless of binding
assert instance_satisfies_meta_instance(meta, inst) is True
def test_find_unsatisfied_uses_binding_not_constraints():
"""find_unsatisfied checks meta_instance_id binding, not just constraint matching."""
meta = _meta_instance()
# Instance matches constraints but is NOT bound to this meta_instance
iid, inst = _instance(node_ids=["node-a"])
topology = _topology("node-a")
result = find_unsatisfied_meta_instances(
{meta.meta_instance_id: meta}, {iid: inst}, topology
)
# Should be unsatisfied because instance.meta_instance_id != meta.meta_instance_id
assert list(result) == [meta]
# =============================================================================
# 5. Concurrent / multi-instance scenarios
# =============================================================================
async def test_health_reconciler_handles_multiple_failing_instances():
"""Multiple instances failing simultaneously should each get their own event."""
meta_a = _meta_instance()
meta_b = _meta_instance()
iid_a, inst_a = _instance(
node_ids=["node-a"], meta_instance_id=meta_a.meta_instance_id
)
iid_b, inst_b = _instance(
node_ids=["node-b"], meta_instance_id=meta_b.meta_instance_id
)
runner_ids_a = list(inst_a.shard_assignments.node_to_runner.values())
runner_ids_b = list(inst_b.shard_assignments.node_to_runner.values())
state = State(
meta_instances={
meta_a.meta_instance_id: meta_a,
meta_b.meta_instance_id: meta_b,
},
instances={iid_a: inst_a, iid_b: inst_b},
runners={
runner_ids_a[0]: RunnerFailed(error_message="OOM"),
runner_ids_b[0]: RunnerFailed(error_message="OOM"),
},
topology=_topology("node-a", "node-b"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 2
# Both should be InstanceRetrying since failures < MAX
assert all(isinstance(e, InstanceRetrying) for e in events)
instance_ids = {e.instance_id for e in events} # type: ignore[union-attr]
assert instance_ids == {iid_a, iid_b}
async def test_health_reconciler_mixed_healthy_and_failing():
"""Only failing instances should produce events; healthy ones should not."""
meta_healthy = _meta_instance()
meta_failing = _meta_instance()
iid_h, inst_h = _instance(
node_ids=["node-a"], meta_instance_id=meta_healthy.meta_instance_id
)
iid_f, inst_f = _instance(
node_ids=["node-b"], meta_instance_id=meta_failing.meta_instance_id
)
runner_ids_h = list(inst_h.shard_assignments.node_to_runner.values())
runner_ids_f = list(inst_f.shard_assignments.node_to_runner.values())
state = State(
meta_instances={
meta_healthy.meta_instance_id: meta_healthy,
meta_failing.meta_instance_id: meta_failing,
},
instances={iid_h: inst_h, iid_f: inst_f},
runners={
runner_ids_h[0]: RunnerReady(),
runner_ids_f[0]: RunnerFailed(error_message="crash"),
},
topology=_topology("node-a", "node-b"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceRetrying)
assert events[0].instance_id == iid_f
async def test_meta_instance_reconciler_empty_state():
"""MetaInstanceReconciler with no meta_instances should produce no events."""
state = State()
reconciler = MetaInstanceReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 0
# =============================================================================
# 6. Placement error tracking
# =============================================================================
def test_placement_failed_sets_error():
"""MetaInstancePlacementFailed should set placement_error on the MetaInstance."""
meta = _meta_instance()
state = State(meta_instances={meta.meta_instance_id: meta})
event = MetaInstancePlacementFailed(
meta_instance_id=meta.meta_instance_id,
reason="Not enough memory",
)
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.placement_error == "Not enough memory"
def test_instance_created_clears_placement_error():
"""InstanceCreated should clear placement_error on the MetaInstance."""
meta = _meta_instance(placement_error="Not enough memory")
_, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(meta_instances={meta.meta_instance_id: meta})
state = apply(state, IndexedEvent(idx=0, event=InstanceCreated(instance=inst)))
mi = state.meta_instances[meta.meta_instance_id]
assert mi.placement_error is None
def test_placement_error_does_not_increment_failures():
"""Placement failures should only set placement_error, not increment consecutive_failures."""
meta = _meta_instance()
state = State(meta_instances={meta.meta_instance_id: meta})
event = MetaInstancePlacementFailed(
meta_instance_id=meta.meta_instance_id,
reason="No resources",
)
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 0
assert mi.placement_error == "No resources"
# =============================================================================
# 7. State serialization roundtrip
# =============================================================================
def test_state_with_meta_instances_serializes():
"""State with meta_instances should serialize and deserialize correctly."""
meta = _meta_instance(consecutive_failures=2, last_failure_error="test")
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
json_str = state.model_dump_json()
restored = State.model_validate_json(json_str)
assert meta.meta_instance_id in restored.meta_instances
mi = restored.meta_instances[meta.meta_instance_id]
assert mi.model_id == meta.model_id
assert mi.consecutive_failures == 2
assert mi.last_failure_error == "test"
assert iid in restored.instances
assert restored.instances[iid].meta_instance_id == meta.meta_instance_id
# =============================================================================
# 8. MetaInstanceReconciler error handling
# =============================================================================
async def test_meta_instance_reconciler_model_load_error_emits_placement_failed(
monkeypatch: "pytest.MonkeyPatch",
):
"""When ModelCard.load raises, reconciler emits MetaInstancePlacementFailed."""
import exo.master.process_managers.meta_instance as mi_mod
meta = _meta_instance()
topo = _topology("node-a")
state = State(
meta_instances={meta.meta_instance_id: meta},
topology=topo,
)
async def _failing_load(_model_id: ModelId) -> ModelCard:
raise RuntimeError("Network error")
monkeypatch.setattr(
mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_failing_load)})
)
reconciler = MetaInstanceReconciler()
events = await reconciler.reconcile(state)
placement_failed = [e for e in events if isinstance(e, MetaInstancePlacementFailed)]
assert len(placement_failed) == 1
assert "Failed to load model card" in placement_failed[0].reason
assert meta.meta_instance_id == placement_failed[0].meta_instance_id
async def test_meta_instance_reconciler_model_load_error_skips_dedup(
monkeypatch: "pytest.MonkeyPatch",
):
"""When ModelCard.load error matches existing placement_error, no duplicate event."""
import exo.master.process_managers.meta_instance as mi_mod
meta = _meta_instance(placement_error="Failed to load model card: Network error")
topo = _topology("node-a")
state = State(
meta_instances={meta.meta_instance_id: meta},
topology=topo,
)
async def _failing_load(_model_id: ModelId) -> ModelCard:
raise RuntimeError("Network error")
monkeypatch.setattr(
mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_failing_load)})
)
reconciler = MetaInstanceReconciler()
events = await reconciler.reconcile(state)
# Error matches existing placement_error, so no duplicate event emitted
assert len(events) == 0
async def test_meta_instance_reconciler_continues_after_error(
monkeypatch: "pytest.MonkeyPatch",
):
"""Reconciler should continue to next meta-instance after one fails to load."""
import exo.master.process_managers.meta_instance as mi_mod
meta_a = _meta_instance(model_id="org/model-a")
meta_b = _meta_instance(model_id="org/model-b")
topo = _topology("node-a")
state = State(
meta_instances={
meta_a.meta_instance_id: meta_a,
meta_b.meta_instance_id: meta_b,
},
topology=topo,
)
call_count = 0
async def _load_second_fails(model_id: ModelId) -> ModelCard:
nonlocal call_count
call_count += 1
raise RuntimeError(f"Cannot load {model_id}")
monkeypatch.setattr(
mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_load_second_fails)})
)
reconciler = MetaInstanceReconciler()
events = await reconciler.reconcile(state)
# Both meta-instances should have been attempted (not short-circuited)
assert call_count == 2
# Both should have placement failed events
placement_failed = [e for e in events if isinstance(e, MetaInstancePlacementFailed)]
assert len(placement_failed) == 2
# =============================================================================
# 8. Cascade delete with task cancellation
# =============================================================================
def test_cascade_delete_cancels_active_tasks():
"""Deleting a MetaInstance should cancel tasks on backing instances.
Regression test: previously, cascade-deleting backing instances via
DeleteMetaInstance did not emit TaskStatusUpdated(Cancelled) for active
tasks, leaving orphaned task references in state.
"""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
task_id = TaskId()
task = LoadModel(task_id=task_id, instance_id=iid, task_status=TaskStatus.Running)
# Build state with meta-instance, backing instance, and active task
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
tasks={task_id: task},
topology=_topology("node-a"),
)
# Simulate the cascade-delete event sequence produced by main.py:
# 1. MetaInstanceDeleted
# 2. TaskStatusUpdated(Cancelled) for active tasks
# 3. InstanceDeleted
idx = 0
state = apply(
state,
IndexedEvent(
idx=idx,
event=MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id),
),
)
idx += 1
state = apply(
state,
IndexedEvent(
idx=idx,
event=TaskStatusUpdated(task_id=task_id, task_status=TaskStatus.Cancelled),
),
)
idx += 1
state = apply(
state,
IndexedEvent(idx=idx, event=InstanceDeleted(instance_id=iid)),
)
# Verify everything is cleaned up
assert len(state.meta_instances) == 0
assert len(state.instances) == 0
assert state.tasks[task_id].task_status == TaskStatus.Cancelled
def test_cascade_delete_skips_completed_tasks():
"""Cascade delete should only cancel Pending/Running tasks, not completed ones."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
running_task_id = TaskId()
completed_task_id = TaskId()
running_task = LoadModel(
task_id=running_task_id, instance_id=iid, task_status=TaskStatus.Running
)
completed_task = LoadModel(
task_id=completed_task_id, instance_id=iid, task_status=TaskStatus.Complete
)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
tasks={running_task_id: running_task, completed_task_id: completed_task},
topology=_topology("node-a"),
)
# Only the running task should be cancelled — we verify the logic pattern
# by checking which tasks are Pending or Running
active_tasks = [
t
for t in state.tasks.values()
if t.instance_id == iid
and t.task_status in (TaskStatus.Pending, TaskStatus.Running)
]
assert len(active_tasks) == 1
assert active_tasks[0].task_id == running_task_id
+3 -3
View File
@@ -239,7 +239,7 @@ def test_get_transition_events_no_change(instance: Instance):
target_instances = {instance_id: instance}
# act
events = get_transition_events(current_instances, target_instances)
events = get_transition_events(current_instances, target_instances, {})
# assert
assert len(events) == 0
@@ -252,7 +252,7 @@ def test_get_transition_events_create_instance(instance: Instance):
target_instances: dict[InstanceId, Instance] = {instance_id: instance}
# act
events = get_transition_events(current_instances, target_instances)
events = get_transition_events(current_instances, target_instances, {})
# assert
assert len(events) == 1
@@ -266,7 +266,7 @@ def test_get_transition_events_delete_instance(instance: Instance):
target_instances: dict[InstanceId, Instance] = {}
# act
events = get_transition_events(current_instances, target_instances)
events = get_transition_events(current_instances, target_instances, {})
# assert
assert len(events) == 1
+6 -6
View File
@@ -3,10 +3,10 @@ import pytest
from exo.master.placement_utils import (
allocate_layers_proportionally,
filter_cycles_by_memory,
get_largest_cycles,
get_mlx_jaccl_coordinators,
get_shard_assignments,
get_shard_assignments_for_pipeline_parallel,
get_smallest_cycles,
)
from exo.master.tests.conftest import (
create_node_memory,
@@ -143,7 +143,7 @@ def test_filter_multiple_cycles_by_memory():
}
def test_get_smallest_cycles():
def test_get_largest_cycles():
# arrange
node_a_id = NodeId()
node_b_id = NodeId()
@@ -175,12 +175,12 @@ def test_get_smallest_cycles():
cycles = [c for c in topology.get_cycles() if len(c) != 1] # ignore singletons
# act
smallest_cycles = get_smallest_cycles(cycles)
largest_cycles = get_largest_cycles(cycles)
# assert
assert len(smallest_cycles) == 1
assert len(smallest_cycles[0]) == 2
assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id}
assert len(largest_cycles) == 1
assert len(largest_cycles[0]) == 3
assert set(n for n in largest_cycles[0]) == {node_a_id, node_b_id, node_c_id}
@pytest.mark.parametrize(
+742
View File
@@ -0,0 +1,742 @@
from exo.master.process_managers.instance_health import InstanceHealthReconciler
from exo.master.reconcile import (
find_unsatisfied_meta_instances,
instance_connections_healthy,
instance_runners_failed,
instance_satisfies_meta_instance,
)
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.topology import Topology
from exo.shared.types.common import Host, MetaInstanceId, NodeId
from exo.shared.types.events import (
IndexedEvent,
InstanceCreated,
InstanceDeleted,
InstanceRetrying,
MetaInstanceCreated,
MetaInstanceDeleted,
)
from exo.shared.types.memory import Memory
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.state import State
from exo.shared.types.topology import Connection, SocketConnection
from exo.shared.types.worker.instances import (
InstanceId,
MlxRingInstance,
)
from exo.shared.types.worker.runners import (
RunnerFailed,
RunnerId,
RunnerLoading,
RunnerReady,
RunnerShutdown,
ShardAssignments,
)
from exo.shared.types.worker.shards import PipelineShardMetadata
def _model_card(model_id: str = "test-org/test-model") -> ModelCard:
return ModelCard(
model_id=ModelId(model_id),
storage_size=Memory.from_kb(1000),
n_layers=10,
hidden_size=30,
supports_tensor=True,
tasks=[ModelTask.TextGeneration],
)
def _topology(*node_ids: str, connect: bool = True) -> Topology:
"""Build a topology with nodes connected in a bidirectional ring with unique IPs.
Node at index ``i`` gets IP ``10.0.0.{i+1}``. Edges go in both directions
between consecutive nodes (including wrap-around).
"""
t = Topology()
nodes = [NodeId(n) for n in node_ids]
for n in nodes:
t.add_node(n)
if connect and len(nodes) > 1:
for i in range(len(nodes)):
j = (i + 1) % len(nodes)
t.add_connection(
Connection(
source=nodes[i],
sink=nodes[j],
edge=SocketConnection(
sink_multiaddr=Multiaddr(
address=f"/ip4/10.0.0.{j + 1}/tcp/50000"
)
),
)
)
t.add_connection(
Connection(
source=nodes[j],
sink=nodes[i],
edge=SocketConnection(
sink_multiaddr=Multiaddr(
address=f"/ip4/10.0.0.{i + 1}/tcp/50000"
)
),
)
)
return t
def _meta_instance(
model_id: str = "test-org/test-model",
*,
min_nodes: int = 1,
node_ids: list[NodeId] | None = None,
meta_instance_id: MetaInstanceId | None = None,
) -> MetaInstance:
return MetaInstance(
meta_instance_id=meta_instance_id or MetaInstanceId(),
model_id=ModelId(model_id),
min_nodes=min_nodes,
node_ids=node_ids,
)
def _instance(
model_id: str = "test-org/test-model",
node_ids: list[str] | None = None,
instance_id: InstanceId | None = None,
meta_instance_id: MetaInstanceId | None = None,
) -> tuple[InstanceId, MlxRingInstance]:
"""Create a test instance with hosts_by_node matching ``_topology()`` IPs."""
iid = instance_id or InstanceId()
nodes = node_ids or ["node-a"]
n = len(nodes)
mc = _model_card(model_id)
ephemeral_port = 50000
node_to_runner = {NodeId(nd): RunnerId() for nd in nodes}
runner_to_shard = {
runner_id: PipelineShardMetadata(
model_card=mc,
device_rank=i,
world_size=n,
start_layer=0,
end_layer=mc.n_layers,
n_layers=mc.n_layers,
)
for i, runner_id in enumerate(node_to_runner.values())
}
# Build hosts_by_node with IPs matching _topology() convention:
# node at index idx has IP 10.0.0.{idx+1}
hosts_by_node: dict[NodeId, list[Host]] = {}
for r, node_str in enumerate(nodes):
hosts: list[Host] = []
for idx in range(n):
if idx == r:
hosts.append(Host(ip="0.0.0.0", port=ephemeral_port))
elif n > 1 and idx in ((r - 1) % n, (r + 1) % n):
hosts.append(Host(ip=f"10.0.0.{idx + 1}", port=ephemeral_port))
else:
hosts.append(Host(ip="198.51.100.1", port=0))
hosts_by_node[NodeId(node_str)] = hosts
return iid, MlxRingInstance(
instance_id=iid,
shard_assignments=ShardAssignments(
model_id=ModelId(model_id),
runner_to_shard=runner_to_shard,
node_to_runner=node_to_runner,
),
hosts_by_node=hosts_by_node,
ephemeral_port=ephemeral_port,
meta_instance_id=meta_instance_id,
)
# --- instance_satisfies_meta_instance (pure constraint matching) ---
def test_satisfies_matching_model():
meta = _meta_instance()
_, inst = _instance(node_ids=["node-a"])
assert instance_satisfies_meta_instance(meta, inst) is True
def test_not_satisfies_wrong_model():
meta = _meta_instance("test-org/model-a")
_, inst = _instance("test-org/model-b")
assert instance_satisfies_meta_instance(meta, inst) is False
def test_not_satisfies_missing_required_node():
meta = _meta_instance(node_ids=[NodeId("node-c")])
_, inst = _instance(node_ids=["node-a", "node-b"])
assert instance_satisfies_meta_instance(meta, inst) is False
def test_not_satisfies_fewer_than_min_nodes():
meta = _meta_instance(min_nodes=3)
_, inst = _instance(node_ids=["node-a", "node-b"])
assert instance_satisfies_meta_instance(meta, inst) is False
def test_satisfies_with_node_ids_specified():
meta = _meta_instance(node_ids=[NodeId("node-a"), NodeId("node-b")], min_nodes=2)
_, inst = _instance(node_ids=["node-a", "node-b", "node-c"])
assert instance_satisfies_meta_instance(meta, inst) is True
# --- instance_connections_healthy ---
def test_healthy_single_node_present():
_, inst = _instance(node_ids=["node-a"])
topology = _topology("node-a")
assert instance_connections_healthy(inst, topology) is True
def test_unhealthy_single_node_missing():
_, inst = _instance(node_ids=["node-a"])
topology = Topology() # empty
assert instance_connections_healthy(inst, topology) is False
def test_healthy_two_node_ring():
_, inst = _instance(node_ids=["node-a", "node-b"])
topology = _topology("node-a", "node-b")
assert instance_connections_healthy(inst, topology) is True
def test_unhealthy_two_node_edge_removed():
"""Nodes present but edge removed — ring broken."""
_, inst = _instance(node_ids=["node-a", "node-b"])
topology = _topology("node-a", "node-b", connect=False)
assert instance_connections_healthy(inst, topology) is False
def test_unhealthy_two_node_ip_changed():
"""Edge exists but with a different IP than instance was configured with."""
_, inst = _instance(node_ids=["node-a", "node-b"])
# Build topology with different IPs than _instance() expects
topology = Topology()
topology.add_node(NodeId("node-a"))
topology.add_node(NodeId("node-b"))
topology.add_connection(
Connection(
source=NodeId("node-a"),
sink=NodeId("node-b"),
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/192.168.99.99/tcp/50000")
),
)
)
topology.add_connection(
Connection(
source=NodeId("node-b"),
sink=NodeId("node-a"),
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/192.168.99.98/tcp/50000")
),
)
)
assert instance_connections_healthy(inst, topology) is False
def test_healthy_three_node_ring():
_, inst = _instance(node_ids=["node-a", "node-b", "node-c"])
topology = _topology("node-a", "node-b", "node-c")
assert instance_connections_healthy(inst, topology) is True
def test_unhealthy_three_node_one_edge_removed():
"""Remove one edge from a three-node ring — instance unhealthy."""
_, inst = _instance(node_ids=["node-a", "node-b", "node-c"])
# Build topology with one direction of one edge missing
topology = Topology()
nodes = [NodeId("node-a"), NodeId("node-b"), NodeId("node-c")]
for n in nodes:
topology.add_node(n)
# Add all edges except node-a → node-b
topology.add_connection(
Connection(
source=nodes[1],
sink=nodes[0],
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/50000")
),
)
)
topology.add_connection(
Connection(
source=nodes[1],
sink=nodes[2],
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.3/tcp/50000")
),
)
)
topology.add_connection(
Connection(
source=nodes[2],
sink=nodes[1],
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.2/tcp/50000")
),
)
)
topology.add_connection(
Connection(
source=nodes[2],
sink=nodes[0],
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/50000")
),
)
)
topology.add_connection(
Connection(
source=nodes[0],
sink=nodes[2],
edge=SocketConnection(
sink_multiaddr=Multiaddr(address="/ip4/10.0.0.3/tcp/50000")
),
)
)
# Missing: node-a → node-b (ip 10.0.0.2)
assert instance_connections_healthy(inst, topology) is False
def test_unhealthy_node_missing_from_topology():
"""Instance has a node that's not in the topology at all."""
_, inst = _instance(node_ids=["node-a", "node-b"])
topology = _topology("node-a") # node-b not present
assert instance_connections_healthy(inst, topology) is False
def test_healthy_extra_nodes_in_topology():
"""Extra nodes in topology don't affect instance health."""
_, inst = _instance(node_ids=["node-a", "node-b"])
topology = _topology("node-a", "node-b", "node-c")
assert instance_connections_healthy(inst, topology) is True
# --- find_unsatisfied_meta_instances ---
def test_unsatisfied_no_meta_instances():
result = find_unsatisfied_meta_instances({}, {}, Topology())
assert list(result) == []
def test_unsatisfied_one_satisfied():
meta = _meta_instance()
id_a, inst_a = _instance(meta_instance_id=meta.meta_instance_id)
topology = _topology("node-a")
result = find_unsatisfied_meta_instances(
{meta.meta_instance_id: meta},
{id_a: inst_a},
topology,
)
assert list(result) == []
def test_unsatisfied_one_not_satisfied():
meta = _meta_instance("test-org/model-x")
id_a, inst_a = _instance("test-org/model-y")
topology = _topology("node-a")
result = find_unsatisfied_meta_instances(
{meta.meta_instance_id: meta}, {id_a: inst_a}, topology
)
assert list(result) == [meta]
def test_unsatisfied_mix():
meta_satisfied = _meta_instance("test-org/model-a")
meta_unsatisfied = _meta_instance("test-org/model-b")
id_a, inst_a = _instance(
"test-org/model-a", meta_instance_id=meta_satisfied.meta_instance_id
)
topology = _topology("node-a")
result = find_unsatisfied_meta_instances(
{
meta_satisfied.meta_instance_id: meta_satisfied,
meta_unsatisfied.meta_instance_id: meta_unsatisfied,
},
{id_a: inst_a},
topology,
)
assert list(result) == [meta_unsatisfied]
def test_unsatisfied_node_disconnect():
meta = _meta_instance()
id_a, inst_a = _instance(
node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id
)
topology = _topology("node-a") # node-b disconnected
result = find_unsatisfied_meta_instances(
{meta.meta_instance_id: meta},
{id_a: inst_a},
topology,
)
assert list(result) == [meta]
def test_unsatisfied_edge_break():
"""Instance exists but its connections broke — meta-instance becomes unsatisfied."""
meta = _meta_instance()
id_a, inst_a = _instance(
node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id
)
topology = _topology("node-a", "node-b", connect=False) # nodes present, no edges
result = find_unsatisfied_meta_instances(
{meta.meta_instance_id: meta},
{id_a: inst_a},
topology,
)
assert list(result) == [meta]
def test_unsatisfied_idempotent():
meta = _meta_instance("test-org/model-x")
topology = _topology("node-a")
meta_instances = {meta.meta_instance_id: meta}
instances: dict[InstanceId, MlxRingInstance] = {}
result_1 = list(
find_unsatisfied_meta_instances(meta_instances, instances, topology)
)
result_2 = list(
find_unsatisfied_meta_instances(meta_instances, instances, topology)
)
assert result_1 == result_2
def test_unsatisfied_exclusive_binding():
"""Two MetaInstances for the same model: one is bound via meta_instance_id, the other is unsatisfied."""
meta_a = _meta_instance("test-org/model-x")
meta_b = _meta_instance("test-org/model-x")
id_inst, inst = _instance(
"test-org/model-x", meta_instance_id=meta_a.meta_instance_id
)
topology = _topology("node-a")
result = find_unsatisfied_meta_instances(
{
meta_a.meta_instance_id: meta_a,
meta_b.meta_instance_id: meta_b,
},
{id_inst: inst},
topology,
)
assert list(result) == [meta_b]
# --- apply handlers ---
def test_apply_meta_instance_created():
state = State()
meta = _meta_instance()
event = MetaInstanceCreated(meta_instance=meta)
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert meta.meta_instance_id in new_state.meta_instances
assert new_state.meta_instances[meta.meta_instance_id] == meta
def test_apply_meta_instance_deleted():
meta = _meta_instance()
state = State(meta_instances={meta.meta_instance_id: meta})
event = MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id)
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert meta.meta_instance_id not in new_state.meta_instances
def test_apply_meta_instance_deleted_clears_failure_info():
meta = _meta_instance().model_copy(
update={"consecutive_failures": 2, "last_failure_error": "OOM"}
)
state = State(meta_instances={meta.meta_instance_id: meta})
event = MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id)
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert meta.meta_instance_id not in new_state.meta_instances
# --- instance_runners_failed ---
def test_runners_failed_all_failed():
"""All runners in RunnerFailed -> instance is failed."""
_, inst = _instance(node_ids=["node-a", "node-b"])
runners = {
rid: RunnerFailed(error_message="OOM")
for rid in inst.shard_assignments.node_to_runner.values()
}
is_failed, error = instance_runners_failed(inst, runners, {})
assert is_failed is True
assert error is not None
assert "OOM" in error
def test_runners_failed_mixed_failed_shutdown():
"""One Failed + one Shutdown = failed."""
_, inst = _instance(node_ids=["node-a", "node-b"])
runner_ids = list(inst.shard_assignments.node_to_runner.values())
runners = {
runner_ids[0]: RunnerFailed(error_message="crash"),
runner_ids[1]: RunnerShutdown(),
}
is_failed, error = instance_runners_failed(inst, runners, {})
assert is_failed is True
assert error is not None
assert "crash" in error
def test_runners_not_failed_all_shutdown():
"""All Shutdown (graceful) = not a failure."""
_, inst = _instance(node_ids=["node-a"])
runners = {
rid: RunnerShutdown() for rid in inst.shard_assignments.node_to_runner.values()
}
is_failed, _ = instance_runners_failed(inst, runners, {})
assert is_failed is False
def test_runners_not_failed_still_active():
"""Some runners still active = not failed yet."""
_, inst = _instance(node_ids=["node-a", "node-b"])
runner_ids = list(inst.shard_assignments.node_to_runner.values())
runners = {
runner_ids[0]: RunnerFailed(error_message="OOM"),
runner_ids[1]: RunnerLoading(),
}
is_failed, _ = instance_runners_failed(inst, runners, {})
assert is_failed is False
def test_runners_not_failed_no_status():
"""Runner not yet reported = not failed."""
_, inst = _instance(node_ids=["node-a"])
is_failed, _ = instance_runners_failed(inst, {}, {})
assert is_failed is False
def test_runners_not_failed_healthy():
"""Runners in Ready state = not failed."""
_, inst = _instance(node_ids=["node-a"])
runners = {
rid: RunnerReady() for rid in inst.shard_assignments.node_to_runner.values()
}
is_failed, _ = instance_runners_failed(inst, runners, {})
assert is_failed is False
# --- failure tracking in apply_instance_deleted ---
def test_apply_instance_deleted_tracks_failure():
"""InstanceDeleted with failure_error increments meta instance failure count."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
event = InstanceDeleted(instance_id=iid, failure_error="Runner OOM")
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 1
assert mi.last_failure_error == "Runner OOM"
def test_apply_instance_deleted_increments_failure():
"""Subsequent failures increment the counter."""
meta = _meta_instance().model_copy(
update={"consecutive_failures": 2, "last_failure_error": "previous error"}
)
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
event = InstanceDeleted(instance_id=iid, failure_error="new error")
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 3
assert mi.last_failure_error == "new error"
def test_apply_instance_deleted_no_failure_no_tracking():
"""InstanceDeleted without failure_error does not track."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
event = InstanceDeleted(instance_id=iid)
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 0
def test_apply_instance_deleted_orphan_no_tracking():
"""InstanceDeleted for orphan instance (no meta_instance_id) does not track."""
iid, inst = _instance(node_ids=["node-a"])
state = State(instances={iid: inst})
event = InstanceDeleted(instance_id=iid, failure_error="crash")
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert len(new_state.meta_instances) == 0
# --- InstanceRetrying ---
def test_apply_instance_retrying_removes_runners():
"""InstanceRetrying removes the instance's runners from state but keeps the instance."""
meta = _meta_instance()
iid, inst = _instance(
node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id
)
runner_ids = list(inst.shard_assignments.node_to_runner.values())
runners = {
runner_ids[0]: RunnerFailed(error_message="OOM"),
runner_ids[1]: RunnerShutdown(),
}
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
runners=runners,
)
event = InstanceRetrying(
instance_id=iid,
meta_instance_id=meta.meta_instance_id,
failure_error="OOM",
)
new_state = apply(state, IndexedEvent(idx=0, event=event))
# Instance still exists
assert iid in new_state.instances
# Runners removed
assert runner_ids[0] not in new_state.runners
assert runner_ids[1] not in new_state.runners
def test_apply_instance_retrying_increments_failure():
"""InstanceRetrying increments consecutive_failures on the MetaInstance."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
event = InstanceRetrying(
instance_id=iid,
meta_instance_id=meta.meta_instance_id,
failure_error="crash",
)
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 1
assert mi.last_failure_error == "crash"
def test_apply_instance_retrying_skips_missing_runners():
"""InstanceRetrying doesn't assert if runners haven't reported yet."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
# No runners in state at all
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
)
event = InstanceRetrying(
instance_id=iid,
meta_instance_id=meta.meta_instance_id,
failure_error="crash",
)
# Should not raise
new_state = apply(state, IndexedEvent(idx=0, event=event))
assert iid in new_state.instances
def test_apply_instance_created_resets_failure_counter():
"""InstanceCreated resets consecutive_failures but preserves last_failure_error."""
meta = _meta_instance().model_copy(
update={"consecutive_failures": 3, "last_failure_error": "old error"}
)
_, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
state = State(meta_instances={meta.meta_instance_id: meta})
event = InstanceCreated(instance=inst)
new_state = apply(state, IndexedEvent(idx=0, event=event))
mi = new_state.meta_instances[meta.meta_instance_id]
assert mi.consecutive_failures == 0
assert mi.last_failure_error == "old error"
assert mi.placement_error is None
# --- InstanceHealthReconciler retry-vs-delete ---
async def test_health_reconciler_retries_when_under_limit():
"""InstanceHealthReconciler emits InstanceRetrying when consecutive_failures < 3."""
meta = _meta_instance()
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
runner_ids = list(inst.shard_assignments.node_to_runner.values())
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
runners={runner_ids[0]: RunnerFailed(error_message="OOM")},
topology=_topology("node-a"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceRetrying)
assert events[0].instance_id == iid
assert events[0].meta_instance_id == meta.meta_instance_id
async def test_health_reconciler_deletes_when_limit_reached():
"""InstanceHealthReconciler emits InstanceDeleted when consecutive_failures >= 3."""
meta = _meta_instance().model_copy(update={"consecutive_failures": 3})
iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id)
runner_ids = list(inst.shard_assignments.node_to_runner.values())
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
runners={runner_ids[0]: RunnerFailed(error_message="OOM")},
topology=_topology("node-a"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceDeleted)
async def test_health_reconciler_deletes_without_meta_instance():
"""Instances without a MetaInstance are deleted immediately on runner failure."""
iid, inst = _instance(node_ids=["node-a"])
runner_ids = list(inst.shard_assignments.node_to_runner.values())
state = State(
instances={iid: inst},
runners={runner_ids[0]: RunnerFailed(error_message="crash")},
topology=_topology("node-a"),
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceDeleted)
async def test_health_reconciler_network_failure_always_deletes():
"""Network failure always triggers InstanceDeleted regardless of retry count."""
meta = _meta_instance()
iid, inst = _instance(
node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id
)
state = State(
meta_instances={meta.meta_instance_id: meta},
instances={iid: inst},
topology=_topology("node-a"), # node-b missing
)
reconciler = InstanceHealthReconciler()
events = await reconciler.reconcile(state)
assert len(events) == 1
assert isinstance(events[0], InstanceDeleted)
assert events[0].failure_error == "Network connection lost"
+9 -5
View File
@@ -101,8 +101,12 @@ class TopicRouter[T: CamelCaseModel]:
class Router:
@classmethod
def create(cls, identity: Keypair) -> "Router":
return cls(handle=NetworkingHandle(identity))
def create(
cls,
identity: Keypair,
bootstrap_peers: list[str] | None = None,
) -> "Router":
return cls(handle=NetworkingHandle(identity, bootstrap_peers or []))
def __init__(self, handle: NetworkingHandle):
self.topic_routers: dict[str, TopicRouter[CamelCaseModel]] = {}
@@ -207,10 +211,10 @@ class Router:
try:
logger.trace(f"Sending message on {topic} with payload {data}")
await self._net.gossipsub_publish(topic, data)
# As a hack, this also catches AllQueuesFull
# Need to fix that ASAP.
except (NoPeersSubscribedToTopicError, AllQueuesFullError):
except NoPeersSubscribedToTopicError:
pass
except AllQueuesFullError:
logger.warning(f"All peer queues full, dropping message on {topic}")
def get_node_id_keypair(
+147 -4
View File
@@ -4,7 +4,7 @@ from datetime import datetime
from loguru import logger
from exo.shared.types.common import NodeId
from exo.shared.types.common import MetaInstanceId, NodeId
from exo.shared.types.events import (
ChunkGenerated,
Event,
@@ -12,6 +12,12 @@ from exo.shared.types.events import (
InputChunkReceived,
InstanceCreated,
InstanceDeleted,
InstanceRetrying,
JacclSideChannelData,
JacclSideChannelGathered,
MetaInstanceCreated,
MetaInstanceDeleted,
MetaInstancePlacementFailed,
NodeDownloadProgress,
NodeGatheredInfo,
NodeTimedOut,
@@ -28,9 +34,11 @@ from exo.shared.types.events import (
TracesCollected,
TracesMerged,
)
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.profiling import (
NodeIdentity,
NodeNetworkInfo,
NodeRdmaCtlStatus,
NodeThunderboltInfo,
ThunderboltBridgeStatus,
)
@@ -47,7 +55,9 @@ from exo.utils.info_gatherer.info_gatherer import (
MemoryUsage,
MiscData,
NodeConfig,
NodeDiskUsage,
NodeNetworkInterfaces,
RdmaCtlStatus,
StaticNodeInformation,
ThunderboltBridgeInfo,
)
@@ -63,12 +73,22 @@ def event_apply(event: Event, state: State) -> State:
| InputChunkReceived()
| TracesCollected()
| TracesMerged()
| JacclSideChannelData()
| JacclSideChannelGathered()
): # Pass-through events that don't modify state
return state
case InstanceCreated():
return apply_instance_created(event, state)
case InstanceDeleted():
return apply_instance_deleted(event, state)
case InstanceRetrying():
return apply_instance_retrying(event, state)
case MetaInstanceCreated():
return apply_meta_instance_created(event, state)
case MetaInstanceDeleted():
return apply_meta_instance_deleted(event, state)
case MetaInstancePlacementFailed():
return apply_meta_instance_placement_failed(event, state)
case NodeTimedOut():
return apply_node_timed_out(event, state)
case NodeDownloadProgress():
@@ -171,20 +191,123 @@ def apply_task_failed(event: TaskFailed, state: State) -> State:
return state.model_copy(update={"tasks": new_tasks})
def _update_meta_instance(
state: State, mid: MetaInstanceId, **fields: object
) -> Mapping[MetaInstanceId, MetaInstance]:
mi = state.meta_instances[mid]
return {**state.meta_instances, mid: mi.model_copy(update=fields)}
def apply_instance_created(event: InstanceCreated, state: State) -> State:
instance = event.instance
new_instances: Mapping[InstanceId, Instance] = {
**state.instances,
instance.instance_id: instance,
}
return state.model_copy(update={"instances": new_instances})
update: dict[str, object] = {"instances": new_instances}
# Reset failure tracking when a new instance is created for a meta-instance
if instance.meta_instance_id and instance.meta_instance_id in state.meta_instances:
mi = state.meta_instances[instance.meta_instance_id]
if mi.placement_error is not None or mi.consecutive_failures > 0:
update["meta_instances"] = _update_meta_instance(
state,
instance.meta_instance_id,
placement_error=None,
consecutive_failures=0,
)
return state.model_copy(update=update)
def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
deleted_instance = state.instances.get(event.instance_id)
new_instances: Mapping[InstanceId, Instance] = {
iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
}
return state.model_copy(update={"instances": new_instances})
update: dict[str, object] = {"instances": new_instances}
# Track failure on the MetaInstance itself
if (
event.failure_error
and deleted_instance
and deleted_instance.meta_instance_id
and deleted_instance.meta_instance_id in state.meta_instances
):
mid = deleted_instance.meta_instance_id
mi = state.meta_instances[mid]
update["meta_instances"] = {
**state.meta_instances,
mid: mi.model_copy(
update={
"consecutive_failures": mi.consecutive_failures + 1,
"last_failure_error": event.failure_error,
}
),
}
return state.model_copy(update=update)
def apply_instance_retrying(event: InstanceRetrying, state: State) -> State:
"""Runners failed but retry limit not reached — remove runners, keep instance."""
instance = state.instances.get(event.instance_id)
if instance is None:
# Instance was already deleted (e.g. cascade from DeleteMetaInstance).
# The InstanceDeleted handler already incremented consecutive_failures
# on the MetaInstance, so skipping here avoids double-counting.
return state
# Remove all runners belonging to this instance from state
runner_ids_to_remove = set(instance.shard_assignments.node_to_runner.values())
new_runners: Mapping[RunnerId, RunnerStatus] = {
rid: rs for rid, rs in state.runners.items() if rid not in runner_ids_to_remove
}
update: dict[str, object] = {"runners": new_runners}
# Increment failure count on the MetaInstance
if event.meta_instance_id in state.meta_instances:
update["meta_instances"] = _update_meta_instance(
state,
event.meta_instance_id,
consecutive_failures=state.meta_instances[
event.meta_instance_id
].consecutive_failures
+ 1,
last_failure_error=event.failure_error,
)
return state.model_copy(update=update)
def apply_meta_instance_created(event: MetaInstanceCreated, state: State) -> State:
new_meta: Mapping[MetaInstanceId, MetaInstance] = {
**state.meta_instances,
event.meta_instance.meta_instance_id: event.meta_instance,
}
return state.model_copy(update={"meta_instances": new_meta})
def apply_meta_instance_deleted(event: MetaInstanceDeleted, state: State) -> State:
new_meta: Mapping[MetaInstanceId, MetaInstance] = {
mid: mi
for mid, mi in state.meta_instances.items()
if mid != event.meta_instance_id
}
return state.model_copy(update={"meta_instances": new_meta})
def apply_meta_instance_placement_failed(
event: MetaInstancePlacementFailed, state: State
) -> State:
if event.meta_instance_id not in state.meta_instances:
return state
return state.model_copy(
update={
"meta_instances": _update_meta_instance(
state, event.meta_instance_id, placement_error=event.reason
)
}
)
def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
@@ -223,6 +346,9 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
node_memory = {
key: value for key, value in state.node_memory.items() if key != event.node_id
}
node_disk = {
key: value for key, value in state.node_disk.items() if key != event.node_id
}
node_system = {
key: value for key, value in state.node_system.items() if key != event.node_id
}
@@ -239,6 +365,9 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
for key, value in state.node_thunderbolt_bridge.items()
if key != event.node_id
}
node_rdma_ctl = {
key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
}
# Only recompute cycles if the leaving node had TB bridge enabled
leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
leaving_node_had_tb_enabled = (
@@ -256,10 +385,12 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
"last_seen": last_seen,
"node_identities": node_identities,
"node_memory": node_memory,
"node_disk": node_disk,
"node_system": node_system,
"node_network": node_network,
"node_thunderbolt": node_thunderbolt,
"node_thunderbolt_bridge": node_thunderbolt_bridge,
"node_rdma_ctl": node_rdma_ctl,
"thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
}
)
@@ -288,6 +419,8 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
case MemoryUsage():
update["node_memory"] = {**state.node_memory, event.node_id: info}
case NodeDiskUsage():
update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage}
case NodeConfig():
pass
case MiscData():
@@ -302,7 +435,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
case StaticNodeInformation():
current_identity = state.node_identities.get(event.node_id, NodeIdentity())
new_identity = current_identity.model_copy(
update={"model_id": info.model, "chip_id": info.chip}
update={
"model_id": info.model,
"chip_id": info.chip,
"os_version": info.os_version,
"os_build_version": info.os_build_version,
}
)
update["node_identities"] = {
**state.node_identities,
@@ -354,6 +492,11 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
new_tb_bridge, state.node_network
)
)
case RdmaCtlStatus():
update["node_rdma_ctl"] = {
**state.node_rdma_ctl,
event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
}
return state.model_copy(update=update)
+3 -1
View File
@@ -43,7 +43,8 @@ DASHBOARD_DIR = (
)
# Log files (data/logs or cache)
EXO_LOG = EXO_CACHE_HOME / "exo.log"
EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
EXO_LOG = EXO_LOG_DIR / "exo.log"
EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
# Identity (config)
@@ -60,6 +61,7 @@ EXO_MAX_CHUNK_SIZE = 512 * 1024
EXO_CUSTOM_MODEL_CARDS_DIR = EXO_DATA_HOME / "custom_model_cards"
EXO_EVENT_LOG_DIR = EXO_DATA_HOME / "event_log"
EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
EXO_TRACING_CACHE_DIR = EXO_CACHE_HOME / "traces"
+23 -1
View File
@@ -1,11 +1,30 @@
import logging
import sys
from collections.abc import Iterator
from pathlib import Path
import zstandard
from hypercorn import Config
from hypercorn.logging import Logger as HypercornLogger
from loguru import logger
_MAX_LOG_ARCHIVES = 5
def _zstd_compress(filepath: str) -> None:
source = Path(filepath)
dest = source.with_suffix(source.suffix + ".zst")
cctx = zstandard.ZstdCompressor()
with open(source, "rb") as f_in, open(dest, "wb") as f_out:
cctx.copy_stream(f_in, f_out)
source.unlink()
def _once_then_never() -> Iterator[bool]:
yield True
while True:
yield False
class InterceptLogger(HypercornLogger):
def __init__(self, config: Config):
@@ -53,13 +72,16 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
enqueue=True,
)
if log_file:
rotate_once = _once_then_never()
logger.add(
log_file,
format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}",
level="INFO",
colorize=False,
enqueue=True,
rotation="1 week",
rotation=lambda _, __: next(rotate_once),
retention=_MAX_LOG_ARCHIVES,
compression=_zstd_compress,
)
+13 -6
View File
@@ -30,11 +30,11 @@ from exo.utils.pydantic_ext import CamelCaseModel
# kinda ugly...
# TODO: load search path from config.toml
_custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
_csp = [Path(RESOURCES_DIR) / "inference_model_cards", _custom_cards_dir]
if EXO_ENABLE_IMAGE_MODELS:
_csp.append(Path(RESOURCES_DIR) / "image_model_cards")
CARD_SEARCH_PATH = _csp
CARD_SEARCH_PATH = [
Path(RESOURCES_DIR) / "inference_model_cards",
Path(RESOURCES_DIR) / "image_model_cards",
_custom_cards_dir,
]
_card_cache: dict[ModelId, "ModelCard"] = {}
@@ -49,10 +49,16 @@ async def _refresh_card_cache():
pass
def _is_image_card(card: "ModelCard") -> bool:
return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
async def get_model_cards() -> list["ModelCard"]:
if len(_card_cache) == 0:
await _refresh_card_cache()
return list(_card_cache.values())
if EXO_ENABLE_IMAGE_MODELS:
return list(_card_cache.values())
return [c for c in _card_cache.values() if not _is_image_card(c)]
class ModelTask(str, Enum):
@@ -183,6 +189,7 @@ class ConfigData(BaseModel):
["MiniMaxM2ForCausalLM"],
["LlamaForCausalLM"],
["GptOssForCausalLM"],
["Step3p5ForCausalLM"],
]
@model_validator(mode="before")
+25 -10
View File
@@ -1,12 +1,12 @@
import time
from collections.abc import Generator
from typing import Annotated, Any, Literal
from uuid import uuid4
from pydantic import BaseModel, Field, field_validator
from pydantic_core import PydanticUseDefault
from pydantic import BaseModel, Field
from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.common import CommandId, MetaInstanceId, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -60,6 +60,7 @@ class ChatCompletionMessageText(BaseModel):
class ToolCallItem(BaseModel):
id: str = Field(default_factory=lambda: str(uuid4()))
name: str
arguments: str
@@ -197,6 +198,7 @@ class ChatCompletionRequest(BaseModel):
top_p: float | None = None
top_k: int | None = None
tools: list[dict[str, Any]] | None = None
enable_thinking: bool | None = None
tool_choice: str | dict[str, Any] | None = None
parallel_tool_calls: bool | None = None
user: str | None = None
@@ -225,13 +227,6 @@ class PlaceInstanceParams(BaseModel):
instance_meta: InstanceMeta = InstanceMeta.MlxRing
min_nodes: int = 1
@field_validator("sharding", "instance_meta", mode="plain")
@classmethod
def use_default(cls, v: object):
if not v or not isinstance(v, (Sharding, InstanceMeta)):
raise PydanticUseDefault()
return v
class CreateInstanceParams(BaseModel):
instance: Instance
@@ -267,6 +262,26 @@ class DeleteInstanceResponse(BaseModel):
instance_id: InstanceId
class CreateMetaInstanceParams(BaseModel):
model_id: ModelId
sharding: Sharding = Sharding.Pipeline
instance_meta: InstanceMeta = InstanceMeta.MlxRing
min_nodes: int = 1
node_ids: list[NodeId] | None = None
class CreateMetaInstanceResponse(BaseModel):
message: str
command_id: CommandId
meta_instance_id: MetaInstanceId
class DeleteMetaInstanceResponse(BaseModel):
message: str
command_id: CommandId
meta_instance_id: MetaInstanceId
class AdvancedImageParams(BaseModel):
seed: Annotated[int, Field(ge=0)] | None = None
num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None
+17 -1
View File
@@ -6,7 +6,8 @@ from exo.shared.types.api import (
ImageGenerationTaskParams,
)
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.common import CommandId, MetaInstanceId, NodeId
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -48,6 +49,18 @@ class DeleteInstance(BaseCommand):
instance_id: InstanceId
class CreateMetaInstance(BaseCommand):
meta_instance: MetaInstance
class DeleteMetaInstance(BaseCommand):
meta_instance_id: MetaInstanceId
class TaskCancelled(BaseCommand):
cancelled_command_id: CommandId
class TaskFinished(BaseCommand):
finished_command_id: CommandId
@@ -89,6 +102,9 @@ Command = (
| PlaceInstance
| CreateInstance
| DeleteInstance
| CreateMetaInstance
| DeleteMetaInstance
| TaskCancelled
| TaskFinished
| SendInputChunk
)
+4
View File
@@ -42,6 +42,10 @@ class CommandId(Id):
pass
class MetaInstanceId(Id):
"""Identifier for a MetaInstance."""
class Host(CamelCaseModel):
ip: str
port: int
+77 -3
View File
@@ -1,11 +1,14 @@
import base64
from collections.abc import Mapping
from datetime import datetime
from typing import final
from typing import Annotated, final
from pydantic import Field
from pydantic import BeforeValidator, Field, PlainSerializer
from exo.shared.topology import Connection
from exo.shared.types.chunks import GenerationChunk, InputImageChunk
from exo.shared.types.common import CommandId, Id, NodeId, SessionId
from exo.shared.types.common import CommandId, Id, MetaInstanceId, NodeId, SessionId
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import DownloadProgress
from exo.shared.types.worker.instances import Instance, InstanceId
@@ -14,6 +17,28 @@ from exo.utils.info_gatherer.info_gatherer import GatheredInfo
from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel
def _decode_base64_bytes(v: bytes | str) -> bytes:
if isinstance(v, bytes):
return v
return base64.b64decode(v)
def _encode_base64_bytes(v: bytes) -> str:
return base64.b64encode(v).decode("ascii")
Base64Bytes = Annotated[
bytes,
BeforeValidator(_decode_base64_bytes),
PlainSerializer(_encode_base64_bytes, return_type=str),
]
"""bytes that serialize to/from base64 strings in JSON.
Needed because TaggedModel's wrap validator converts JSON→Python validation
context, which breaks strict-mode bytes deserialization from JSON strings.
"""
class EventId(Id):
"""
Newtype around `ID`
@@ -66,6 +91,30 @@ class InstanceCreated(BaseEvent):
class InstanceDeleted(BaseEvent):
instance_id: InstanceId
failure_error: str | None = None
class MetaInstanceCreated(BaseEvent):
meta_instance: MetaInstance
class MetaInstanceDeleted(BaseEvent):
meta_instance_id: MetaInstanceId
@final
class MetaInstancePlacementFailed(BaseEvent):
meta_instance_id: MetaInstanceId
reason: str
@final
class InstanceRetrying(BaseEvent):
"""Runners failed but retry count is below the limit — restart runners, keep instance."""
instance_id: InstanceId
meta_instance_id: MetaInstanceId
failure_error: str
class RunnerStatusUpdated(BaseEvent):
@@ -132,6 +181,25 @@ class TracesMerged(BaseEvent):
traces: list[TraceEventData]
@final
class JacclSideChannelData(BaseEvent):
"""A runner's local contribution to a JACCL SideChannel all_gather round."""
instance_id: InstanceId
runner_id: RunnerId
sequence: int
data: Base64Bytes
@final
class JacclSideChannelGathered(BaseEvent):
"""Gathered result of a JACCL SideChannel all_gather round."""
instance_id: InstanceId
sequence: int
gathered_data: Mapping[RunnerId, Base64Bytes]
Event = (
TestEvent
| TaskCreated
@@ -141,6 +209,10 @@ Event = (
| TaskAcknowledged
| InstanceCreated
| InstanceDeleted
| InstanceRetrying
| MetaInstanceCreated
| MetaInstanceDeleted
| MetaInstancePlacementFailed
| RunnerStatusUpdated
| RunnerDeleted
| NodeTimedOut
@@ -152,6 +224,8 @@ Event = (
| TopologyEdgeDeleted
| TracesCollected
| TracesMerged
| JacclSideChannelData
| JacclSideChannelGathered
)
+25
View File
@@ -0,0 +1,25 @@
from typing import final
from pydantic import Field
from exo.shared.models.model_cards import ModelId
from exo.shared.types.common import MetaInstanceId, NodeId
from exo.shared.types.worker.instances import InstanceMeta
from exo.shared.types.worker.shards import Sharding
from exo.utils.pydantic_ext import FrozenModel
@final
class MetaInstance(FrozenModel):
"""Declarative constraint: ensure an instance matching these parameters always exists."""
meta_instance_id: MetaInstanceId = Field(default_factory=MetaInstanceId)
model_id: ModelId
sharding: Sharding = Sharding.Pipeline
instance_meta: InstanceMeta = InstanceMeta.MlxRing
min_nodes: int = 1
node_ids: list[NodeId] | None = None
# Failure tracking
placement_error: str | None = None
consecutive_failures: int = 0
last_failure_error: str | None = None
+26
View File
@@ -1,4 +1,6 @@
import shutil
from collections.abc import Sequence
from pathlib import Path
from typing import Literal, Self
import psutil
@@ -38,6 +40,22 @@ class MemoryUsage(CamelCaseModel):
)
class DiskUsage(CamelCaseModel):
"""Disk space usage for the models directory."""
total: Memory
available: Memory
@classmethod
def from_path(cls, path: Path) -> Self:
"""Get disk usage stats for the partition containing path."""
total, _used, free = shutil.disk_usage(path)
return cls(
total=Memory.from_bytes(total),
available=Memory.from_bytes(free),
)
class SystemPerformanceProfile(CamelCaseModel):
# TODO: flops_fp16: float
@@ -63,6 +81,8 @@ class NodeIdentity(CamelCaseModel):
model_id: str = "Unknown"
chip_id: str = "Unknown"
friendly_name: str = "Unknown"
os_version: str = "Unknown"
os_build_version: str = "Unknown"
class NodeNetworkInfo(CamelCaseModel):
@@ -77,6 +97,12 @@ class NodeThunderboltInfo(CamelCaseModel):
interfaces: Sequence[ThunderboltIdentifier] = []
class NodeRdmaCtlStatus(CamelCaseModel):
"""Whether RDMA is enabled on this node (via rdma_ctl)."""
enabled: bool
class ThunderboltBridgeStatus(CamelCaseModel):
"""Whether the Thunderbolt Bridge network service is enabled on this node."""
+7 -1
View File
@@ -6,11 +6,14 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator
from pydantic.alias_generators import to_camel
from exo.shared.topology import Topology, TopologySnapshot
from exo.shared.types.common import NodeId
from exo.shared.types.common import MetaInstanceId, NodeId
from exo.shared.types.meta_instance import MetaInstance
from exo.shared.types.profiling import (
DiskUsage,
MemoryUsage,
NodeIdentity,
NodeNetworkInfo,
NodeRdmaCtlStatus,
NodeThunderboltInfo,
SystemPerformanceProfile,
ThunderboltBridgeStatus,
@@ -39,6 +42,7 @@ class State(CamelCaseModel):
arbitrary_types_allowed=True,
)
instances: Mapping[InstanceId, Instance] = {}
meta_instances: Mapping[MetaInstanceId, MetaInstance] = {}
runners: Mapping[RunnerId, RunnerStatus] = {}
downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
tasks: Mapping[TaskId, Task] = {}
@@ -49,10 +53,12 @@ class State(CamelCaseModel):
# Granular node state mappings (update independently at different frequencies)
node_identities: Mapping[NodeId, NodeIdentity] = {}
node_memory: Mapping[NodeId, MemoryUsage] = {}
node_disk: Mapping[NodeId, DiskUsage] = {}
node_system: Mapping[NodeId, SystemPerformanceProfile] = {}
node_network: Mapping[NodeId, NodeNetworkInfo] = {}
node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
# Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
+7
View File
@@ -24,6 +24,7 @@ class TaskStatus(str, Enum):
Complete = "Complete"
TimedOut = "TimedOut"
Failed = "Failed"
Cancelled = "Cancelled"
class BaseTask(TaggedModel):
@@ -60,6 +61,11 @@ class TextGeneration(BaseTask): # emitted by Master
error_message: str | None = Field(default=None)
class CancelTask(BaseTask):
cancelled_task_id: TaskId
runner_id: RunnerId
class ImageGeneration(BaseTask): # emitted by Master
command_id: CommandId
task_params: ImageGenerationTaskParams
@@ -87,6 +93,7 @@ Task = (
| LoadModel
| StartWarmup
| TextGeneration
| CancelTask
| ImageGeneration
| ImageEdits
| Shutdown

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