Merge remote-tracking branch 'origin/main' into alexcheema/speculative-decoding

# Conflicts:
#	src/exo/shared/types/commands.py
This commit is contained in:
Alex Cheema
2026-02-13 05:46:52 -08:00
135 changed files with 6753 additions and 2292 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"
+6
View File
@@ -32,3 +32,9 @@ dashboard/.svelte-kit/
# host config snapshots
hosts_*.json
.swp
# bench files
bench/**/*.json
# tmp
tmp/models
@@ -0,0 +1,7 @@
"""
This type stub file was generated by pyright.
"""
from mflux.models.flux.variants.kontext.flux_kontext import Flux1Kontext
__all__ = ["Flux1Kontext"]
@@ -0,0 +1,49 @@
"""
This type stub file was generated by pyright.
"""
from pathlib import Path
from typing import Any
from mlx import nn
from mflux.models.common.config.model_config import ModelConfig
from mflux.models.flux.model.flux_text_encoder.clip_encoder.clip_encoder import (
CLIPEncoder,
)
from mflux.models.flux.model.flux_text_encoder.t5_encoder.t5_encoder import T5Encoder
from mflux.models.flux.model.flux_transformer.transformer import Transformer
from mflux.models.flux.model.flux_vae.vae import VAE
from mflux.utils.generated_image import GeneratedImage
class Flux1Kontext(nn.Module):
vae: VAE
transformer: Transformer
t5_text_encoder: T5Encoder
clip_text_encoder: CLIPEncoder
bits: int | None
lora_paths: list[str] | None
lora_scales: list[float] | None
prompt_cache: dict[str, Any]
tokenizers: dict[str, Any]
def __init__(
self,
quantize: int | None = ...,
model_path: str | None = ...,
lora_paths: list[str] | None = ...,
lora_scales: list[float] | None = ...,
model_config: ModelConfig = ...,
) -> None: ...
def generate_image(
self,
seed: int,
prompt: str,
num_inference_steps: int = ...,
height: int = ...,
width: int = ...,
guidance: float = ...,
image_path: Path | str | None = ...,
image_strength: float | None = ...,
scheduler: str = ...,
) -> GeneratedImage: ...
@@ -0,0 +1,16 @@
"""
This type stub file was generated by pyright.
"""
import mlx.core as mx
from mflux.models.flux.model.flux_vae.vae import VAE
class KontextUtil:
@staticmethod
def create_image_conditioning_latents(
vae: VAE,
height: int,
width: int,
image_path: str,
) -> tuple[mx.array, mx.array]: ...
+3 -3
View File
@@ -1139,7 +1139,7 @@ class array:
) -> array:
"""See :func:`flatten`."""
def reshape(self, *shape, stream: Stream | Device | None = ...) -> array:
def reshape(self, *shape: int, stream: Stream | Device | None = ...) -> array:
"""
Equivalent to :func:`reshape` but the shape can be passed either as a
:obj:`tuple` or as separate arguments.
@@ -1222,7 +1222,7 @@ class array:
) -> array:
"""See :func:`swapaxes`."""
def transpose(self, *axes, stream: Stream | Device | None = ...) -> array:
def transpose(self, *axes: int, stream: Stream | Device | None = ...) -> array:
"""
Equivalent to :func:`transpose` but the axes can be passed either as
a tuple or as separate arguments.
@@ -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:
@@ -30,6 +30,9 @@ class Conv1d(Module):
bias (bool, optional): If ``True`` add a learnable bias to the output.
Default: ``True``
"""
weight: mx.array
groups: int
def __init__(
self,
in_channels: int,
+5 -1
View File
@@ -11,7 +11,10 @@ import mlx.core as mx
class Cache(Protocol):
keys: mx.array
values: mx.array
def update_and_fetch(self, keys: mx.array, values: mx.array) -> None: ...
offset: int
def update_and_fetch(
self, keys: mx.array, values: mx.array
) -> tuple[mx.array, mx.array]: ...
@property
def state(self) -> tuple[mx.array, mx.array]: ...
@state.setter
@@ -87,6 +90,7 @@ def create_attention_mask(
class _BaseCache(Cache):
keys: mx.array
values: mx.array
offset: int
@property
def state(self) -> tuple[mx.array, mx.array]: ...
@state.setter
+5 -1
View File
@@ -5,6 +5,7 @@ from typing import Any, Dict, Optional
import mlx.core as mx
import mlx.nn as nn
from mlx_lm.models.mla import MultiLinear
from .base import BaseModelArgs
from .switch_layers import SwitchGLU
@@ -60,7 +61,10 @@ class DeepseekV3Attention(nn.Module):
q_b_proj: nn.Linear
kv_a_proj_with_mqa: nn.Linear
kv_a_layernorm: nn.RMSNorm
kv_b_proj: nn.Linear
# kv_b_proj: nn.Linear
embed_q: MultiLinear
unembed_out: MultiLinear
o_proj: nn.Linear
rope: Any
+153
View File
@@ -0,0 +1,153 @@
from dataclasses import dataclass
from typing import Any, Dict, 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
vocab_size: int
hidden_size: int
intermediate_size: int
max_position_embeddings: int
moe_intermediate_size: int
norm_topk_prob: bool
num_attention_heads: int
n_group: int
head_dim: int
topk_group: int
n_shared_experts: int
n_routed_experts: int
routed_scaling_factor: float
num_experts_per_tok: int
first_k_dense_replace: int
num_hidden_layers: int
num_key_value_heads: int
rms_norm_eps: float
rope_theta: float
rope_scaling: Optional[Dict[str, Any]]
use_qk_norm: bool
tie_word_embeddings: bool
attention_bias: bool
partial_rotary_factor: float
scoring_func: str
topk_method: str
class Attention(nn.Module):
n_heads: int
n_kv_heads: int
scale: float
q_proj: nn.Linear
k_proj: nn.Linear
v_proj: nn.Linear
o_proj: nn.Linear
use_qk_norm: bool
q_norm: nn.RMSNorm
k_norm: nn.RMSNorm
rope: nn.RoPE
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class MLP(nn.Module):
config: ModelArgs
hidden_size: int
intermediate_size: int
gate_proj: nn.Linear
up_proj: nn.Linear
down_proj: nn.Linear
def __init__(
self,
config: ModelArgs,
hidden_size: Optional[int] = None,
intermediate_size: Optional[int] = None,
) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class MoEGate(nn.Module):
config: ModelArgs
top_k: int
norm_topk_prob: bool
n_routed_experts: int
routed_scaling_factor: float
n_group: int
topk_group: int
weight: mx.array
e_score_correction_bias: mx.array
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
class MoE(nn.Module):
config: ModelArgs
num_experts_per_tok: int
switch_mlp: SwitchGLU
gate: MoEGate
shared_experts: MLP
sharding_group: Optional[mx.distributed.Group]
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class DecoderLayer(nn.Module):
self_attn: Attention
mlp: MLP | MoE
input_layernorm: nn.RMSNorm
post_attention_layernorm: nn.RMSNorm
def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class LanguageModel(nn.Module):
vocab_size: int
embed_tokens: nn.Embedding
layers: list[DecoderLayer]
norm: nn.RMSNorm
pipeline_rank: int
pipeline_size: int
start_idx: int
end_idx: Optional[int]
num_layers: int
def __init__(self, config: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
@property
def pipeline_layers(self) -> list[DecoderLayer]: ...
class Model(nn.Module):
args: ModelArgs
model_type: str
model: LanguageModel
lm_head: nn.Linear
def __init__(self, config: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[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[DecoderLayer]: ...
@property
def cast_predicate(self) -> Any: ...
+114
View File
@@ -0,0 +1,114 @@
"""Type stubs for mlx_lm.models.qwen3_next"""
from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .switch_layers import SwitchGLU
class Qwen3NextMLP(nn.Module):
gate_proj: nn.Linear
down_proj: nn.Linear
up_proj: nn.Linear
def __init__(self, dim: int, hidden_dim: int) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class Qwen3NextGatedDeltaNet(nn.Module):
hidden_size: int
num_v_heads: int
num_k_heads: int
head_k_dim: int
head_v_dim: int
key_dim: int
value_dim: int
conv_kernel_size: int
conv_dim: int
conv1d: nn.Conv1d
in_proj_qkvz: nn.Linear
in_proj_ba: nn.Linear
dt_bias: mx.array
A_log: mx.array
out_proj: nn.Linear
def __init__(self, config: Any) -> None: ...
def __call__(
self,
inputs: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class Qwen3NextAttention(nn.Module):
num_attention_heads: int
num_key_value_heads: int
head_dim: int
scale: float
q_proj: nn.Linear
k_proj: nn.Linear
v_proj: nn.Linear
o_proj: nn.Linear
def __init__(self, args: Any) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class Qwen3NextSparseMoeBlock(nn.Module):
norm_topk_prob: bool
num_experts: int
top_k: int
gate: nn.Linear
switch_mlp: SwitchGLU
shared_expert: Qwen3NextMLP
shared_expert_gate: nn.Linear
def __init__(self, args: Any) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class Qwen3NextDecoderLayer(nn.Module):
is_linear: bool
linear_attn: Qwen3NextGatedDeltaNet
self_attn: Qwen3NextAttention
input_layernorm: nn.RMSNorm
post_attention_layernorm: nn.RMSNorm
mlp: Qwen3NextMLP | Qwen3NextSparseMoeBlock
def __init__(self, args: Any, layer_idx: int) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class Qwen3NextModel(nn.Module):
embed_tokens: nn.Embedding
layers: list[Qwen3NextDecoderLayer]
norm: nn.RMSNorm
def __init__(self, args: Any) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
class Model(nn.Module):
model_type: str
model: Qwen3NextModel
lm_head: nn.Linear
def __init__(self, args: Any) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
@property
def layers(self) -> list[Qwen3NextDecoderLayer]: ...
+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: ...
+4
View File
@@ -113,6 +113,10 @@ class TokenizerWrapper:
bos_token: str | None
vocab_size: int
all_special_tokens: list[str]
think_start: str | None
think_end: str | None
think_start_id: int | None
think_end_id: int | None
def __init__(
self,
+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"
-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 [] }
+93 -43
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
@@ -388,6 +393,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",
@@ -431,7 +496,12 @@ def main() -> int:
ap.add_argument(
"--skip-pipeline-jaccl",
action="store_true",
help="Pipeline jaccl is often pointless, skip by default",
help="Skip pipeline+jaccl placements, as it's often pointless.",
)
ap.add_argument(
"--skip-tensor-ring",
action="store_true",
help="Skip tensor+ring placements, as it's so slow.",
)
ap.add_argument(
"--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
@@ -450,6 +520,7 @@ def main() -> int:
default="bench/results.json",
help="Write raw per-run results JSON to this path.",
)
ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
ap.add_argument(
"--dry-run", action="store_true", help="List selected placements and exit."
)
@@ -458,6 +529,12 @@ 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).",
)
args = ap.parse_args()
pp_list = parse_int_list(args.pp)
@@ -481,11 +558,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")
@@ -497,44 +569,20 @@ 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
selected = fetch_and_filter_placements(client, full_model_id, args)
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()
if not selected and args.settle_timeout > 0:
backoff = _SETTLE_INITIAL_BACKOFF_S
deadline = time.monotonic() + args.settle_timeout
while not selected and time.monotonic() < deadline:
remaining = 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.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.")
@@ -652,7 +700,9 @@ def main() -> int:
time.sleep(5)
if args.json_out:
if args.stdout:
json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
elif args.json_out:
with open(args.json_out, "w", encoding="utf-8") as f:
json.dump(all_rows, f, indent=2, ensure_ascii=False)
logger.debug(f"\nWrote results JSON: {args.json_out}")
+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
+46 -3
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) &&
@@ -282,7 +293,11 @@
// 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 +535,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 +643,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());
@@ -389,10 +393,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 +475,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"
@@ -789,3 +827,8 @@
</button>
{/if}
</div>
<ImageLightbox
src={expandedImageSrc}
onclose={() => (expandedImageSrc = null)}
/>
@@ -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"
@@ -14,6 +14,7 @@
isAdding: boolean;
onAdd: () => void;
onSelect: () => void;
downloadedOnNodes?: string[];
};
let {
@@ -22,6 +23,7 @@
isAdding,
onAdd,
onSelect,
downloadedOnNodes = [],
}: HuggingFaceResultItemProps = $props();
function formatNumber(num: number): string {
@@ -45,6 +47,28 @@
<span class="text-sm font-mono text-white truncate" title={model.id}
>{modelName}</span
>
{#if downloadedOnNodes.length > 0}
<span
class="flex-shrink-0"
title={`Downloaded on ${downloadedOnNodes.join(", ")}`}
>
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
</span>
{/if}
{#if isAdded}
<span
class="px-1.5 py-0.5 text-[10px] font-mono bg-green-500/20 text-green-400 rounded"
@@ -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}
@@ -64,6 +64,8 @@
"1024x1024",
"1024x768",
"768x1024",
"1024x1365",
"1365x1024",
];
const qualityOptions: ImageGenerationParams["quality"][] = [
@@ -148,6 +150,15 @@
setImageGenerationParams({ guidance: null });
}
function handleNumSyncStepsChange(event: Event) {
const value = parseInt((event.target as HTMLInputElement).value, 10);
setImageGenerationParams({ numSyncSteps: value });
}
function clearNumSyncSteps() {
setImageGenerationParams({ numSyncSteps: null });
}
function handleReset() {
resetImageGenerationParams();
showAdvanced = false;
@@ -157,7 +168,8 @@
params.seed !== null ||
params.numInferenceSteps !== null ||
params.guidance !== null ||
(params.negativePrompt !== null && params.negativePrompt.trim() !== ""),
(params.negativePrompt !== null && params.negativePrompt.trim() !== "") ||
params.numSyncSteps !== null,
);
</script>
@@ -578,7 +590,50 @@
</div>
</div>
<!-- Row 3: Negative Prompt -->
<!-- Row 3: Sync Steps -->
<div class="flex items-center gap-1.5">
<span
class="text-xs text-exo-light-gray uppercase tracking-wider whitespace-nowrap"
>SYNC STEPS:</span
>
<div class="flex items-center gap-2 flex-1 max-w-xs">
<input
type="range"
min="1"
max="100"
value={params.numSyncSteps ?? 1}
oninput={handleNumSyncStepsChange}
class="flex-1 h-1 bg-exo-medium-gray/50 rounded appearance-none cursor-pointer accent-exo-yellow"
/>
<span class="text-xs font-mono text-exo-yellow w-8 text-right">
{params.numSyncSteps ?? "--"}
</span>
{#if params.numSyncSteps !== null}
<button
type="button"
onclick={clearNumSyncSteps}
class="text-exo-light-gray hover:text-exo-yellow transition-colors"
title="Clear"
>
<svg
class="w-3 h-3"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
stroke-width="2"
d="M6 18L18 6M6 6l12 12"
/>
</svg>
</button>
{/if}
</div>
</div>
<!-- Row 4: Negative Prompt -->
<div class="flex flex-col gap-1.5">
<span class="text-xs text-exo-light-gray uppercase tracking-wider"
>NEGATIVE PROMPT:</span
+10 -3
View File
@@ -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 });
@@ -5,6 +5,7 @@
interface FilterState {
capabilities: string[];
sizeRange: { min: number; max: number } | null;
downloadedOnly: boolean;
}
type ModelFilterPopoverProps = {
@@ -23,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
@@ -141,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>
@@ -148,6 +190,36 @@
</div>
</div>
<!-- Downloaded only -->
<div>
<h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
<button
type="button"
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
? 'bg-green-500/20 text-green-400 border border-green-500/30'
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
onclick={() =>
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
>
<svg
class="w-3.5 h-3.5 inline-block"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
<span class="ml-1">Downloaded</span>
</button>
</div>
<!-- Size range -->
<div>
<h4 class="text-xs font-mono text-white/50 mb-2">Model Size</h4>
@@ -21,16 +21,26 @@
hasMultipleVariants: boolean;
}
type DownloadAvailability = {
available: boolean;
nodeNames: string[];
nodeIds: string[];
};
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
type ModelPickerGroupProps = {
group: ModelGroup;
isExpanded: boolean;
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 {
@@ -39,12 +49,25 @@
isFavorite,
selectedModelId,
canModelFit,
getModelFitStatus,
onToggleExpand,
onSelectModel,
onToggleFavorite,
onShowInfo,
downloadStatusMap,
launchedAt,
}: ModelPickerGroupProps = $props();
// Group-level download status: show if any variant is downloaded
const groupDownloadStatus = $derived.by(() => {
if (!downloadStatusMap || downloadStatusMap.size === 0) return undefined;
// Return the first available entry (prefer "available" ones)
for (const avail of downloadStatusMap.values()) {
if (avail.available) return avail;
}
return downloadStatusMap.values().next().value;
});
// Format storage size
function formatSize(mb: number | undefined): string {
if (!mb) return "";
@@ -54,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(
@@ -182,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}
@@ -193,15 +284,65 @@
<!-- 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}
<!-- Variant count -->
<!-- Variant count with size range -->
{#if group.hasMultipleVariants}
<span class="text-xs font-mono text-white/30 flex-shrink-0">
{group.variants.length} variants
{@const sizes = group.variants
.map((v) => v.storage_size_megabytes || 0)
.filter((s) => s > 0)
.sort((a, b) => a - b)}
<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
class="flex-shrink-0"
title={groupDownloadStatus.available
? `Ready — downloaded on ${groupDownloadStatus.nodeNames.join(", ")}`
: `Downloaded on ${groupDownloadStatus.nodeNames.join(", ")} (may need more nodes)`}
>
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
</span>
{/if}
@@ -277,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
@@ -301,10 +443,41 @@
</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>
<!-- Download indicator for this variant -->
{#if downloadStatusMap?.get(variant.id)}
{@const variantDl = downloadStatusMap.get(variant.id)}
{#if variantDl}
<span
class="flex-shrink-0"
title={`Downloaded on ${variantDl.nodeNames.join(", ")}`}
>
<svg
class="w-3.5 h-3.5"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
</span>
{/if}
{/if}
<!-- Check mark if selected -->
{#if isSelected}
<svg
@@ -5,6 +5,8 @@
import ModelPickerGroup from "./ModelPickerGroup.svelte";
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;
@@ -33,6 +35,7 @@
interface FilterState {
capabilities: string[];
sizeRange: { min: number; max: number } | null;
downloadedOnly: boolean;
}
interface HuggingFaceModel {
@@ -44,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;
@@ -58,6 +66,15 @@
onDeleteModel: (modelId: string) => Promise<void>;
totalMemoryGB: number;
usedMemoryGB: number;
downloadsData?: Record<string, unknown[]>;
topologyNodes?: Record<
string,
{
friendly_name?: string;
system_info?: { model_id?: string };
macmon_info?: { memory?: { ram_total?: number } };
}
>;
};
let {
@@ -65,8 +82,11 @@
models,
selectedModelId,
favorites,
recentModelIds = [],
hasRecents: hasRecentsTab = false,
existingModelIds,
canModelFit,
getModelFitStatus,
onSelect,
onClose,
onToggleFavorite,
@@ -74,6 +94,8 @@
onDeleteModel,
totalMemoryGB,
usedMemoryGB,
downloadsData,
topologyNodes,
}: ModelPickerModalProps = $props();
// Local state
@@ -81,9 +103,75 @@
let selectedFamily = $state<string | null>(null);
let expandedGroups = $state<Set<string>>(new Set());
let showFilters = $state(false);
let filters = $state<FilterState>({ capabilities: [], sizeRange: null });
let filters = $state<FilterState>({
capabilities: [],
sizeRange: null,
downloadedOnly: false,
});
let infoGroup = $state<ModelGroup | null>(null);
// Download availability per model group
type DownloadAvailability = {
available: boolean;
nodeNames: string[];
nodeIds: string[];
};
function getNodeName(nodeId: string): string {
const node = topologyNodes?.[nodeId];
return (
node?.friendly_name || node?.system_info?.model_id || nodeId.slice(0, 8)
);
}
const modelDownloadAvailability = $derived.by(() => {
const result = new Map<string, DownloadAvailability>();
if (!downloadsData || !topologyNodes) return result;
for (const model of models) {
const nodeIds = getNodesWithModelDownloaded(downloadsData, model.id);
if (nodeIds.length === 0) continue;
// Sum total RAM across nodes that have the model
let totalRamBytes = 0;
for (const nodeId of nodeIds) {
const ramTotal = topologyNodes[nodeId]?.macmon_info?.memory?.ram_total;
if (typeof ramTotal === "number") totalRamBytes += ramTotal;
}
const modelSizeBytes = (model.storage_size_megabytes || 0) * 1024 * 1024;
result.set(model.id, {
available: modelSizeBytes > 0 && totalRamBytes >= modelSizeBytes,
nodeNames: nodeIds.map(getNodeName),
nodeIds,
});
}
return result;
});
// Aggregate download availability per group (available if ANY variant is available)
function getGroupDownloadAvailability(
group: ModelGroup,
): DownloadAvailability | undefined {
for (const variant of group.variants) {
const avail = modelDownloadAvailability.get(variant.id);
if (avail && avail.nodeIds.length > 0) return avail;
}
return undefined;
}
// Get per-variant download map for a group
function getVariantDownloadMap(
group: ModelGroup,
): Map<string, DownloadAvailability> {
const map = new Map<string, DownloadAvailability>();
for (const variant of group.variants) {
const avail = modelDownloadAvailability.get(variant.id);
if (avail && avail.nodeIds.length > 0) map.set(variant.id, avail);
}
return map;
}
// HuggingFace Hub state
let hfSearchQuery = $state("");
let hfSearchResults = $state<HuggingFaceModel[]>([]);
@@ -95,15 +183,12 @@
let manualModelId = $state("");
let addModelError = $state<string | null>(null);
// Reset state when modal opens
// Reset transient state when modal opens, but preserve tab selection
$effect(() => {
if (isOpen) {
searchQuery = "";
selectedFamily = null;
expandedGroups = new Set();
showFilters = false;
hfSearchQuery = "";
hfSearchResults = [];
manualModelId = "";
addModelError = null;
}
@@ -287,6 +372,8 @@
"deepseek",
"gpt-oss",
"llama",
"flux",
"qwen-image",
];
return Array.from(families).sort((a, b) => {
const aIdx = familyOrder.indexOf(a);
@@ -305,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);
}
@@ -339,13 +430,33 @@
});
}
// Sort: models that fit first, then by size (largest first)
result.sort((a, b) => {
const aFits = a.variants.some((v) => canModelFit(v.id));
const bFits = b.variants.some((v) => canModelFit(v.id));
// Filter to downloaded models only
if (filters.downloadedOnly) {
result = result.filter((g) =>
g.variants.some((v) => {
const avail = modelDownloadAvailability.get(v.id);
return avail && avail.nodeIds.length > 0;
}),
);
}
if (aFits && !bFits) return -1;
if (!aFits && bFits) return 1;
// Sort: fits-now first, then fits-cluster-capacity, then too-large
result.sort((a, b) => {
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;
};
const aRank = getGroupFitRank(a);
const bRank = getGroupFitRank(b);
if (aRank !== bRank) return aRank - bRank;
return (
(b.smallestVariant.storage_size_megabytes || 0) -
@@ -359,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)) {
@@ -385,11 +538,13 @@
}
function clearFilters() {
filters = { capabilities: [], sizeRange: null };
filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
}
const hasActiveFilters = $derived(
filters.capabilities.length > 0 || filters.sizeRange !== null,
filters.capabilities.length > 0 ||
filters.sizeRange !== null ||
filters.downloadedOnly,
);
</script>
@@ -514,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">
@@ -536,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"
@@ -576,6 +732,12 @@
isAdding={addingModelId === model.id}
onAdd={() => handleAddModel(model.id)}
onSelect={() => handleSelectHfModel(model.id)}
downloadedOnNodes={downloadsData
? getNodesWithModelDownloaded(
downloadsData,
model.id,
).map(getNodeName)
: []}
/>
{/each}
{/if}
@@ -615,6 +777,44 @@
</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}
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"
@@ -646,10 +846,12 @@
isFavorite={favorites.has(group.id)}
{selectedModelId}
{canModelFit}
{getModelFitStatus}
onToggleExpand={() => toggleGroupExpanded(group.id)}
onSelectModel={handleSelect}
{onToggleFavorite}
onShowInfo={(g) => (infoGroup = g)}
downloadStatusMap={getVariantDownloadMap(group)}
/>
{/each}
{/if}
@@ -667,6 +869,11 @@
>{cap}</span
>
{/each}
{#if filters.downloadedOnly}
<span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
>Downloaded</span
>
{/if}
{#if filters.sizeRange}
<span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
{filters.sizeRange.min}GB - {filters.sizeRange.max}GB
@@ -742,6 +949,40 @@
</div>
</div>
{/if}
{#if getGroupDownloadAvailability(infoGroup)?.nodeNames?.length}
{@const infoDownload = getGroupDownloadAvailability(infoGroup)}
{#if infoDownload}
<div class="mt-3 pt-3 border-t border-exo-yellow/10">
<div class="flex items-center gap-2 mb-1">
<svg
class="w-3.5 h-3.5"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path
class="text-white/40"
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
/>
<path class="text-green-400" d="m9 13 2 2 4-4" />
</svg>
<span class="text-white/40">Downloaded on:</span>
</div>
<div class="flex flex-wrap gap-1 mt-1">
{#each infoDownload.nodeNames as nodeName}
<span
class="px-1.5 py-0.5 bg-green-500/10 text-green-400/80 border border-green-500/20 rounded text-[10px]"
>
{nodeName}
</span>
{/each}
</div>
</div>
{/if}
{/if}
</div>
</div>
{/if}
@@ -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})` : ""}`,
);
}
}
});
+128 -23
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 {
@@ -78,6 +81,8 @@ interface RawNodeIdentity {
modelId?: string;
chipId?: string;
friendlyName?: string;
osVersion?: string;
osBuildVersion?: string;
}
interface RawMemoryUsage {
@@ -227,6 +232,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,
@@ -280,6 +298,7 @@ export interface Conversation {
modelId: string | null;
sharding: string | null;
instanceType: string | null;
enableThinking: boolean | null;
}
const STORAGE_KEY = "exo-conversations";
@@ -288,7 +307,14 @@ const IMAGE_PARAMS_STORAGE_KEY = "exo-image-generation-params";
// Image generation params interface matching backend API
export interface ImageGenerationParams {
// Basic params
size: "512x512" | "768x768" | "1024x1024" | "1024x768" | "768x1024";
size:
| "512x512"
| "768x768"
| "1024x1024"
| "1024x768"
| "768x1024"
| "1024x1365"
| "1365x1024";
quality: "low" | "medium" | "high";
outputFormat: "png" | "jpeg";
numImages: number;
@@ -300,6 +326,7 @@ export interface ImageGenerationParams {
numInferenceSteps: number | null;
guidance: number | null;
negativePrompt: string | null;
numSyncSteps: number | null;
// Edit mode params
inputFidelity: "low" | "high";
}
@@ -321,6 +348,7 @@ const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = {
numInferenceSteps: null,
guidance: null,
negativePrompt: null,
numSyncSteps: null,
inputFidelity: "low",
};
@@ -418,6 +446,7 @@ function transformTopology(
},
last_macmon_update: Date.now() / 1000,
friendly_name: identity?.friendlyName,
os_version: identity?.osVersion,
};
}
@@ -430,6 +459,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) {
@@ -437,10 +468,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,
});
}
}
}
@@ -483,12 +527,32 @@ 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);
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,
@@ -544,6 +608,7 @@ class AppStore {
modelId: conversation.modelId ?? null,
sharding: conversation.sharding ?? null,
instanceType: conversation.instanceType ?? null,
enableThinking: conversation.enableThinking ?? null,
}));
}
} catch (error) {
@@ -733,6 +798,7 @@ class AppStore {
modelId: derivedModelId,
sharding: derivedSharding,
instanceType: derivedInstanceType,
enableThinking: null,
};
this.conversations.unshift(conversation);
@@ -758,6 +824,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;
@@ -1199,6 +1266,15 @@ class AppStore {
if (data.downloads) {
this.downloads = data.downloads;
}
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
@@ -1862,6 +1938,11 @@ class AppStore {
}
}
/**
* Whether thinking is enabled for the current conversation
*/
thinkingEnabled = $state(true);
/**
* Selected model for chat (can be set by the UI)
*/
@@ -2040,6 +2121,7 @@ class AppStore {
textContent?: string;
preview?: string;
}[],
enableThinking?: boolean | null,
): Promise<void> {
if ((!content.trim() && (!files || files.length === 0)) || this.isLoading)
return;
@@ -2187,6 +2269,9 @@ class AppStore {
stream: true,
logprobs: true,
top_logprobs: 5,
...(enableThinking != null && {
enable_thinking: enableThinking,
}),
}),
});
@@ -2398,7 +2483,9 @@ class AppStore {
params.seed !== null ||
params.numInferenceSteps !== null ||
params.guidance !== null ||
(params.negativePrompt !== null && params.negativePrompt.trim() !== "");
(params.negativePrompt !== null &&
params.negativePrompt.trim() !== "") ||
params.numSyncSteps !== null;
const requestBody: Record<string, unknown> = {
model,
@@ -2423,6 +2510,9 @@ class AppStore {
params.negativePrompt.trim() !== "" && {
negative_prompt: params.negativePrompt,
}),
...(params.numSyncSteps !== null && {
num_sync_steps: params.numSyncSteps,
}),
};
}
@@ -2672,29 +2762,19 @@ class AppStore {
formData.append("input_fidelity", params.inputFidelity);
// Advanced params
if (params.seed !== null) {
formData.append(
"advanced_params",
JSON.stringify({
seed: params.seed,
...(params.numInferenceSteps !== null && {
num_inference_steps: params.numInferenceSteps,
}),
...(params.guidance !== null && { guidance: params.guidance }),
...(params.negativePrompt !== null &&
params.negativePrompt.trim() !== "" && {
negative_prompt: params.negativePrompt,
}),
}),
);
} else if (
const hasAdvancedParams =
params.seed !== null ||
params.numInferenceSteps !== null ||
params.guidance !== null ||
(params.negativePrompt !== null && params.negativePrompt.trim() !== "")
) {
(params.negativePrompt !== null &&
params.negativePrompt.trim() !== "") ||
params.numSyncSteps !== null;
if (hasAdvancedParams) {
formData.append(
"advanced_params",
JSON.stringify({
...(params.seed !== null && { seed: params.seed }),
...(params.numInferenceSteps !== null && {
num_inference_steps: params.numInferenceSteps,
}),
@@ -2703,6 +2783,9 @@ class AppStore {
params.negativePrompt.trim() !== "" && {
negative_prompt: params.negativePrompt,
}),
...(params.numSyncSteps !== null && {
num_sync_steps: params.numSyncSteps,
}),
}),
);
}
@@ -2847,6 +2930,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
*/
@@ -2953,12 +3048,14 @@ export const topologyData = () => appStore.topologyData;
export const instances = () => appStore.instances;
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();
@@ -2974,7 +3071,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 = (
@@ -3017,6 +3115,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;
@@ -3033,7 +3133,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();
+152
View File
@@ -0,0 +1,152 @@
/**
* Shared utilities for parsing and querying download state.
*
* The download state from `/state` is shaped as:
* Record<NodeId, Array<TaggedDownloadEntry>>
*
* Each entry is a tagged union object like:
* { "DownloadCompleted": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } }
*/
/** Unwrap one level of tagged-union envelope, returning [tag, payload]. */
function unwrapTagged(
obj: Record<string, unknown>,
): [string, Record<string, unknown>] | null {
const keys = Object.keys(obj);
if (keys.length !== 1) return null;
const tag = keys[0];
const payload = obj[tag];
if (!payload || typeof payload !== "object") return null;
return [tag, payload as Record<string, unknown>];
}
/** Extract the model ID string from a download entry's nested shard_metadata. */
export function extractModelIdFromDownload(
downloadPayload: Record<string, unknown>,
): string | null {
const shardMetadata =
downloadPayload.shard_metadata ?? downloadPayload.shardMetadata;
if (!shardMetadata || typeof shardMetadata !== "object") return null;
const unwrapped = unwrapTagged(shardMetadata as Record<string, unknown>);
if (!unwrapped) return null;
const [, shardData] = unwrapped;
const modelMeta = shardData.model_card ?? shardData.modelCard;
if (!modelMeta || typeof modelMeta !== "object") return null;
const meta = modelMeta as Record<string, unknown>;
return (meta.model_id as string) ?? (meta.modelId as string) ?? null;
}
/** Extract the shard_metadata object from a download entry payload. */
export function extractShardMetadata(
downloadPayload: Record<string, unknown>,
): Record<string, unknown> | null {
const shardMetadata =
downloadPayload.shard_metadata ?? downloadPayload.shardMetadata;
if (!shardMetadata || typeof shardMetadata !== "object") return null;
return shardMetadata as Record<string, unknown>;
}
/** Get the download tag (DownloadCompleted, DownloadOngoing, etc.) from a wrapped entry. */
export function getDownloadTag(
entry: unknown,
): [string, Record<string, unknown>] | null {
if (!entry || typeof entry !== "object") return null;
return unwrapTagged(entry as Record<string, unknown>);
}
/**
* Iterate over all download entries for a given node, yielding [tag, payload, modelId].
*/
function* iterNodeDownloads(
nodeDownloads: unknown[],
): Generator<[string, Record<string, unknown>, string]> {
for (const entry of nodeDownloads) {
const tagged = getDownloadTag(entry);
if (!tagged) continue;
const [tag, payload] = tagged;
const modelId = extractModelIdFromDownload(payload);
if (!modelId) continue;
yield [tag, payload, modelId];
}
}
/** Check if a specific model is fully downloaded (DownloadCompleted) on a specific node. */
export function isModelDownloadedOnNode(
downloadsData: Record<string, unknown[]>,
nodeId: string,
modelId: string,
): boolean {
const nodeDownloads = downloadsData[nodeId];
if (!Array.isArray(nodeDownloads)) return false;
for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) {
if (tag === "DownloadCompleted" && entryModelId === modelId) return true;
}
return false;
}
/** Get all node IDs where a model is fully downloaded (DownloadCompleted). */
export function getNodesWithModelDownloaded(
downloadsData: Record<string, unknown[]>,
modelId: string,
): string[] {
const result: string[] = [];
for (const nodeId of Object.keys(downloadsData)) {
if (isModelDownloadedOnNode(downloadsData, nodeId, modelId)) {
result.push(nodeId);
}
}
return result;
}
/**
* Find shard metadata for a model from any download entry across all nodes.
* Returns the first match found (completed entries are preferred).
*/
export function getShardMetadataForModel(
downloadsData: Record<string, unknown[]>,
modelId: string,
): Record<string, unknown> | null {
let fallback: Record<string, unknown> | null = null;
for (const nodeDownloads of Object.values(downloadsData)) {
if (!Array.isArray(nodeDownloads)) continue;
for (const [tag, payload, entryModelId] of iterNodeDownloads(
nodeDownloads,
)) {
if (entryModelId !== modelId) continue;
const shard = extractShardMetadata(payload);
if (!shard) continue;
if (tag === "DownloadCompleted") return shard;
if (!fallback) fallback = shard;
}
}
return fallback;
}
/**
* Get the download status tag for a specific model on a specific node.
* Returns the "best" status: DownloadCompleted > DownloadOngoing > others.
*/
export function getModelDownloadStatus(
downloadsData: Record<string, unknown[]>,
nodeId: string,
modelId: string,
): string | null {
const nodeDownloads = downloadsData[nodeId];
if (!Array.isArray(nodeDownloads)) return null;
let best: string | null = null;
for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) {
if (entryModelId !== modelId) continue;
if (tag === "DownloadCompleted") return tag;
if (tag === "DownloadOngoing") best = tag;
else if (!best) best = tag;
}
return best;
}
+450 -134
View File
@@ -12,6 +12,11 @@
toggleFavorite,
getFavoritesSet,
} from "$lib/stores/favorites.svelte";
import {
hasRecents,
getRecentModelIds,
recordRecentLaunch,
} from "$lib/stores/recents.svelte";
import {
hasStartedChat,
isTopologyMinimized,
@@ -37,8 +42,11 @@
toggleTopologyOnlyMode,
chatSidebarVisible,
toggleChatSidebarVisible,
nodeThunderbolt,
nodeRdmaCtl,
thunderboltBridgeCycles,
nodeThunderboltBridge,
nodeIdentities,
type DownloadProgress,
type PlacementPreview,
} from "$lib/stores/app.svelte";
@@ -62,8 +70,50 @@
const sidebarVisible = $derived(chatSidebarVisible());
const tbBridgeCycles = $derived(thunderboltBridgeCycles());
const tbBridgeData = $derived(nodeThunderboltBridge());
const identitiesData = $derived(nodeIdentities());
const tbIdentifiers = $derived(nodeThunderbolt());
const rdmaCtlData = $derived(nodeRdmaCtl());
const nodeFilter = $derived(previewNodeFilter());
// Detect macOS version mismatches across cluster nodes
const macosVersionMismatch = $derived.by(() => {
if (!identitiesData) return null;
const entries = Object.entries(identitiesData);
// Filter to macOS nodes (version starts with a digit, e.g. "15.3")
const macosNodes = entries.filter(([_, id]) => {
const v = id.osVersion;
return v && v !== "Unknown" && /^\d/.test(v);
});
if (macosNodes.length < 2) return null;
// Compare on buildVersion for precise mismatch detection
const buildVersions = new Set(
macosNodes.map(([_, id]) => id.osBuildVersion ?? id.osVersion),
);
if (buildVersions.size <= 1) return null;
return macosNodes.map(([nodeId, id]) => ({
nodeId,
friendlyName: getNodeName(nodeId),
version: id.osVersion!,
buildVersion: id.osBuildVersion ?? "Unknown",
}));
});
// Detect TB5 nodes where RDMA is not enabled
const tb5WithoutRdma = $derived.by(() => {
const rdmaCtl = rdmaCtlData;
if (!rdmaCtl) return false;
const ids = tbIdentifiers;
if (!ids) return false;
// Find nodes with TB5 hardware (any TB interface)
const tb5NodeIds = Object.entries(ids)
.filter(([_, node]) => node.interfaces.length > 0)
.map(([id]) => id);
if (tb5NodeIds.length < 2) return false;
// At least one TB5 node has RDMA disabled
return tb5NodeIds.some((id) => rdmaCtl[id]?.enabled !== true);
});
let tb5InfoDismissed = $state(false);
// Helper to get friendly node name from node ID
function getNodeName(nodeId: string): string {
const node = data?.nodes?.[nodeId];
@@ -96,6 +146,12 @@
}
}
// Warning icon SVG path (reused across warning snippets)
const warningIconPath =
"M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z";
const infoIconPath =
"M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z";
let mounted = $state(false);
// Instance launch state
@@ -113,6 +169,10 @@
capabilities?: string[];
}>
>([]);
type ModelMemoryFitStatus =
| "fits_now"
| "fits_cluster_capacity"
| "too_large";
// Model tasks lookup for ChatForm - maps both short IDs and full HuggingFace IDs
const modelTasks = $derived(() => {
@@ -130,6 +190,19 @@
return tasks;
});
const modelCapabilities = $derived(() => {
const caps: Record<string, string[]> = {};
for (const model of models) {
if (model.capabilities && model.capabilities.length > 0) {
caps[model.id] = model.capabilities;
if (model.hugging_face_id) {
caps[model.hugging_face_id] = model.capabilities;
}
}
}
return caps;
});
// Helper to check if a model supports image generation
function modelSupportsImageGeneration(modelId: string): boolean {
const model = models.find(
@@ -228,6 +301,10 @@
// Favorites state (reactive)
const favoritesSet = $derived(getFavoritesSet());
// Recent models state (reactive)
const recentModelIds = $derived(getRecentModelIds());
const showRecentsTab = $derived(hasRecents());
// Slider dragging state
let isDraggingSlider = $state(false);
let sliderTrackElement: HTMLDivElement | null = $state(null);
@@ -450,14 +527,41 @@
);
});
// Calculate total memory in the cluster (in GB)
const clusterTotalMemoryGB = $derived(() => {
if (!data) return 0;
return (
Object.values(data.nodes).reduce((acc, n) => {
const total =
n.macmon_info?.memory?.ram_total ?? n.system_info?.memory ?? 0;
return acc + total;
}, 0) /
(1024 * 1024 * 1024)
);
});
function getModelMemoryFitStatus(model: {
id: string;
name?: string;
storage_size_megabytes?: number;
}): ModelMemoryFitStatus {
const modelSizeGB = getModelSizeGB(model);
if (modelSizeGB <= availableMemoryGB()) {
return "fits_now";
}
if (modelSizeGB <= clusterTotalMemoryGB()) {
return "fits_cluster_capacity";
}
return "too_large";
}
// Check if a model has enough memory to run
function hasEnoughMemory(model: {
id: string;
name?: string;
storage_size_megabytes?: number;
}): boolean {
const modelSizeGB = getModelSizeGB(model);
return modelSizeGB <= availableMemoryGB();
return getModelMemoryFitStatus(model) === "fits_now";
}
// Sorted models for dropdown - biggest first, unrunnable at the end
@@ -630,6 +734,9 @@
// Always auto-select the newly launched model so the user chats to what they just launched
setSelectedChatModel(modelId);
// Record the launch in recent models history
recordRecentLaunch(modelId);
// Scroll to the bottom of instances container to show the new instance
// Use multiple attempts to ensure DOM has updated with the new instance
const scrollToBottom = () => {
@@ -1657,6 +1764,249 @@
});
</script>
{#snippet clusterWarnings()}
{#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)}
<div class="absolute top-4 left-4 flex flex-col gap-2 z-40">
{#if tbBridgeCycles.length > 0}
{@const cycle = tbBridgeCycles[0]}
{@const serviceName = getTbBridgeServiceName(cycle)}
{@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
<div class="group relative" role="alert">
<div
class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
>
<svg
class="w-5 h-5 text-yellow-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={warningIconPath}
/>
</svg>
<span class="text-sm font-mono text-yellow-200">
THUNDERBOLT BRIDGE CYCLE DETECTED
</span>
</div>
<!-- Tooltip on hover -->
<div
class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
>
<p class="text-xs text-white/80 mb-2">
A network routing cycle was detected between nodes connected via
Thunderbolt Bridge. This can cause connectivity issues.
</p>
<p class="text-xs text-white/60 mb-2">
<span class="text-yellow-300">Affected nodes:</span>
{cycle.map(getNodeName).join(" → ")}
</p>
<p class="text-xs text-white/60 mb-1">
<span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
Bridge on one of the affected nodes:
</p>
<button
type="button"
onclick={() => copyToClipboard(disableCmd)}
class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
title="Click to copy"
>
<span class="flex-1">{disableCmd}</span>
<svg
class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
{#if copiedCommand}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M5 13l4 4L19 7"
/>
{:else}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
/>
{/if}
</svg>
</button>
</div>
</div>
{/if}
{#if macosVersionMismatch}
<div class="group relative" role="alert">
<div
class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
>
<svg
class="w-5 h-5 text-yellow-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={warningIconPath}
/>
</svg>
<span class="text-sm font-mono text-yellow-200">
INCOMPATIBLE macOS VERSIONS
</span>
</div>
<!-- Tooltip on hover -->
<div
class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
>
<p class="text-xs text-white/80 mb-2">
Nodes in this cluster are running different macOS versions. This
may cause inference compatibility issues.
</p>
<div class="text-xs text-white/60 mb-2">
<span class="text-yellow-300">Node versions:</span>
{#each macosVersionMismatch as node}
<div class="ml-2">
{node.friendlyName} — macOS {node.version} ({node.buildVersion})
</div>
{/each}
</div>
<p class="text-xs text-white/60">
<span class="text-yellow-300">Suggested action:</span> Update all nodes
to the same macOS version for best compatibility.
</p>
</div>
</div>
{/if}
{#if tb5WithoutRdma && !tb5InfoDismissed}
<div
class="flex items-center gap-2 px-3 py-2 rounded border border-blue-400/50 bg-blue-400/10 backdrop-blur-sm"
role="status"
>
<svg
class="w-5 h-5 text-blue-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={infoIconPath}
/>
</svg>
<span class="text-sm font-mono text-blue-200"> RDMA AVAILABLE </span>
<button
type="button"
onclick={() => (tb5InfoDismissed = true)}
class="ml-1 text-blue-300/60 hover:text-blue-200 transition-colors cursor-pointer"
title="Dismiss"
>
<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="M6 18L18 6M6 6l12 12"
/>
</svg>
</button>
</div>
{/if}
</div>
{/if}
{/snippet}
{#snippet clusterWarningsCompact()}
{#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)}
<div class="absolute top-2 left-2 flex flex-col gap-1">
{#if tbBridgeCycles.length > 0}
<div
class="flex items-center gap-1.5 px-2 py-1 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm"
title="Thunderbolt Bridge cycle detected"
>
<svg
class="w-3.5 h-3.5 text-yellow-400"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={warningIconPath}
/>
</svg>
<span class="text-[10px] font-mono text-yellow-200">TB CYCLE</span>
</div>
{/if}
{#if macosVersionMismatch}
<div
class="flex items-center gap-1.5 px-2 py-1 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm"
title="Incompatible macOS versions detected"
>
<svg
class="w-3.5 h-3.5 text-yellow-400"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={warningIconPath}
/>
</svg>
<span class="text-[10px] font-mono text-yellow-200"
>macOS MISMATCH</span
>
</div>
{/if}
{#if tb5WithoutRdma && !tb5InfoDismissed}
<div
class="flex items-center gap-1.5 px-2 py-1 rounded border border-blue-400/50 bg-blue-400/10 backdrop-blur-sm"
title="Thunderbolt 5 detected — RDMA can be enabled for better performance"
>
<svg
class="w-3.5 h-3.5 text-blue-400"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d={infoIconPath}
/>
</svg>
<span class="text-[10px] font-mono text-blue-200">RDMA AVAILABLE</span
>
</div>
{/if}
</div>
{/if}
{/snippet}
<!-- Global event listeners for slider dragging -->
<svelte:window
onmousemove={handleSliderMouseMove}
@@ -1724,17 +2074,40 @@
onNodeClick={togglePreviewNodeFilter}
/>
<!-- Thunderbolt Bridge Cycle Warning -->
{#if tbBridgeCycles.length > 0}
{@const cycle = tbBridgeCycles[0]}
{@const serviceName = getTbBridgeServiceName(cycle)}
{@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
<div class="absolute top-4 left-4 group" role="alert">
<div
class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
{@render clusterWarnings()}
<!-- TB5 RDMA Available Info -->
{#if tb5WithoutRdma && !tb5InfoDismissed}
<div
class="absolute left-4 flex items-center gap-2 px-3 py-2 rounded border border-blue-400/50 bg-blue-400/10 backdrop-blur-sm"
class:top-16={tbBridgeCycles.length > 0}
class:top-4={tbBridgeCycles.length === 0}
role="status"
>
<svg
class="w-5 h-5 text-blue-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<span class="text-sm font-mono text-blue-200">
RDMA AVAILABLE
</span>
<button
type="button"
onclick={() => (tb5InfoDismissed = true)}
class="ml-1 text-blue-300/60 hover:text-blue-200 transition-colors cursor-pointer"
title="Dismiss"
>
<svg
class="w-5 h-5 text-yellow-400 flex-shrink-0"
class="w-4 h-4"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
@@ -1743,60 +2116,10 @@
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
d="M6 18L18 6M6 6l12 12"
/>
</svg>
<span class="text-sm font-mono text-yellow-200">
THUNDERBOLT BRIDGE CYCLE DETECTED
</span>
</div>
<!-- Tooltip on hover -->
<div
class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
>
<p class="text-xs text-white/80 mb-2">
A network routing cycle was detected between nodes connected
via Thunderbolt Bridge. This can cause connectivity issues.
</p>
<p class="text-xs text-white/60 mb-2">
<span class="text-yellow-300">Affected nodes:</span>
{cycle.map(getNodeName).join(" → ")}
</p>
<p class="text-xs text-white/60 mb-1">
<span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
Bridge on one of the affected nodes:
</p>
<button
type="button"
onclick={() => copyToClipboard(disableCmd)}
class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
title="Click to copy"
>
<span class="flex-1">{disableCmd}</span>
<svg
class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
{#if copiedCommand}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M5 13l4 4L19 7"
/>
{:else}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
/>
{/if}
</svg>
</button>
</div>
</button>
</div>
{/if}
@@ -1843,17 +2166,21 @@
onNodeClick={togglePreviewNodeFilter}
/>
<!-- Thunderbolt Bridge Cycle Warning -->
{#if tbBridgeCycles.length > 0}
{@const cycle = tbBridgeCycles[0]}
{@const serviceName = getTbBridgeServiceName(cycle)}
{@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
<div class="absolute top-4 left-4 group" role="alert">
{@render clusterWarnings()}
<!-- TB5 RDMA Available Info -->
{#if tb5WithoutRdma && !tb5InfoDismissed}
<div
class="absolute left-4 group"
class:top-16={tbBridgeCycles.length > 0}
class:top-4={tbBridgeCycles.length === 0}
role="status"
>
<div
class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
class="flex items-center gap-2 px-3 py-2 rounded border border-blue-400/50 bg-blue-400/10 backdrop-blur-sm"
>
<svg
class="w-5 h-5 text-yellow-400 flex-shrink-0"
class="w-5 h-5 text-blue-400 flex-shrink-0"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
@@ -1862,60 +2189,62 @@
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
d="M13 16h-1v-4h-1m1-4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
/>
</svg>
<span class="text-sm font-mono text-yellow-200">
THUNDERBOLT BRIDGE CYCLE DETECTED
<span class="text-sm font-mono text-blue-200">
RDMA AVAILABLE
</span>
</div>
<!-- Tooltip on hover -->
<div
class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
>
<p class="text-xs text-white/80 mb-2">
A network routing cycle was detected between nodes connected
via Thunderbolt Bridge. This can cause connectivity issues.
</p>
<p class="text-xs text-white/60 mb-2">
<span class="text-yellow-300">Affected nodes:</span>
{cycle.map(getNodeName).join(" → ")}
</p>
<p class="text-xs text-white/60 mb-1">
<span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
Bridge on one of the affected nodes:
</p>
<button
type="button"
onclick={() => copyToClipboard(disableCmd)}
class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
title="Click to copy"
onclick={() => (tb5InfoDismissed = true)}
class="ml-1 text-blue-300/60 hover:text-blue-200 transition-colors cursor-pointer"
title="Dismiss"
>
<span class="flex-1">{disableCmd}</span>
<svg
class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
class="w-4 h-4"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
{#if copiedCommand}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M5 13l4 4L19 7"
/>
{:else}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
/>
{/if}
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M6 18L18 6M6 6l12 12"
/>
</svg>
</button>
</div>
<!-- Tooltip on hover -->
<div
class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-blue-400/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
>
<p class="text-xs text-white/80 mb-2">
Thunderbolt 5 hardware detected on multiple nodes. Enable
RDMA for significantly faster inter-node communication.
</p>
<p class="text-xs text-white/60 mb-1.5">
<span class="text-blue-300">To enable:</span>
</p>
<ol
class="text-xs text-white/60 list-decimal list-inside space-y-0.5 mb-1.5"
>
<li>Connect nodes with TB5 cables</li>
<li>Boot to Recovery (hold power 10s → Options)</li>
<li>
Run
<code class="text-blue-300 bg-blue-400/10 px-1 rounded"
>rdma_ctl enable</code
>
</li>
<li>Reboot</li>
</ol>
<p class="text-xs text-white/40">
Requires macOS 26.2+, TB5 cables, and matching OS versions.
</p>
</div>
</div>
{/if}
@@ -1954,6 +2283,7 @@
showHelperText={false}
showModelSelector={true}
modelTasks={modelTasks()}
modelCapabilities={modelCapabilities()}
/>
</div>
</div>
@@ -2733,6 +3063,7 @@
placeholder="Ask anything"
showModelSelector={true}
modelTasks={modelTasks()}
modelCapabilities={modelCapabilities()}
/>
</div>
</div>
@@ -2773,30 +3104,7 @@
onNodeClick={togglePreviewNodeFilter}
/>
<!-- Thunderbolt Bridge Cycle Warning (compact) -->
{#if tbBridgeCycles.length > 0}
<div
class="absolute top-2 left-2 flex items-center gap-1.5 px-2 py-1 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm"
title="Thunderbolt Bridge cycle detected - click to view details"
>
<svg
class="w-3.5 h-3.5 text-yellow-400"
fill="none"
viewBox="0 0 24 24"
stroke="currentColor"
stroke-width="2"
>
<path
stroke-linecap="round"
stroke-linejoin="round"
d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
/>
</svg>
<span class="text-[10px] font-mono text-yellow-200"
>TB CYCLE</span
>
</div>
{/if}
{@render clusterWarningsCompact()}
</div>
</button>
@@ -3252,11 +3560,17 @@
{models}
{selectedModelId}
favorites={favoritesSet}
{recentModelIds}
hasRecents={showRecentsTab}
existingModelIds={new Set(models.map((m) => m.id))}
canModelFit={(modelId) => {
const model = models.find((m) => m.id === modelId);
return model ? hasEnoughMemory(model) : false;
}}
getModelFitStatus={(modelId): ModelMemoryFitStatus => {
const model = models.find((m) => m.id === modelId);
return model ? getModelMemoryFitStatus(model) : "too_large";
}}
onSelect={handleModelPickerSelect}
onClose={() => (isModelPickerOpen = false)}
onToggleFavorite={toggleFavorite}
@@ -3264,4 +3578,6 @@
onDeleteModel={deleteCustomModel}
totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)}
usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)}
{downloadsData}
topologyNodes={data?.nodes}
/>
+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>
+5 -1
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";
@@ -118,9 +121,10 @@
{
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
mlx = pkgs.callPackage ./nix/mlx.nix {
metal-toolchain = self'.packages.metal-toolchain;
inherit (self'.packages) metal-toolchain;
inherit uvLockMlxVersion;
};
default = self'.packages.exo;
}
);
+1 -1
View File
@@ -20,7 +20,7 @@ sync-clean:
rust-rebuild:
cargo run --bin stub_gen
just sync-clean
uv sync --reinstall-package exo_pyo3_bindings
build-dashboard:
#!/usr/bin/env bash
+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.4"; 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}")
+9 -7
View File
@@ -17,22 +17,22 @@ dependencies = [
"loguru>=0.7.3",
"exo_pyo3_bindings", # rust bindings
"anyio==4.11.0",
"mlx==0.30.4; sys_platform == 'darwin'",
"mlx[cpu]==0.30.4; sys_platform == 'linux'",
"mlx-lm",
"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",
"openai-harmony>=0.0.8",
"httpx>=0.28.1",
"tomlkit>=0.14.0",
"pillow>=11.0,<12.0", # compatibility with mflux
"mflux==0.15.4",
"mflux==0.15.5",
"python-multipart>=0.0.21",
"msgspec>=0.19.0",
"zstandard>=0.23.0",
]
[project.scripts]
exo-master = "exo.master.main:main"
exo-worker = "exo.worker.main:main"
exo = "exo.main:main"
# dependencies only required for development
@@ -59,11 +59,12 @@ dev = [
[tool.uv.workspace]
members = [
"rust/exo_pyo3_bindings",
"bench",
]
[tool.uv.sources]
exo_pyo3_bindings = { workspace = true }
mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
# Uncomment to use local mlx/mlx-lm development versions:
# mlx = { path = "/Users/Shared/mlx", editable=true }
# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
@@ -105,6 +106,7 @@ root = "src"
# supported platforms for this project
[tool.uv]
required-version = ">=0.8.6"
prerelease = "allow"
environments = [
"sys_platform == 'darwin'",
+40 -12
View File
@@ -59,6 +59,32 @@
}
);
mkPythonScript = name: path: pkgs.writeShellApplication {
inherit name;
runtimeInputs = [ exoVenv ];
runtimeEnv = {
EXO_DASHBOARD_DIR = self'.packages.dashboard;
EXO_RESOURCES_DIR = inputs.self + /resources;
};
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 ];
text = ''exec python ${path} "$@"'';
};
exoPackage = pkgs.runCommand "exo"
{
nativeBuildInputs = [ pkgs.makeWrapper ];
@@ -66,28 +92,30 @@
''
mkdir -p $out/bin
# Create wrapper scripts
for script in exo exo-master exo-worker; do
makeWrapper ${exoVenv}/bin/$script $out/bin/$script \
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
--set EXO_RESOURCES_DIR ${inputs.self + "/resources"} \
${lib.optionalString pkgs.stdenv.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
done
# Create wrapper script
makeWrapper ${exoVenv}/bin/exo $out/bin/exo \
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
'';
in
{
# Python package only available on macOS (requires MLX/Metal)
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
exo = exoPackage;
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
exo-test-env = testVenv;
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
{
exo = exoPackage;
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
exo-test-env = testVenv;
} // {
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);
};
checks = {
# Ruff linting (works on all platforms)
lint = pkgs.runCommand "ruff-lint" { } ''
export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
${pkgs.ruff}/bin/ruff check ${inputs.self}/
${pkgs.ruff}/bin/ruff check ${inputs.self}
touch $out
'';
};
@@ -0,0 +1,49 @@
model_id = "exolabs/FLUX.1-Kontext-dev-4bit"
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
[[components]]
component_name = "text_encoder"
component_path = "text_encoder/"
n_layers = 12
can_shard = false
[components.storage_size]
in_bytes = 0
[[components]]
component_name = "text_encoder_2"
component_path = "text_encoder_2/"
n_layers = 24
can_shard = false
safetensors_index_filename = "model.safetensors.index.json"
[components.storage_size]
in_bytes = 9524621312
[[components]]
component_name = "transformer"
component_path = "transformer/"
n_layers = 57
can_shard = true
safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json"
[components.storage_size]
in_bytes = 5950704160
[[components]]
component_name = "vae"
component_path = "vae/"
can_shard = false
[components.storage_size]
in_bytes = 0
@@ -0,0 +1,49 @@
model_id = "exolabs/FLUX.1-Kontext-dev-8bit"
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
[[components]]
component_name = "text_encoder"
component_path = "text_encoder/"
n_layers = 12
can_shard = false
[components.storage_size]
in_bytes = 0
[[components]]
component_name = "text_encoder_2"
component_path = "text_encoder_2/"
n_layers = 24
can_shard = false
safetensors_index_filename = "model.safetensors.index.json"
[components.storage_size]
in_bytes = 9524621312
[[components]]
component_name = "transformer"
component_path = "transformer/"
n_layers = 57
can_shard = true
safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json"
[components.storage_size]
in_bytes = 11901408320
[[components]]
component_name = "vae"
component_path = "vae/"
can_shard = false
[components.storage_size]
in_bytes = 0
@@ -0,0 +1,49 @@
model_id = "exolabs/FLUX.1-Kontext-dev"
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
[[components]]
component_name = "text_encoder"
component_path = "text_encoder/"
n_layers = 12
can_shard = false
[components.storage_size]
in_bytes = 0
[[components]]
component_name = "text_encoder_2"
component_path = "text_encoder_2/"
n_layers = 24
can_shard = false
safetensors_index_filename = "model.safetensors.index.json"
[components.storage_size]
in_bytes = 9524621312
[[components]]
component_name = "transformer"
component_path = "transformer/"
n_layers = 57
can_shard = true
safetensors_index_filename = "diffusion_pytorch_model.safetensors.index.json"
[components.storage_size]
in_bytes = 23802816640
[[components]]
component_name = "vae"
component_path = "vae/"
can_shard = false
[components.storage_size]
in_bytes = 0
@@ -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,8 @@
model_id = "mlx-community/Qwen3-Coder-Next-4bit"
n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
[storage_size]
in_bytes = 45644286500
@@ -0,0 +1,8 @@
model_id = "mlx-community/Qwen3-Coder-Next-5bit"
n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
[storage_size]
in_bytes = 57657697020
@@ -0,0 +1,8 @@
model_id = "mlx-community/Qwen3-Coder-Next-6bit"
n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
[storage_size]
in_bytes = 68899327465
@@ -0,0 +1,8 @@
model_id = "mlx-community/Qwen3-Coder-Next-8bit"
n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
[storage_size]
in_bytes = 89357758772
@@ -0,0 +1,8 @@
model_id = "mlx-community/Qwen3-Coder-Next-bf16"
n_layers = 48
hidden_size = 2048
supports_tensor = true
tasks = ["TextGeneration"]
[storage_size]
in_bytes = 157548627945
@@ -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
+3 -2
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;
@@ -532,7 +531,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)
-2
View File
@@ -106,7 +106,6 @@ mod behaviour {
use crate::{alias, discovery};
use libp2p::swarm::NetworkBehaviour;
use libp2p::{gossipsub, identity};
use std::time::Duration;
/// Behavior of the Swarm which composes all desired behaviors:
/// Right now its just [`discovery::Behaviour`] and [`gossipsub::Behaviour`].
@@ -134,7 +133,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()
}
}
+59 -45
View File
@@ -16,6 +16,7 @@ from exo.download.download_utils import (
from exo.download.shard_downloader import ShardDownloader
from exo.shared.models.model_cards import ModelId
from exo.shared.types.commands import (
CancelDownload,
DeleteDownload,
ForwarderDownloadCommand,
StartDownload,
@@ -53,11 +54,51 @@ class DownloadCoordinator:
# Internal event channel for forwarding (initialized in __post_init__)
event_sender: Sender[Event] = field(init=False)
event_receiver: Receiver[Event] = field(init=False)
_tg: TaskGroup = field(init=False)
_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._tg = anyio.create_task_group()
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")
@@ -108,16 +149,23 @@ class DownloadCoordinator:
await self._start_download(shard)
case DeleteDownload(model_id=model_id):
await self._delete_download(model_id)
case CancelDownload(model_id=model_id):
await self._cancel_download(model_id)
async def _cancel_download(self, model_id: ModelId) -> None:
if model_id in self.active_downloads and model_id in self.download_status:
logger.info(f"Cancelling download for {model_id}")
self.active_downloads.pop(model_id).cancel()
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
@@ -162,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)
@@ -276,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,
+97 -3
View File
@@ -158,6 +158,78 @@ async def seed_models(seed_dir: str | Path):
logger.error(traceback.format_exc())
async def _build_file_list_from_local_directory(
model_id: ModelId,
recursive: bool = False,
) -> list[FileListEntry] | None:
"""Build a file list from locally existing model files.
We can only figure out the files we need from safetensors index, so
a local directory must contain a *.safetensors.index.json and
safetensors listed there.
"""
model_dir = (await ensure_models_dir()) / model_id.normalize()
if not await aios.path.exists(model_dir):
return None
def _scan() -> list[FileListEntry] | None:
index_files = list(model_dir.glob("**/*.safetensors.index.json"))
if not index_files:
return None
entries_by_path: dict[str, FileListEntry] = {}
if recursive:
for dirpath, _, filenames in os.walk(model_dir):
for filename in filenames:
if filename.endswith(".partial"):
continue
full_path = Path(dirpath) / filename
rel_path = str(full_path.relative_to(model_dir))
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=full_path.stat().st_size,
)
else:
for item in model_dir.iterdir():
if item.is_file() and not item.name.endswith(".partial"):
entries_by_path[item.name] = FileListEntry(
type="file",
path=item.name,
size=item.stat().st_size,
)
# Add expected weight files from index that haven't been downloaded yet
for index_file in index_files:
try:
index_data = ModelSafetensorsIndex.model_validate_json(
index_file.read_text()
)
relative_dir = index_file.parent.relative_to(model_dir)
for filename in set(index_data.weight_map.values()):
rel_path = (
str(relative_dir / filename)
if relative_dir != Path(".")
else filename
)
if rel_path not in entries_by_path:
entries_by_path[rel_path] = FileListEntry(
type="file",
path=rel_path,
size=None,
)
except Exception:
continue
return list(entries_by_path.values())
file_list = await asyncio.to_thread(_scan)
if not file_list:
return None
return file_list
_fetched_file_lists_this_session: set[str] = set()
@@ -183,6 +255,14 @@ async def fetch_file_list_with_cache(
if await aios.path.exists(cache_file):
async with aiofiles.open(cache_file, "r") as f:
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
local_file_list = await _build_file_list_from_local_directory(
model_id, recursive
)
if local_file_list is not None:
logger.warning(
f"No internet and no cached file list for {model_id} - using local file list"
)
return local_file_list
raise FileNotFoundError(
f"No internet connection and no cached file list for {model_id}"
)
@@ -203,10 +283,18 @@ async def fetch_file_list_with_cache(
except Exception as e:
if await aios.path.exists(cache_file):
logger.warning(
f"Failed to fetch file list for {model_id}, using cached data: {e}"
f"No internet and no cached file list for {model_id} - using local file list"
)
async with aiofiles.open(cache_file, "r") as f:
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
local_file_list = await _build_file_list_from_local_directory(
model_id, recursive
)
if local_file_list is not None:
logger.warning(
f"Failed to fetch file list for {model_id} and no cache exists, "
)
return local_file_list
raise FileNotFoundError(f"Failed to fetch file list for {model_id}: {e}") from e
@@ -369,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
@@ -378,10 +468,14 @@ async def download_file_with_retry(
logger.error(traceback.format_exc())
await asyncio.sleep(2.0**attempt)
except Exception as e:
on_connection_lost()
if attempt == n_attempts - 1:
on_connection_lost()
raise e
break
logger.error(
f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
)
logger.error(traceback.format_exc())
await asyncio.sleep(2.0**attempt)
raise Exception(
f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
)
+5 -1
View File
@@ -195,6 +195,10 @@ class ResumableShardDownloader(ShardDownloader):
self, shard: ShardMetadata
) -> RepoDownloadProgress:
_, progress = await download_shard(
shard, self.on_progress_wrapper, skip_download=True
shard,
self.on_progress_wrapper,
skip_download=True,
skip_internet=not self.internet_connection,
on_connection_lost=lambda: self.set_internet_connection(False),
)
return progress
+8 -3
View File
@@ -27,7 +27,6 @@ from exo.utils.pydantic_ext import CamelCaseModel
from exo.worker.main import Worker
# I marked this as a dataclass as I want trivial constructors.
@dataclass
class Node:
router: Router
@@ -106,6 +105,7 @@ class Node:
global_event_sender=router.sender(topics.GLOBAL_EVENTS),
local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
command_receiver=router.receiver(topics.COMMANDS),
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
)
er_send, er_recv = channel[ElectionResult]()
@@ -136,7 +136,6 @@ class Node:
async def run(self):
async with self._tg as tg:
signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
tg.start_soon(self.router.run)
tg.start_soon(self.election.run)
if self.download_coordinator:
@@ -148,6 +147,8 @@ class Node:
if self.api:
tg.start_soon(self.api.run)
tg.start_soon(self._elect_loop)
signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
def shutdown(self):
# if this is our second call to shutdown, just sys.exit
@@ -188,6 +189,9 @@ class Node:
global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
command_receiver=self.router.receiver(topics.COMMANDS),
download_command_sender=self.router.sender(
topics.DOWNLOAD_COMMANDS
),
)
self._tg.start_soon(self.master.run)
elif (
@@ -247,7 +251,8 @@ 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")
# TODO: Refactor the current verbosity system
+3 -3
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,
@@ -80,6 +79,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,
@@ -141,7 +141,7 @@ async def generate_chat_stream(
if isinstance(chunk, ToolCallChunk):
tool_call_deltas = [
ToolCall(
id=str(uuid4()),
id=tool.id,
index=i,
function=tool,
)
@@ -207,7 +207,7 @@ async def collect_chat_response(
if isinstance(chunk, ToolCallChunk):
tool_calls.extend(
ToolCall(
id=str(uuid4()),
id=tool.id,
index=i,
function=tool,
)
+2 -3
View File
@@ -3,7 +3,6 @@
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.chunks import ErrorChunk, TokenChunk, ToolCallChunk
@@ -179,7 +178,7 @@ async def collect_claude_response(
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]
)
@@ -264,7 +263,7 @@ async def generate_claude_stream(
# 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
+4 -5
View File
@@ -3,7 +3,6 @@
from collections.abc import AsyncGenerator
from itertools import count
from typing import Any
from uuid import uuid4
from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
from exo.shared.types.common import CommandId
@@ -140,8 +139,8 @@ async def collect_responses_response(
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,
)
@@ -246,8 +245,8 @@ async def generate_responses_stream(
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(
+56 -17
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,
@@ -146,6 +148,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 +179,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 +227,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]()
@@ -289,7 +295,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)
@@ -386,7 +392,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):
@@ -689,6 +700,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:
@@ -1320,29 +1347,41 @@ class API:
]
async def run(self):
shutdown_ev = anyio.Event()
try:
async with create_task_group() as tg:
self._tg = tg
logger.info("Starting API")
tg.start_soon(self._apply_state)
tg.start_soon(self._pause_on_new_election)
tg.start_soon(self._cleanup_expired_images)
print_startup_banner(self.port)
tg.start_soon(self.run_api, shutdown_ev)
try:
await anyio.sleep_forever()
finally:
with anyio.CancelScope(shield=True):
shutdown_ev.set()
finally:
self._event_log.close()
self.command_sender.close()
self.global_event_receiver.close()
async def run_api(self, ev: anyio.Event):
cfg = Config()
cfg.bind = f"0.0.0.0:{self.port}"
cfg.bind = [f"0.0.0.0:{self.port}"]
# nb: shared.logging needs updating if any of this changes
cfg.accesslog = None
cfg.errorlog = "-"
cfg.logger_class = InterceptLogger
async with create_task_group() as tg:
self._tg = tg
logger.info("Starting API")
tg.start_soon(self._apply_state)
tg.start_soon(self._pause_on_new_election)
tg.start_soon(self._cleanup_expired_images)
print_startup_banner(self.port)
with anyio.CancelScope(shield=True):
await serve(
cast(ASGIFramework, self.app),
cfg,
shutdown_trigger=lambda: anyio.sleep_forever(),
shutdown_trigger=ev.wait,
)
self.command_sender.close()
self.global_event_receiver.close()
async def _apply_state(self):
with self.global_event_receiver as events:
async for f_event in events:
+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()
+35 -21
View File
@@ -4,18 +4,21 @@ 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,
delete_instance,
get_transition_events,
place_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.types.commands import (
CreateInstance,
DeleteInstance,
ForwarderCommand,
ForwarderDownloadCommand,
ImageEdits,
ImageGeneration,
PlaceInstance,
@@ -68,12 +71,9 @@ class Master:
session_id: SessionId,
*,
command_receiver: Receiver[ForwarderCommand],
# Receiving indexed events from the forwarder to be applied to state
# Ideally these would be WorkerForwarderEvents but type system says no :(
local_event_receiver: Receiver[ForwarderEvent],
# Send events to the forwarder to be indexed (usually from command processing)
# Ideally these would be MasterForwarderEvents but type system says no :(
global_event_sender: Sender[ForwarderEvent],
download_command_sender: Sender[ForwarderDownloadCommand],
):
self.state = State()
self._tg: TaskGroup = anyio.create_task_group()
@@ -83,6 +83,7 @@ class Master:
self.command_receiver = command_receiver
self.local_event_receiver = local_event_receiver
self.global_event_sender = global_event_sender
self.download_command_sender = download_command_sender
send, recv = channel[Event]()
self.event_sender: Sender[Event] = send
self._loopback_event_receiver: Receiver[Event] = recv
@@ -90,24 +91,26 @@ class Master:
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]] = {}
async def run(self):
logger.info("Starting 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)
self.global_event_sender.close()
self.local_event_receiver.close()
self.command_receiver.close()
self._loopback_event_sender.close()
self._loopback_event_receiver.close()
try:
async with self._tg as tg:
tg.start_soon(self._event_processor)
tg.start_soon(self._command_processor)
tg.start_soon(self._loopback_processor)
tg.start_soon(self._plan)
finally:
self._event_log.close()
self.global_event_sender.close()
self.local_event_receiver.close()
self.command_receiver.close()
self._loopback_event_sender.close()
self._loopback_event_receiver.close()
async def shutdown(self):
logger.info("Stopping Master")
@@ -280,6 +283,14 @@ class Master:
transition_events = get_transition_events(
self.state.instances, placement
)
for cmd in cancel_unnecessary_downloads(
placement, self.state.downloads
):
await self.download_command_sender.send(
ForwarderDownloadCommand(
origin=self.node_id, command=cmd
)
)
generated_events.extend(transition_events)
case PlaceInstance():
placement = place_instance(
@@ -332,10 +343,13 @@ class Master:
]
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)
except ValueError as e:
+32
View File
@@ -15,14 +15,20 @@ from exo.master.placement_utils import (
from exo.shared.models.model_cards import ModelId
from exo.shared.topology import Topology
from exo.shared.types.commands import (
CancelDownload,
CreateInstance,
DeleteInstance,
DownloadCommand,
PlaceInstance,
)
from exo.shared.types.common import NodeId
from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.worker.downloads import (
DownloadOngoing,
DownloadProgress,
)
from exo.shared.types.worker.instances import (
Instance,
InstanceId,
@@ -206,3 +212,29 @@ def get_transition_events(
)
return events
def cancel_unnecessary_downloads(
instances: Mapping[InstanceId, Instance],
download_status: Mapping[NodeId, Sequence[DownloadProgress]],
) -> Sequence[DownloadCommand]:
commands: list[DownloadCommand] = []
currently_downloading = [
(k, v.shard_metadata.model_card.model_id)
for k, vs in download_status.items()
for v in vs
if isinstance(v, (DownloadOngoing))
]
active_models = set(
(
node_id,
instance.shard_assignments.runner_to_shard[runner_id].model_card.model_id,
)
for instance in instances.values()
for node_id, runner_id in instance.shard_assignments.node_to_runner.items()
)
for pair in currently_downloading:
if pair not in active_models:
commands.append(CancelDownload(target_node_id=pair[0], model_id=pair[1]))
return commands
+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()
+3
View File
@@ -11,6 +11,7 @@ from exo.shared.models.model_cards import ModelCard, ModelTask
from exo.shared.types.commands import (
CommandId,
ForwarderCommand,
ForwarderDownloadCommand,
PlaceInstance,
TextGeneration,
)
@@ -47,6 +48,7 @@ async def test_master():
ge_sender, global_event_receiver = channel[ForwarderEvent]()
command_sender, co_receiver = channel[ForwarderCommand]()
local_event_sender, le_receiver = channel[ForwarderEvent]()
fcds, _fcdr = channel[ForwarderDownloadCommand]()
all_events: list[IndexedEvent] = []
@@ -67,6 +69,7 @@ async def test_master():
global_event_sender=ge_sender,
local_event_receiver=le_receiver,
command_receiver=co_receiver,
download_command_sender=fcds,
)
logger.info("run the master")
async with anyio.create_task_group() as tg:
+21 -17
View File
@@ -9,6 +9,7 @@ from anyio import (
BrokenResourceError,
ClosedResourceError,
create_task_group,
move_on_after,
sleep_forever,
)
from anyio.abc import TaskGroup
@@ -146,18 +147,21 @@ class Router:
async def run(self):
logger.debug("Starting Router")
async with create_task_group() as tg:
self._tg = tg
for topic in self.topic_routers:
router = self.topic_routers[topic]
tg.start_soon(router.run)
tg.start_soon(self._networking_recv)
tg.start_soon(self._networking_recv_connection_messages)
tg.start_soon(self._networking_publish)
# Router only shuts down if you cancel it.
await sleep_forever()
for topic in self.topic_routers:
await self._networking_unsubscribe(str(topic))
try:
async with create_task_group() as tg:
self._tg = tg
for topic in self.topic_routers:
router = self.topic_routers[topic]
tg.start_soon(router.run)
tg.start_soon(self._networking_recv)
tg.start_soon(self._networking_recv_connection_messages)
tg.start_soon(self._networking_publish)
# Router only shuts down if you cancel it.
await sleep_forever()
finally:
with move_on_after(1, shield=True):
for topic in self.topic_routers:
await self._networking_unsubscribe(str(topic))
async def shutdown(self):
logger.debug("Shutting down Router")
@@ -166,12 +170,12 @@ class Router:
self._tg.cancel_scope.cancel()
async def _networking_subscribe(self, topic: str):
logger.info(f"Subscribing to {topic}")
await self._net.gossipsub_subscribe(topic)
logger.info(f"Subscribed to {topic}")
async def _networking_unsubscribe(self, topic: str):
logger.info(f"Unsubscribing from {topic}")
await self._net.gossipsub_unsubscribe(topic)
logger.info(f"Unsubscribed from {topic}")
async def _networking_recv(self):
while True:
@@ -203,10 +207,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(
+24 -1
View File
@@ -32,6 +32,7 @@ from exo.shared.types.events import (
from exo.shared.types.profiling import (
NodeIdentity,
NodeNetworkInfo,
NodeRdmaCtlStatus,
NodeThunderboltInfo,
ThunderboltBridgeStatus,
)
@@ -48,7 +49,9 @@ from exo.utils.info_gatherer.info_gatherer import (
MemoryUsage,
MiscData,
NodeConfig,
NodeDiskUsage,
NodeNetworkInterfaces,
RdmaCtlStatus,
StaticNodeInformation,
ThunderboltBridgeInfo,
)
@@ -245,6 +248,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
}
@@ -261,6 +267,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 = (
@@ -278,10 +287,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,
}
)
@@ -310,6 +321,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():
@@ -324,7 +337,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,
@@ -376,6 +394,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"
+22 -21
View File
@@ -86,28 +86,29 @@ class Election:
async def run(self):
logger.info("Starting Election")
async with create_task_group() as tg:
self._tg = tg
tg.start_soon(self._election_receiver)
tg.start_soon(self._connection_receiver)
tg.start_soon(self._command_counter)
try:
async with create_task_group() as tg:
self._tg = tg
tg.start_soon(self._election_receiver)
tg.start_soon(self._connection_receiver)
tg.start_soon(self._command_counter)
# And start an election immediately, that instantly resolves
candidates: list[ElectionMessage] = []
logger.debug("Starting initial campaign")
self._candidates = candidates
await self._campaign(candidates, campaign_timeout=0.0)
logger.debug("Initial campaign finished")
# Cancel and wait for the last election to end
if self._campaign_cancel_scope is not None:
logger.debug("Cancelling campaign")
self._campaign_cancel_scope.cancel()
if self._campaign_done is not None:
logger.debug("Waiting for campaign to finish")
await self._campaign_done.wait()
logger.debug("Campaign cancelled and finished")
logger.info("Election finished")
# And start an election immediately, that instantly resolves
candidates: list[ElectionMessage] = []
logger.debug("Starting initial campaign")
self._candidates = candidates
await self._campaign(candidates, campaign_timeout=0.0)
logger.debug("Initial campaign finished")
finally:
# Cancel and wait for the last election to end
if self._campaign_cancel_scope is not None:
logger.debug("Cancelling campaign")
self._campaign_cancel_scope.cancel()
if self._campaign_done is not None:
logger.debug("Waiting for campaign to finish")
await self._campaign_done.wait()
logger.debug("Campaign cancelled and finished")
logger.info("Election shutdown")
async def elect(self, em: ElectionMessage) -> None:
logger.debug(f"Electing: {em}")
+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")
+4
View File
@@ -1,6 +1,7 @@
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
@@ -60,6 +61,7 @@ class ChatCompletionMessageText(BaseModel):
class ToolCallItem(BaseModel):
id: str = Field(default_factory=lambda: str(uuid4()))
name: str
arguments: str
@@ -197,6 +199,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
@@ -272,6 +275,7 @@ class AdvancedImageParams(BaseModel):
num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None
guidance: Annotated[float, Field(ge=1.0, le=20.0)] | None = None
negative_prompt: str | None = None
num_sync_steps: Annotated[int, Field(ge=1, le=100)] | None = None
class ImageGenerationTaskParams(BaseModel):
+6 -1
View File
@@ -80,7 +80,12 @@ class SetInstanceDraftModel(BaseCommand):
num_draft_tokens: int = 4
DownloadCommand = StartDownload | DeleteDownload
class CancelDownload(BaseCommand):
target_node_id: NodeId
model_id: ModelId
DownloadCommand = StartDownload | DeleteDownload | CancelDownload
Command = (

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