Compare commits
68 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| acb4818057 | |||
| b26b50053d | |||
| ebc7e6c100 | |||
| 54ad960ffb | |||
| 81c5d9379f | |||
| cd718ed726 | |||
| 3536161f15 | |||
| a6aa07ed83 | |||
| 131ad0ff36 | |||
| d01636100a | |||
| 79c8dbeacd | |||
| 0096159728 | |||
| 7a36d3968d | |||
| eee3432738 | |||
| e8c3a873a6 | |||
| afab3095b0 | |||
| b9d40e8e35 | |||
| 3a4d635d0c | |||
| 8485805042 | |||
| 4de8f801c7 | |||
| 5777bf3c39 | |||
| 886192f1e6 | |||
| d914acd64e | |||
| 37296c8249 | |||
| 28817d3ee3 | |||
| 0e1b9501c3 | |||
| f0d4ccbeb3 | |||
| 858dc808df | |||
| 635118ef24 | |||
| dc0bb5e13b | |||
| 152a27ea5d | |||
| db36bd5ac6 | |||
| 639243aa09 | |||
| db73c4fd5d | |||
| eaed92952c | |||
| ba611f9cd0 | |||
| eab3e0b456 | |||
| c4e874e97d | |||
| e23c3a3026 | |||
| 190e63e56d | |||
| 7660893538 | |||
| 9a2d2a4a7c | |||
| bea64fe85e | |||
| 14526d281a | |||
| 73e50df827 | |||
| 2b417f28ff | |||
| b65982ddd7 | |||
| 2fe689315b | |||
| 644c5573ce | |||
| 12c3015f52 | |||
| 365dd68d9a | |||
| d3d129581e | |||
| c90a0cec78 | |||
| e8c1337168 | |||
| 7024ddcf3e | |||
| dc89ba662a | |||
| 5cd96b5060 | |||
| 05986f77aa | |||
| dab7ed4821 | |||
| 2261014715 | |||
| 61d2a2b6cf | |||
| 0ff99a2c40 | |||
| fbb80e1cc9 | |||
| 8d94eab6c6 | |||
| f370452d7e | |||
| a4c2aa2b87 | |||
| 7312c535b4 | |||
| 18717023ad |
@@ -215,6 +215,22 @@ class StreamContext:
|
||||
traceback: object | None = ...,
|
||||
) -> None: ...
|
||||
|
||||
def device_info() -> dict[str, str | int]:
|
||||
"""
|
||||
Get information about the GPU device and system settings.
|
||||
|
||||
Currently returns:
|
||||
|
||||
* ``architecture``
|
||||
* ``max_buffer_size``
|
||||
* ``max_recommended_working_set_size``
|
||||
* ``memory_size``
|
||||
* ``resource_limit``
|
||||
|
||||
Returns:
|
||||
dict: A dictionary with string keys and string or integer values.
|
||||
"""
|
||||
|
||||
def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
|
||||
"""
|
||||
Element-wise absolute value.
|
||||
|
||||
@@ -73,9 +73,11 @@ class GenerationResponse:
|
||||
finish_reason: Optional[str] = ...
|
||||
|
||||
def maybe_quantize_kv_cache(
|
||||
prompt_cache, quantized_kv_start, kv_group_size, kv_bits
|
||||
): # -> None:
|
||||
...
|
||||
prompt_cache: Any,
|
||||
quantized_kv_start: int | None,
|
||||
kv_group_size: int | None,
|
||||
kv_bits: int | None,
|
||||
) -> None: ...
|
||||
def generate_step(
|
||||
prompt: mx.array,
|
||||
model: nn.Module,
|
||||
@@ -252,7 +254,14 @@ class BatchResponse:
|
||||
|
||||
texts: List[str]
|
||||
stats: BatchStats
|
||||
caches: Optional[List[List[Any]]]
|
||||
|
||||
def _left_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
|
||||
def _right_pad_prompts(prompts: Any, max_length: Optional[int] = ...) -> mx.array: ...
|
||||
def _make_cache(
|
||||
model: Any, left_padding: Any, max_kv_size: Optional[int]
|
||||
) -> List[Any]: ...
|
||||
def _merge_caches(caches: Any) -> List[Any]: ...
|
||||
@dataclass
|
||||
class Batch:
|
||||
uids: List[int]
|
||||
@@ -261,39 +270,71 @@ class Batch:
|
||||
max_tokens: List[int]
|
||||
num_tokens: List[int]
|
||||
cache: List[Any]
|
||||
def __len__(self): # -> int:
|
||||
...
|
||||
def filter(self, keep_idx: List[int]): # -> None:
|
||||
...
|
||||
def extend(self, other): # -> None:
|
||||
...
|
||||
samplers: List[Any]
|
||||
logits_processors: List[Any]
|
||||
tokens: List[mx.array]
|
||||
def __len__(self) -> int: ...
|
||||
def filter(self, keep_idx: List[int]) -> None: ...
|
||||
def extend(self, other: "Batch") -> None: ...
|
||||
def extract_cache(self, idx: int) -> List[Any]: ...
|
||||
|
||||
class BatchGenerator:
|
||||
model: Any
|
||||
max_kv_size: Optional[int]
|
||||
prefill_step_size: int
|
||||
unprocessed_prompts: List[Any]
|
||||
active_batch: Optional[Batch]
|
||||
prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
|
||||
_stats: BatchStats
|
||||
|
||||
@dataclass
|
||||
class Response:
|
||||
uid: int
|
||||
token: int
|
||||
logprobs: mx.array
|
||||
finish_reason: Optional[str]
|
||||
prompt_cache: Any
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model,
|
||||
model: Any,
|
||||
max_tokens: int = ...,
|
||||
stop_tokens: Optional[set] = ...,
|
||||
stop_tokens: Optional[set[int]] = ...,
|
||||
sampler: Optional[Callable[[mx.array], mx.array]] = ...,
|
||||
logits_processors: Optional[
|
||||
List[Callable[[mx.array, mx.array], mx.array]]
|
||||
] = ...,
|
||||
completion_batch_size: int = ...,
|
||||
prefill_batch_size: int = ...,
|
||||
prefill_step_size: int = ...,
|
||||
prompt_progress_callback: Optional[
|
||||
Callable[[List[Tuple[int, int, int]]], None]
|
||||
] = ...,
|
||||
max_kv_size: Optional[int] = ...,
|
||||
) -> None: ...
|
||||
def close(self) -> None: ...
|
||||
def insert(
|
||||
self, prompts, max_tokens: Union[List[int], int, None] = ...
|
||||
): # -> list[Any]:
|
||||
...
|
||||
def stats(self): # -> BatchStats:
|
||||
...
|
||||
def next(self): # -> list[Any]:
|
||||
...
|
||||
self,
|
||||
prompts: Any,
|
||||
max_tokens: Union[List[int], int, None] = ...,
|
||||
caches: Any = ...,
|
||||
samplers: Optional[List[Any]] = ...,
|
||||
logits_processors: Optional[List[Any]] = ...,
|
||||
) -> List[int]: ...
|
||||
def remove(
|
||||
self, uids: List[int], return_prompt_caches: bool = ...
|
||||
) -> Optional[dict[int, List[Any]]]: ...
|
||||
def stats(self) -> BatchStats: ...
|
||||
def next(self) -> List[Response]: ...
|
||||
def _process_prompts(self, prompts: List[Any]) -> Batch: ...
|
||||
def _step(
|
||||
self,
|
||||
input_tokens: mx.array,
|
||||
prompt_cache: List[Any],
|
||||
samplers: Optional[List[Any]],
|
||||
logits_processors: Optional[List[Any]],
|
||||
tokens: List[mx.array],
|
||||
) -> Tuple[mx.array, List[mx.array]]: ...
|
||||
|
||||
def batch_generate(
|
||||
model,
|
||||
|
||||
@@ -16,7 +16,7 @@ class Cache(Protocol):
|
||||
self, keys: mx.array, values: mx.array
|
||||
) -> tuple[mx.array, mx.array]: ...
|
||||
@property
|
||||
def state(self) -> tuple[mx.array, mx.array]: ...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
|
||||
@@ -92,13 +92,14 @@ class _BaseCache(Cache):
|
||||
values: mx.array
|
||||
offset: int
|
||||
@property
|
||||
def state(self) -> tuple[mx.array, mx.array]: ...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
@property
|
||||
def meta_state(self) -> Literal[""]: ...
|
||||
@meta_state.setter
|
||||
def meta_state(self, v) -> None: ...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def is_trimmable(self) -> Literal[False]: ...
|
||||
@classmethod
|
||||
def from_state(cls, state, meta_state) -> Self: ...
|
||||
@@ -114,15 +115,13 @@ class ConcatenateKVCache(_BaseCache):
|
||||
def update_and_fetch(self, keys, values): # -> tuple[Any | array, Any | array]:
|
||||
...
|
||||
@property
|
||||
def state(self): # -> tuple[Any | array | None, Any | array | None]:
|
||||
...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
@@ -132,10 +131,7 @@ class QuantizedKVCache(_BaseCache):
|
||||
def update_and_fetch(self, keys, values): # -> Any:
|
||||
...
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
): # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
|
||||
...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -147,8 +143,7 @@ class QuantizedKVCache(_BaseCache):
|
||||
...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
|
||||
@@ -160,22 +155,30 @@ class KVCache(_BaseCache):
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
) -> tuple[array, array]: ...
|
||||
) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v) -> None: ...
|
||||
def is_trimmable(self): # -> Literal[True]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
) -> QuantizedKVCache: ...
|
||||
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
|
||||
...
|
||||
def make_mask(
|
||||
self, *args: Any, **kwargs: Any
|
||||
) -> mx.array | Literal["causal"] | None: ...
|
||||
|
||||
class RotatingKVCache(_BaseCache):
|
||||
step = ...
|
||||
keys: mx.array | None
|
||||
values: mx.array | None
|
||||
keep: int
|
||||
max_size: int
|
||||
_idx: int
|
||||
def __init__(self, max_size, keep=...) -> None: ...
|
||||
def _trim(
|
||||
self, trim_size: int, v: mx.array, append: mx.array | None = ...
|
||||
) -> mx.array: ...
|
||||
def update_and_fetch(
|
||||
self, keys, values
|
||||
): # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
|
||||
@@ -183,8 +186,7 @@ class RotatingKVCache(_BaseCache):
|
||||
@property
|
||||
def state(
|
||||
self,
|
||||
): # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
|
||||
...
|
||||
) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -196,8 +198,7 @@ class RotatingKVCache(_BaseCache):
|
||||
...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
def to_quantized(
|
||||
self, group_size: int = ..., bits: int = ...
|
||||
) -> QuantizedKVCache: ...
|
||||
@@ -212,8 +213,7 @@ class ArraysCache(_BaseCache):
|
||||
...
|
||||
def __getitem__(self, idx): ...
|
||||
@property
|
||||
def state(self): # -> list[Any | array] | list[array]:
|
||||
...
|
||||
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
@@ -227,8 +227,7 @@ class ArraysCache(_BaseCache):
|
||||
In-place extend this cache with the other cache.
|
||||
"""
|
||||
|
||||
def make_mask(self, N: int): # -> array | None:
|
||||
...
|
||||
def make_mask(self, N: int) -> mx.array | None: ...
|
||||
|
||||
class MambaCache(ArraysCache):
|
||||
def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
|
||||
@@ -239,8 +238,7 @@ class ChunkedKVCache(KVCache):
|
||||
...
|
||||
def update_and_fetch(self, keys, values): # -> tuple[array, array]:
|
||||
...
|
||||
def trim(self, n): # -> int:
|
||||
...
|
||||
def trim(self, n: int) -> int: ...
|
||||
@property
|
||||
def meta_state(self): # -> tuple[str, ...]:
|
||||
...
|
||||
@@ -253,10 +251,9 @@ class CacheList(_BaseCache):
|
||||
def __getitem__(self, idx): ...
|
||||
def is_trimmable(self): # -> bool:
|
||||
...
|
||||
def trim(self, n): ...
|
||||
def trim(self, n: int) -> int: ...
|
||||
@property
|
||||
def state(self): # -> list[Any]:
|
||||
...
|
||||
def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
|
||||
@state.setter
|
||||
def state(self, v): # -> None:
|
||||
...
|
||||
|
||||
@@ -0,0 +1,153 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .qwen3_next import (
|
||||
Qwen3NextAttention as Attention,
|
||||
Qwen3NextMLP as MLP,
|
||||
Qwen3NextRMSNormGated as RMSNormGated,
|
||||
Qwen3NextSparseMoeBlock,
|
||||
)
|
||||
|
||||
SparseMoeBlock = Qwen3NextSparseMoeBlock
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
@dataclass
|
||||
class TextModelArgs:
|
||||
model_type: str
|
||||
hidden_size: int
|
||||
intermediate_size: int
|
||||
num_hidden_layers: int
|
||||
num_attention_heads: int
|
||||
rms_norm_eps: float
|
||||
vocab_size: int
|
||||
num_key_value_heads: int
|
||||
max_position_embeddings: int
|
||||
linear_num_value_heads: int
|
||||
linear_num_key_heads: int
|
||||
linear_key_head_dim: int
|
||||
linear_value_head_dim: int
|
||||
linear_conv_kernel_dim: int
|
||||
tie_word_embeddings: bool
|
||||
attention_bias: bool
|
||||
head_dim: Optional[int]
|
||||
full_attention_interval: int
|
||||
num_experts: int
|
||||
num_experts_per_tok: int
|
||||
decoder_sparse_step: int
|
||||
shared_expert_intermediate_size: int
|
||||
moe_intermediate_size: int
|
||||
norm_topk_prob: bool
|
||||
rope_parameters: Optional[dict[str, Any]]
|
||||
partial_rotary_factor: float
|
||||
rope_theta: float
|
||||
rope_scaling: Optional[dict[str, Any]]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> TextModelArgs: ...
|
||||
def __post_init__(self) -> None: ...
|
||||
|
||||
class GatedDeltaNet(nn.Module):
|
||||
hidden_size: int
|
||||
num_v_heads: int
|
||||
num_k_heads: int
|
||||
head_k_dim: int
|
||||
head_v_dim: int
|
||||
key_dim: int
|
||||
value_dim: int
|
||||
conv_kernel_size: int
|
||||
conv_dim: int
|
||||
conv1d: nn.Conv1d
|
||||
in_proj_qkv: nn.Linear
|
||||
in_proj_z: nn.Linear
|
||||
in_proj_b: nn.Linear
|
||||
in_proj_a: nn.Linear
|
||||
dt_bias: mx.array
|
||||
A_log: mx.array
|
||||
norm: RMSNormGated
|
||||
out_proj: nn.Linear
|
||||
|
||||
def __init__(self, config: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class DecoderLayer(nn.Module):
|
||||
is_linear: bool
|
||||
linear_attn: GatedDeltaNet
|
||||
self_attn: Attention
|
||||
input_layernorm: nn.RMSNorm
|
||||
post_attention_layernorm: nn.RMSNorm
|
||||
mlp: MLP | SparseMoeBlock
|
||||
|
||||
def __init__(self, args: TextModelArgs, layer_idx: int) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
x: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
cache: Optional[Any] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class Qwen3_5TextModel(nn.Module):
|
||||
embed_tokens: nn.Embedding
|
||||
layers: list[DecoderLayer]
|
||||
norm: nn.RMSNorm
|
||||
ssm_idx: int
|
||||
fa_idx: int
|
||||
|
||||
def __init__(self, args: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
|
||||
class TextModel(nn.Module):
|
||||
args: TextModelArgs
|
||||
model_type: str
|
||||
model: Qwen3_5TextModel
|
||||
lm_head: nn.Linear
|
||||
|
||||
def __init__(self, args: TextModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
text_config: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
|
||||
class Model(nn.Module):
|
||||
args: ModelArgs
|
||||
model_type: str
|
||||
language_model: TextModel
|
||||
|
||||
def __init__(self, args: ModelArgs) -> None: ...
|
||||
def __call__(
|
||||
self,
|
||||
inputs: mx.array,
|
||||
cache: Optional[Any] = None,
|
||||
input_embeddings: Optional[mx.array] = None,
|
||||
) -> mx.array: ...
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@property
|
||||
def layers(self) -> list[DecoderLayer]: ...
|
||||
def make_cache(self) -> list[ArraysCache | KVCache]: ...
|
||||
@@ -0,0 +1,19 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import mlx.core as mx
|
||||
import mlx.nn as nn
|
||||
|
||||
from .cache import ArraysCache, KVCache
|
||||
from .qwen3_5 import DecoderLayer, Model as Qwen3_5Model, TextModel
|
||||
|
||||
@dataclass
|
||||
class ModelArgs:
|
||||
model_type: str
|
||||
text_config: dict[str, Any]
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
|
||||
|
||||
class Model(Qwen3_5Model):
|
||||
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
|
||||
@@ -7,6 +7,15 @@ import mlx.nn as nn
|
||||
|
||||
from .switch_layers import SwitchGLU
|
||||
|
||||
class Qwen3NextRMSNormGated(nn.Module):
|
||||
eps: float
|
||||
weight: mx.array
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = ...) -> None: ...
|
||||
def __call__(
|
||||
self, hidden_states: mx.array, gate: mx.array | None = None
|
||||
) -> mx.array: ...
|
||||
|
||||
class Qwen3NextMLP(nn.Module):
|
||||
gate_proj: nn.Linear
|
||||
down_proj: nn.Linear
|
||||
|
||||
@@ -48,7 +48,7 @@ def make_logits_processors(
|
||||
logit_bias: Optional[Dict[int, float]] = ...,
|
||||
repetition_penalty: Optional[float] = ...,
|
||||
repetition_context_size: Optional[int] = ...,
|
||||
): # -> list[Any]:
|
||||
) -> list[Callable[[mx.array, mx.array], mx.array]]:
|
||||
"""
|
||||
Make logits processors for use with ``generate_step``.
|
||||
|
||||
|
||||
+113
@@ -39,6 +39,119 @@ Write pure functions where possible. When adding new code, prefer Rust unless th
|
||||
|
||||
Run `nix fmt` to auto-format your code before submitting.
|
||||
|
||||
## Model Cards
|
||||
|
||||
EXO uses TOML-based model cards to define model metadata and capabilities. Model cards are stored in:
|
||||
- `resources/inference_model_cards/` for text generation models
|
||||
- `resources/image_model_cards/` for image generation models
|
||||
- `~/.exo/custom_model_cards/` for user-added custom models
|
||||
|
||||
### Adding a Model Card
|
||||
|
||||
To add a new model, create a TOML file with the following structure:
|
||||
|
||||
```toml
|
||||
model_id = "mlx-community/Llama-3.2-1B-Instruct-4bit"
|
||||
n_layers = 16
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "llama"
|
||||
quantization = "4bit"
|
||||
base_model = "Llama 3.2 1B"
|
||||
capabilities = ["text"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 729808896
|
||||
```
|
||||
|
||||
### Required Fields
|
||||
|
||||
- `model_id`: Hugging Face model identifier
|
||||
- `n_layers`: Number of transformer layers
|
||||
- `hidden_size`: Hidden dimension size
|
||||
- `supports_tensor`: Whether the model supports tensor parallelism
|
||||
- `tasks`: List of supported tasks (`TextGeneration`, `TextToImage`, `ImageToImage`)
|
||||
- `family`: Model family (e.g., "llama", "deepseek", "qwen")
|
||||
- `quantization`: Quantization level (e.g., "4bit", "8bit", "bf16")
|
||||
- `base_model`: Human-readable base model name
|
||||
- `capabilities`: List of capabilities (e.g., `["text"]`, `["text", "thinking"]`)
|
||||
|
||||
### Optional Fields
|
||||
|
||||
- `components`: For multi-component models (like image models with separate text encoders and transformers)
|
||||
- `uses_cfg`: Whether the model uses classifier-free guidance (for image models)
|
||||
- `trust_remote_code`: Whether to allow remote code execution (defaults to `false` for security)
|
||||
|
||||
### Capabilities
|
||||
|
||||
The `capabilities` field defines what the model can do:
|
||||
- `text`: Standard text generation
|
||||
- `thinking`: Model supports chain-of-thought reasoning
|
||||
- `thinking_toggle`: Thinking can be enabled/disabled via `enable_thinking` parameter
|
||||
- `image_edit`: Model supports image-to-image editing (FLUX.1-Kontext)
|
||||
|
||||
### Security Note
|
||||
|
||||
By default, `trust_remote_code` is set to `false` for security. Only enable it if the model explicitly requires remote code execution from the Hugging Face hub.
|
||||
|
||||
## API Adapters
|
||||
|
||||
EXO supports multiple API formats through an adapter pattern. Adapters convert API-specific request formats to the internal `TextGenerationTaskParams` format and convert internal token chunks back to API-specific responses.
|
||||
|
||||
### Adapter Architecture
|
||||
|
||||
All adapters live in `src/exo/master/adapters/` and follow the same pattern:
|
||||
|
||||
1. Convert API-specific requests to `TextGenerationTaskParams`
|
||||
2. Handle both streaming and non-streaming response generation
|
||||
3. Convert internal `TokenChunk` objects to API-specific formats
|
||||
4. Manage error handling and edge cases
|
||||
|
||||
### Existing Adapters
|
||||
|
||||
- `chat_completions.py`: OpenAI Chat Completions API
|
||||
- `claude.py`: Anthropic Claude Messages API
|
||||
- `responses.py`: OpenAI Responses API
|
||||
- `ollama.py`: Ollama API (for OpenWebUI compatibility)
|
||||
|
||||
### Adding a New API Adapter
|
||||
|
||||
To add support for a new API format:
|
||||
|
||||
1. Create a new adapter file in `src/exo/master/adapters/`
|
||||
2. Implement a request conversion function:
|
||||
```python
|
||||
def your_api_request_to_text_generation(
|
||||
request: YourAPIRequest,
|
||||
) -> TextGenerationTaskParams:
|
||||
# Convert API request to internal format
|
||||
pass
|
||||
```
|
||||
3. Implement streaming response generation:
|
||||
```python
|
||||
async def generate_your_api_stream(
|
||||
command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None],
|
||||
) -> AsyncGenerator[str, None]:
|
||||
# Convert internal chunks to API-specific streaming format
|
||||
pass
|
||||
```
|
||||
4. Implement non-streaming response collection:
|
||||
```python
|
||||
async def collect_your_api_response(
|
||||
command_id: CommandId,
|
||||
chunk_stream: AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None],
|
||||
) -> AsyncGenerator[str]:
|
||||
# Collect all chunks and return single response
|
||||
pass
|
||||
```
|
||||
5. Register the adapter endpoints in `src/exo/master/api.py`
|
||||
|
||||
The adapter pattern keeps API-specific logic isolated from core inference systems. Internal systems (worker, runner, event sourcing) only see `TextGenerationTaskParams` and `TokenChunk` objects - no API-specific types cross the adapter boundary.
|
||||
|
||||
For detailed API documentation, see [docs/api.md](docs/api.md).
|
||||
|
||||
## Testing
|
||||
|
||||
EXO relies heavily on manual testing at this point in the project, but this is evolving. Before submitting a change, test both before and after to demonstrate how your change improves behavior. Do the best you can with the hardware you have available - if you need help testing, ask and we'll do our best to assist. Add automated tests where possible - we're actively working to substantially improve our automated testing story.
|
||||
|
||||
Generated
+24
@@ -216,6 +216,28 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.111",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-trait"
|
||||
version = "0.1.89"
|
||||
@@ -2759,6 +2781,7 @@ dependencies = [
|
||||
name = "networking"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"async-stream",
|
||||
"delegate",
|
||||
"either",
|
||||
"extend",
|
||||
@@ -2767,6 +2790,7 @@ dependencies = [
|
||||
"keccak-const",
|
||||
"libp2p",
|
||||
"log",
|
||||
"pin-project",
|
||||
"tokio",
|
||||
"tracing-subscriber",
|
||||
"util",
|
||||
|
||||
+2
-5
@@ -1,10 +1,6 @@
|
||||
[workspace]
|
||||
resolver = "3"
|
||||
members = [
|
||||
"rust/networking",
|
||||
"rust/exo_pyo3_bindings",
|
||||
"rust/util",
|
||||
]
|
||||
members = ["rust/networking", "rust/exo_pyo3_bindings", "rust/util"]
|
||||
|
||||
[workspace.package]
|
||||
version = "0.0.1"
|
||||
@@ -34,6 +30,7 @@ delegate = "0.13"
|
||||
keccak-const = "0.2"
|
||||
|
||||
# Async dependencies
|
||||
async-stream = "0.3"
|
||||
tokio = "1.46"
|
||||
futures-lite = "2.6.1"
|
||||
futures-timer = "3.0"
|
||||
|
||||
@@ -26,6 +26,8 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
|
||||
- **Topology-Aware Auto Parallel**: exo figures out the best way to split your model across all available devices based on a realtime view of your device topology. It takes into account device resources and network latency/bandwidth between each link.
|
||||
- **Tensor Parallelism**: exo supports sharding models, for up to 1.8x speedup on 2 devices and 3.2x speedup on 4 devices.
|
||||
- **MLX Support**: exo uses [MLX](https://github.com/ml-explore/mlx) as an inference backend and [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) for distributed communication.
|
||||
- **Multiple API Compatibility**: Compatible with OpenAI Chat Completions API, Claude Messages API, OpenAI Responses API, and Ollama API - use your existing tools and clients.
|
||||
- **Custom Model Support**: Load custom models from HuggingFace hub to expand the range of available models.
|
||||
|
||||
## Dashboard
|
||||
|
||||
@@ -196,6 +198,8 @@ exo follows the [XDG Base Directory Specification](https://specifications.freede
|
||||
- **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
|
||||
- **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
|
||||
- **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
|
||||
- **Log files**: `~/.cache/exo/exo_log/` (with automatic log rotation)
|
||||
- **Custom model cards**: `~/.local/share/exo/custom_model_cards/`
|
||||
|
||||
You can override these locations by setting the corresponding XDG environment variables.
|
||||
|
||||
@@ -275,8 +279,47 @@ After that, RDMA will be enabled in macOS and exo will take care of the rest.
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables
|
||||
|
||||
exo supports several environment variables for configuration:
|
||||
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `EXO_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None |
|
||||
| `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
|
||||
| `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
|
||||
| `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
|
||||
| `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
|
||||
| `EXO_FAST_SYNCH` | Control MLX_METAL_FAST_SYNCH behavior (for JACCL backend) | Auto |
|
||||
| `EXO_TRACING_ENABLED` | Enable distributed tracing for performance analysis | `false` |
|
||||
|
||||
**Example usage:**
|
||||
|
||||
```bash
|
||||
# Use pre-downloaded models from NFS mount
|
||||
EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-models uv run exo
|
||||
|
||||
# Run in offline mode
|
||||
EXO_OFFLINE=true uv run exo
|
||||
|
||||
# Enable image models
|
||||
EXO_ENABLE_IMAGE_MODELS=true uv run exo
|
||||
|
||||
# Use custom namespace for cluster isolation
|
||||
EXO_LIBP2P_NAMESPACE=my-dev-cluster uv run exo
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
### Using the API
|
||||
|
||||
exo provides multiple API-compatible interfaces for maximum compatibility with existing tools:
|
||||
|
||||
- **OpenAI Chat Completions API** - Compatible with OpenAI clients
|
||||
- **Claude Messages API** - Compatible with Anthropic's Claude format
|
||||
- **OpenAI Responses API** - Compatible with OpenAI's Responses format
|
||||
- **Ollama API** - Compatible with Ollama and tools like OpenWebUI
|
||||
|
||||
If you prefer to interact with exo via the API, here is an example creating an instance of a small model (`mlx-community/Llama-3.2-1B-Instruct-4bit`), sending a chat completions request and deleting the instance.
|
||||
|
||||
---
|
||||
@@ -366,14 +409,85 @@ When you're done, delete the instance by its ID (find it via `/state` or `/insta
|
||||
curl -X DELETE http://localhost:52415/instance/YOUR_INSTANCE_ID
|
||||
```
|
||||
|
||||
### Claude Messages API Compatibility
|
||||
|
||||
Use the Claude Messages API format with the `/v1/messages` endpoint:
|
||||
|
||||
```bash
|
||||
curl -N -X POST http://localhost:52415/v1/messages \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"max_tokens": 1024,
|
||||
"stream": true
|
||||
}'
|
||||
```
|
||||
|
||||
### OpenAI Responses API Compatibility
|
||||
|
||||
Use the OpenAI Responses API format with the `/v1/responses` endpoint:
|
||||
|
||||
```bash
|
||||
curl -N -X POST http://localhost:52415/v1/responses \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"stream": true
|
||||
}'
|
||||
```
|
||||
|
||||
### Ollama API Compatibility
|
||||
|
||||
exo supports Ollama API endpoints for compatibility with tools like OpenWebUI:
|
||||
|
||||
```bash
|
||||
# Ollama chat
|
||||
curl -X POST http://localhost:52415/ollama/api/chat \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "mlx-community/Llama-3.2-1B-Instruct-4bit",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"}
|
||||
],
|
||||
"stream": false
|
||||
}'
|
||||
|
||||
# List models (Ollama format)
|
||||
curl http://localhost:52415/ollama/api/tags
|
||||
```
|
||||
|
||||
### Custom Model Loading from HuggingFace
|
||||
|
||||
You can add custom models from the HuggingFace hub:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:52415/models/add \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model_id": "mlx-community/my-custom-model"
|
||||
}'
|
||||
```
|
||||
|
||||
**Security Note:**
|
||||
|
||||
Custom models requiring `trust_remote_code` in their configuration must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code execution. Models are fetched from HuggingFace and stored locally as custom model cards.
|
||||
|
||||
**Other useful API endpoints*:**
|
||||
|
||||
- List all models: `curl http://localhost:52415/models`
|
||||
- List downloaded models only: `curl http://localhost:52415/models?status=downloaded`
|
||||
- Search HuggingFace: `curl "http://localhost:52415/models/search?query=llama&limit=10"`
|
||||
- Inspect instance IDs and deployment state: `curl http://localhost:52415/state`
|
||||
|
||||
For further details, see:
|
||||
|
||||
- API basic documentation in [docs/api.md](docs/api.md).
|
||||
- API documentation in [docs/api.md](docs/api.md).
|
||||
- API types and endpoints in [src/exo/master/api.py](src/exo/master/api.py).
|
||||
|
||||
---
|
||||
@@ -432,4 +546,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
|
||||
See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
|
||||
|
||||
+215
-148
@@ -15,14 +15,23 @@ struct ContentView: View {
|
||||
@EnvironmentObject private var localNetworkChecker: LocalNetworkChecker
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var settingsWindowController: SettingsWindowController
|
||||
@State private var focusedNode: NodeViewModel?
|
||||
@State private var deletingInstanceIDs: Set<String> = []
|
||||
@State private var showAllNodes = false
|
||||
@State private var showAllInstances = false
|
||||
@State private var baseURLCopied = false
|
||||
@State private var showAdvanced = false
|
||||
@State private var showDebugInfo = false
|
||||
@State private var bugReportInFlight = false
|
||||
@State private var bugReportMessage: String?
|
||||
private enum BugReportPhase: Equatable {
|
||||
case idle
|
||||
case prompting
|
||||
case sending(String)
|
||||
case success(String)
|
||||
case failure(String)
|
||||
}
|
||||
@State private var bugReportPhase: BugReportPhase = .idle
|
||||
@State private var bugReportUserDescription: String = ""
|
||||
@State private var uninstallInProgress = false
|
||||
@State private var pendingNamespace: String = ""
|
||||
@State private var pendingHFToken: String = ""
|
||||
@@ -258,139 +267,79 @@ struct ContentView: View {
|
||||
VStack(alignment: .leading, spacing: 0) {
|
||||
if controller.status != .stopped {
|
||||
dashboardButton
|
||||
baseURLRow
|
||||
Divider()
|
||||
.padding(.vertical, 8)
|
||||
} else {
|
||||
Divider()
|
||||
.padding(.vertical, 4)
|
||||
}
|
||||
advancedSection
|
||||
.padding(.bottom, 8)
|
||||
controlButton(title: "Quit", tint: .secondary) {
|
||||
HoverButton(
|
||||
title: "Settings",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "gear"
|
||||
) {
|
||||
settingsWindowController.open(
|
||||
controller: controller,
|
||||
updater: updater,
|
||||
networkStatusService: networkStatusService,
|
||||
thunderboltBridgeService: thunderboltBridgeService,
|
||||
stateService: stateService
|
||||
)
|
||||
}
|
||||
HoverButton(
|
||||
title: "Check for Updates",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "arrow.triangle.2.circlepath"
|
||||
) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
.padding(.bottom, 8)
|
||||
HoverButton(title: "Quit", tint: .secondary) {
|
||||
controller.stop()
|
||||
NSApplication.shared.terminate(nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var advancedSection: some View {
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
HStack {
|
||||
Text("Advanced")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
Spacer()
|
||||
collapseButton(isExpanded: $showAdvanced)
|
||||
}
|
||||
.animation(nil, value: showAdvanced)
|
||||
if showAdvanced {
|
||||
VStack(alignment: .leading, spacing: 8) {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("Cluster Namespace")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
HStack {
|
||||
TextField("optional", text: $pendingNamespace)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingNamespace = controller.customNamespace
|
||||
}
|
||||
Button("Save & Restart") {
|
||||
controller.customNamespace = pendingNamespace
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingNamespace == controller.customNamespace)
|
||||
}
|
||||
}
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text("HuggingFace Token")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
HStack {
|
||||
SecureField("optional", text: $pendingHFToken)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingHFToken = controller.hfToken
|
||||
}
|
||||
Button("Save & Restart") {
|
||||
controller.hfToken = pendingHFToken
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingHFToken == controller.hfToken)
|
||||
}
|
||||
}
|
||||
Divider()
|
||||
HStack {
|
||||
Toggle(
|
||||
"Enable Image Models (experimental)", isOn: $pendingEnableImageModels
|
||||
)
|
||||
.toggleStyle(.switch)
|
||||
.font(.caption2)
|
||||
.onAppear {
|
||||
pendingEnableImageModels = controller.enableImageModels
|
||||
}
|
||||
|
||||
Spacer()
|
||||
|
||||
Button("Save & Restart") {
|
||||
controller.enableImageModels = pendingEnableImageModels
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.disabled(pendingEnableImageModels == controller.enableImageModels)
|
||||
}
|
||||
HoverButton(title: "Check for Updates", small: true) {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
debugSection
|
||||
HoverButton(title: "Uninstall", tint: .red, small: true) {
|
||||
showUninstallConfirmationAlert()
|
||||
}
|
||||
.disabled(uninstallInProgress)
|
||||
}
|
||||
.transition(.opacity)
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.25), value: showAdvanced)
|
||||
}
|
||||
|
||||
private func controlButton(title: String, tint: Color = .primary, action: @escaping () -> Void)
|
||||
-> some View
|
||||
{
|
||||
HoverButton(title: title, tint: tint, trailingSystemImage: nil, action: action)
|
||||
}
|
||||
|
||||
private var dashboardButton: some View {
|
||||
Button {
|
||||
HoverButton(
|
||||
title: "Web Dashboard",
|
||||
tint: .primary,
|
||||
trailingSystemImage: "arrow.up.right"
|
||||
) {
|
||||
guard let url = URL(string: "http://localhost:52415/") else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
} label: {
|
||||
HStack {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
.imageScale(.small)
|
||||
Text("Dashboard")
|
||||
.fontWeight(.medium)
|
||||
Spacer()
|
||||
}
|
||||
.padding(.vertical, 8)
|
||||
.padding(.horizontal, 10)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 8, style: .continuous)
|
||||
.fill(Color(red: 1.0, green: 0.87, blue: 0.0).opacity(0.2))
|
||||
)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.padding(.bottom, 4)
|
||||
}
|
||||
|
||||
private var baseURLRow: some View {
|
||||
HStack(spacing: 6) {
|
||||
Image(systemName: "link")
|
||||
.imageScale(.small)
|
||||
.foregroundColor(.secondary)
|
||||
Text("localhost:52415/v1")
|
||||
.font(.system(.caption, design: .monospaced))
|
||||
.foregroundColor(.primary)
|
||||
Spacer()
|
||||
Button {
|
||||
NSPasteboard.general.clearContents()
|
||||
NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string)
|
||||
baseURLCopied = true
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 2) {
|
||||
baseURLCopied = false
|
||||
}
|
||||
} label: {
|
||||
Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc")
|
||||
.imageScale(.small)
|
||||
.foregroundColor(baseURLCopied ? .green : .secondary)
|
||||
.contentTransition(.symbolEffect(.replace))
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.help("Copy API base URL")
|
||||
}
|
||||
.padding(.vertical, 4)
|
||||
.padding(.horizontal, 8)
|
||||
}
|
||||
|
||||
private func collapseButton(isExpanded: Binding<Bool>) -> some View {
|
||||
@@ -611,39 +560,115 @@ struct ContentView: View {
|
||||
}
|
||||
|
||||
private var sendBugReportButton: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Button {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
} label: {
|
||||
HStack {
|
||||
if bugReportInFlight {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
switch bugReportPhase {
|
||||
case .idle:
|
||||
Button {
|
||||
bugReportPhase = .prompting
|
||||
bugReportUserDescription = ""
|
||||
} label: {
|
||||
HStack {
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
}
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
.padding(.vertical, 6)
|
||||
.padding(.horizontal, 8)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color.accentColor.opacity(0.12))
|
||||
)
|
||||
}
|
||||
.padding(.vertical, 6)
|
||||
.padding(.horizontal, 8)
|
||||
.buttonStyle(.plain)
|
||||
|
||||
case .prompting:
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
Text("What's the issue? (optional)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
TextEditor(text: $bugReportUserDescription)
|
||||
.font(.caption2)
|
||||
.frame(height: 60)
|
||||
.overlay(
|
||||
RoundedRectangle(cornerRadius: 4)
|
||||
.stroke(Color.secondary.opacity(0.3), lineWidth: 1)
|
||||
)
|
||||
HStack(spacing: 8) {
|
||||
Button("Send") {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.borderedProminent)
|
||||
.controlSize(.small)
|
||||
Button("Cancel") {
|
||||
bugReportPhase = .idle
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.bordered)
|
||||
.controlSize(.small)
|
||||
}
|
||||
}
|
||||
.padding(8)
|
||||
.background(
|
||||
RoundedRectangle(cornerRadius: 6)
|
||||
.fill(Color.accentColor.opacity(0.12))
|
||||
.fill(Color.accentColor.opacity(0.06))
|
||||
)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
.disabled(bugReportInFlight)
|
||||
|
||||
if let message = bugReportMessage {
|
||||
Text(message)
|
||||
case .sending(let message):
|
||||
HStack(spacing: 6) {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
case .success(let message):
|
||||
VStack(alignment: .leading, spacing: 6) {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
Button {
|
||||
openGitHubIssue()
|
||||
} label: {
|
||||
HStack(spacing: 4) {
|
||||
Image(systemName: "arrow.up.right.square")
|
||||
.imageScale(.small)
|
||||
Text("Create GitHub Issue")
|
||||
.font(.caption2)
|
||||
}
|
||||
}
|
||||
.buttonStyle(.bordered)
|
||||
.controlSize(.small)
|
||||
Button("Done") {
|
||||
bugReportPhase = .idle
|
||||
bugReportUserDescription = ""
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.plain)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
|
||||
case .failure(let message):
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.red)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
Button("Dismiss") {
|
||||
bugReportPhase = .idle
|
||||
}
|
||||
.font(.caption2)
|
||||
.buttonStyle(.plain)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
}
|
||||
}
|
||||
.animation(.easeInOut(duration: 0.2), value: bugReportPhase)
|
||||
}
|
||||
|
||||
private var processToggleBinding: Binding<Bool> {
|
||||
@@ -687,16 +712,58 @@ struct ContentView: View {
|
||||
}
|
||||
|
||||
private func sendBugReport() async {
|
||||
bugReportInFlight = true
|
||||
bugReportMessage = "Collecting logs..."
|
||||
bugReportPhase = .sending("Collecting logs...")
|
||||
let service = BugReportService()
|
||||
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
do {
|
||||
let outcome = try await service.sendReport(isManual: true)
|
||||
bugReportMessage = outcome.message
|
||||
let outcome = try await service.sendReport(
|
||||
isManual: true,
|
||||
userDescription: description.isEmpty ? nil : description
|
||||
)
|
||||
if outcome.success {
|
||||
bugReportPhase = .success(outcome.message)
|
||||
} else {
|
||||
bugReportPhase = .failure(outcome.message)
|
||||
}
|
||||
} catch {
|
||||
bugReportMessage = error.localizedDescription
|
||||
bugReportPhase = .failure(error.localizedDescription)
|
||||
}
|
||||
}
|
||||
|
||||
private func openGitHubIssue() {
|
||||
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
|
||||
var bodyParts: [String] = []
|
||||
bodyParts.append("## Describe the bug")
|
||||
bodyParts.append("")
|
||||
if !description.isEmpty {
|
||||
bodyParts.append(description)
|
||||
} else {
|
||||
bodyParts.append("A clear and concise description of what the bug is.")
|
||||
}
|
||||
bodyParts.append("")
|
||||
bodyParts.append("## Environment")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
|
||||
bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("## Additional context")
|
||||
bodyParts.append("")
|
||||
bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
|
||||
|
||||
let body = bodyParts.joined(separator: "\n")
|
||||
|
||||
var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
|
||||
components.queryItems = [
|
||||
URLQueryItem(name: "template", value: "bug_report.md"),
|
||||
URLQueryItem(name: "title", value: "[BUG] "),
|
||||
URLQueryItem(name: "body", value: body),
|
||||
URLQueryItem(name: "labels", value: "bug"),
|
||||
]
|
||||
|
||||
if let url = components.url {
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
bugReportInFlight = false
|
||||
}
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
|
||||
@@ -21,7 +21,9 @@ struct EXOApp: App {
|
||||
@StateObject private var localNetworkChecker: LocalNetworkChecker
|
||||
@StateObject private var updater: SparkleUpdater
|
||||
@StateObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@StateObject private var settingsWindowController: SettingsWindowController
|
||||
private let terminationObserver: TerminationObserver
|
||||
private let firstLaunchPopout = FirstLaunchPopout()
|
||||
private let ciContext = CIContext(options: nil)
|
||||
|
||||
init() {
|
||||
@@ -43,12 +45,13 @@ struct EXOApp: App {
|
||||
_updater = StateObject(wrappedValue: updater)
|
||||
let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
|
||||
_thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
|
||||
_settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
|
||||
enableLaunchAtLoginIfNeeded()
|
||||
// Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
|
||||
NetworkSetupHelper.promptAndInstallIfNeeded()
|
||||
// Check local network access periodically (warning disappears when user grants permission)
|
||||
localNetwork.startPeriodicChecking(interval: 10)
|
||||
controller.scheduleLaunch(after: 15)
|
||||
controller.scheduleLaunch(after: 5)
|
||||
service.startPolling()
|
||||
networkStatus.startPolling()
|
||||
}
|
||||
@@ -62,8 +65,19 @@ struct EXOApp: App {
|
||||
.environmentObject(localNetworkChecker)
|
||||
.environmentObject(updater)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(settingsWindowController)
|
||||
} label: {
|
||||
menuBarIcon
|
||||
.onReceive(controller.$isFirstLaunchReady) { ready in
|
||||
if ready {
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 3.0) {
|
||||
self.firstLaunchPopout.onComplete = { [weak controller] in
|
||||
controller?.markOnboardingCompleted()
|
||||
}
|
||||
self.firstLaunchPopout.show()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
.menuBarExtraStyle(.window)
|
||||
}
|
||||
|
||||
@@ -5,6 +5,8 @@ import Foundation
|
||||
private let customNamespaceKey = "EXOCustomNamespace"
|
||||
private let hfTokenKey = "EXOHFToken"
|
||||
private let enableImageModelsKey = "EXOEnableImageModels"
|
||||
private let offlineModeKey = "EXOOfflineMode"
|
||||
private let onboardingCompletedKey = "EXOOnboardingCompleted"
|
||||
|
||||
@MainActor
|
||||
final class ExoProcessController: ObservableObject {
|
||||
@@ -59,6 +61,17 @@ final class ExoProcessController: ObservableObject {
|
||||
UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
|
||||
}
|
||||
}
|
||||
@Published var offlineMode: Bool = {
|
||||
return UserDefaults.standard.bool(forKey: offlineModeKey)
|
||||
}()
|
||||
{
|
||||
didSet {
|
||||
UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
|
||||
}
|
||||
}
|
||||
|
||||
/// Fires once when EXO transitions to `.running` for the very first time (fresh install).
|
||||
@Published private(set) var isFirstLaunchReady = false
|
||||
|
||||
private var process: Process?
|
||||
private var runtimeDirectoryURL: URL?
|
||||
@@ -113,6 +126,9 @@ final class ExoProcessController: ObservableObject {
|
||||
try child.run()
|
||||
process = child
|
||||
status = .running
|
||||
|
||||
// Show welcome popout on every launch
|
||||
isFirstLaunchReady = true
|
||||
} catch {
|
||||
process = nil
|
||||
status = .failed(message: "Launch error")
|
||||
@@ -164,6 +180,17 @@ final class ExoProcessController: ObservableObject {
|
||||
launch()
|
||||
}
|
||||
|
||||
/// Mark onboarding as completed (user interacted with the welcome popout).
|
||||
func markOnboardingCompleted() {
|
||||
UserDefaults.standard.set(true, forKey: onboardingCompletedKey)
|
||||
}
|
||||
|
||||
/// Reset onboarding so the welcome popout appears on next launch.
|
||||
func resetOnboarding() {
|
||||
UserDefaults.standard.removeObject(forKey: onboardingCompletedKey)
|
||||
isFirstLaunchReady = false
|
||||
}
|
||||
|
||||
func scheduleLaunch(after seconds: TimeInterval) {
|
||||
cancelPendingLaunch()
|
||||
let start = max(1, Int(ceil(seconds)))
|
||||
@@ -249,6 +276,9 @@ final class ExoProcessController: ObservableObject {
|
||||
if enableImageModels {
|
||||
environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
|
||||
}
|
||||
if offlineMode {
|
||||
environment["EXO_OFFLINE"] = "true"
|
||||
}
|
||||
|
||||
var paths: [String] = []
|
||||
if let existing = environment["PATH"], !existing.isEmpty {
|
||||
|
||||
@@ -38,7 +38,8 @@ struct BugReportService {
|
||||
func sendReport(
|
||||
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
|
||||
now: Date = Date(),
|
||||
isManual: Bool = false
|
||||
isManual: Bool = false,
|
||||
userDescription: String? = nil
|
||||
) async throws -> BugReportOutcome {
|
||||
let timestamp = Self.runTimestampString(now)
|
||||
let dayPrefix = Self.dayPrefixString(now)
|
||||
@@ -60,7 +61,8 @@ struct BugReportService {
|
||||
ifconfig: ifconfigText,
|
||||
debugInfo: debugInfo,
|
||||
isManual: isManual,
|
||||
clusterTbBridgeStatus: clusterTbBridgeStatus
|
||||
clusterTbBridgeStatus: clusterTbBridgeStatus,
|
||||
userDescription: userDescription
|
||||
)
|
||||
|
||||
let eventLogFiles = readAllEventLogs()
|
||||
@@ -306,7 +308,8 @@ struct BugReportService {
|
||||
ifconfig: String,
|
||||
debugInfo: DebugInfo,
|
||||
isManual: Bool,
|
||||
clusterTbBridgeStatus: [[String: Any]]?
|
||||
clusterTbBridgeStatus: [[String: Any]]?,
|
||||
userDescription: String? = nil
|
||||
) -> Data? {
|
||||
let system = readSystemMetadata()
|
||||
let exo = readExoMetadata()
|
||||
@@ -323,6 +326,9 @@ struct BugReportService {
|
||||
if let tbStatus = clusterTbBridgeStatus {
|
||||
payload["cluster_thunderbolt_bridge"] = tbStatus
|
||||
}
|
||||
if let desc = userDescription, !desc.isEmpty {
|
||||
payload["user_description"] = desc
|
||||
}
|
||||
return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// A popover callout anchored to the menu bar icon on every launch,
|
||||
/// pointing the user to the web dashboard with an arrow connecting to the icon.
|
||||
@MainActor
|
||||
final class FirstLaunchPopout {
|
||||
private var popover: NSPopover?
|
||||
private var countdownTask: Task<Void, Never>?
|
||||
private static let dashboardURL = "http://localhost:52415/"
|
||||
|
||||
/// Called when the user completes onboarding (clicks Open Dashboard or dismisses).
|
||||
var onComplete: (() -> Void)?
|
||||
|
||||
func show() {
|
||||
guard popover == nil else { return }
|
||||
|
||||
// The status bar button may not exist yet on first launch; retry generously.
|
||||
showWithRetry(attemptsRemaining: 15)
|
||||
}
|
||||
|
||||
private func showWithRetry(attemptsRemaining: Int) {
|
||||
guard attemptsRemaining > 0 else {
|
||||
// Exhausted retries — fall back to just opening the dashboard directly.
|
||||
openDashboard()
|
||||
onComplete?()
|
||||
return
|
||||
}
|
||||
|
||||
guard let button = Self.findStatusItemButton() else {
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { [weak self] in
|
||||
self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
let pop = NSPopover()
|
||||
pop.behavior = .applicationDefined
|
||||
pop.animates = true
|
||||
pop.contentSize = NSSize(width: 280, height: 120)
|
||||
pop.contentViewController = NSHostingController(
|
||||
rootView: WelcomeCalloutView(
|
||||
countdownDuration: 10,
|
||||
onDismiss: { [weak self] in
|
||||
self?.onComplete?()
|
||||
self?.dismiss()
|
||||
},
|
||||
onOpen: { [weak self] in
|
||||
self?.openDashboard()
|
||||
self?.onComplete?()
|
||||
self?.dismiss()
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
self.popover = pop
|
||||
pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY)
|
||||
|
||||
// Auto-open dashboard after 10s then dismiss
|
||||
countdownTask = Task {
|
||||
try? await Task.sleep(nanoseconds: 10_000_000_000)
|
||||
if !Task.isCancelled {
|
||||
openDashboard()
|
||||
onComplete?()
|
||||
dismiss()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func dismiss() {
|
||||
countdownTask?.cancel()
|
||||
countdownTask = nil
|
||||
guard let pop = popover else { return }
|
||||
popover = nil
|
||||
pop.performClose(nil)
|
||||
}
|
||||
|
||||
private func openDashboard() {
|
||||
guard let url = URL(string: Self.dashboardURL) else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
|
||||
/// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra.
|
||||
/// Walks the view hierarchy to find the actual button rather than the content view.
|
||||
private static func findStatusItemButton() -> NSView? {
|
||||
for window in NSApp.windows {
|
||||
let className = NSStringFromClass(type(of: window))
|
||||
// Match NSStatusBarWindow or any internal SwiftUI status bar window
|
||||
guard className.contains("StatusBar") || className.contains("MenuBarExtra") else {
|
||||
continue
|
||||
}
|
||||
if let content = window.contentView {
|
||||
if let button = findButton(in: content) {
|
||||
return button
|
||||
}
|
||||
// Fall back to the content view itself if it has a non-zero frame
|
||||
if content.frame.width > 0 {
|
||||
return content
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
/// Recursively searches the view hierarchy for an NSStatusBarButton.
|
||||
private static func findButton(in view: NSView) -> NSView? {
|
||||
let className = NSStringFromClass(type(of: view))
|
||||
if className.contains("StatusBarButton") || className.contains("StatusItem") {
|
||||
return view
|
||||
}
|
||||
for subview in view.subviews {
|
||||
if let found = findButton(in: subview) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimal welcome callout — friendly pointer, not a wall of text.
|
||||
/// Rendered inside the NSPopover which provides its own chrome and arrow.
|
||||
private struct WelcomeCalloutView: View {
|
||||
let countdownDuration: Int
|
||||
let onDismiss: () -> Void
|
||||
let onOpen: () -> Void
|
||||
@State private var countdown: Int
|
||||
@State private var timerTask: Task<Void, Never>?
|
||||
|
||||
init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) {
|
||||
self.countdownDuration = countdownDuration
|
||||
self.onDismiss = onDismiss
|
||||
self.onOpen = onOpen
|
||||
self._countdown = State(initialValue: countdownDuration)
|
||||
}
|
||||
|
||||
var body: some View {
|
||||
VStack(alignment: .leading, spacing: 10) {
|
||||
HStack(alignment: .top) {
|
||||
Text("EXO is running")
|
||||
.font(.system(.headline, design: .rounded))
|
||||
.fontWeight(.semibold)
|
||||
.foregroundColor(.primary)
|
||||
Spacer()
|
||||
Button {
|
||||
onDismiss()
|
||||
} label: {
|
||||
Image(systemName: "xmark.circle.fill")
|
||||
.font(.system(size: 14))
|
||||
.foregroundStyle(.tertiary)
|
||||
}
|
||||
.buttonStyle(.plain)
|
||||
}
|
||||
|
||||
Text("Run your first model here:")
|
||||
.font(.system(.subheadline, design: .default))
|
||||
.foregroundColor(.secondary)
|
||||
|
||||
HStack {
|
||||
Button {
|
||||
onOpen()
|
||||
} label: {
|
||||
Label("Open Dashboard", systemImage: "arrow.up.right.square")
|
||||
.font(.system(.caption, design: .default))
|
||||
.fontWeight(.medium)
|
||||
}
|
||||
.buttonStyle(.borderedProminent)
|
||||
.tint(.accentColor)
|
||||
.controlSize(.small)
|
||||
|
||||
Spacer()
|
||||
|
||||
if countdown > 0 {
|
||||
Text("Launching in \(countdown) secs...")
|
||||
.font(.system(.caption2, design: .default))
|
||||
.foregroundColor(.secondary.opacity(0.6))
|
||||
.monospacedDigit()
|
||||
}
|
||||
}
|
||||
}
|
||||
.padding(14)
|
||||
.onAppear {
|
||||
startCountdown()
|
||||
}
|
||||
.onDisappear {
|
||||
timerTask?.cancel()
|
||||
timerTask = nil
|
||||
}
|
||||
}
|
||||
|
||||
private func startCountdown() {
|
||||
timerTask = Task {
|
||||
while countdown > 0 {
|
||||
try? await Task.sleep(nanoseconds: 1_000_000_000)
|
||||
if !Task.isCancelled {
|
||||
countdown -= 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Native macOS Settings window following Apple HIG.
|
||||
/// Organized into General, Model, Advanced, and About sections.
|
||||
struct SettingsView: View {
|
||||
@EnvironmentObject private var controller: ExoProcessController
|
||||
@EnvironmentObject private var updater: SparkleUpdater
|
||||
@EnvironmentObject private var networkStatusService: NetworkStatusService
|
||||
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
|
||||
@EnvironmentObject private var stateService: ClusterStateService
|
||||
|
||||
@State private var pendingNamespace: String = ""
|
||||
@State private var pendingHFToken: String = ""
|
||||
@State private var pendingEnableImageModels = false
|
||||
@State private var pendingOfflineMode = false
|
||||
@State private var needsRestart = false
|
||||
@State private var bugReportInFlight = false
|
||||
@State private var bugReportMessage: String?
|
||||
@State private var uninstallInProgress = false
|
||||
|
||||
var body: some View {
|
||||
TabView {
|
||||
generalTab
|
||||
.tabItem {
|
||||
Label("General", systemImage: "gear")
|
||||
}
|
||||
modelTab
|
||||
.tabItem {
|
||||
Label("Model", systemImage: "cube")
|
||||
}
|
||||
advancedTab
|
||||
.tabItem {
|
||||
Label("Advanced", systemImage: "wrench.and.screwdriver")
|
||||
}
|
||||
aboutTab
|
||||
.tabItem {
|
||||
Label("About", systemImage: "info.circle")
|
||||
}
|
||||
}
|
||||
.frame(width: 450, height: 400)
|
||||
.onAppear {
|
||||
pendingNamespace = controller.customNamespace
|
||||
pendingHFToken = controller.hfToken
|
||||
pendingEnableImageModels = controller.enableImageModels
|
||||
pendingOfflineMode = controller.offlineMode
|
||||
needsRestart = false
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - General Tab
|
||||
|
||||
private var generalTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
LabeledContent("Cluster Namespace") {
|
||||
TextField("default", text: $pendingNamespace)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 200)
|
||||
}
|
||||
Text("Nodes with the same namespace form a cluster. Leave empty for default.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
LabeledContent("HuggingFace Token") {
|
||||
SecureField("optional", text: $pendingHFToken)
|
||||
.textFieldStyle(.roundedBorder)
|
||||
.frame(width: 200)
|
||||
}
|
||||
Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
Toggle("Offline Mode", isOn: $pendingOfflineMode)
|
||||
Text("Skip internet checks and use only locally available models.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyGeneralSettings()
|
||||
}
|
||||
.disabled(!hasGeneralChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Model Tab
|
||||
|
||||
private var modelTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels)
|
||||
Text("Allow text-to-image and image-to-image models in the model picker.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
|
||||
Section {
|
||||
HStack {
|
||||
Spacer()
|
||||
Button("Save & Restart") {
|
||||
applyModelSettings()
|
||||
}
|
||||
.disabled(!hasModelChanges)
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Advanced Tab
|
||||
|
||||
private var advancedTab: some View {
|
||||
Form {
|
||||
Section("Onboarding") {
|
||||
HStack {
|
||||
VStack(alignment: .leading) {
|
||||
Text("Reset Onboarding")
|
||||
Text("Opens the dashboard and resets the onboarding wizard.")
|
||||
.font(.caption)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
Spacer()
|
||||
Button("Reset") {
|
||||
guard let url = URL(string: "http://localhost:52415/?reset-onboarding")
|
||||
else { return }
|
||||
NSWorkspace.shared.open(url)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Section("Debug Info") {
|
||||
LabeledContent("Thunderbolt Bridge") {
|
||||
Text(thunderboltStatusText)
|
||||
.foregroundColor(thunderboltStatusColor)
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
clusterThunderboltBridgeView
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
interfaceIpList
|
||||
}
|
||||
|
||||
VStack(alignment: .leading, spacing: 2) {
|
||||
rdmaStatusView
|
||||
}
|
||||
|
||||
sendBugReportButton
|
||||
}
|
||||
|
||||
Section("Danger Zone") {
|
||||
Button(role: .destructive) {
|
||||
showUninstallConfirmationAlert()
|
||||
} label: {
|
||||
HStack {
|
||||
Text("Uninstall EXO")
|
||||
Spacer()
|
||||
Image(systemName: "trash")
|
||||
.imageScale(.small)
|
||||
}
|
||||
}
|
||||
.disabled(uninstallInProgress)
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - About Tab
|
||||
|
||||
private var aboutTab: some View {
|
||||
Form {
|
||||
Section {
|
||||
LabeledContent("Version") {
|
||||
Text(buildTag)
|
||||
.textSelection(.enabled)
|
||||
}
|
||||
LabeledContent("Commit") {
|
||||
Text(buildCommit)
|
||||
.font(.system(.body, design: .monospaced))
|
||||
.textSelection(.enabled)
|
||||
}
|
||||
}
|
||||
|
||||
Section {
|
||||
Button("Check for Updates") {
|
||||
updater.checkForUpdates()
|
||||
}
|
||||
}
|
||||
}
|
||||
.formStyle(.grouped)
|
||||
.padding()
|
||||
}
|
||||
|
||||
// MARK: - Debug Info Views (moved from ContentView)
|
||||
|
||||
private var thunderboltStatusText: String {
|
||||
switch networkStatusService.status.thunderboltBridgeState {
|
||||
case .some(.disabled):
|
||||
return "Disabled"
|
||||
case .some(.deleted):
|
||||
return "Deleted"
|
||||
case .some(.enabled):
|
||||
return "Enabled"
|
||||
case nil:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
private var thunderboltStatusColor: Color {
|
||||
switch networkStatusService.status.thunderboltBridgeState {
|
||||
case .some(.disabled), .some(.deleted):
|
||||
return .green
|
||||
case .some(.enabled):
|
||||
return .red
|
||||
case nil:
|
||||
return .secondary
|
||||
}
|
||||
}
|
||||
|
||||
private var clusterThunderboltBridgeView: some View {
|
||||
let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
|
||||
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
if bridgeStatuses.isEmpty {
|
||||
Text("Cluster TB Bridge: No data")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
Text("Cluster TB Bridge Status:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
|
||||
if let status = bridgeStatuses[nodeId] {
|
||||
let nodeName =
|
||||
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
|
||||
let isLocal = nodeId == localNodeId
|
||||
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
|
||||
let statusText =
|
||||
!status.exists
|
||||
? "N/A"
|
||||
: (status.enabled ? "Enabled" : "Disabled")
|
||||
let color: Color =
|
||||
!status.exists
|
||||
? .secondary
|
||||
: (status.enabled ? .red : .green)
|
||||
Text("\(prefix) \(statusText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(color)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var interfaceIpList: some View {
|
||||
let statuses = networkStatusService.status.interfaceStatuses
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
Text("Interfaces (en0–en7):")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
if statuses.isEmpty {
|
||||
Text(" Unknown")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
ForEach(statuses, id: \.interfaceName) { status in
|
||||
let ipText = status.ipAddress ?? "No IP"
|
||||
Text(" \(status.interfaceName): \(ipText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(status.ipAddress == nil ? .red : .green)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var rdmaStatusView: some View {
|
||||
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
|
||||
let localNodeId = stateService.localNodeId
|
||||
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
|
||||
let localDevices = networkStatusService.status.localRdmaDevices
|
||||
let localPorts = networkStatusService.status.localRdmaActivePorts
|
||||
|
||||
return VStack(alignment: .leading, spacing: 1) {
|
||||
if rdmaStatuses.isEmpty {
|
||||
Text("Cluster RDMA: No data")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
} else {
|
||||
Text("Cluster RDMA Status:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
|
||||
if let status = rdmaStatuses[nodeId] {
|
||||
let nodeName =
|
||||
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
|
||||
let isLocal = nodeId == localNodeId
|
||||
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
|
||||
let statusText = status.enabled ? "Enabled" : "Disabled"
|
||||
let color: Color = status.enabled ? .green : .orange
|
||||
Text("\(prefix) \(statusText)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(color)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !localDevices.isEmpty {
|
||||
Text(" Local Devices: \(localDevices.joined(separator: ", "))")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
}
|
||||
if !localPorts.isEmpty {
|
||||
Text(" Local Active Ports:")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
ForEach(localPorts, id: \.device) { port in
|
||||
Text(" \(port.device) port \(port.port): \(port.state)")
|
||||
.font(.caption2)
|
||||
.foregroundColor(.green)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private var sendBugReportButton: some View {
|
||||
VStack(alignment: .leading, spacing: 4) {
|
||||
Button {
|
||||
Task {
|
||||
await sendBugReport()
|
||||
}
|
||||
} label: {
|
||||
HStack {
|
||||
if bugReportInFlight {
|
||||
ProgressView()
|
||||
.scaleEffect(0.6)
|
||||
}
|
||||
Text("Send Bug Report")
|
||||
.font(.caption)
|
||||
.fontWeight(.semibold)
|
||||
Spacer()
|
||||
}
|
||||
}
|
||||
.disabled(bugReportInFlight)
|
||||
|
||||
if let message = bugReportMessage {
|
||||
Text(message)
|
||||
.font(.caption2)
|
||||
.foregroundColor(.secondary)
|
||||
.fixedSize(horizontal: false, vertical: true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Actions
|
||||
|
||||
private func sendBugReport() async {
|
||||
bugReportInFlight = true
|
||||
bugReportMessage = "Collecting logs..."
|
||||
let service = BugReportService()
|
||||
do {
|
||||
let outcome = try await service.sendReport(isManual: true)
|
||||
bugReportMessage = outcome.message
|
||||
} catch {
|
||||
bugReportMessage = error.localizedDescription
|
||||
}
|
||||
bugReportInFlight = false
|
||||
}
|
||||
|
||||
private func showUninstallConfirmationAlert() {
|
||||
let alert = NSAlert()
|
||||
alert.messageText = "Uninstall EXO"
|
||||
alert.informativeText = """
|
||||
This will remove EXO and all its system components:
|
||||
|
||||
• Network configuration daemon
|
||||
• Launch at login registration
|
||||
• EXO network location
|
||||
|
||||
The app will be moved to Trash.
|
||||
"""
|
||||
alert.alertStyle = .warning
|
||||
alert.addButton(withTitle: "Uninstall")
|
||||
alert.addButton(withTitle: "Cancel")
|
||||
|
||||
if let uninstallButton = alert.buttons.first {
|
||||
uninstallButton.hasDestructiveAction = true
|
||||
}
|
||||
|
||||
let response = alert.runModal()
|
||||
if response == .alertFirstButtonReturn {
|
||||
performUninstall()
|
||||
}
|
||||
}
|
||||
|
||||
private func performUninstall() {
|
||||
uninstallInProgress = true
|
||||
|
||||
controller.cancelPendingLaunch()
|
||||
controller.stop()
|
||||
stateService.stopPolling()
|
||||
|
||||
DispatchQueue.global(qos: .utility).async {
|
||||
do {
|
||||
try NetworkSetupHelper.uninstall()
|
||||
|
||||
DispatchQueue.main.async {
|
||||
LaunchAtLoginHelper.disable()
|
||||
self.moveAppToTrash()
|
||||
|
||||
DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
|
||||
NSApplication.shared.terminate(nil)
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
DispatchQueue.main.async {
|
||||
let errorAlert = NSAlert()
|
||||
errorAlert.messageText = "Uninstall Failed"
|
||||
errorAlert.informativeText = error.localizedDescription
|
||||
errorAlert.alertStyle = .critical
|
||||
errorAlert.addButton(withTitle: "OK")
|
||||
errorAlert.runModal()
|
||||
self.uninstallInProgress = false
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func moveAppToTrash() {
|
||||
guard let appURL = Bundle.main.bundleURL as URL? else { return }
|
||||
do {
|
||||
try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
|
||||
} catch {
|
||||
// If we can't trash the app, that's OK - user can do it manually
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Helpers
|
||||
|
||||
private var hasGeneralChanges: Bool {
|
||||
pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
|
||||
|| pendingOfflineMode != controller.offlineMode
|
||||
}
|
||||
|
||||
private var hasModelChanges: Bool {
|
||||
pendingEnableImageModels != controller.enableImageModels
|
||||
}
|
||||
|
||||
private func applyGeneralSettings() {
|
||||
controller.customNamespace = pendingNamespace
|
||||
controller.hfToken = pendingHFToken
|
||||
controller.offlineMode = pendingOfflineMode
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func applyModelSettings() {
|
||||
controller.enableImageModels = pendingEnableImageModels
|
||||
restartIfRunning()
|
||||
}
|
||||
|
||||
private func restartIfRunning() {
|
||||
if controller.status == .running || controller.status == .starting {
|
||||
controller.restart()
|
||||
}
|
||||
}
|
||||
|
||||
private var buildTag: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
|
||||
}
|
||||
|
||||
private var buildCommit: String {
|
||||
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
import AppKit
|
||||
import SwiftUI
|
||||
|
||||
/// Manages a standalone native macOS Settings window.
|
||||
/// Ensures only one instance exists and brings it to front on repeated opens.
|
||||
@MainActor
|
||||
final class SettingsWindowController: ObservableObject {
|
||||
private var window: NSWindow?
|
||||
|
||||
func open(
|
||||
controller: ExoProcessController,
|
||||
updater: SparkleUpdater,
|
||||
networkStatusService: NetworkStatusService,
|
||||
thunderboltBridgeService: ThunderboltBridgeService,
|
||||
stateService: ClusterStateService
|
||||
) {
|
||||
if let existing = window, existing.isVisible {
|
||||
existing.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate(ignoringOtherApps: true)
|
||||
return
|
||||
}
|
||||
|
||||
let settingsView = SettingsView()
|
||||
.environmentObject(controller)
|
||||
.environmentObject(updater)
|
||||
.environmentObject(networkStatusService)
|
||||
.environmentObject(thunderboltBridgeService)
|
||||
.environmentObject(stateService)
|
||||
|
||||
let hostingView = NSHostingView(rootView: settingsView)
|
||||
|
||||
let newWindow = NSWindow(
|
||||
contentRect: NSRect(x: 0, y: 0, width: 450, height: 400),
|
||||
styleMask: [.titled, .closable],
|
||||
backing: .buffered,
|
||||
defer: false
|
||||
)
|
||||
newWindow.title = "EXO Settings"
|
||||
newWindow.contentView = hostingView
|
||||
newWindow.center()
|
||||
newWindow.isReleasedWhenClosed = false
|
||||
newWindow.makeKeyAndOrderFront(nil)
|
||||
NSApp.activate(ignoringOtherApps: true)
|
||||
|
||||
window = newWindow
|
||||
}
|
||||
}
|
||||
+1
-3
@@ -2,6 +2,4 @@
|
||||
#
|
||||
# Lists the suite files to include. Each file defines benchmarks
|
||||
# with shared constraints, topology, and default args.
|
||||
include = [
|
||||
"single-m3-ultra.toml",
|
||||
]
|
||||
include = ["single-m3-ultra.toml"]
|
||||
|
||||
@@ -0,0 +1,292 @@
|
||||
# Model evaluation configurations for exo_eval.
|
||||
#
|
||||
# Each [[model]] entry uses `patterns` — a list of substrings matched
|
||||
# against the model_id. First matching entry wins.
|
||||
#
|
||||
# Required fields:
|
||||
# name, patterns, reasoning
|
||||
#
|
||||
# Optional per-model overrides (CLI flags take priority over these):
|
||||
# temperature, top_p, max_tokens, reasoning_effort
|
||||
#
|
||||
# Fallback defaults (when no per-model config):
|
||||
# reasoning: temperature=1.0, max_tokens=131072, reasoning_effort="high"
|
||||
# non-reasoning: temperature=0.0, max_tokens=16384
|
||||
#
|
||||
# All per-model values below are sourced from official model cards,
|
||||
# generation_config.json files, and vendor documentation.
|
||||
|
||||
# ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards (Qwen/Qwen3.5-*)
|
||||
# 35B-A3B thinking general: temp=1.0, top_p=0.95, top_k=20
|
||||
# 397B thinking: temp=0.6, top_p=0.95, top_k=20
|
||||
# Non-thinking: temp=0.7, top_p=0.8, top_k=20
|
||||
# max_tokens: 32768 general, 81920 for complex math/code
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 2B"
|
||||
patterns = ["Qwen3.5-2B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 81920
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 9B"
|
||||
patterns = ["Qwen3.5-9B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 81920
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 27B"
|
||||
patterns = ["Qwen3.5-27B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 81920
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 35B A3B"
|
||||
patterns = ["Qwen3.5-35B-A3B"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
max_tokens = 81920
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 122B A10B"
|
||||
patterns = ["Qwen3.5-122B-A10B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 81920
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3.5 397B A17B"
|
||||
patterns = ["Qwen3.5-397B-A17B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 81920
|
||||
|
||||
# ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards (Qwen/Qwen3-*)
|
||||
# Thinking: temp=0.6, top_p=0.95, top_k=20
|
||||
# Non-thinking: temp=0.7, top_p=0.8, top_k=20
|
||||
# max_tokens: 32768 general, 38912 for complex math/code
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 0.6B"
|
||||
patterns = ["Qwen3-0.6B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 30B A3B"
|
||||
patterns = ["Qwen3-30B-A3B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 235B A22B"
|
||||
patterns = ["Qwen3-235B-A22B"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Next 80B Thinking"
|
||||
patterns = ["Qwen3-Next-80B-A3B-Thinking"]
|
||||
reasoning = true
|
||||
temperature = 0.6
|
||||
top_p = 0.95
|
||||
max_tokens = 38912
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Next 80B Instruct"
|
||||
patterns = ["Qwen3-Next-80B-A3B-Instruct"]
|
||||
reasoning = false
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
max_tokens = 16384
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Coder 480B"
|
||||
patterns = ["Qwen3-Coder-480B"]
|
||||
reasoning = false
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
max_tokens = 16384
|
||||
|
||||
[[model]]
|
||||
name = "Qwen3 Coder Next"
|
||||
patterns = ["Qwen3-Coder-Next"]
|
||||
reasoning = false
|
||||
temperature = 0.7
|
||||
top_p = 0.8
|
||||
max_tokens = 16384
|
||||
|
||||
# ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
|
||||
# Source: OpenAI GitHub README + HuggingFace discussion #21
|
||||
# temp=1.0, top_p=1.0, NO top_k, NO repetition_penalty
|
||||
# reasoning_effort supported: low/medium/high
|
||||
|
||||
[[model]]
|
||||
name = "GPT-OSS 20B"
|
||||
patterns = ["gpt-oss-20b"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 1.0
|
||||
|
||||
[[model]]
|
||||
name = "GPT-OSS 120B"
|
||||
patterns = ["gpt-oss-120b"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 1.0
|
||||
|
||||
# ─── DeepSeek ────────────────────────────────────────────────────────
|
||||
# Source: https://api-docs.deepseek.com/quick_start/parameter_settings
|
||||
# Coding/Math: temp=0.0, General: temp=1.3, Creative: temp=1.5
|
||||
# NOTE: DeepSeek API applies nonlinear temp mapping. These are API values.
|
||||
# When running model directly: API temp 1.0 = model temp 0.3
|
||||
# We use temp=0.0 for eval (coding/math focus).
|
||||
|
||||
[[model]]
|
||||
name = "DeepSeek V3.1"
|
||||
patterns = ["DeepSeek-V3.1"]
|
||||
reasoning = true
|
||||
temperature = 0.0
|
||||
|
||||
# ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
|
||||
# Source: HuggingFace model cards + generation_config.json + docs.z.ai
|
||||
# GLM 4.5+: temp=1.0, top_p=0.95
|
||||
# Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7
|
||||
|
||||
[[model]]
|
||||
name = "GLM-5"
|
||||
patterns = ["GLM-5"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
max_tokens = 131072
|
||||
|
||||
[[model]]
|
||||
name = "GLM 4.5 Air"
|
||||
patterns = ["GLM-4.5-Air"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
|
||||
[[model]]
|
||||
name = "GLM 4.7"
|
||||
patterns = ["GLM-4.7-"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
max_tokens = 131072
|
||||
# Note: matches both GLM-4.7 and GLM-4.7-Flash
|
||||
|
||||
# ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards (moonshotai/Kimi-K2-*)
|
||||
# K2-Instruct: temp=0.6
|
||||
# K2-Thinking: temp=1.0, max_length=262144
|
||||
# K2.5: thinking temp=1.0, top_p=0.95; instant temp=0.6, top_p=0.95
|
||||
|
||||
[[model]]
|
||||
name = "Kimi K2 Thinking"
|
||||
patterns = ["Kimi-K2-Thinking"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
max_tokens = 131072
|
||||
|
||||
[[model]]
|
||||
name = "Kimi K2.5"
|
||||
patterns = ["Kimi-K2.5"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
max_tokens = 131072
|
||||
|
||||
[[model]]
|
||||
name = "Kimi K2 Instruct"
|
||||
patterns = ["Kimi-K2-Instruct"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
|
||||
# ─── MiniMax ─────────────────────────────────────────────────────────
|
||||
# Source: HuggingFace model cards + generation_config.json
|
||||
# All models: temp=1.0, top_p=0.95, top_k=40
|
||||
|
||||
[[model]]
|
||||
name = "MiniMax M2.5"
|
||||
patterns = ["MiniMax-M2.5"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
|
||||
[[model]]
|
||||
name = "MiniMax M2.1"
|
||||
patterns = ["MiniMax-M2.1"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
|
||||
# ─── Step (StepFun) ─────────────────────────────────────────────────
|
||||
# Source: HuggingFace model card (stepfun-ai/Step-3.5-Flash)
|
||||
# Reasoning: temp=1.0, top_p=0.95
|
||||
# General chat: temp=0.6, top_p=0.95
|
||||
# We use reasoning settings for eval.
|
||||
|
||||
[[model]]
|
||||
name = "Step 3.5 Flash"
|
||||
patterns = ["Step-3.5-Flash"]
|
||||
reasoning = true
|
||||
temperature = 1.0
|
||||
top_p = 0.95
|
||||
|
||||
# ─── Llama (Meta) ───────────────────────────────────────────────────
|
||||
# Source: generation_config.json + meta-llama/llama-models generation.py
|
||||
# All variants: temp=0.6, top_p=0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.2 1B"
|
||||
patterns = ["Llama-3.2-1B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.2 3B"
|
||||
patterns = ["Llama-3.2-3B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.1 8B"
|
||||
patterns = ["Llama-3.1-8B", "Meta-Llama-3.1-8B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.1 70B"
|
||||
patterns = ["Llama-3.1-70B", "Meta-Llama-3.1-70B"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
|
||||
[[model]]
|
||||
name = "Llama 3.3 70B"
|
||||
patterns = ["Llama-3.3-70B", "llama-3.3-70b"]
|
||||
reasoning = false
|
||||
temperature = 0.6
|
||||
top_p = 0.9
|
||||
+130
-46
@@ -24,6 +24,7 @@ import json
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
from statistics import mean
|
||||
from typing import Any
|
||||
@@ -75,7 +76,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
|
||||
model_path = Path(
|
||||
snapshot_download(
|
||||
model_id,
|
||||
allow_patterns=["*.json", "*.py", "*.tiktoken"],
|
||||
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
|
||||
)
|
||||
)
|
||||
|
||||
@@ -252,6 +253,12 @@ def main() -> int:
|
||||
ap.add_argument(
|
||||
"--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
|
||||
)
|
||||
ap.add_argument(
|
||||
"--concurrency",
|
||||
nargs="+",
|
||||
default=["1"],
|
||||
help="Concurrency levels (ints). Accepts commas. E.g. --concurrency 1,2,4,8. Default 1.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--warmup",
|
||||
type=int,
|
||||
@@ -282,6 +289,10 @@ def main() -> int:
|
||||
if args.repeat <= 0:
|
||||
logger.error("--repeat must be >= 1")
|
||||
return 2
|
||||
concurrency_list = parse_int_list(args.concurrency)
|
||||
if not concurrency_list or any(c <= 0 for c in concurrency_list):
|
||||
logger.error("--concurrency values must be >= 1")
|
||||
return 2
|
||||
|
||||
# Log pairing mode
|
||||
use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
|
||||
@@ -293,7 +304,9 @@ def main() -> int:
|
||||
logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs")
|
||||
|
||||
client = ExoClient(args.host, args.port, timeout_s=args.timeout)
|
||||
short_id, full_model_id = resolve_model_short_id(client, args.model)
|
||||
short_id, full_model_id = resolve_model_short_id(
|
||||
client, args.model, force_download=args.force_download
|
||||
)
|
||||
|
||||
tokenizer = load_tokenizer_for_bench(full_model_id)
|
||||
if tokenizer is None:
|
||||
@@ -392,52 +405,123 @@ def main() -> int:
|
||||
pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
|
||||
|
||||
for pp, tg in pp_tg_pairs:
|
||||
runs: list[dict[str, Any]] = []
|
||||
for r in range(args.repeat):
|
||||
time.sleep(3)
|
||||
try:
|
||||
row, actual_pp_tokens = run_one_completion(
|
||||
client, full_model_id, pp, tg, prompt_sizer
|
||||
for concurrency in concurrency_list:
|
||||
logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---")
|
||||
runs: list[dict[str, Any]] = []
|
||||
for r in range(args.repeat):
|
||||
time.sleep(3)
|
||||
|
||||
if concurrency <= 1:
|
||||
# Sequential: single request
|
||||
try:
|
||||
row, actual_pp_tokens = run_one_completion(
|
||||
client, full_model_id, pp, tg, prompt_sizer
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
continue
|
||||
row.update(
|
||||
{
|
||||
"model_short_id": short_id,
|
||||
"model_id": full_model_id,
|
||||
"placement_sharding": sharding,
|
||||
"placement_instance_meta": instance_meta,
|
||||
"placement_nodes": n_nodes,
|
||||
"instance_id": instance_id,
|
||||
"pp_tokens": actual_pp_tokens,
|
||||
"tg": tg,
|
||||
"repeat_index": r,
|
||||
"concurrency": 1,
|
||||
**(
|
||||
{"download_duration_s": download_duration_s}
|
||||
if download_duration_s is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
runs.append(row)
|
||||
all_rows.append(row)
|
||||
else:
|
||||
# Concurrent: fire N requests in parallel
|
||||
# Each thread gets its own ExoClient (separate HTTP connection)
|
||||
batch_results: list[tuple[dict[str, Any], int]] = []
|
||||
batch_errors = 0
|
||||
|
||||
def _run_concurrent(
|
||||
idx: int, *, _pp: int = pp, _tg: int = tg
|
||||
) -> tuple[dict[str, Any], int]:
|
||||
c = ExoClient(
|
||||
args.host, args.port, timeout_s=args.timeout
|
||||
)
|
||||
return run_one_completion(
|
||||
c, full_model_id, _pp, _tg, prompt_sizer
|
||||
)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=concurrency) as pool:
|
||||
futures = {
|
||||
pool.submit(_run_concurrent, i): i
|
||||
for i in range(concurrency)
|
||||
}
|
||||
for fut in as_completed(futures):
|
||||
try:
|
||||
batch_results.append(fut.result())
|
||||
except Exception as e:
|
||||
logger.error(f"Concurrent request failed: {e}")
|
||||
batch_errors += 1
|
||||
|
||||
for idx, (row, actual_pp_tokens) in enumerate(
|
||||
batch_results
|
||||
):
|
||||
row.update(
|
||||
{
|
||||
"model_short_id": short_id,
|
||||
"model_id": full_model_id,
|
||||
"placement_sharding": sharding,
|
||||
"placement_instance_meta": instance_meta,
|
||||
"placement_nodes": n_nodes,
|
||||
"instance_id": instance_id,
|
||||
"pp_tokens": actual_pp_tokens,
|
||||
"tg": tg,
|
||||
"repeat_index": r,
|
||||
"concurrency": concurrency,
|
||||
"concurrent_index": idx,
|
||||
**(
|
||||
{"download_duration_s": download_duration_s}
|
||||
if download_duration_s is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
runs.append(row)
|
||||
all_rows.append(row)
|
||||
|
||||
if batch_results:
|
||||
agg_gen_tps = sum(
|
||||
x["stats"]["generation_tps"]
|
||||
for x, _ in batch_results
|
||||
if x["stats"]["generation_tps"] > 0
|
||||
)
|
||||
logger.info(
|
||||
f"[concurrent {concurrency}x] "
|
||||
f"agg_gen_tps={agg_gen_tps:.2f} "
|
||||
f"errors={batch_errors}"
|
||||
)
|
||||
|
||||
if runs:
|
||||
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
|
||||
gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
|
||||
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
|
||||
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
|
||||
peak = mean(
|
||||
x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(e)
|
||||
continue
|
||||
row.update(
|
||||
{
|
||||
"model_short_id": short_id,
|
||||
"model_id": full_model_id,
|
||||
"placement_sharding": sharding,
|
||||
"placement_instance_meta": instance_meta,
|
||||
"placement_nodes": n_nodes,
|
||||
"instance_id": instance_id,
|
||||
"pp_tokens": actual_pp_tokens,
|
||||
"tg": tg,
|
||||
"repeat_index": r,
|
||||
**(
|
||||
{"download_duration_s": download_duration_s}
|
||||
if download_duration_s is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
)
|
||||
runs.append(row)
|
||||
all_rows.append(row)
|
||||
|
||||
if runs:
|
||||
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
|
||||
gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
|
||||
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
|
||||
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
|
||||
peak = mean(
|
||||
x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
|
||||
f"prompt_tokens={ptok} gen_tokens={gtok} "
|
||||
f"peak_memory={format_peak_memory(peak)}\n"
|
||||
)
|
||||
time.sleep(2)
|
||||
logger.info(
|
||||
f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
|
||||
f"prompt_tokens={ptok} gen_tokens={gtok} "
|
||||
f"peak_memory={format_peak_memory(peak)}\n"
|
||||
)
|
||||
time.sleep(2)
|
||||
finally:
|
||||
try:
|
||||
client.request_json("DELETE", f"/instance/{instance_id}")
|
||||
|
||||
+1499
File diff suppressed because it is too large
Load Diff
+19
-2
@@ -165,7 +165,9 @@ def wait_for_instance_gone(
|
||||
raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
|
||||
|
||||
|
||||
def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
|
||||
def resolve_model_short_id(
|
||||
client: ExoClient, model_arg: str, *, force_download: bool = False
|
||||
) -> tuple[str, str]:
|
||||
models = client.request_json("GET", "/models") or {}
|
||||
data = models.get("data") or []
|
||||
|
||||
@@ -181,6 +183,16 @@ def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]
|
||||
full_id = str(m["hugging_face_id"])
|
||||
return short_id, full_id
|
||||
|
||||
if force_download and "/" in model_arg:
|
||||
logger.info(f"Model not in /models, adding from HuggingFace: {model_arg}")
|
||||
result = client.request_json(
|
||||
"POST", "/models/add", body={"model_id": model_arg}
|
||||
)
|
||||
if result:
|
||||
short_id = str(result.get("name") or model_arg.rsplit("/", 1)[-1])
|
||||
full_id = str(result.get("hugging_face_id") or model_arg)
|
||||
return short_id, full_id
|
||||
|
||||
raise ValueError(f"Model not found in /models: {model_arg}")
|
||||
|
||||
|
||||
@@ -371,7 +383,7 @@ def run_planning_phase(
|
||||
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
|
||||
"modelId"
|
||||
],
|
||||
p["DownloadCompleted"]["totalBytes"]["inBytes"],
|
||||
p["DownloadCompleted"]["total"]["inBytes"],
|
||||
)
|
||||
for p in node_downloads
|
||||
if "DownloadCompleted" in p
|
||||
@@ -445,6 +457,11 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
|
||||
"--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
|
||||
)
|
||||
ap.add_argument("--model", required=True, help="Model short id or huggingface id")
|
||||
ap.add_argument(
|
||||
"--force-download",
|
||||
action="store_true",
|
||||
help="If model not in /models, add it from HuggingFace via exo and download.",
|
||||
)
|
||||
ap.add_argument(
|
||||
"--max-nodes",
|
||||
type=int,
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
# type: ignore
|
||||
import argparse
|
||||
import asyncio
|
||||
import sys
|
||||
import termios
|
||||
import time
|
||||
import tty
|
||||
|
||||
import aiohttp
|
||||
|
||||
NUM_REQUESTS = 10
|
||||
BASE_URL = ""
|
||||
|
||||
QUESTIONS = [
|
||||
"What is the capital of Australia?",
|
||||
"How many bones are in the human body?",
|
||||
"What year did World War II end?",
|
||||
"What is the speed of light in meters per second?",
|
||||
"Who wrote Romeo and Juliet?",
|
||||
"What is the chemical formula for water?",
|
||||
"How many planets are in our solar system?",
|
||||
"What is the largest ocean on Earth?",
|
||||
"Who painted the Mona Lisa?",
|
||||
"What is the boiling point of water in Celsius?",
|
||||
]
|
||||
|
||||
|
||||
def write(s: str) -> None:
|
||||
sys.stdout.write(s)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model picker (same style as exo_eval)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def fetch_models() -> list[str]:
|
||||
import json
|
||||
import urllib.request
|
||||
|
||||
with urllib.request.urlopen(f"{BASE_URL}/state") as resp:
|
||||
data = json.loads(resp.read())
|
||||
model_ids: set[str] = set()
|
||||
for instance in data.get("instances", {}).values():
|
||||
for variant in instance.values():
|
||||
sa = variant.get("shardAssignments", {})
|
||||
model_id = sa.get("modelId")
|
||||
if model_id:
|
||||
model_ids.add(model_id)
|
||||
return sorted(model_ids)
|
||||
|
||||
|
||||
def pick_model() -> str | None:
|
||||
models = fetch_models()
|
||||
if not models:
|
||||
print("No models found.")
|
||||
return None
|
||||
|
||||
cursor = 0
|
||||
total_lines = len(models) + 4
|
||||
|
||||
def render(first: bool = False) -> None:
|
||||
if not first:
|
||||
write(f"\033[{total_lines}A")
|
||||
write("\033[J")
|
||||
write("\033[1mSelect model\033[0m (up/down, enter confirm, q quit)\r\n\r\n")
|
||||
for i, model in enumerate(models):
|
||||
line = f" {'>' if i == cursor else ' '} {model}"
|
||||
write(f"\033[7m{line}\033[0m\r\n" if i == cursor else f"{line}\r\n")
|
||||
write("\r\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
fd = sys.stdin.fileno()
|
||||
old = termios.tcgetattr(fd)
|
||||
try:
|
||||
tty.setraw(fd)
|
||||
write("\033[?25l")
|
||||
render(first=True)
|
||||
while True:
|
||||
ch = sys.stdin.read(1)
|
||||
if ch in ("q", "\x03"):
|
||||
write("\033[?25h\033[0m\r\n")
|
||||
return None
|
||||
elif ch in ("\r", "\n"):
|
||||
break
|
||||
elif ch == "\x1b":
|
||||
seq = sys.stdin.read(2)
|
||||
if seq == "[A":
|
||||
cursor = (cursor - 1) % len(models)
|
||||
elif seq == "[B":
|
||||
cursor = (cursor + 1) % len(models)
|
||||
render()
|
||||
finally:
|
||||
termios.tcsetattr(fd, termios.TCSADRAIN, old)
|
||||
write(f"\033[{total_lines}A\033[J") # clear picker UI
|
||||
write("\033[?25h\033[0m")
|
||||
sys.stdout.flush()
|
||||
|
||||
return models[cursor]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parallel requests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
statuses: list[str] = []
|
||||
times: list[str] = []
|
||||
previews: list[str] = []
|
||||
tokens: list[str] = []
|
||||
full_responses: list[dict | None] = []
|
||||
total_lines = 0
|
||||
start_time: float = 0
|
||||
selected_model: str = ""
|
||||
|
||||
|
||||
def render_progress(first: bool = False) -> None:
|
||||
if not first:
|
||||
write(f"\033[{total_lines}A")
|
||||
write("\033[J")
|
||||
elapsed = time.monotonic() - start_time if start_time else 0
|
||||
done = sum(1 for s in statuses if s == "done")
|
||||
write(
|
||||
f"\033[1m{selected_model}\033[0m [{done}/{NUM_REQUESTS}] {elapsed:.1f}s\r\n\r\n"
|
||||
)
|
||||
|
||||
for i in range(NUM_REQUESTS):
|
||||
q = QUESTIONS[i % len(QUESTIONS)]
|
||||
status = statuses[i]
|
||||
if status == "pending":
|
||||
color = "\033[33m" # yellow
|
||||
elif status == "running":
|
||||
color = "\033[36m" # cyan
|
||||
elif status == "done":
|
||||
color = "\033[32m" # green
|
||||
else:
|
||||
color = "\033[31m" # red
|
||||
write(
|
||||
f" {i:>2} {color}{status:<8}\033[0m {times[i]:>6} {tokens[i]:>5}tok {q[:40]:<40} {previews[i][:50]}\r\n"
|
||||
)
|
||||
|
||||
write("\r\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
async def send_request(
|
||||
session: aiohttp.ClientSession, i: int, lock: asyncio.Lock
|
||||
) -> None:
|
||||
payload = {
|
||||
"model": selected_model,
|
||||
"messages": [{"role": "user", "content": QUESTIONS[i % len(QUESTIONS)]}],
|
||||
"max_tokens": 1024,
|
||||
}
|
||||
statuses[i] = "running"
|
||||
async with lock:
|
||||
render_progress()
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
async with session.post(
|
||||
f"{BASE_URL}/v1/chat/completions", json=payload
|
||||
) as resp:
|
||||
data = await resp.json()
|
||||
elapsed = time.monotonic() - t0
|
||||
full_responses[i] = data
|
||||
times[i] = f"{elapsed:.1f}s"
|
||||
if resp.status == 200:
|
||||
choice = data["choices"][0]
|
||||
msg = choice["message"]
|
||||
content = msg.get("content", "")
|
||||
previews[i] = content[:50].replace("\n", " ") or "(empty)"
|
||||
if "usage" in data:
|
||||
tokens[i] = str(data["usage"].get("total_tokens", ""))
|
||||
statuses[i] = "done"
|
||||
else:
|
||||
statuses[i] = f"err:{resp.status}"
|
||||
previews[i] = str(data.get("error", {}).get("message", ""))[:50]
|
||||
except Exception as e:
|
||||
elapsed = time.monotonic() - t0
|
||||
times[i] = f"{elapsed:.1f}s"
|
||||
statuses[i] = "error"
|
||||
previews[i] = str(e)[:50]
|
||||
async with lock:
|
||||
render_progress()
|
||||
|
||||
|
||||
async def run_requests(print_stdout: bool = False) -> None:
|
||||
global start_time, total_lines, statuses, times, previews, tokens, full_responses
|
||||
|
||||
statuses = ["pending"] * NUM_REQUESTS
|
||||
times = ["-"] * NUM_REQUESTS
|
||||
previews = ["-"] * NUM_REQUESTS
|
||||
tokens = ["-"] * NUM_REQUESTS
|
||||
full_responses = [None] * NUM_REQUESTS
|
||||
total_lines = NUM_REQUESTS + 4
|
||||
|
||||
write("\033[?25l") # hide cursor
|
||||
start_time = time.monotonic()
|
||||
render_progress(first=True)
|
||||
lock = asyncio.Lock()
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [send_request(session, i, lock) for i in range(NUM_REQUESTS)]
|
||||
await asyncio.gather(*tasks)
|
||||
total = time.monotonic() - start_time
|
||||
write(
|
||||
f"\033[1m=== All {NUM_REQUESTS} requests done in {total:.1f}s ===\033[0m\r\n\r\n"
|
||||
)
|
||||
|
||||
if print_stdout:
|
||||
for i in range(NUM_REQUESTS):
|
||||
data = full_responses[i]
|
||||
if not data or "choices" not in data:
|
||||
continue
|
||||
choice = data["choices"][0]
|
||||
msg = choice["message"]
|
||||
q = QUESTIONS[i % len(QUESTIONS)]
|
||||
write(f"\033[1m--- #{i}: {q} ---\033[0m\r\n")
|
||||
if msg.get("reasoning_content"):
|
||||
write(f"\033[2m[Thinking]: {msg['reasoning_content']}\033[0m\r\n")
|
||||
write(f"{msg.get('content', '')}\r\n")
|
||||
if "usage" in data:
|
||||
u = data["usage"]
|
||||
write(
|
||||
f"\033[2m[Usage: prompt={u.get('prompt_tokens')}, "
|
||||
f"completion={u.get('completion_tokens')}, "
|
||||
f"total={u.get('total_tokens')}]\033[0m\r\n"
|
||||
)
|
||||
write("\r\n")
|
||||
finally:
|
||||
write("\033[?25h") # show cursor
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
global selected_model, BASE_URL
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Send parallel requests to an exo cluster"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host", required=True, help="Hostname of the exo node (e.g. s1)"
|
||||
)
|
||||
parser.add_argument("--port", type=int, default=52415, help="Port (default: 52415)")
|
||||
parser.add_argument(
|
||||
"--stdout", action="store_true", help="Print full responses after completion"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
BASE_URL = f"http://{args.host}:{args.port}"
|
||||
model = pick_model()
|
||||
if not model:
|
||||
return
|
||||
selected_model = model
|
||||
asyncio.run(run_requests(print_stdout=args.stdout))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+12
-6
@@ -4,12 +4,18 @@ version = "0.1.0"
|
||||
description = "Benchmarking tool for exo distributed inference"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = [
|
||||
"httpx>=0.27.0",
|
||||
"loguru>=0.7.3",
|
||||
"transformers>=5.0.0",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"tiktoken>=0.12.0",
|
||||
"jinja2>=3.1.0",
|
||||
"httpx>=0.27.0",
|
||||
"loguru>=0.7.3",
|
||||
"transformers>=5.0.0",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"tiktoken>=0.12.0",
|
||||
"jinja2>=3.1.0",
|
||||
"protobuf>=5.29.0",
|
||||
"datasets>=2.0.0",
|
||||
"math-verify>=0.7.0",
|
||||
"lm-eval[api,math]>=0.4.0",
|
||||
"human-eval>=1.0.3",
|
||||
"numpy>=1.24.0",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
|
||||
@@ -2,10 +2,10 @@
|
||||
#
|
||||
# Shared constraints applied to ALL benchmarks in this file.
|
||||
constraints = [
|
||||
"All(MacOsBuild(=25D125))",
|
||||
"Hosts(=1)",
|
||||
"All(Chip(m3_ultra))",
|
||||
"All(GpuCores(=80))",
|
||||
"All(MacOsBuild(=25D125))",
|
||||
"Hosts(=1)",
|
||||
"All(Chip(m3_ultra))",
|
||||
"All(GpuCores(=80))",
|
||||
]
|
||||
|
||||
[topology]
|
||||
|
||||
Vendored
Vendored
+593
@@ -0,0 +1,593 @@
|
||||
# type: ignore
|
||||
# Vendored from LiveCodeBench (https://github.com/LiveCodeBench/LiveCodeBench)
|
||||
# File: lcb_runner/evaluation/testing_util.py
|
||||
# License: MIT
|
||||
# Vendored 2026-03-07 — do not modify without updating from upstream.
|
||||
|
||||
import ast
|
||||
import faulthandler
|
||||
import json
|
||||
import platform
|
||||
|
||||
# to run the solution files we're using a timing based approach
|
||||
import signal
|
||||
import sys
|
||||
import time
|
||||
|
||||
# used for debugging to time steps
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
from enum import Enum
|
||||
from io import StringIO
|
||||
|
||||
# from pyext import RuntimeModule
|
||||
from types import ModuleType
|
||||
|
||||
# used for testing the code that reads from input
|
||||
from unittest.mock import mock_open, patch
|
||||
|
||||
import_string = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(50000)\n"
|
||||
|
||||
|
||||
def truncatefn(s, length=300):
|
||||
if isinstance(s, str):
|
||||
pass
|
||||
else:
|
||||
s = str(s)
|
||||
if len(s) <= length:
|
||||
return s
|
||||
|
||||
return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :]
|
||||
|
||||
|
||||
class CODE_TYPE(Enum):
|
||||
call_based = 0
|
||||
standard_input = 1
|
||||
|
||||
|
||||
# stuff for setting up signal timer
|
||||
class TimeoutException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def timeout_handler(signum, frame):
|
||||
print("timeout occured: alarm went off")
|
||||
raise TimeoutException
|
||||
|
||||
|
||||
# used to capture stdout as a list
|
||||
# from https://stackoverflow.com/a/16571630/6416660
|
||||
# alternative use redirect_stdout() from contextlib
|
||||
class Capturing(list):
|
||||
def __enter__(self):
|
||||
self._stdout = sys.stdout
|
||||
sys.stdout = self._stringio = StringIO()
|
||||
# Make closing the StringIO a no-op
|
||||
self._stringio.close = lambda x: 1
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.append(self._stringio.getvalue())
|
||||
del self._stringio # free up some memory
|
||||
sys.stdout = self._stdout
|
||||
|
||||
|
||||
# Custom mock for sys.stdin that supports buffer attribute
|
||||
class MockStdinWithBuffer:
|
||||
def __init__(self, inputs: str):
|
||||
self.inputs = inputs
|
||||
self._stringio = StringIO(inputs)
|
||||
self.buffer = MockBuffer(inputs)
|
||||
|
||||
def read(self, *args):
|
||||
return self.inputs
|
||||
|
||||
def readline(self, *args):
|
||||
return self._stringio.readline(*args)
|
||||
|
||||
def readlines(self, *args):
|
||||
return self.inputs.split("\n")
|
||||
|
||||
def __getattr__(self, name):
|
||||
# Delegate other attributes to StringIO
|
||||
return getattr(self._stringio, name)
|
||||
|
||||
|
||||
class MockBuffer:
|
||||
def __init__(self, inputs: str):
|
||||
self.inputs = inputs.encode("utf-8") # Convert to bytes
|
||||
|
||||
def read(self, *args):
|
||||
# Return as byte strings that can be split
|
||||
return self.inputs
|
||||
|
||||
def readline(self, *args):
|
||||
return self.inputs.split(b"\n")[0] + b"\n"
|
||||
|
||||
|
||||
def clean_if_name(code: str) -> str:
|
||||
try:
|
||||
astree = ast.parse(code)
|
||||
last_block = astree.body[-1]
|
||||
if isinstance(last_block, ast.If):
|
||||
condition = last_block.test
|
||||
if ast.unparse(condition).strip() == "__name__ == '__main__'":
|
||||
code = (
|
||||
ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body) # type: ignore
|
||||
)
|
||||
except:
|
||||
pass
|
||||
|
||||
return code
|
||||
|
||||
|
||||
def make_function(code: str) -> str:
|
||||
try:
|
||||
import_stmts = []
|
||||
all_other_stmts = []
|
||||
astree = ast.parse(code)
|
||||
for stmt in astree.body:
|
||||
if isinstance(stmt, (ast.Import, ast.ImportFrom)):
|
||||
import_stmts.append(stmt)
|
||||
else:
|
||||
all_other_stmts.append(stmt)
|
||||
|
||||
function_ast = ast.FunctionDef(
|
||||
name="wrapped_function",
|
||||
args=ast.arguments(
|
||||
posonlyargs=[], args=[], kwonlyargs=[], kw_defaults=[], defaults=[]
|
||||
),
|
||||
body=all_other_stmts,
|
||||
decorator_list=[],
|
||||
lineno=-1,
|
||||
)
|
||||
main_code = (
|
||||
import_string
|
||||
+ "\n"
|
||||
+ ast.unparse(import_stmts)
|
||||
+ "\n"
|
||||
+ ast.unparse(function_ast)
|
||||
)
|
||||
return main_code
|
||||
except Exception:
|
||||
return code
|
||||
|
||||
|
||||
def call_method(method, inputs):
|
||||
if isinstance(inputs, list):
|
||||
inputs = "\n".join(inputs)
|
||||
|
||||
inputs_line_iterator = iter(inputs.split("\n"))
|
||||
|
||||
# Create custom stdin mock with buffer support
|
||||
mock_stdin = MockStdinWithBuffer(inputs)
|
||||
|
||||
# sys.setrecursionlimit(10000)
|
||||
|
||||
# @patch('builtins.input', side_effect=inputs.split("\n"))
|
||||
@patch("builtins.open", mock_open(read_data=inputs))
|
||||
@patch("sys.stdin", mock_stdin) # Use our custom mock instead of StringIO
|
||||
@patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator))
|
||||
@patch("sys.stdin.readlines", lambda *args: inputs.split("\n"))
|
||||
@patch("sys.stdin.read", lambda *args: inputs)
|
||||
# @patch('sys.stdout.write', print)
|
||||
def _inner_call_method(_method):
|
||||
try:
|
||||
return _method()
|
||||
except SystemExit:
|
||||
pass
|
||||
finally:
|
||||
pass
|
||||
|
||||
return _inner_call_method(method)
|
||||
|
||||
|
||||
def get_function(compiled_sol, fn_name: str): # type: ignore
|
||||
try:
|
||||
assert hasattr(compiled_sol, fn_name)
|
||||
return getattr(compiled_sol, fn_name)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def compile_code(code: str, timeout: int):
|
||||
signal.alarm(timeout)
|
||||
try:
|
||||
tmp_sol = ModuleType("tmp_sol", "")
|
||||
exec(code, tmp_sol.__dict__)
|
||||
if "class Solution" in code:
|
||||
# leetcode wraps solutions in `Solution`
|
||||
# this is a hack to check if it is leetcode solution or not
|
||||
# currently livecodebench only supports LeetCode but
|
||||
# else condition allows future extensibility to other platforms
|
||||
compiled_sol = tmp_sol.Solution()
|
||||
else:
|
||||
# do nothing in the other case since function is accesible
|
||||
compiled_sol = tmp_sol
|
||||
|
||||
assert compiled_sol is not None
|
||||
finally:
|
||||
signal.alarm(0)
|
||||
|
||||
return compiled_sol
|
||||
|
||||
|
||||
def convert_line_to_decimals(line: str) -> tuple[bool, list[Decimal]]:
|
||||
try:
|
||||
decimal_line = [Decimal(elem) for elem in line.split()]
|
||||
except:
|
||||
return False, []
|
||||
return True, decimal_line
|
||||
|
||||
|
||||
def get_stripped_lines(val: str):
|
||||
## you don't want empty lines to add empty list after splitlines!
|
||||
val = val.strip()
|
||||
|
||||
return [val_line.strip() for val_line in val.split("\n")]
|
||||
|
||||
|
||||
def grade_call_based(
|
||||
code: str, all_inputs: list, all_outputs: list, fn_name: str, timeout: int
|
||||
):
|
||||
# call-based clean up logic
|
||||
# need to wrap in try-catch logic after to catch the correct errors, but for now this is fine.
|
||||
code = import_string + "\n\n" + code
|
||||
compiled_sol = compile_code(code, timeout)
|
||||
|
||||
if compiled_sol is None:
|
||||
return
|
||||
|
||||
method = get_function(compiled_sol, fn_name)
|
||||
|
||||
if method is None:
|
||||
return
|
||||
|
||||
all_inputs = [
|
||||
[json.loads(line) for line in inputs.split("\n")] for inputs in all_inputs
|
||||
]
|
||||
|
||||
all_outputs = [json.loads(output) for output in all_outputs]
|
||||
|
||||
total_execution = 0
|
||||
all_results = []
|
||||
for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)):
|
||||
signal.alarm(timeout)
|
||||
faulthandler.enable()
|
||||
try:
|
||||
# can lock here so time is useful
|
||||
start = time.time()
|
||||
prediction = method(*gt_inp)
|
||||
total_execution += time.time() - start
|
||||
signal.alarm(0)
|
||||
|
||||
# don't penalize model if it produces tuples instead of lists
|
||||
# ground truth sequences are not tuples
|
||||
if isinstance(prediction, tuple):
|
||||
prediction = list(prediction)
|
||||
|
||||
tmp_result = prediction == gt_out
|
||||
|
||||
# handle floating point comparisons
|
||||
|
||||
all_results.append(tmp_result)
|
||||
|
||||
if not tmp_result:
|
||||
return all_results, {
|
||||
"output": truncatefn(prediction),
|
||||
"inputs": truncatefn(gt_inp),
|
||||
"expected": truncatefn(gt_out),
|
||||
"error_code": -2,
|
||||
"error_message": "Wrong Answer",
|
||||
}
|
||||
except Exception as e:
|
||||
signal.alarm(0)
|
||||
if "timeoutexception" in repr(e).lower():
|
||||
all_results.append(-3)
|
||||
return all_results, {
|
||||
"error": repr(e),
|
||||
"error_code": -3,
|
||||
"error_message": "Time Limit Exceeded",
|
||||
"inputs": truncatefn(gt_inp),
|
||||
"expected": truncatefn(gt_out),
|
||||
}
|
||||
else:
|
||||
all_results.append(-4)
|
||||
return all_results, {
|
||||
"error": repr(e),
|
||||
"error_code": -4,
|
||||
"error_message": "Runtime Error",
|
||||
"inputs": truncatefn(gt_inp),
|
||||
"expected": truncatefn(gt_out),
|
||||
}
|
||||
|
||||
finally:
|
||||
signal.alarm(0)
|
||||
faulthandler.disable()
|
||||
|
||||
return all_results, {"execution time": total_execution}
|
||||
|
||||
|
||||
def grade_stdio(
|
||||
code: str,
|
||||
all_inputs: list,
|
||||
all_outputs: list,
|
||||
timeout: int,
|
||||
):
|
||||
## runtime doesn't interact well with __name__ == '__main__'
|
||||
code = clean_if_name(code)
|
||||
|
||||
## we wrap the given code inside another function
|
||||
code = make_function(code)
|
||||
|
||||
compiled_sol = compile_code(code, timeout)
|
||||
if compiled_sol is None:
|
||||
return
|
||||
|
||||
method = get_function(compiled_sol, "wrapped_function")
|
||||
|
||||
if method is None:
|
||||
return
|
||||
|
||||
all_results = []
|
||||
total_execution_time = 0
|
||||
for idx, (gt_inp, gt_out) in enumerate(zip(all_inputs, all_outputs)):
|
||||
signal.alarm(timeout)
|
||||
faulthandler.enable()
|
||||
|
||||
signal.alarm(timeout)
|
||||
with Capturing() as captured_output:
|
||||
try:
|
||||
start = time.time()
|
||||
call_method(method, gt_inp)
|
||||
total_execution_time += time.time() - start
|
||||
# reset the alarm
|
||||
signal.alarm(0)
|
||||
except Exception as e:
|
||||
signal.alarm(0)
|
||||
if "timeoutexception" in repr(e).lower():
|
||||
all_results.append(-3)
|
||||
return all_results, {
|
||||
"error": repr(e),
|
||||
"error_code": -3,
|
||||
"error_message": "Time Limit Exceeded",
|
||||
"inputs": truncatefn(gt_inp),
|
||||
"expected": truncatefn(gt_out),
|
||||
}
|
||||
else:
|
||||
all_results.append(-4)
|
||||
return all_results, {
|
||||
"error": repr(e),
|
||||
"error_code": -4,
|
||||
"error_message": "Runtime Error",
|
||||
"inputs": truncatefn(gt_inp),
|
||||
"expected": truncatefn(gt_out),
|
||||
}
|
||||
|
||||
finally:
|
||||
signal.alarm(0)
|
||||
faulthandler.disable()
|
||||
|
||||
prediction = captured_output[0]
|
||||
|
||||
stripped_prediction_lines = get_stripped_lines(prediction)
|
||||
stripped_gt_out_lines = get_stripped_lines(gt_out)
|
||||
|
||||
## WA happens in multiple circumstances
|
||||
## so cache the return to make it clean!
|
||||
WA_send_args = {
|
||||
"output": truncatefn(prediction),
|
||||
"inputs": truncatefn(gt_inp),
|
||||
"expected": truncatefn(gt_out),
|
||||
"error_code": -2,
|
||||
}
|
||||
|
||||
if len(stripped_prediction_lines) != len(stripped_gt_out_lines):
|
||||
all_results.append(-2)
|
||||
WA_send_args["error_message"] = "Wrong answer: mismatched output length"
|
||||
return all_results, WA_send_args
|
||||
|
||||
for output_line_idx, (
|
||||
stripped_prediction_line,
|
||||
stripped_gt_out_line,
|
||||
) in enumerate(zip(stripped_prediction_lines, stripped_gt_out_lines)):
|
||||
WA_send_args["error_message"] = (
|
||||
f"Wrong answer at {output_line_idx=}: {truncatefn(stripped_prediction_line)} != {truncatefn(stripped_gt_out_line)}"
|
||||
)
|
||||
|
||||
## CASE 1: exact match
|
||||
if stripped_prediction_line == stripped_gt_out_line:
|
||||
continue
|
||||
|
||||
## CASE 2: element-wise comparision
|
||||
## if there are floating elements
|
||||
## use `decimal` library for good floating point comparision
|
||||
## otherwise gotcha: np.isclose(50000000000000000, 50000000000000001) = True
|
||||
## note that we should always be able to convert to decimals
|
||||
|
||||
success, decimal_prediction_line = convert_line_to_decimals(
|
||||
stripped_prediction_line
|
||||
)
|
||||
if not success:
|
||||
all_results.append(-2)
|
||||
return all_results, WA_send_args
|
||||
success, decimal_gtout_line = convert_line_to_decimals(stripped_gt_out_line)
|
||||
if not success:
|
||||
all_results.append(-2)
|
||||
return all_results, WA_send_args
|
||||
|
||||
if decimal_prediction_line == decimal_gtout_line:
|
||||
continue
|
||||
|
||||
all_results.append(-2)
|
||||
return all_results, WA_send_args
|
||||
all_results.append(True)
|
||||
|
||||
return all_results, {"execution time": total_execution_time}
|
||||
|
||||
|
||||
def run_test(sample, test=None, debug=False, timeout=6):
|
||||
"""
|
||||
if test(generated_code) is not None it'll try to run the code.
|
||||
otherwise it'll just return an input and output pair.
|
||||
"""
|
||||
signal.signal(signal.SIGALRM, timeout_handler)
|
||||
|
||||
# Disable functionalities that can make destructive changes to the test.
|
||||
# max memory is set to 4GB
|
||||
reliability_guard()
|
||||
|
||||
if debug:
|
||||
print(f"start = {datetime.now().time()}")
|
||||
|
||||
try:
|
||||
in_outs = json.loads(sample["input_output"])
|
||||
except ValueError as e:
|
||||
raise e
|
||||
in_outs = None
|
||||
|
||||
if in_outs:
|
||||
if in_outs.get("fn_name") is None:
|
||||
which_type = CODE_TYPE.standard_input # Standard input
|
||||
method_name = None
|
||||
|
||||
else:
|
||||
which_type = CODE_TYPE.call_based # Call-based
|
||||
method_name = in_outs["fn_name"]
|
||||
|
||||
if debug:
|
||||
print(f"loaded input_output = {datetime.now().time()}")
|
||||
|
||||
if test is None:
|
||||
assert False, "should not happen: test code is none"
|
||||
return in_outs, {"error": "no test code provided"}
|
||||
elif test is not None:
|
||||
results = []
|
||||
sol = import_string
|
||||
if debug:
|
||||
print(f"loading test code = {datetime.now().time()}")
|
||||
|
||||
if which_type == CODE_TYPE.call_based:
|
||||
signal.alarm(timeout)
|
||||
try:
|
||||
results, metadata = grade_call_based(
|
||||
code=test,
|
||||
all_inputs=in_outs["inputs"],
|
||||
all_outputs=in_outs["outputs"],
|
||||
fn_name=method_name,
|
||||
timeout=timeout,
|
||||
)
|
||||
return results, metadata
|
||||
except Exception as e:
|
||||
return [-4], {
|
||||
"error_code": -4,
|
||||
"error_message": f"Error during testing: {e}",
|
||||
}
|
||||
finally:
|
||||
signal.alarm(0)
|
||||
elif which_type == CODE_TYPE.standard_input:
|
||||
# sol
|
||||
# if code has if __name__ == "__main__": then remove it
|
||||
|
||||
signal.alarm(timeout)
|
||||
try:
|
||||
results, metadata = grade_stdio(
|
||||
code=test,
|
||||
all_inputs=in_outs["inputs"],
|
||||
all_outputs=in_outs["outputs"],
|
||||
timeout=timeout,
|
||||
)
|
||||
return results, metadata
|
||||
except Exception as e:
|
||||
return [-4], {
|
||||
"error_code": -4,
|
||||
"error_message": f"Error during testing: {e}",
|
||||
}
|
||||
finally:
|
||||
signal.alarm(0)
|
||||
|
||||
|
||||
def reliability_guard(maximum_memory_bytes=None):
|
||||
"""
|
||||
This disables various destructive functions and prevents the generated code
|
||||
from interfering with the test (e.g. fork bomb, killing other processes,
|
||||
removing filesystem files, etc.)
|
||||
WARNING
|
||||
This function is NOT a security sandbox. Untrusted code, including, model-
|
||||
generated code, should not be blindly executed outside of one. See the
|
||||
Codex paper for more information about OpenAI's code sandbox, and proceed
|
||||
with caution.
|
||||
"""
|
||||
|
||||
if maximum_memory_bytes is not None:
|
||||
import resource
|
||||
|
||||
resource.setrlimit(
|
||||
resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)
|
||||
)
|
||||
resource.setrlimit(
|
||||
resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)
|
||||
)
|
||||
if not platform.uname().system == "Darwin":
|
||||
resource.setrlimit(
|
||||
resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)
|
||||
)
|
||||
|
||||
faulthandler.disable()
|
||||
|
||||
import builtins
|
||||
|
||||
# builtins.exit = None
|
||||
builtins.quit = None
|
||||
|
||||
import os
|
||||
|
||||
os.environ["OMP_NUM_THREADS"] = "1"
|
||||
|
||||
os.kill = None
|
||||
os.system = None
|
||||
os.putenv = None
|
||||
os.remove = None
|
||||
os.removedirs = None
|
||||
os.rmdir = None
|
||||
os.fchdir = None
|
||||
os.setuid = None
|
||||
os.fork = None
|
||||
os.forkpty = None
|
||||
os.killpg = None
|
||||
os.rename = None
|
||||
os.renames = None
|
||||
os.truncate = None
|
||||
os.replace = None
|
||||
os.unlink = None
|
||||
os.fchmod = None
|
||||
os.fchown = None
|
||||
os.chmod = None
|
||||
os.chown = None
|
||||
os.chroot = None
|
||||
os.fchdir = None
|
||||
os.lchflags = None
|
||||
os.lchmod = None
|
||||
os.lchown = None
|
||||
os.getcwd = None
|
||||
os.chdir = None
|
||||
|
||||
import shutil
|
||||
|
||||
shutil.rmtree = None
|
||||
shutil.move = None
|
||||
shutil.chown = None
|
||||
|
||||
import subprocess
|
||||
|
||||
subprocess.Popen = None
|
||||
|
||||
__builtins__["help"] = None
|
||||
|
||||
import sys
|
||||
|
||||
sys.modules["ipdb"] = None
|
||||
sys.modules["joblib"] = None
|
||||
sys.modules["resource"] = None
|
||||
sys.modules["psutil"] = None
|
||||
sys.modules["tkinter"] = None
|
||||
@@ -202,6 +202,23 @@
|
||||
filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5));
|
||||
}
|
||||
|
||||
/* Onboarding step 2: connection line between devices */
|
||||
.onboarding-connection-line {
|
||||
stroke: oklch(0.85 0.18 85 / 0.5);
|
||||
stroke-width: 1.5px;
|
||||
stroke-dasharray: 6, 6;
|
||||
animation: flowAnimation 1s linear infinite;
|
||||
filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4));
|
||||
}
|
||||
|
||||
/* Onboarding step 4: red connection line for disconnect */
|
||||
.onboarding-connection-line-red {
|
||||
stroke: rgba(220, 38, 38, 0.7);
|
||||
stroke-width: 1.5px;
|
||||
stroke-dasharray: 6, 6;
|
||||
filter: drop-shadow(0 0 2px rgba(220, 38, 38, 0.3));
|
||||
}
|
||||
|
||||
.graph-link-active {
|
||||
stroke: oklch(0.85 0.18 85 / 0.8);
|
||||
stroke-width: 2px;
|
||||
@@ -320,3 +337,51 @@ input:focus, textarea:focus {
|
||||
transform: translate(400px, 400px);
|
||||
}
|
||||
}
|
||||
|
||||
/* Onboarding smooth fade-in */
|
||||
@keyframes onb-fade-in {
|
||||
from { opacity: 0; transform: translateY(8px); }
|
||||
to { opacity: 1; transform: none; }
|
||||
}
|
||||
|
||||
/* Opacity-only fade-in — no transform, safe for containers with fixed-position children */
|
||||
@keyframes onb-fade-opacity {
|
||||
from { opacity: 0; }
|
||||
to { opacity: 1; }
|
||||
}
|
||||
|
||||
/* Respect reduced motion preference */
|
||||
@media (prefers-reduced-motion: reduce) {
|
||||
.shooting-star,
|
||||
.shooting-star::before {
|
||||
animation: none !important;
|
||||
opacity: 0 !important;
|
||||
}
|
||||
.graph-link {
|
||||
animation: none;
|
||||
}
|
||||
.status-pulse {
|
||||
animation: none;
|
||||
}
|
||||
.cursor-blink {
|
||||
animation: none;
|
||||
}
|
||||
.onboarding-connection-line {
|
||||
animation: none;
|
||||
}
|
||||
.onboarding-connection-line-red {
|
||||
animation: none;
|
||||
}
|
||||
[style*="onb-fade-in"],
|
||||
[style*="onb-fade-opacity"] {
|
||||
animation: none !important;
|
||||
opacity: 1 !important;
|
||||
}
|
||||
*,
|
||||
*::before,
|
||||
*::after {
|
||||
transition-duration: 0.01ms !important;
|
||||
animation-duration: 0.01ms !important;
|
||||
animation-iteration-count: 1 !important;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,8 +7,6 @@
|
||||
editingImage,
|
||||
clearEditingImage,
|
||||
selectedChatModel,
|
||||
setSelectedChatModel,
|
||||
instances,
|
||||
ttftMs,
|
||||
tps,
|
||||
totalTokens,
|
||||
@@ -29,6 +27,19 @@
|
||||
showModelSelector?: boolean;
|
||||
modelTasks?: Record<string, string[]>;
|
||||
modelCapabilities?: Record<string, string[]>;
|
||||
onSend?: () => void;
|
||||
onAutoSend?: (
|
||||
content: string,
|
||||
files?: {
|
||||
id: string;
|
||||
name: string;
|
||||
type: string;
|
||||
textContent?: string;
|
||||
preview?: string;
|
||||
}[],
|
||||
) => void;
|
||||
onOpenModelPicker?: () => void;
|
||||
modelDisplayOverride?: string;
|
||||
}
|
||||
|
||||
let {
|
||||
@@ -39,6 +50,10 @@
|
||||
showModelSelector = false,
|
||||
modelTasks = {},
|
||||
modelCapabilities = {},
|
||||
onSend,
|
||||
onAutoSend,
|
||||
onOpenModelPicker,
|
||||
modelDisplayOverride,
|
||||
}: Props = $props();
|
||||
|
||||
let message = $state("");
|
||||
@@ -49,27 +64,12 @@
|
||||
const thinkingEnabled = $derived(thinkingEnabledStore());
|
||||
let loading = $derived(isLoading());
|
||||
const currentModel = $derived(selectedChatModel());
|
||||
const instanceData = $derived(instances());
|
||||
const currentTtft = $derived(ttftMs());
|
||||
const currentTps = $derived(tps());
|
||||
const currentTokens = $derived(totalTokens());
|
||||
const currentEditingImage = $derived(editingImage());
|
||||
const isEditMode = $derived(currentEditingImage !== null);
|
||||
|
||||
// Custom dropdown state
|
||||
let isModelDropdownOpen = $state(false);
|
||||
let dropdownButtonRef: HTMLButtonElement | undefined = $state();
|
||||
let dropdownPosition = $derived(() => {
|
||||
if (!dropdownButtonRef || !isModelDropdownOpen)
|
||||
return { top: 0, left: 0, width: 0 };
|
||||
const rect = dropdownButtonRef.getBoundingClientRect();
|
||||
return {
|
||||
top: rect.top,
|
||||
left: rect.left,
|
||||
width: rect.width,
|
||||
};
|
||||
});
|
||||
|
||||
// Accept all supported file types
|
||||
const acceptString = getAcceptString(["image", "text", "pdf"]);
|
||||
|
||||
@@ -122,73 +122,14 @@
|
||||
uploadedFiles.length > 0),
|
||||
);
|
||||
|
||||
// Extract available models from running instances
|
||||
const availableModels = $derived(() => {
|
||||
const models: Array<{ id: string; label: string; isImageModel: boolean }> =
|
||||
[];
|
||||
for (const [, instance] of Object.entries(instanceData)) {
|
||||
const modelId = getInstanceModelId(instance);
|
||||
if (
|
||||
modelId &&
|
||||
modelId !== "Unknown" &&
|
||||
!models.some((m) => m.id === modelId)
|
||||
) {
|
||||
models.push({
|
||||
id: modelId,
|
||||
label: modelId.split("/").pop() || modelId,
|
||||
isImageModel: modelSupportsImageGeneration(modelId),
|
||||
});
|
||||
}
|
||||
}
|
||||
return models;
|
||||
});
|
||||
|
||||
// Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
|
||||
let previousModelIds: Set<string> = new Set();
|
||||
|
||||
// Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
|
||||
$effect(() => {
|
||||
const models = availableModels();
|
||||
const currentModelIds = new Set(models.map((m) => m.id));
|
||||
|
||||
if (models.length > 0) {
|
||||
// Find newly added models (in current but not in previous)
|
||||
const newModels = models.filter((m) => !previousModelIds.has(m.id));
|
||||
|
||||
// If no model selected, select the first available
|
||||
if (!currentModel) {
|
||||
setSelectedChatModel(models[0].id);
|
||||
}
|
||||
// If current model is stale (no longer has a running instance), reset to first available
|
||||
else if (!models.some((m) => m.id === currentModel)) {
|
||||
setSelectedChatModel(models[0].id);
|
||||
}
|
||||
// If a new model was just added, select it
|
||||
else if (newModels.length > 0 && previousModelIds.size > 0) {
|
||||
setSelectedChatModel(newModels[0].id);
|
||||
}
|
||||
} else {
|
||||
// No instances running - clear the selected model
|
||||
if (currentModel) {
|
||||
setSelectedChatModel("");
|
||||
}
|
||||
}
|
||||
|
||||
// Update previous model IDs for next comparison
|
||||
previousModelIds = currentModelIds;
|
||||
});
|
||||
|
||||
function getInstanceModelId(instanceWrapped: unknown): string {
|
||||
if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
|
||||
const keys = Object.keys(instanceWrapped as Record<string, unknown>);
|
||||
if (keys.length === 1) {
|
||||
const instance = (instanceWrapped as Record<string, unknown>)[
|
||||
keys[0]
|
||||
] as { shardAssignments?: { modelId?: string } };
|
||||
return instance?.shardAssignments?.modelId || "";
|
||||
}
|
||||
return "";
|
||||
}
|
||||
// Short label for the currently selected model
|
||||
const currentModelLabel = $derived(
|
||||
currentModel
|
||||
? currentModel.split("/").pop() || currentModel
|
||||
: modelDisplayOverride
|
||||
? modelDisplayOverride.split("/").pop() || modelDisplayOverride
|
||||
: "",
|
||||
);
|
||||
|
||||
async function handleFiles(files: File[]) {
|
||||
if (files.length === 0) return;
|
||||
@@ -275,6 +216,15 @@
|
||||
uploadedFiles = [];
|
||||
resetTextareaHeight();
|
||||
|
||||
// When onAutoSend is provided, the parent controls all send logic
|
||||
// (including launching non-running models before sending)
|
||||
if (onAutoSend) {
|
||||
onAutoSend(content, files);
|
||||
onSend?.();
|
||||
setTimeout(() => textareaRef?.focus(), 10);
|
||||
return;
|
||||
}
|
||||
|
||||
// Use image editing if in edit mode
|
||||
if (isEditMode && currentEditingImage && content) {
|
||||
editImage(content, currentEditingImage.imageDataUrl);
|
||||
@@ -306,6 +256,8 @@
|
||||
);
|
||||
}
|
||||
|
||||
onSend?.();
|
||||
|
||||
// Refocus the textarea after sending
|
||||
setTimeout(() => textareaRef?.focus(), 10);
|
||||
}
|
||||
@@ -416,7 +368,7 @@
|
||||
{/if}
|
||||
|
||||
<!-- Model selector (when enabled) -->
|
||||
{#if showModelSelector && availableModels().length > 0}
|
||||
{#if showModelSelector}
|
||||
<div
|
||||
class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
|
||||
>
|
||||
@@ -425,33 +377,22 @@
|
||||
class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
|
||||
>MODEL:</span
|
||||
>
|
||||
<!-- Custom dropdown -->
|
||||
<!-- Model button — opens the full model picker -->
|
||||
<div class="relative flex-1 max-w-xs">
|
||||
<button
|
||||
bind:this={dropdownButtonRef}
|
||||
type="button"
|
||||
onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
|
||||
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
|
||||
? 'border-exo-yellow/70'
|
||||
: ''}"
|
||||
onclick={() => onOpenModelPicker?.()}
|
||||
class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
|
||||
>
|
||||
{#if availableModels().find((m) => m.id === currentModel)}
|
||||
<span class="text-exo-yellow truncate"
|
||||
>{availableModels().find((m) => m.id === currentModel)
|
||||
?.label}</span
|
||||
>
|
||||
{:else if availableModels().length > 0}
|
||||
<span class="text-exo-yellow truncate"
|
||||
>{availableModels()[0].label}</span
|
||||
{#if currentModelLabel}
|
||||
<span class="text-exo-yellow truncate">{currentModelLabel}</span
|
||||
>
|
||||
{:else}
|
||||
<span class="text-exo-light-gray/50">— SELECT MODEL —</span>
|
||||
{/if}
|
||||
</button>
|
||||
<div
|
||||
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
|
||||
? 'rotate-180'
|
||||
: ''}"
|
||||
class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
|
||||
>
|
||||
<svg
|
||||
class="w-3 h-3 text-exo-yellow/60"
|
||||
@@ -468,78 +409,6 @@
|
||||
</svg>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{#if isModelDropdownOpen}
|
||||
<!-- Backdrop to close dropdown -->
|
||||
<button
|
||||
type="button"
|
||||
class="fixed inset-0 z-[9998] cursor-default"
|
||||
onclick={() => (isModelDropdownOpen = false)}
|
||||
aria-label="Close dropdown"
|
||||
></button>
|
||||
|
||||
<!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
|
||||
<div
|
||||
class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
|
||||
style="bottom: calc(100vh - {dropdownPosition()
|
||||
.top}px + 4px); left: {dropdownPosition()
|
||||
.left}px; width: {dropdownPosition().width}px;"
|
||||
>
|
||||
<div class="py-1">
|
||||
{#each availableModels() as model}
|
||||
<button
|
||||
type="button"
|
||||
onclick={() => {
|
||||
setSelectedChatModel(model.id);
|
||||
isModelDropdownOpen = false;
|
||||
}}
|
||||
class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
|
||||
model.id
|
||||
? 'bg-transparent text-exo-yellow'
|
||||
: 'text-exo-light-gray hover:text-exo-yellow'}"
|
||||
>
|
||||
{#if currentModel === model.id}
|
||||
<svg
|
||||
class="w-3 h-3 flex-shrink-0"
|
||||
fill="currentColor"
|
||||
viewBox="0 0 20 20"
|
||||
>
|
||||
<path
|
||||
fill-rule="evenodd"
|
||||
d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
|
||||
clip-rule="evenodd"
|
||||
/>
|
||||
</svg>
|
||||
{:else}
|
||||
<span class="w-3"></span>
|
||||
{/if}
|
||||
{#if model.isImageModel}
|
||||
<svg
|
||||
class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
aria-label="Image generation model"
|
||||
>
|
||||
<rect
|
||||
x="3"
|
||||
y="3"
|
||||
width="18"
|
||||
height="18"
|
||||
rx="2"
|
||||
ry="2"
|
||||
/>
|
||||
<circle cx="8.5" cy="8.5" r="1.5" />
|
||||
<polyline points="21 15 16 10 5 21" />
|
||||
</svg>
|
||||
{/if}
|
||||
<span class="truncate flex-1">{model.label}</span>
|
||||
</button>
|
||||
{/each}
|
||||
</div>
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
<!-- Thinking toggle -->
|
||||
{#if modelSupportsThinking()}
|
||||
|
||||
@@ -807,8 +807,8 @@
|
||||
>
|
||||
AWAITING INPUT
|
||||
</p>
|
||||
<p class="text-sm sm:text-xs text-exo-light-gray tracking-wider mt-1">
|
||||
ENTER A QUERY TO BEGIN
|
||||
<p class="text-xs text-white/30 tracking-wider mt-1.5 font-mono">
|
||||
Type a message below · Shift+Enter for newline
|
||||
</p>
|
||||
</div>
|
||||
{/if}
|
||||
@@ -823,6 +823,7 @@
|
||||
onclick={scrollToBottom}
|
||||
class="sticky bottom-4 left-1/2 -translate-x-1/2 w-10 h-10 rounded-full bg-exo-dark-gray/90 border border-exo-medium-gray/50 flex items-center justify-center text-exo-light-gray hover:text-exo-yellow hover:border-exo-yellow/50 transition-all shadow-lg cursor-pointer z-10"
|
||||
title="Scroll to bottom"
|
||||
aria-label="Scroll to bottom of messages"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5"
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
<script lang="ts" module>
|
||||
export interface ChatModelInfo {
|
||||
id: string;
|
||||
name: string;
|
||||
base_model: string;
|
||||
storage_size_megabytes: number;
|
||||
capabilities: string[];
|
||||
family: string;
|
||||
quantization: string;
|
||||
}
|
||||
|
||||
// Auto mode tier list (for when user just starts typing)
|
||||
export const AUTO_TIERS: string[][] = [
|
||||
// Tier 1 (frontier)
|
||||
["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
|
||||
// Tier 2 (excellent)
|
||||
[
|
||||
"Kimi K2",
|
||||
"Qwen3 235B",
|
||||
"MiniMax M2.5",
|
||||
"Step 3.5 Flash",
|
||||
"Qwen3 Next 80B",
|
||||
],
|
||||
// Tier 3 (great)
|
||||
[
|
||||
"GLM 4.7",
|
||||
"MiniMax M2.1",
|
||||
"Qwen3 Coder 480B",
|
||||
"GLM 4.5 Air",
|
||||
"Llama 3.3 70B",
|
||||
],
|
||||
// Tier 4 (good)
|
||||
["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
|
||||
// Tier 5 (small/fast)
|
||||
[
|
||||
"Llama 3.1 8B",
|
||||
"GPT-OSS 20B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
"Llama 3.2 1B",
|
||||
],
|
||||
];
|
||||
|
||||
/** Return the tier index (0 = best) for a base_model name. */
|
||||
export function getAutoTierIndex(baseModel: string): number {
|
||||
for (let i = 0; i < AUTO_TIERS.length; i++) {
|
||||
if (AUTO_TIERS[i].includes(baseModel)) return i;
|
||||
}
|
||||
return AUTO_TIERS.length; // not in any tier → lowest priority
|
||||
}
|
||||
|
||||
/** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
|
||||
export function pickAutoModel(
|
||||
modelList: ChatModelInfo[],
|
||||
memoryGB: number,
|
||||
): ChatModelInfo | null {
|
||||
for (const tier of AUTO_TIERS) {
|
||||
const candidates: ChatModelInfo[] = [];
|
||||
for (const baseModel of tier) {
|
||||
const variants = modelList
|
||||
.filter(
|
||||
(m) =>
|
||||
m.base_model === baseModel &&
|
||||
(m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
|
||||
(m.storage_size_megabytes || 0) > 0,
|
||||
)
|
||||
.sort(
|
||||
(a, b) =>
|
||||
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
|
||||
);
|
||||
if (variants[0]) candidates.push(variants[0]);
|
||||
}
|
||||
if (candidates.length > 0) {
|
||||
candidates.sort(
|
||||
(a, b) =>
|
||||
(b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
|
||||
);
|
||||
return candidates[0];
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
</script>
|
||||
|
||||
<script lang="ts">
|
||||
interface CategoryRecommendation {
|
||||
category: "coding" | "writing" | "agentic" | "biggest";
|
||||
label: string;
|
||||
model: ChatModelInfo | null;
|
||||
tooltip: string;
|
||||
}
|
||||
|
||||
interface Props {
|
||||
models: ChatModelInfo[];
|
||||
clusterLabel: string;
|
||||
totalMemoryGB: number;
|
||||
onSelect: (modelId: string, category: string) => void;
|
||||
onAddModel: () => void;
|
||||
class?: string;
|
||||
}
|
||||
|
||||
let {
|
||||
models,
|
||||
clusterLabel,
|
||||
totalMemoryGB,
|
||||
onSelect,
|
||||
onAddModel,
|
||||
class: className = "",
|
||||
}: Props = $props();
|
||||
|
||||
// --- Hardcoded Rankings ---
|
||||
const CODING_RANKING = [
|
||||
"Qwen3 Coder Next",
|
||||
"Qwen3 Coder 480B",
|
||||
"Qwen3 30B",
|
||||
"GPT-OSS 20B",
|
||||
"Llama 3.1 8B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
];
|
||||
|
||||
const WRITING_RANKING = [
|
||||
"Kimi K2.5",
|
||||
"Kimi K2",
|
||||
"Qwen3 Next 80B",
|
||||
"Llama 3.3 70B",
|
||||
"MiniMax M2.5",
|
||||
"GLM 4.5 Air",
|
||||
"GLM 4.7 Flash",
|
||||
"GPT-OSS 20B",
|
||||
"Llama 3.1 8B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
];
|
||||
|
||||
const AGENTIC_RANKING = [
|
||||
"DeepSeek V3.1",
|
||||
"GLM-5",
|
||||
"Qwen3 235B",
|
||||
"Step 3.5 Flash",
|
||||
"GLM 4.7",
|
||||
"MiniMax M2.1",
|
||||
"GPT-OSS 120B",
|
||||
"Llama 3.3 70B",
|
||||
"Llama 3.1 70B",
|
||||
"GLM 4.7 Flash",
|
||||
"GPT-OSS 20B",
|
||||
"Qwen3 30B",
|
||||
"Llama 3.1 8B",
|
||||
"Llama 3.2 3B",
|
||||
"Qwen3 0.6B",
|
||||
];
|
||||
|
||||
function getModelSizeGB(m: ChatModelInfo): number {
|
||||
return (m.storage_size_megabytes || 0) / 1024;
|
||||
}
|
||||
|
||||
function fitsInMemory(m: ChatModelInfo): boolean {
|
||||
return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
|
||||
}
|
||||
|
||||
/** For a given base_model name, find the biggest quant variant that fits in memory. */
|
||||
function pickBestVariant(baseModel: string): ChatModelInfo | null {
|
||||
const variants = models
|
||||
.filter((m) => m.base_model === baseModel && fitsInMemory(m))
|
||||
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
|
||||
return variants[0] ?? null;
|
||||
}
|
||||
|
||||
/** Walk a ranked list of base_model names, return the first that has a fitting variant. */
|
||||
function pickFromRanking(ranking: string[]): ChatModelInfo | null {
|
||||
for (const baseModel of ranking) {
|
||||
const pick = pickBestVariant(baseModel);
|
||||
if (pick) return pick;
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
/** Pick the single biggest model that fits. */
|
||||
function pickBiggest(): ChatModelInfo | null {
|
||||
const fitting = models
|
||||
.filter((m) => fitsInMemory(m))
|
||||
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
|
||||
return fitting[0] ?? null;
|
||||
}
|
||||
|
||||
const recommendations = $derived.by((): CategoryRecommendation[] => {
|
||||
return [
|
||||
{
|
||||
category: "coding",
|
||||
label: "Best for Coding",
|
||||
model: pickFromRanking(CODING_RANKING),
|
||||
tooltip:
|
||||
"Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
|
||||
},
|
||||
{
|
||||
category: "writing",
|
||||
label: "Best for Writing",
|
||||
model: pickFromRanking(WRITING_RANKING),
|
||||
tooltip: "Ranked by creative writing quality and instruction following",
|
||||
},
|
||||
{
|
||||
category: "agentic",
|
||||
label: "Best Agentic",
|
||||
model: pickFromRanking(AGENTIC_RANKING),
|
||||
tooltip: "Ranked by reasoning, planning, and tool-use capability",
|
||||
},
|
||||
{
|
||||
category: "biggest",
|
||||
label: "Biggest",
|
||||
model: pickBiggest(),
|
||||
tooltip: "Largest model that fits in your available memory",
|
||||
},
|
||||
];
|
||||
});
|
||||
|
||||
function formatSize(mb: number): string {
|
||||
const gb = mb / 1024;
|
||||
if (gb >= 100) return `${Math.round(gb)} GB`;
|
||||
return `${gb.toFixed(1)} GB`;
|
||||
}
|
||||
|
||||
// Category icons (SVG paths)
|
||||
const categoryIcons: Record<string, string> = {
|
||||
coding:
|
||||
"M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
|
||||
writing:
|
||||
"M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
|
||||
agentic:
|
||||
"M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
|
||||
biggest:
|
||||
"M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
|
||||
};
|
||||
|
||||
let hoveredTooltip = $state<string | null>(null);
|
||||
let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
|
||||
|
||||
function showTooltip(category: string, e: MouseEvent | FocusEvent) {
|
||||
hoveredTooltip = category;
|
||||
const target = e.currentTarget as HTMLElement;
|
||||
const rect = target.getBoundingClientRect();
|
||||
tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
|
||||
}
|
||||
|
||||
function hideTooltip() {
|
||||
hoveredTooltip = null;
|
||||
tooltipAnchor = null;
|
||||
}
|
||||
</script>
|
||||
|
||||
<div class="flex flex-col items-center justify-center gap-6 {className}">
|
||||
<!-- Header -->
|
||||
<div class="text-center">
|
||||
<p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
|
||||
Recommended for your
|
||||
</p>
|
||||
<p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
|
||||
</div>
|
||||
|
||||
<!-- Category Cards Grid -->
|
||||
<div class="grid grid-cols-2 gap-3 w-full max-w-md">
|
||||
{#each recommendations as rec}
|
||||
{#if rec.model}
|
||||
<button
|
||||
type="button"
|
||||
onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
|
||||
class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
|
||||
>
|
||||
<!-- Category icon + label -->
|
||||
<div class="flex items-center gap-2 w-full">
|
||||
<svg
|
||||
class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="1.5"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d={categoryIcons[rec.category]}
|
||||
/>
|
||||
</svg>
|
||||
<span
|
||||
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
|
||||
>
|
||||
{rec.label}
|
||||
</span>
|
||||
<!-- Info tooltip -->
|
||||
<div class="ml-auto flex-shrink-0">
|
||||
<span
|
||||
role="button"
|
||||
tabindex="-1"
|
||||
class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
|
||||
onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
|
||||
onmouseleave={() => hideTooltip()}
|
||||
onclick={(e: MouseEvent) => {
|
||||
e.stopPropagation();
|
||||
if (hoveredTooltip === rec.category) {
|
||||
hideTooltip();
|
||||
} else {
|
||||
showTooltip(rec.category, e);
|
||||
}
|
||||
}}
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<circle cx="12" cy="12" r="10" />
|
||||
<path d="M12 16v-4m0-4h.01" />
|
||||
</svg>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Model name + size -->
|
||||
<div class="w-full">
|
||||
<p class="text-sm text-white font-mono truncate">
|
||||
{rec.model.base_model}
|
||||
</p>
|
||||
<p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
|
||||
{formatSize(rec.model.storage_size_megabytes)}
|
||||
{#if rec.model.quantization}
|
||||
<span class="text-exo-light-gray/40"
|
||||
>· {rec.model.quantization}</span
|
||||
>
|
||||
{/if}
|
||||
</p>
|
||||
</div>
|
||||
</button>
|
||||
{:else}
|
||||
<!-- No model fits for this category -->
|
||||
<div
|
||||
class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
|
||||
>
|
||||
<div class="flex items-center gap-2">
|
||||
<svg
|
||||
class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="1.5"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d={categoryIcons[rec.category]}
|
||||
/>
|
||||
</svg>
|
||||
<span
|
||||
class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
|
||||
>{rec.label}</span
|
||||
>
|
||||
</div>
|
||||
<p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
|
||||
</div>
|
||||
{/if}
|
||||
{/each}
|
||||
</div>
|
||||
|
||||
<!-- Add Model Button -->
|
||||
<button
|
||||
type="button"
|
||||
onclick={onAddModel}
|
||||
class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
|
||||
</svg>
|
||||
Add Model
|
||||
</button>
|
||||
|
||||
<!-- Auto hint -->
|
||||
<p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
|
||||
Or just start typing — we'll pick the best model automatically
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
|
||||
{#if hoveredTooltip && tooltipAnchor}
|
||||
{@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
|
||||
{#if rec}
|
||||
<div
|
||||
class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
|
||||
style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
|
||||
8}px; transform: translate(-50%, -100%);"
|
||||
>
|
||||
{rec.tooltip}
|
||||
</div>
|
||||
{/if}
|
||||
{/if}
|
||||
@@ -2,7 +2,6 @@
|
||||
import {
|
||||
conversations,
|
||||
activeConversationId,
|
||||
createConversation,
|
||||
loadConversation,
|
||||
deleteConversation,
|
||||
deleteAllConversations,
|
||||
@@ -17,9 +16,21 @@
|
||||
|
||||
interface Props {
|
||||
class?: string;
|
||||
onNewChat?: () => void;
|
||||
onSelectConversation?: () => void;
|
||||
isMobileDrawer?: boolean;
|
||||
isOpen?: boolean;
|
||||
onClose?: () => void;
|
||||
}
|
||||
|
||||
let { class: className = "" }: Props = $props();
|
||||
let {
|
||||
class: className = "",
|
||||
onNewChat,
|
||||
onSelectConversation,
|
||||
isMobileDrawer = false,
|
||||
isOpen = false,
|
||||
onClose,
|
||||
}: Props = $props();
|
||||
|
||||
const conversationList = $derived(conversations());
|
||||
const activeId = $derived(activeConversationId());
|
||||
@@ -42,11 +53,16 @@
|
||||
);
|
||||
|
||||
function handleNewChat() {
|
||||
createConversation();
|
||||
onNewChat?.();
|
||||
}
|
||||
|
||||
function handleSelectConversation(id: string) {
|
||||
onSelectConversation?.();
|
||||
loadConversation(id);
|
||||
// Close mobile drawer when selecting a conversation
|
||||
if (isMobileDrawer && isOpen) {
|
||||
onClose?.();
|
||||
}
|
||||
}
|
||||
|
||||
function handleStartEdit(id: string, name: string, event: MouseEvent) {
|
||||
@@ -185,11 +201,7 @@
|
||||
|
||||
let instanceType: string | null = null;
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (
|
||||
instanceTag === "MlxIbvInstance" ||
|
||||
instanceTag === "MlxJacclInstance"
|
||||
)
|
||||
instanceType = "MLX RDMA";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
@@ -250,9 +262,7 @@
|
||||
}
|
||||
</script>
|
||||
|
||||
<aside
|
||||
class="flex flex-col h-full bg-exo-dark-gray border-r border-exo-yellow/10 {className}"
|
||||
>
|
||||
{#snippet sidebarContent()}
|
||||
<!-- Header -->
|
||||
<div class="p-4">
|
||||
<button
|
||||
@@ -307,7 +317,7 @@
|
||||
<div class="py-2">
|
||||
<div class="px-4 py-2">
|
||||
<span
|
||||
class="text-sm text-white/70 font-mono tracking-wider uppercase"
|
||||
class="text-xs text-exo-light-gray font-mono tracking-wider uppercase"
|
||||
>
|
||||
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"}
|
||||
</span>
|
||||
@@ -376,39 +386,37 @@
|
||||
onkeydown={(e) =>
|
||||
e.key === "Enter" &&
|
||||
handleSelectConversation(conversation.id)}
|
||||
class="group w-full flex items-center justify-between p-2 rounded mb-1 transition-all text-left cursor-pointer
|
||||
class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer
|
||||
{activeId === conversation.id
|
||||
? 'bg-transparent border border-exo-yellow/30'
|
||||
: 'hover:border-exo-yellow/20 border border-transparent'}"
|
||||
? 'bg-exo-yellow/5 border border-exo-yellow/30'
|
||||
: 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}"
|
||||
>
|
||||
<div class="flex-1 min-w-0 pr-2">
|
||||
<div
|
||||
class="text-sm truncate {activeId === conversation.id
|
||||
class="text-sm font-medium truncate {activeId ===
|
||||
conversation.id
|
||||
? 'text-exo-yellow'
|
||||
: 'text-white/90'}"
|
||||
: 'text-white'}"
|
||||
>
|
||||
{conversation.name}
|
||||
</div>
|
||||
<div class="text-sm text-white/50 mt-0.5">
|
||||
<div class="text-xs text-white/60 mt-0.5">
|
||||
{formatDate(conversation.updatedAt)}
|
||||
</div>
|
||||
<div class="text-sm text-white/70 truncate">
|
||||
<div class="text-xs text-exo-light-gray truncate">
|
||||
{info.modelLabel}
|
||||
</div>
|
||||
<div class="text-xs text-white/60 font-mono">
|
||||
Strategy: <span class="text-white/80"
|
||||
>{info.strategyLabel}</span
|
||||
>
|
||||
</div>
|
||||
{#if stats}
|
||||
<div class="text-xs text-white/60 font-mono mt-1">
|
||||
{#if stats.ttftMs}<span class="text-white/40">TTFT</span>
|
||||
{stats.ttftMs.toFixed(
|
||||
0,
|
||||
)}ms{/if}{#if stats.ttftMs && stats.tps}<span
|
||||
class="text-white/30 mx-1.5">•</span
|
||||
>{/if}{#if stats.tps}{stats.tps.toFixed(1)}
|
||||
<span class="text-white/40">tok/s</span>{/if}
|
||||
<div class="text-xs text-white/70 font-mono mt-1">
|
||||
{#if stats.ttftMs}<span class="text-white/50">TTFT</span>
|
||||
<span class="text-exo-yellow/80"
|
||||
>{stats.ttftMs.toFixed(0)}ms</span
|
||||
>{/if}{#if stats.ttftMs && stats.tps}<span
|
||||
class="text-white/30 mx-1.5">·</span
|
||||
>{/if}{#if stats.tps}<span class="text-exo-yellow/80"
|
||||
>{stats.tps.toFixed(1)}</span
|
||||
>
|
||||
<span class="text-white/50">tok/s</span>{/if}
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
@@ -591,4 +599,30 @@
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</aside>
|
||||
{/snippet}
|
||||
|
||||
{#if isMobileDrawer}
|
||||
<!-- Mobile drawer with overlay -->
|
||||
{#if isOpen}
|
||||
<!-- Overlay backdrop -->
|
||||
<button
|
||||
type="button"
|
||||
class="fixed inset-0 bg-black/60 backdrop-blur-sm z-40 md:hidden"
|
||||
onclick={() => onClose?.()}
|
||||
aria-label="Close sidebar"
|
||||
></button>
|
||||
<!-- Drawer panel -->
|
||||
<aside
|
||||
class="fixed left-0 top-0 bottom-0 w-72 bg-exo-dark-gray border-r border-exo-yellow/10 z-50 flex flex-col md:hidden"
|
||||
>
|
||||
{@render sidebarContent()}
|
||||
</aside>
|
||||
{/if}
|
||||
{:else}
|
||||
<!-- Desktop sidebar -->
|
||||
<aside
|
||||
class="flex flex-col h-full bg-exo-dark-gray border-r border-exo-yellow/10 {className}"
|
||||
>
|
||||
{@render sidebarContent()}
|
||||
</aside>
|
||||
{/if}
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
<script lang="ts">
|
||||
import { isConnected } from "$lib/stores/app.svelte";
|
||||
import { slide } from "svelte/transition";
|
||||
|
||||
const connected = $derived(isConnected());
|
||||
</script>
|
||||
|
||||
{#if !connected}
|
||||
<div
|
||||
transition:slide={{ duration: 200 }}
|
||||
class="relative z-50 flex items-center justify-center gap-2 px-4 py-2 bg-red-950/80 border-b border-red-500/30"
|
||||
role="alert"
|
||||
aria-live="assertive"
|
||||
>
|
||||
<div class="w-2 h-2 bg-red-500 rounded-full animate-pulse"></div>
|
||||
<span class="text-xs font-mono text-red-300 tracking-wider uppercase">
|
||||
Connection lost — Reconnecting to backend…
|
||||
</span>
|
||||
</div>
|
||||
{/if}
|
||||
@@ -0,0 +1,277 @@
|
||||
<script lang="ts">
|
||||
/**
|
||||
* DeviceIcon — renders a device icon as an SVG <g> element.
|
||||
* Uses the exact same proportional math as TopologyGraph.svelte
|
||||
* so that devices look identical in both the topology view and
|
||||
* the onboarding animation.
|
||||
*
|
||||
* Must be placed inside an <svg> element.
|
||||
*/
|
||||
|
||||
interface Props {
|
||||
/** "macbook pro" | "mac studio" | "mac mini" etc. */
|
||||
deviceType: string;
|
||||
/** Center X coordinate in SVG space */
|
||||
cx: number;
|
||||
/** Center Y coordinate in SVG space */
|
||||
cy: number;
|
||||
/** Base sizing factor (equivalent to TopologyGraph's nodeRadius) */
|
||||
size?: number;
|
||||
/** RAM usage 0–100 */
|
||||
ramPercent?: number;
|
||||
/** Unique id suffix for clip-path ids */
|
||||
uid?: string;
|
||||
}
|
||||
|
||||
let {
|
||||
deviceType,
|
||||
cx,
|
||||
cy,
|
||||
size = 60,
|
||||
ramPercent = 60,
|
||||
uid = "dev",
|
||||
}: Props = $props();
|
||||
|
||||
// Apple logo path — same constant used by TopologyGraph
|
||||
const APPLE_LOGO_PATH =
|
||||
"M788.1 340.9c-5.8 4.5-108.2 62.2-108.2 190.5 0 148.4 130.3 200.9 134.2 202.2-.6 3.2-20.7 71.9-68.7 141.9-42.8 61.6-87.5 123.1-155.5 123.1s-85.5-39.5-164-39.5c-76.5 0-103.7 40.8-165.9 40.8s-105.6-57-155.5-127C46.7 790.7 0 663 0 541.8c0-194.4 126.4-297.5 250.8-297.5 66.1 0 121.2 43.4 162.7 43.4 39.5 0 101.1-46 176.3-46 28.5 0 130.9 2.6 198.3 99.2zm-234-181.5c31.1-36.9 53.1-88.1 53.1-139.3 0-7.1-.6-14.3-1.9-20.1-50.6 1.9-110.8 33.7-147.1 75.8-28.5 32.4-55.1 83.6-55.1 135.5 0 7.8 1.3 15.6 1.9 18.1 3.2.6 8.4 1.3 13.6 1.3 45.4 0 102.5-30.4 135.5-71.3z";
|
||||
const LOGO_NATIVE_WIDTH = 814;
|
||||
const LOGO_NATIVE_HEIGHT = 1000;
|
||||
|
||||
const wireColor = "rgba(179,179,179,0.8)";
|
||||
const strokeWidth = 1.5;
|
||||
|
||||
const modelLower = $derived(deviceType.toLowerCase());
|
||||
|
||||
// ── Mac Studio dimensions (same ratios as TopologyGraph) ──
|
||||
const studioW = $derived(size * 1.25);
|
||||
const studioH = $derived(size * 0.85);
|
||||
const studioX = $derived(cx - studioW / 2);
|
||||
const studioY = $derived(cy - studioH / 2);
|
||||
const studioCorner = 4;
|
||||
const studioTopH = $derived(studioH * 0.15);
|
||||
|
||||
// Studio front panel details
|
||||
const studioSlotH = $derived(studioH * 0.14);
|
||||
const studioVSlotW = $derived(studioW * 0.05);
|
||||
const studioVSlotY = $derived(
|
||||
studioY + studioTopH + (studioH - studioTopH) * 0.6,
|
||||
);
|
||||
const studioVSlot1X = $derived(studioX + studioW * 0.18);
|
||||
const studioVSlot2X = $derived(studioX + studioW * 0.28);
|
||||
const studioHSlotW = $derived(studioW * 0.2);
|
||||
const studioHSlotX = $derived(studioX + studioW * 0.5 - studioHSlotW / 2);
|
||||
|
||||
// Studio memory fill
|
||||
const studioMemTotalH = $derived(studioH - studioTopH);
|
||||
const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
|
||||
|
||||
// ── MacBook dimensions (same ratios as TopologyGraph) ──
|
||||
const mbW = $derived((size * 1.6 * 0.85) / 1.15);
|
||||
const mbH = $derived(size * 0.85);
|
||||
const mbX = $derived(cx - mbW / 2);
|
||||
const mbY = $derived(cy - mbH / 2);
|
||||
|
||||
const mbScreenH = $derived(mbH * 0.7);
|
||||
const mbBaseH = $derived(mbH * 0.3);
|
||||
const mbScreenW = $derived(mbW * 0.85);
|
||||
const mbScreenX = $derived(cx - mbScreenW / 2);
|
||||
const mbBezel = 3;
|
||||
|
||||
// MacBook memory fill
|
||||
const mbMemTotalH = $derived(mbScreenH - mbBezel * 2);
|
||||
const mbMemH = $derived((ramPercent / 100) * mbMemTotalH);
|
||||
|
||||
// Apple logo sizing
|
||||
const mbLogoTargetH = $derived(mbScreenH * 0.22);
|
||||
const mbLogoScale = $derived(mbLogoTargetH / LOGO_NATIVE_HEIGHT);
|
||||
const mbLogoX = $derived(cx - (LOGO_NATIVE_WIDTH * mbLogoScale) / 2);
|
||||
const mbLogoY = $derived(
|
||||
mbY + mbScreenH / 2 - (LOGO_NATIVE_HEIGHT * mbLogoScale) / 2,
|
||||
);
|
||||
|
||||
// MacBook base (trapezoidal)
|
||||
const mbBaseY = $derived(mbY + mbScreenH);
|
||||
const mbBaseTopW = $derived(mbScreenW);
|
||||
const mbBaseBottomW = $derived(mbW);
|
||||
const mbBaseTopX = $derived(cx - mbBaseTopW / 2);
|
||||
const mbBaseBottomX = $derived(cx - mbBaseBottomW / 2);
|
||||
|
||||
// Keyboard
|
||||
const mbKbX = $derived(mbBaseTopX + 6);
|
||||
const mbKbY = $derived(mbBaseY + 3);
|
||||
const mbKbW = $derived(mbBaseTopW - 12);
|
||||
const mbKbH = $derived(mbBaseH * 0.55);
|
||||
|
||||
// Trackpad
|
||||
const mbTpW = $derived(mbBaseTopW * 0.4);
|
||||
const mbTpX = $derived(cx - mbTpW / 2);
|
||||
const mbTpY = $derived(mbBaseY + mbKbH + 5);
|
||||
const mbTpH = $derived(mbBaseH * 0.3);
|
||||
|
||||
// Clip IDs
|
||||
const screenClipId = $derived(`di-screen-${uid}`);
|
||||
const studioClipId = $derived(`di-studio-${uid}`);
|
||||
</script>
|
||||
|
||||
{#if modelLower === "mac studio" || modelLower === "mac mini"}
|
||||
<!-- Mac Studio / Mac Mini -->
|
||||
<defs>
|
||||
<clipPath id={studioClipId}>
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY + studioTopH}
|
||||
width={studioW}
|
||||
height={studioH - studioTopH}
|
||||
rx={studioCorner - 1}
|
||||
/>
|
||||
</clipPath>
|
||||
</defs>
|
||||
|
||||
<!-- Main body -->
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY}
|
||||
width={studioW}
|
||||
height={studioH}
|
||||
rx={studioCorner}
|
||||
fill="#1a1a1a"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
|
||||
<!-- Memory fill -->
|
||||
{#if ramPercent > 0}
|
||||
<rect
|
||||
x={studioX}
|
||||
y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
|
||||
width={studioW}
|
||||
height={studioMemH}
|
||||
fill="rgba(255,215,0,0.75)"
|
||||
clip-path="url(#{studioClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Top surface divider -->
|
||||
<line
|
||||
x1={studioX}
|
||||
y1={studioY + studioTopH}
|
||||
x2={studioX + studioW}
|
||||
y2={studioY + studioTopH}
|
||||
stroke="rgba(179,179,179,0.3)"
|
||||
stroke-width="0.5"
|
||||
/>
|
||||
|
||||
<!-- Front panel: vertical slots -->
|
||||
<rect
|
||||
x={studioVSlot1X - studioVSlotW / 2}
|
||||
y={studioVSlotY}
|
||||
width={studioVSlotW}
|
||||
height={studioSlotH}
|
||||
fill="rgba(0,0,0,0.35)"
|
||||
rx="1.5"
|
||||
/>
|
||||
<rect
|
||||
x={studioVSlot2X - studioVSlotW / 2}
|
||||
y={studioVSlotY}
|
||||
width={studioVSlotW}
|
||||
height={studioSlotH}
|
||||
fill="rgba(0,0,0,0.35)"
|
||||
rx="1.5"
|
||||
/>
|
||||
|
||||
<!-- Horizontal slot (SD card) -->
|
||||
<rect
|
||||
x={studioHSlotX}
|
||||
y={studioVSlotY}
|
||||
width={studioHSlotW}
|
||||
height={studioSlotH * 0.6}
|
||||
fill="rgba(0,0,0,0.35)"
|
||||
rx="1"
|
||||
/>
|
||||
{:else}
|
||||
<!-- MacBook Pro -->
|
||||
<defs>
|
||||
<clipPath id={screenClipId}>
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbScreenH - mbBezel * 2}
|
||||
rx="2"
|
||||
/>
|
||||
</clipPath>
|
||||
</defs>
|
||||
|
||||
<!-- Screen outer frame -->
|
||||
<rect
|
||||
x={mbScreenX}
|
||||
y={mbY}
|
||||
width={mbScreenW}
|
||||
height={mbScreenH}
|
||||
rx="3"
|
||||
fill="#1a1a1a"
|
||||
stroke={wireColor}
|
||||
stroke-width={strokeWidth}
|
||||
/>
|
||||
|
||||
<!-- Screen inner (dark) -->
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbScreenH - mbBezel * 2}
|
||||
rx="2"
|
||||
fill="#0a0a12"
|
||||
/>
|
||||
|
||||
<!-- Memory fill on screen -->
|
||||
{#if ramPercent > 0}
|
||||
<rect
|
||||
x={mbScreenX + mbBezel}
|
||||
y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
|
||||
width={mbScreenW - mbBezel * 2}
|
||||
height={mbMemH}
|
||||
fill="rgba(255,215,0,0.85)"
|
||||
clip-path="url(#{screenClipId})"
|
||||
/>
|
||||
{/if}
|
||||
|
||||
<!-- Apple logo -->
|
||||
<path
|
||||
d={APPLE_LOGO_PATH}
|
||||
transform="translate({mbLogoX}, {mbLogoY}) scale({mbLogoScale})"
|
||||
fill="#FFFFFF"
|
||||
opacity="0.9"
|
||||
/>
|
||||
|
||||
<!-- Keyboard base (trapezoidal) -->
|
||||
<path
|
||||
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
|
||||
mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
|
||||
mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
|
||||
fill="#2c2c2c"
|
||||
stroke={wireColor}
|
||||
stroke-width="1"
|
||||
/>
|
||||
|
||||
<!-- Keyboard area -->
|
||||
<rect
|
||||
x={mbKbX}
|
||||
y={mbKbY}
|
||||
width={mbKbW}
|
||||
height={mbKbH}
|
||||
fill="rgba(0,0,0,0.2)"
|
||||
rx="2"
|
||||
/>
|
||||
|
||||
<!-- Trackpad -->
|
||||
<rect
|
||||
x={mbTpX}
|
||||
y={mbTpY}
|
||||
width={mbTpW}
|
||||
height={mbTpH}
|
||||
fill="rgba(255,255,255,0.08)"
|
||||
rx="2"
|
||||
/>
|
||||
{/if}
|
||||
@@ -41,13 +41,13 @@
|
||||
</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] overflow-y-auto scrollbar-hide"
|
||||
class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[72px] sm:min-w-[64px] overflow-y-auto scrollbar-hide"
|
||||
>
|
||||
<!-- All models (no filter) -->
|
||||
<button
|
||||
type="button"
|
||||
onclick={() => onSelect(null)}
|
||||
class="group flex flex-col items-center justify-center p-2 rounded transition-all duration-200 cursor-pointer {selectedFamily ===
|
||||
class="group flex flex-col items-center justify-center p-2 sm:p-2 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
|
||||
null
|
||||
? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
|
||||
: 'hover:bg-white/5 border-l-2 border-transparent'}"
|
||||
|
||||
@@ -6,6 +6,16 @@
|
||||
export let showSidebarToggle = false;
|
||||
export let sidebarVisible = true;
|
||||
export let onToggleSidebar: (() => void) | null = null;
|
||||
export let showMobileMenuToggle = false;
|
||||
export let mobileMenuOpen = false;
|
||||
export let onToggleMobileMenu: (() => void) | null = null;
|
||||
export let showMobileRightToggle = false;
|
||||
export let mobileRightOpen = false;
|
||||
export let onToggleMobileRight: (() => void) | null = null;
|
||||
export let downloadProgress: {
|
||||
count: number;
|
||||
percentage: number;
|
||||
} | null = null;
|
||||
|
||||
function handleHome(): void {
|
||||
if (onHome) {
|
||||
@@ -23,45 +33,96 @@
|
||||
onToggleSidebar();
|
||||
}
|
||||
}
|
||||
|
||||
function handleToggleMobileMenu(): void {
|
||||
if (onToggleMobileMenu) {
|
||||
onToggleMobileMenu();
|
||||
}
|
||||
}
|
||||
|
||||
function handleToggleMobileRight(): void {
|
||||
if (onToggleMobileRight) {
|
||||
onToggleMobileRight();
|
||||
}
|
||||
}
|
||||
</script>
|
||||
|
||||
<header
|
||||
class="relative z-20 flex items-center justify-center px-6 pt-8 pb-4 bg-exo-dark-gray"
|
||||
class="relative z-20 flex items-center justify-center px-4 md:px-6 pt-4 md:pt-8 pb-3 md:pb-4 bg-exo-dark-gray"
|
||||
>
|
||||
<!-- Left: Sidebar Toggle -->
|
||||
{#if showSidebarToggle}
|
||||
<div class="absolute left-6 top-1/2 -translate-y-1/2">
|
||||
<button
|
||||
onclick={handleToggleSidebar}
|
||||
class="p-2 rounded border border-exo-medium-gray/40 hover:border-exo-yellow/50 transition-colors cursor-pointer"
|
||||
title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
|
||||
<!-- Left: Sidebar Toggle (desktop) or Mobile Sidebar Toggle (mobile) -->
|
||||
<div
|
||||
class="absolute left-4 md:left-6 top-1/2 -translate-y-1/2 flex items-center gap-2"
|
||||
>
|
||||
<!-- Mobile sidebar toggle -->
|
||||
<button
|
||||
onclick={handleToggleMobileMenu}
|
||||
class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer md:hidden"
|
||||
title={mobileMenuOpen ? "Hide sidebar" : "Show sidebar"}
|
||||
aria-label={mobileMenuOpen
|
||||
? "Hide conversation sidebar"
|
||||
: "Show conversation sidebar"}
|
||||
aria-pressed={mobileMenuOpen}
|
||||
>
|
||||
<svg
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
class="w-5 h-5 {mobileMenuOpen
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-light-gray'}"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5 {sidebarVisible
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-medium-gray'}"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
{#if sidebarVisible}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
|
||||
/>
|
||||
{:else}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M13 5l7 7-7 7M5 5l7 7-7 7"
|
||||
/>
|
||||
{/if}
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
{/if}
|
||||
{#if mobileMenuOpen}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
|
||||
></path>
|
||||
{:else}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M13 5l7 7-7 7M5 5l7 7-7 7"
|
||||
></path>
|
||||
{/if}
|
||||
</svg>
|
||||
</button>
|
||||
<!-- Desktop sidebar toggle -->
|
||||
<button
|
||||
onclick={handleToggleSidebar}
|
||||
class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer hidden md:block"
|
||||
title={sidebarVisible ? "Hide sidebar" : "Show sidebar"}
|
||||
aria-label={sidebarVisible
|
||||
? "Hide conversation sidebar"
|
||||
: "Show conversation sidebar"}
|
||||
aria-pressed={sidebarVisible}
|
||||
>
|
||||
<svg
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
class="w-5 h-5 {sidebarVisible
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-light-gray'}"
|
||||
>
|
||||
{#if sidebarVisible}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
|
||||
></path>
|
||||
{:else}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M13 5l7 7-7 7M5 5l7 7-7 7"
|
||||
></path>
|
||||
{/if}
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<!-- Center: Logo (clickable to go home) -->
|
||||
<button
|
||||
@@ -75,18 +136,55 @@
|
||||
<img
|
||||
src="/exo-logo.png"
|
||||
alt="EXO"
|
||||
class="h-18 drop-shadow-[0_0_20px_rgba(255,215,0,0.5)]"
|
||||
class="h-12 md:h-18 drop-shadow-[0_0_4px_rgba(255,215,0,0.3)]"
|
||||
/>
|
||||
</button>
|
||||
|
||||
<!-- Right: Home + Downloads -->
|
||||
<div
|
||||
class="absolute right-6 top-1/2 -translate-y-1/2 flex items-center gap-4"
|
||||
<!-- Right: Home + Downloads + Mobile Right Toggle -->
|
||||
<nav
|
||||
class="absolute right-4 md:right-6 top-1/2 -translate-y-1/2 flex items-center gap-2 md:gap-4"
|
||||
aria-label="Main navigation"
|
||||
>
|
||||
<!-- Mobile right sidebar toggle (instances/models) - only show when not in chat mode -->
|
||||
{#if showMobileRightToggle}
|
||||
<button
|
||||
onclick={handleToggleMobileRight}
|
||||
class="p-2 rounded border border-exo-light-gray/30 hover:border-exo-yellow/50 hover:bg-exo-medium-gray/30 transition-colors cursor-pointer md:hidden"
|
||||
title={mobileRightOpen ? "Hide instances" : "Show instances"}
|
||||
aria-label={mobileRightOpen
|
||||
? "Hide instances panel"
|
||||
: "Show instances panel"}
|
||||
aria-pressed={mobileRightOpen}
|
||||
>
|
||||
<svg
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
class="w-5 h-5 {mobileRightOpen
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-light-gray'}"
|
||||
>
|
||||
{#if mobileRightOpen}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M13 5l7 7-7 7M5 5l7 7-7 7"
|
||||
></path>
|
||||
{:else}
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M11 19l-7-7 7-7m8 14l-7-7 7-7"
|
||||
></path>
|
||||
{/if}
|
||||
</svg>
|
||||
</button>
|
||||
{/if}
|
||||
{#if showHome}
|
||||
<button
|
||||
onclick={handleHome}
|
||||
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
|
||||
class="flex text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase items-center gap-2 cursor-pointer"
|
||||
title="Back to topology view"
|
||||
>
|
||||
<svg
|
||||
@@ -102,28 +200,64 @@
|
||||
d="M3 12l2-2m0 0l7-7 7 7M5 10v10a1 1 0 001 1h3m10-11l2 2m-2-2v10a1 1 0 01-1 1h-3m-6 0a1 1 0 001-1v-4a1 1 0 011-1h2a1 1 0 011 1v4a1 1 0 001 1m-6 0h6"
|
||||
/>
|
||||
</svg>
|
||||
Home
|
||||
<span class="hidden sm:inline">Home</span>
|
||||
</button>
|
||||
{/if}
|
||||
<a
|
||||
href="/#/downloads"
|
||||
class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
|
||||
class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
|
||||
title="View downloads overview"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M12 3v12" />
|
||||
<path d="M7 12l5 5 5-5" />
|
||||
<path d="M5 21h14" />
|
||||
</svg>
|
||||
Downloads
|
||||
{#if downloadProgress}
|
||||
<!-- Compact download progress indicator -->
|
||||
<div class="relative w-4 h-4 flex-shrink-0">
|
||||
<svg class="w-4 h-4 -rotate-90" viewBox="0 0 20 20">
|
||||
<circle
|
||||
cx="10"
|
||||
cy="10"
|
||||
r="8"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
opacity="0.2"
|
||||
/>
|
||||
<circle
|
||||
cx="10"
|
||||
cy="10"
|
||||
r="8"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-dasharray={2 * Math.PI * 8}
|
||||
stroke-dashoffset={2 *
|
||||
Math.PI *
|
||||
8 *
|
||||
(1 - downloadProgress.percentage / 100)}
|
||||
class="text-blue-400 transition-all duration-300"
|
||||
/>
|
||||
</svg>
|
||||
<div
|
||||
class="absolute inset-0 flex items-center justify-center text-[6px] font-mono text-blue-400"
|
||||
>
|
||||
{downloadProgress.count}
|
||||
</div>
|
||||
</div>
|
||||
{:else}
|
||||
<svg
|
||||
class="w-4 h-4"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M12 3v12" />
|
||||
<path d="M7 12l5 5 5-5" />
|
||||
<path d="M5 21h14" />
|
||||
</svg>
|
||||
{/if}
|
||||
<span class="hidden sm:inline">Downloads</span>
|
||||
</a>
|
||||
</div>
|
||||
</nav>
|
||||
</header>
|
||||
|
||||
@@ -36,8 +36,12 @@
|
||||
return num.toString();
|
||||
}
|
||||
|
||||
// Extract model name from full ID (e.g., "mlx-community/Llama-3.2-1B" -> "Llama-3.2-1B")
|
||||
const modelName = $derived(model.id.split("/").pop() || model.id);
|
||||
// Show short name for mlx-community models, full ID for everything else
|
||||
const modelName = $derived(
|
||||
model.author === "mlx-community"
|
||||
? model.id.split("/").pop() || model.id
|
||||
: model.id,
|
||||
);
|
||||
</script>
|
||||
|
||||
<div
|
||||
|
||||
@@ -507,9 +507,29 @@
|
||||
});
|
||||
|
||||
$effect(() => {
|
||||
if (containerRef && processedHtml) {
|
||||
setupCopyButtons();
|
||||
if (!containerRef || !browser) return;
|
||||
|
||||
function handleDelegatedClick(event: MouseEvent) {
|
||||
const codeBtn = (event.target as HTMLElement).closest(
|
||||
".copy-code-btn",
|
||||
) as HTMLButtonElement | null;
|
||||
if (codeBtn) {
|
||||
handleCopyClick({ currentTarget: codeBtn } as unknown as Event);
|
||||
return;
|
||||
}
|
||||
const mathBtn = (event.target as HTMLElement).closest(
|
||||
".copy-math-btn",
|
||||
) as HTMLButtonElement | null;
|
||||
if (mathBtn) {
|
||||
handleMathCopyClick({ currentTarget: mathBtn } as unknown as Event);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
containerRef.addEventListener("click", handleDelegatedClick);
|
||||
return () => {
|
||||
containerRef?.removeEventListener("click", handleDelegatedClick);
|
||||
};
|
||||
});
|
||||
</script>
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@
|
||||
} | null;
|
||||
nodes?: Record<string, NodeInfo>;
|
||||
sharding?: "Pipeline" | "Tensor";
|
||||
runtime?: "MlxRing" | "MlxIbv" | "MlxJaccl";
|
||||
runtime?: "MlxRing" | "MlxJaccl";
|
||||
onLaunch?: () => void;
|
||||
tags?: string[];
|
||||
apiPreview?: PlacementPreview | null;
|
||||
@@ -348,7 +348,7 @@
|
||||
// Debug mode state
|
||||
const isDebugMode = $derived(debugMode());
|
||||
const topology = $derived(topologyData());
|
||||
const isRdma = $derived(runtime === "MlxIbv" || runtime === "MlxJaccl");
|
||||
const isRdma = $derived(runtime === "MlxJaccl");
|
||||
|
||||
// Get interface name for an IP from node data
|
||||
function getInterfaceForIp(nodeId: string, ip?: string): string | null {
|
||||
@@ -567,20 +567,46 @@
|
||||
<div class="flex items-center gap-1.5 mb-2">
|
||||
<span
|
||||
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
|
||||
title={sharding === "Pipeline"
|
||||
? "Pipeline: splits model into sequential stages across devices. Lower network overhead."
|
||||
: "Tensor: splits each layer across devices. Best with high-bandwidth connections (Thunderbolt)."}
|
||||
>
|
||||
{sharding}
|
||||
</span>
|
||||
<span
|
||||
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
|
||||
title={runtime === "MlxRing"
|
||||
? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
|
||||
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
|
||||
>
|
||||
{runtime === "MlxRing"
|
||||
? "MLX Ring"
|
||||
: runtime === "MlxIbv" || runtime === "MlxJaccl"
|
||||
: runtime === "MlxJaccl"
|
||||
? "MLX RDMA"
|
||||
: runtime}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<!-- Download Status -->
|
||||
{#if isDownloading && progress}
|
||||
<div class="mb-2 space-y-1">
|
||||
<div class="flex items-center justify-between text-xs font-mono">
|
||||
<span class="text-blue-400 tracking-wider uppercase">Downloading</span
|
||||
>
|
||||
<span class="text-white/60"
|
||||
>{percentage.toFixed(1)}% · {formatSpeed(progress.speed)}
|
||||
· {formatEta(progress.etaMs)}</span
|
||||
>
|
||||
</div>
|
||||
<div class="h-1 bg-exo-medium-gray/30 rounded overflow-hidden">
|
||||
<div
|
||||
class="h-full bg-blue-500/70 transition-all duration-300"
|
||||
style="width: {percentage}%"
|
||||
></div>
|
||||
</div>
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
<!-- Mini Topology Preview -->
|
||||
{#if placementPreview().nodes.length > 0}
|
||||
{@const preview = placementPreview()}
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
capabilities: string[];
|
||||
sizeRange: { min: number; max: number } | null;
|
||||
downloadedOnly: boolean;
|
||||
readyOnly: boolean;
|
||||
}
|
||||
|
||||
type ModelFilterPopoverProps = {
|
||||
@@ -190,34 +191,58 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Downloaded only -->
|
||||
<!-- Availability filters -->
|
||||
<div>
|
||||
<h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
|
||||
<button
|
||||
type="button"
|
||||
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
|
||||
? 'bg-green-500/20 text-green-400 border border-green-500/30'
|
||||
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
|
||||
onclick={() =>
|
||||
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 inline-block"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
<div class="flex flex-wrap gap-1.5">
|
||||
<button
|
||||
type="button"
|
||||
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
|
||||
? 'bg-green-500/20 text-green-400 border border-green-500/30'
|
||||
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
|
||||
onclick={() =>
|
||||
onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
|
||||
>
|
||||
<path
|
||||
class="text-white/40"
|
||||
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
|
||||
/>
|
||||
<path class="text-green-400" d="m9 13 2 2 4-4" />
|
||||
</svg>
|
||||
<span class="ml-1">Downloaded</span>
|
||||
</button>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 inline-block"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path
|
||||
class="text-white/40"
|
||||
d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
|
||||
/>
|
||||
<path class="text-green-400" d="m9 13 2 2 4-4" />
|
||||
</svg>
|
||||
<span class="ml-1">Downloaded</span>
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
|
||||
? 'bg-green-500/20 text-green-400 border border-green-500/30'
|
||||
: 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
|
||||
onclick={() =>
|
||||
onChange({ ...filters, readyOnly: !filters.readyOnly })}
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 inline-block"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<circle cx="12" cy="12" r="10" />
|
||||
<path d="m9 12 2 2 4-4" />
|
||||
</svg>
|
||||
<span class="ml-1">Ready</span>
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Size range -->
|
||||
|
||||
@@ -32,6 +32,7 @@
|
||||
group: ModelGroup;
|
||||
isExpanded: boolean;
|
||||
isFavorite: boolean;
|
||||
isHighlighted?: boolean;
|
||||
selectedModelId: string | null;
|
||||
canModelFit: (id: string) => boolean;
|
||||
getModelFitStatus: (id: string) => ModelFitStatus;
|
||||
@@ -41,12 +42,14 @@
|
||||
onShowInfo: (group: ModelGroup) => void;
|
||||
downloadStatusMap?: Map<string, DownloadAvailability>;
|
||||
launchedAt?: number;
|
||||
instanceStatuses?: Record<string, { status: string; statusClass: string }>;
|
||||
};
|
||||
|
||||
let {
|
||||
group,
|
||||
isExpanded,
|
||||
isFavorite,
|
||||
isHighlighted = false,
|
||||
selectedModelId,
|
||||
canModelFit,
|
||||
getModelFitStatus,
|
||||
@@ -56,6 +59,7 @@
|
||||
onShowInfo,
|
||||
downloadStatusMap,
|
||||
launchedAt,
|
||||
instanceStatuses = {},
|
||||
}: ModelPickerGroupProps = $props();
|
||||
|
||||
// Group-level download status: show if any variant is downloaded
|
||||
@@ -92,6 +96,12 @@
|
||||
const anyVariantFits = $derived(
|
||||
group.variants.some((v) => canModelFit(v.id)),
|
||||
);
|
||||
// Check if any variant has an active instance (ready, loading, downloading)
|
||||
const anyVariantHasInstance = $derived(
|
||||
instanceStatuses
|
||||
? group.variants.some((v) => instanceStatuses[v.id] != null)
|
||||
: false,
|
||||
);
|
||||
const groupFitStatus = $derived.by((): ModelFitStatus => {
|
||||
let hasClusterCapacityOnly = false;
|
||||
for (const variant of group.variants) {
|
||||
@@ -122,16 +132,36 @@
|
||||
!group.hasMultipleVariants &&
|
||||
group.variants.some((v) => v.id === selectedModelId),
|
||||
);
|
||||
|
||||
// Group-level instance status: show the "best" status across all variants
|
||||
const groupInstanceStatus = $derived.by(() => {
|
||||
if (!instanceStatuses) return null;
|
||||
const readyStatuses = ["READY", "LOADED", "RUNNING"];
|
||||
const loadingStatuses = ["LOADING", "WARMING UP"];
|
||||
let bestStatus: { status: string; statusClass: string } | null = null;
|
||||
for (const variant of group.variants) {
|
||||
const s = instanceStatuses[variant.id];
|
||||
if (!s) continue;
|
||||
if (readyStatuses.includes(s.status)) return s; // Ready is best
|
||||
if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
|
||||
bestStatus = s;
|
||||
}
|
||||
}
|
||||
return bestStatus;
|
||||
});
|
||||
</script>
|
||||
|
||||
<div
|
||||
class="border-b border-white/5 last:border-b-0 {!anyVariantFits
|
||||
data-model-ids={group.variants.map((v) => v.id).join(" ")}
|
||||
class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
|
||||
!anyVariantHasInstance
|
||||
? 'opacity-50'
|
||||
: ''}"
|
||||
: ''} {isHighlighted ? 'model-just-added' : ''}"
|
||||
>
|
||||
<!-- Main row -->
|
||||
<div
|
||||
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
|
||||
class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
|
||||
anyVariantHasInstance
|
||||
? 'hover:bg-white/5 cursor-pointer'
|
||||
: 'cursor-not-allowed'} {isMainSelected
|
||||
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
|
||||
@@ -141,7 +171,7 @@
|
||||
onToggleExpand();
|
||||
} else {
|
||||
const modelId = group.variants[0]?.id;
|
||||
if (modelId && canModelFit(modelId)) {
|
||||
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
|
||||
onSelectModel(modelId);
|
||||
}
|
||||
}
|
||||
@@ -155,7 +185,7 @@
|
||||
onToggleExpand();
|
||||
} else {
|
||||
const modelId = group.variants[0]?.id;
|
||||
if (modelId && canModelFit(modelId)) {
|
||||
if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
|
||||
onSelectModel(modelId);
|
||||
}
|
||||
}
|
||||
@@ -346,6 +376,47 @@
|
||||
</span>
|
||||
{/if}
|
||||
|
||||
<!-- Instance status badge -->
|
||||
{#if groupInstanceStatus}
|
||||
{#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
|
||||
<span class="flex-shrink-0" title="Running">
|
||||
<svg
|
||||
class="w-3 h-3 text-green-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if groupInstanceStatus.status === "DOWNLOADING"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Downloading">
|
||||
<svg
|
||||
class="w-3.5 h-3.5 text-blue-400"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
|
||||
<polyline points="7 10 12 15 17 10" />
|
||||
<line x1="12" y1="15" x2="12" y2="3" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Loading">
|
||||
<svg
|
||||
class="w-3 h-3 text-yellow-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
<!-- Check mark if selected (single-variant) -->
|
||||
{#if isMainSelected}
|
||||
<svg
|
||||
@@ -420,20 +491,30 @@
|
||||
{#each group.variants as variant}
|
||||
{@const fitStatus = getModelFitStatus(variant.id)}
|
||||
{@const modelCanFit = canModelFit(variant.id)}
|
||||
{@const variantHasInstance = instanceStatuses[variant.id] != null}
|
||||
{@const isSelected = selectedModelId === variant.id}
|
||||
<button
|
||||
type="button"
|
||||
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
|
||||
<div
|
||||
class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
|
||||
!variantHasInstance
|
||||
? 'opacity-50 cursor-not-allowed'
|
||||
: 'cursor-pointer'} {isSelected
|
||||
? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
|
||||
: 'border-l-2 border-transparent'}"
|
||||
disabled={!modelCanFit}
|
||||
role="button"
|
||||
tabindex="0"
|
||||
onclick={() => {
|
||||
if (modelCanFit) {
|
||||
if (modelCanFit || variantHasInstance) {
|
||||
onSelectModel(variant.id);
|
||||
}
|
||||
}}
|
||||
onkeydown={(e) => {
|
||||
if (e.key === "Enter" || e.key === " ") {
|
||||
e.preventDefault();
|
||||
if (modelCanFit) {
|
||||
onSelectModel(variant.id);
|
||||
}
|
||||
}
|
||||
}}
|
||||
>
|
||||
<!-- Quantization badge -->
|
||||
<span
|
||||
@@ -478,6 +559,48 @@
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
<!-- Instance status badge -->
|
||||
{#if instanceStatuses[variant.id]}
|
||||
{@const instStatus = instanceStatuses[variant.id]}
|
||||
{#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
|
||||
<span class="flex-shrink-0" title="Running">
|
||||
<svg
|
||||
class="w-3 h-3 text-green-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if instStatus.status === "DOWNLOADING"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Downloading">
|
||||
<svg
|
||||
class="w-3.5 h-3.5 text-blue-400"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
>
|
||||
<path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
|
||||
<polyline points="7 10 12 15 17 10" />
|
||||
<line x1="12" y1="15" x2="12" y2="3" />
|
||||
</svg>
|
||||
</span>
|
||||
{:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
|
||||
<span class="flex-shrink-0 animate-pulse" title="Loading">
|
||||
<svg
|
||||
class="w-3 h-3 text-yellow-400"
|
||||
viewBox="0 0 12 12"
|
||||
fill="currentColor"
|
||||
>
|
||||
<circle cx="6" cy="6" r="5" />
|
||||
</svg>
|
||||
</span>
|
||||
{/if}
|
||||
{/if}
|
||||
|
||||
<!-- Check mark if selected -->
|
||||
{#if isSelected}
|
||||
<svg
|
||||
@@ -490,8 +613,55 @@
|
||||
/>
|
||||
</svg>
|
||||
{/if}
|
||||
</button>
|
||||
|
||||
<!-- Info button -->
|
||||
<button
|
||||
type="button"
|
||||
class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
|
||||
onclick={(e) => {
|
||||
e.stopPropagation();
|
||||
onShowInfo({
|
||||
id: variant.id,
|
||||
name: variant.name || variant.id,
|
||||
capabilities: group.capabilities,
|
||||
family: group.family,
|
||||
variants: [variant],
|
||||
smallestVariant: variant,
|
||||
hasMultipleVariants: false,
|
||||
});
|
||||
}}
|
||||
title="View variant details"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4 text-white/30 hover:text-white/50"
|
||||
viewBox="0 0 24 24"
|
||||
fill="currentColor"
|
||||
>
|
||||
<path
|
||||
d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
|
||||
/>
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
{/each}
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
|
||||
<style>
|
||||
.model-just-added {
|
||||
animation: highlightFade 4s ease-out forwards;
|
||||
}
|
||||
|
||||
@keyframes highlightFade {
|
||||
0%,
|
||||
40% {
|
||||
background-color: rgba(20, 83, 45, 0.25);
|
||||
box-shadow: inset 0 0 0 1px rgba(74, 222, 128, 0.4);
|
||||
}
|
||||
100% {
|
||||
background-color: transparent;
|
||||
box-shadow: none;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
<script lang="ts">
|
||||
import { tick } from "svelte";
|
||||
import { fade, fly } from "svelte/transition";
|
||||
import { cubicOut } from "svelte/easing";
|
||||
import FamilySidebar from "./FamilySidebar.svelte";
|
||||
@@ -7,6 +8,7 @@
|
||||
import HuggingFaceResultItem from "./HuggingFaceResultItem.svelte";
|
||||
import { getNodesWithModelDownloaded } from "$lib/utils/downloads";
|
||||
import { getRecentEntries } from "$lib/stores/recents.svelte";
|
||||
import { addToast } from "$lib/stores/toast.svelte";
|
||||
|
||||
interface ModelInfo {
|
||||
id: string;
|
||||
@@ -36,6 +38,7 @@
|
||||
capabilities: string[];
|
||||
sizeRange: { min: number; max: number } | null;
|
||||
downloadedOnly: boolean;
|
||||
readyOnly: boolean;
|
||||
}
|
||||
|
||||
interface HuggingFaceModel {
|
||||
@@ -49,6 +52,11 @@
|
||||
|
||||
type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
|
||||
|
||||
export type InstanceStatus = {
|
||||
status: string;
|
||||
statusClass: string;
|
||||
};
|
||||
|
||||
type ModelPickerModalProps = {
|
||||
isOpen: boolean;
|
||||
models: ModelInfo[];
|
||||
@@ -75,6 +83,7 @@
|
||||
macmon_info?: { memory?: { ram_total?: number } };
|
||||
}
|
||||
>;
|
||||
instanceStatuses?: Record<string, InstanceStatus>;
|
||||
};
|
||||
|
||||
let {
|
||||
@@ -96,6 +105,7 @@
|
||||
usedMemoryGB,
|
||||
downloadsData,
|
||||
topologyNodes,
|
||||
instanceStatuses = {},
|
||||
}: ModelPickerModalProps = $props();
|
||||
|
||||
// Local state
|
||||
@@ -107,6 +117,7 @@
|
||||
capabilities: [],
|
||||
sizeRange: null,
|
||||
downloadedOnly: false,
|
||||
readyOnly: false,
|
||||
});
|
||||
let infoGroup = $state<ModelGroup | null>(null);
|
||||
|
||||
@@ -182,6 +193,13 @@
|
||||
let hfSearchDebounceTimer: ReturnType<typeof setTimeout> | null = null;
|
||||
let manualModelId = $state("");
|
||||
let addModelError = $state<string | null>(null);
|
||||
let justAddedModelId = $state<string | null>(null);
|
||||
let justAddedTimer: ReturnType<typeof setTimeout> | null = null;
|
||||
|
||||
// Inline HuggingFace search in main search bar
|
||||
let mainSearchHfResults = $state<HuggingFaceModel[]>([]);
|
||||
let mainSearchHfLoading = $state(false);
|
||||
let mainSearchDebounceTimer: ReturnType<typeof setTimeout> | null = null;
|
||||
|
||||
// Reset transient state when modal opens, but preserve tab selection
|
||||
$effect(() => {
|
||||
@@ -191,6 +209,11 @@
|
||||
showFilters = false;
|
||||
manualModelId = "";
|
||||
addModelError = null;
|
||||
justAddedModelId = null;
|
||||
if (justAddedTimer) {
|
||||
clearTimeout(justAddedTimer);
|
||||
justAddedTimer = null;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
@@ -205,6 +228,50 @@
|
||||
}
|
||||
});
|
||||
|
||||
// Inline HuggingFace search when local search returns no results
|
||||
$effect(() => {
|
||||
const query = searchQuery.trim();
|
||||
const noLocalResults = filteredGroups.length === 0;
|
||||
|
||||
if (mainSearchDebounceTimer) {
|
||||
clearTimeout(mainSearchDebounceTimer);
|
||||
mainSearchDebounceTimer = null;
|
||||
}
|
||||
|
||||
if (
|
||||
selectedFamily === "huggingface" ||
|
||||
selectedFamily === "recents" ||
|
||||
selectedFamily === "favorites" ||
|
||||
query.length < 2 ||
|
||||
!noLocalResults
|
||||
) {
|
||||
mainSearchHfResults = [];
|
||||
mainSearchHfLoading = false;
|
||||
return;
|
||||
}
|
||||
|
||||
mainSearchHfLoading = true;
|
||||
mainSearchDebounceTimer = setTimeout(async () => {
|
||||
try {
|
||||
const response = await fetch(
|
||||
`/models/search?query=${encodeURIComponent(query)}&limit=10`,
|
||||
);
|
||||
if (response.ok) {
|
||||
const results: HuggingFaceModel[] = await response.json();
|
||||
mainSearchHfResults = results.filter(
|
||||
(r) => !existingModelIds.has(r.id),
|
||||
);
|
||||
} else {
|
||||
mainSearchHfResults = [];
|
||||
}
|
||||
} catch {
|
||||
mainSearchHfResults = [];
|
||||
} finally {
|
||||
mainSearchHfLoading = false;
|
||||
}
|
||||
}, 500);
|
||||
});
|
||||
|
||||
async function fetchTrendingModels() {
|
||||
hfIsLoadingTrending = true;
|
||||
try {
|
||||
@@ -265,6 +332,24 @@
|
||||
addModelError = null;
|
||||
try {
|
||||
await onAddModel(modelId);
|
||||
// Success: show toast, switch to All Models, highlight the model
|
||||
const shortName = modelId.split("/").pop() || modelId;
|
||||
addToast({ type: "success", message: `Added ${shortName}` });
|
||||
justAddedModelId = modelId;
|
||||
selectedFamily = null;
|
||||
searchQuery = "";
|
||||
// Scroll to the newly added model after DOM update
|
||||
await tick();
|
||||
const el = document.querySelector(
|
||||
`[data-model-ids~="${CSS.escape(modelId)}"]`,
|
||||
);
|
||||
el?.scrollIntoView({ behavior: "smooth", block: "center" });
|
||||
// Clear highlight after 4 seconds
|
||||
if (justAddedTimer) clearTimeout(justAddedTimer);
|
||||
justAddedTimer = setTimeout(() => {
|
||||
justAddedModelId = null;
|
||||
justAddedTimer = null;
|
||||
}, 4000);
|
||||
} catch (error) {
|
||||
addModelError =
|
||||
error instanceof Error ? error.message : "Failed to add model";
|
||||
@@ -440,6 +525,16 @@
|
||||
);
|
||||
}
|
||||
|
||||
// Filter to ready/running models only
|
||||
if (filters.readyOnly) {
|
||||
result = result.filter((g) =>
|
||||
g.variants.some((v) => {
|
||||
const s = instanceStatuses[v.id];
|
||||
return s && s.statusClass === "ready";
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
// Sort: fits-now first, then fits-cluster-capacity, then too-large
|
||||
result.sort((a, b) => {
|
||||
const getGroupFitRank = (group: ModelGroup): number => {
|
||||
@@ -512,6 +607,18 @@
|
||||
);
|
||||
});
|
||||
|
||||
// Split filtered groups into recommended (fits_now) and others for visual separation
|
||||
const recommendedGroups = $derived(
|
||||
filteredGroups.filter((g) =>
|
||||
g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
|
||||
),
|
||||
);
|
||||
const otherGroups = $derived(
|
||||
filteredGroups.filter(
|
||||
(g) => !g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"),
|
||||
),
|
||||
);
|
||||
|
||||
function toggleGroupExpanded(groupId: string) {
|
||||
const next = new Set(expandedGroups);
|
||||
if (next.has(groupId)) {
|
||||
@@ -538,13 +645,19 @@
|
||||
}
|
||||
|
||||
function clearFilters() {
|
||||
filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
|
||||
filters = {
|
||||
capabilities: [],
|
||||
sizeRange: null,
|
||||
downloadedOnly: false,
|
||||
readyOnly: false,
|
||||
};
|
||||
}
|
||||
|
||||
const hasActiveFilters = $derived(
|
||||
filters.capabilities.length > 0 ||
|
||||
filters.sizeRange !== null ||
|
||||
filters.downloadedOnly,
|
||||
filters.downloadedOnly ||
|
||||
filters.readyOnly,
|
||||
);
|
||||
</script>
|
||||
|
||||
@@ -804,6 +917,8 @@
|
||||
{group}
|
||||
isExpanded={expandedGroups.has(group.id)}
|
||||
isFavorite={favorites.has(group.id)}
|
||||
isHighlighted={justAddedModelId !== null &&
|
||||
group.variants.some((v) => v.id === justAddedModelId)}
|
||||
{selectedModelId}
|
||||
{canModelFit}
|
||||
{getModelFitStatus}
|
||||
@@ -813,6 +928,7 @@
|
||||
onShowInfo={(g) => (infoGroup = g)}
|
||||
downloadStatusMap={getVariantDownloadMap(group)}
|
||||
launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
|
||||
{instanceStatuses}
|
||||
/>
|
||||
{/each}
|
||||
{/if}
|
||||
@@ -840,11 +956,40 @@
|
||||
{/if}
|
||||
</div>
|
||||
{:else}
|
||||
{#each filteredGroups as group}
|
||||
<!-- Recommended for your cluster -->
|
||||
{#if recommendedGroups.length > 0 && otherGroups.length > 0 && !searchQuery.trim()}
|
||||
<div
|
||||
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-green-950/60 border-b border-green-500/20 backdrop-blur-sm"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5 text-green-400 flex-shrink-0"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z"
|
||||
/>
|
||||
</svg>
|
||||
<span
|
||||
class="text-xs font-mono text-green-400 tracking-wider uppercase"
|
||||
>Recommended for your cluster</span
|
||||
>
|
||||
<span class="text-xs font-mono text-green-400/50"
|
||||
>— fits in available memory</span
|
||||
>
|
||||
</div>
|
||||
{/if}
|
||||
{#each recommendedGroups as group}
|
||||
<ModelPickerGroup
|
||||
{group}
|
||||
isExpanded={expandedGroups.has(group.id)}
|
||||
isFavorite={favorites.has(group.id)}
|
||||
isHighlighted={justAddedModelId !== null &&
|
||||
group.variants.some((v) => v.id === justAddedModelId)}
|
||||
{selectedModelId}
|
||||
{canModelFit}
|
||||
{getModelFitStatus}
|
||||
@@ -853,8 +998,87 @@
|
||||
{onToggleFavorite}
|
||||
onShowInfo={(g) => (infoGroup = g)}
|
||||
downloadStatusMap={getVariantDownloadMap(group)}
|
||||
{instanceStatuses}
|
||||
/>
|
||||
{/each}
|
||||
<!-- Other models -->
|
||||
{#if otherGroups.length > 0 && recommendedGroups.length > 0 && !searchQuery.trim()}
|
||||
<div
|
||||
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-exo-dark-gray/80 border-y border-exo-medium-gray/20 backdrop-blur-sm"
|
||||
>
|
||||
<span
|
||||
class="text-xs font-mono text-white/40 tracking-wider uppercase"
|
||||
>Other models</span
|
||||
>
|
||||
</div>
|
||||
{/if}
|
||||
{#each otherGroups as group}
|
||||
<ModelPickerGroup
|
||||
{group}
|
||||
isExpanded={expandedGroups.has(group.id)}
|
||||
isFavorite={favorites.has(group.id)}
|
||||
isHighlighted={justAddedModelId !== null &&
|
||||
group.variants.some((v) => v.id === justAddedModelId)}
|
||||
{selectedModelId}
|
||||
{canModelFit}
|
||||
{getModelFitStatus}
|
||||
onToggleExpand={() => toggleGroupExpanded(group.id)}
|
||||
onSelectModel={handleSelect}
|
||||
{onToggleFavorite}
|
||||
onShowInfo={(g) => (infoGroup = g)}
|
||||
downloadStatusMap={getVariantDownloadMap(group)}
|
||||
{instanceStatuses}
|
||||
/>
|
||||
{/each}
|
||||
<!-- Inline HuggingFace search results (shown when no local results match) -->
|
||||
{#if filteredGroups.length === 0 && searchQuery.trim().length >= 2 && selectedFamily !== "huggingface" && selectedFamily !== "recents" && selectedFamily !== "favorites"}
|
||||
{#if mainSearchHfLoading}
|
||||
<div
|
||||
class="flex items-center gap-2 px-3 py-2 border-t border-orange-400/20 bg-orange-950/20"
|
||||
>
|
||||
<span
|
||||
class="w-4 h-4 border-2 border-orange-400 border-t-transparent rounded-full animate-spin"
|
||||
></span>
|
||||
<span class="text-xs font-mono text-orange-400/60"
|
||||
>Searching HuggingFace...</span
|
||||
>
|
||||
</div>
|
||||
{:else if mainSearchHfResults.length > 0}
|
||||
<div
|
||||
class="sticky top-0 z-10 flex items-center gap-2 px-3 py-2 bg-orange-950/30 border-y border-orange-400/20 backdrop-blur-sm"
|
||||
>
|
||||
<span
|
||||
class="text-xs font-mono text-orange-400 tracking-wider uppercase"
|
||||
>From HuggingFace</span
|
||||
>
|
||||
</div>
|
||||
{#each mainSearchHfResults as model}
|
||||
<HuggingFaceResultItem
|
||||
{model}
|
||||
isAdded={existingModelIds.has(model.id)}
|
||||
isAdding={addingModelId === model.id}
|
||||
onAdd={() => handleAddModel(model.id)}
|
||||
onSelect={() => handleSelectHfModel(model.id)}
|
||||
downloadedOnNodes={downloadsData
|
||||
? getNodesWithModelDownloaded(downloadsData, model.id).map(
|
||||
getNodeName,
|
||||
)
|
||||
: []}
|
||||
/>
|
||||
{/each}
|
||||
<button
|
||||
type="button"
|
||||
class="w-full px-3 py-2 text-xs font-mono text-orange-400/60 hover:text-orange-400 hover:bg-orange-500/10 transition-colors text-center"
|
||||
onclick={() => {
|
||||
hfSearchQuery = searchQuery;
|
||||
searchHuggingFace(searchQuery);
|
||||
selectedFamily = "huggingface";
|
||||
}}
|
||||
>
|
||||
See all results on Hub
|
||||
</button>
|
||||
{/if}
|
||||
{/if}
|
||||
{/if}
|
||||
</div>
|
||||
</div>
|
||||
@@ -875,6 +1099,11 @@
|
||||
>Downloaded</span
|
||||
>
|
||||
{/if}
|
||||
{#if filters.readyOnly}
|
||||
<span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
|
||||
>Ready</span
|
||||
>
|
||||
{/if}
|
||||
{#if filters.sizeRange}
|
||||
<span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
|
||||
{filters.sizeRange.min}GB - {filters.sizeRange.max}GB
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
<script lang="ts">
|
||||
import { toasts, dismissToast, type Toast } from "$lib/stores/toast.svelte";
|
||||
import { fly, fade } from "svelte/transition";
|
||||
import { flip } from "svelte/animate";
|
||||
|
||||
const items = $derived(toasts());
|
||||
|
||||
const typeStyles: Record<
|
||||
Toast["type"],
|
||||
{ border: string; icon: string; iconColor: string }
|
||||
> = {
|
||||
success: {
|
||||
border: "border-l-green-500",
|
||||
icon: "M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z",
|
||||
iconColor: "text-green-400",
|
||||
},
|
||||
error: {
|
||||
border: "border-l-red-500",
|
||||
icon: "M12 9v3.75m9-.75a9 9 0 11-18 0 9 9 0 0118 0zm-9 3.75h.008v.008H12v-.008z",
|
||||
iconColor: "text-red-400",
|
||||
},
|
||||
warning: {
|
||||
border: "border-l-yellow-500",
|
||||
icon: "M12 9v3.75m-9.303 3.376c-.866 1.5.217 3.374 1.948 3.374h14.71c1.73 0 2.813-1.874 1.948-3.374L13.949 3.378c-.866-1.5-3.032-1.5-3.898 0L2.697 16.126z",
|
||||
iconColor: "text-yellow-400",
|
||||
},
|
||||
info: {
|
||||
border: "border-l-blue-500",
|
||||
icon: "M11.25 11.25l.041-.02a.75.75 0 011.063.852l-.708 2.836a.75.75 0 001.063.853l.041-.021M21 12a9 9 0 11-18 0 9 9 0 0118 0zm-9-3.75h.008v.008H12V8.25z",
|
||||
iconColor: "text-blue-400",
|
||||
},
|
||||
};
|
||||
</script>
|
||||
|
||||
{#if items.length > 0}
|
||||
<div
|
||||
class="fixed bottom-6 right-6 z-[9999] flex flex-col gap-2 pointer-events-none"
|
||||
role="log"
|
||||
aria-live="polite"
|
||||
aria-label="Notifications"
|
||||
>
|
||||
{#each items as toast (toast.id)}
|
||||
{@const style = typeStyles[toast.type]}
|
||||
<div
|
||||
class="pointer-events-auto max-w-sm w-80 bg-exo-dark-gray/95 backdrop-blur-sm border border-exo-medium-gray/60 border-l-[3px] {style.border} rounded shadow-lg shadow-black/40"
|
||||
in:fly={{ x: 80, duration: 250 }}
|
||||
out:fade={{ duration: 150 }}
|
||||
animate:flip={{ duration: 200 }}
|
||||
role="alert"
|
||||
>
|
||||
<div class="flex items-start gap-3 px-4 py-3">
|
||||
<!-- Icon -->
|
||||
<svg
|
||||
class="w-5 h-5 flex-shrink-0 mt-0.5 {style.iconColor}"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke-width="1.5"
|
||||
stroke="currentColor"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d={style.icon}
|
||||
/>
|
||||
</svg>
|
||||
|
||||
<!-- Message -->
|
||||
<p class="flex-1 text-sm text-white/90 font-mono leading-snug">
|
||||
{toast.message}
|
||||
</p>
|
||||
|
||||
<!-- Dismiss button -->
|
||||
<button
|
||||
onclick={() => dismissToast(toast.id)}
|
||||
class="flex-shrink-0 p-0.5 text-white/40 hover:text-white/80 transition-colors cursor-pointer"
|
||||
aria-label="Dismiss notification"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4"
|
||||
fill="none"
|
||||
viewBox="0 0 24 24"
|
||||
stroke-width="2"
|
||||
stroke="currentColor"
|
||||
>
|
||||
<path
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
d="M6 18L18 6M6 6l12 12"
|
||||
/>
|
||||
</svg>
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<!-- Auto-dismiss progress bar -->
|
||||
{#if toast.duration > 0}
|
||||
<div class="h-0.5 bg-white/5 rounded-b overflow-hidden">
|
||||
<div
|
||||
class="h-full {style.border.replace('border-l-', 'bg-')}/60"
|
||||
style="animation: shrink {toast.duration}ms linear forwards"
|
||||
></div>
|
||||
</div>
|
||||
{/if}
|
||||
</div>
|
||||
{/each}
|
||||
</div>
|
||||
{/if}
|
||||
|
||||
<style>
|
||||
@keyframes shrink {
|
||||
from {
|
||||
width: 100%;
|
||||
}
|
||||
to {
|
||||
width: 0%;
|
||||
}
|
||||
}
|
||||
</style>
|
||||
@@ -12,3 +12,4 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
|
||||
export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
|
||||
export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
|
||||
export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
|
||||
export { default as ChatModelSelector } from "./ChatModelSelector.svelte";
|
||||
|
||||
@@ -168,7 +168,7 @@ export interface ModelDownloadStatus {
|
||||
export interface PlacementPreview {
|
||||
model_id: string;
|
||||
sharding: "Pipeline" | "Tensor";
|
||||
instance_meta: "MlxRing" | "MlxIbv" | "MlxJaccl";
|
||||
instance_meta: "MlxRing" | "MlxJaccl";
|
||||
instance: unknown | null;
|
||||
memory_delta_by_node: Record<string, number> | null;
|
||||
error: string | null;
|
||||
@@ -219,7 +219,6 @@ interface RawStateResponse {
|
||||
string,
|
||||
{
|
||||
MlxRingInstance?: Instance;
|
||||
MlxIbvInstance?: Instance;
|
||||
MlxJacclInstance?: Instance;
|
||||
}
|
||||
>;
|
||||
@@ -581,6 +580,8 @@ class AppStore {
|
||||
debugMode = $state(false);
|
||||
topologyOnlyMode = $state(false);
|
||||
chatSidebarVisible = $state(true); // Shown by default
|
||||
mobileChatSidebarOpen = $state(false); // Mobile drawer state
|
||||
mobileRightSidebarOpen = $state(false); // Mobile right drawer state
|
||||
|
||||
// Image generation params
|
||||
imageGenerationParams = $state<ImageGenerationParams>({
|
||||
@@ -590,6 +591,12 @@ class AppStore {
|
||||
// Image editing state
|
||||
editingImage = $state<EditingImage | null>(null);
|
||||
|
||||
/** True when the backend is reachable. */
|
||||
isConnected = $state<boolean>(true);
|
||||
/** Number of consecutive fetch failures. */
|
||||
private consecutiveFailures = 0;
|
||||
private static readonly CONNECTION_LOST_THRESHOLD = 3;
|
||||
|
||||
private fetchInterval: ReturnType<typeof setInterval> | null = null;
|
||||
private previewsInterval: ReturnType<typeof setInterval> | null = null;
|
||||
private lastConversationPersistTs = 0;
|
||||
@@ -842,6 +849,10 @@ class AppStore {
|
||||
this.thinkingEnabled = conversation.enableThinking ?? true;
|
||||
this.refreshConversationModelFromInstances();
|
||||
|
||||
// Sync global selection to the loaded conversation's model so reactive
|
||||
// effects in +page.svelte can determine the correct chat launch state.
|
||||
this.selectedChatModel = conversation.modelId || "";
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -912,11 +923,7 @@ class AppStore {
|
||||
|
||||
let instanceType: string | null = null;
|
||||
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
|
||||
else if (
|
||||
instanceTag === "MlxIbvInstance" ||
|
||||
instanceTag === "MlxJacclInstance"
|
||||
)
|
||||
instanceType = "MLX RDMA";
|
||||
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
|
||||
|
||||
let sharding: string | null = null;
|
||||
const inst = instance as {
|
||||
@@ -1240,6 +1247,30 @@ class AppStore {
|
||||
this.saveChatSidebarVisibleToStorage();
|
||||
}
|
||||
|
||||
getMobileChatSidebarOpen(): boolean {
|
||||
return this.mobileChatSidebarOpen;
|
||||
}
|
||||
|
||||
setMobileChatSidebarOpen(open: boolean) {
|
||||
this.mobileChatSidebarOpen = open;
|
||||
}
|
||||
|
||||
toggleMobileChatSidebar() {
|
||||
this.mobileChatSidebarOpen = !this.mobileChatSidebarOpen;
|
||||
}
|
||||
|
||||
getMobileRightSidebarOpen(): boolean {
|
||||
return this.mobileRightSidebarOpen;
|
||||
}
|
||||
|
||||
setMobileRightSidebarOpen(open: boolean) {
|
||||
this.mobileRightSidebarOpen = open;
|
||||
}
|
||||
|
||||
toggleMobileRightSidebar() {
|
||||
this.mobileRightSidebarOpen = !this.mobileRightSidebarOpen;
|
||||
}
|
||||
|
||||
startPolling() {
|
||||
this.fetchState();
|
||||
this.fetchInterval = setInterval(() => this.fetchState(), 1000);
|
||||
@@ -1295,7 +1326,19 @@ class AppStore {
|
||||
// Thunderbolt bridge status per node
|
||||
this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
|
||||
this.lastUpdate = Date.now();
|
||||
// Connection recovered
|
||||
if (!this.isConnected) {
|
||||
this.isConnected = true;
|
||||
}
|
||||
this.consecutiveFailures = 0;
|
||||
} catch (error) {
|
||||
this.consecutiveFailures++;
|
||||
if (
|
||||
this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD &&
|
||||
this.isConnected
|
||||
) {
|
||||
this.isConnected = false;
|
||||
}
|
||||
console.error("Error fetching state:", error);
|
||||
}
|
||||
}
|
||||
@@ -1836,7 +1879,7 @@ class AppStore {
|
||||
assistantMessage.id,
|
||||
(msg) => {
|
||||
msg.content =
|
||||
"Error: No model available. Please launch an instance first.";
|
||||
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.";
|
||||
},
|
||||
);
|
||||
this.syncActiveMessagesIfNeeded(targetConversationId);
|
||||
@@ -2307,7 +2350,7 @@ class AppStore {
|
||||
const modelToUse = this.getModelForRequest();
|
||||
if (!modelToUse) {
|
||||
throw new Error(
|
||||
"No model selected and no running instances available. Please launch an instance first.",
|
||||
"No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.",
|
||||
);
|
||||
}
|
||||
|
||||
@@ -3141,6 +3184,23 @@ class AppStore {
|
||||
return (await response.json()) as TraceStatsResponse;
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete traces by task IDs
|
||||
*/
|
||||
async deleteTraces(
|
||||
taskIds: string[],
|
||||
): Promise<{ deleted: string[]; notFound: string[] }> {
|
||||
const response = await fetch("/v1/traces/delete", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ taskIds }),
|
||||
});
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to delete traces: ${response.status}`);
|
||||
}
|
||||
return await response.json();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the URL for the raw trace file (for Perfetto)
|
||||
*/
|
||||
@@ -3248,8 +3308,23 @@ export const toggleChatSidebarVisible = () =>
|
||||
appStore.toggleChatSidebarVisible();
|
||||
export const setChatSidebarVisible = (visible: boolean) =>
|
||||
appStore.setChatSidebarVisible(visible);
|
||||
|
||||
// Mobile sidebar state
|
||||
export const mobileChatSidebarOpen = () => appStore.mobileChatSidebarOpen;
|
||||
export const toggleMobileChatSidebar = () => appStore.toggleMobileChatSidebar();
|
||||
export const setMobileChatSidebarOpen = (open: boolean) =>
|
||||
appStore.setMobileChatSidebarOpen(open);
|
||||
export const mobileRightSidebarOpen = () => appStore.mobileRightSidebarOpen;
|
||||
export const toggleMobileRightSidebar = () =>
|
||||
appStore.toggleMobileRightSidebar();
|
||||
export const setMobileRightSidebarOpen = (open: boolean) =>
|
||||
appStore.setMobileRightSidebarOpen(open);
|
||||
|
||||
export const refreshState = () => appStore.fetchState();
|
||||
|
||||
// Connection status
|
||||
export const isConnected = () => appStore.isConnected;
|
||||
|
||||
// Node identities (for OS version mismatch detection)
|
||||
export const nodeIdentities = () => appStore.nodeIdentities;
|
||||
|
||||
@@ -3281,3 +3356,5 @@ export const fetchTraceStats = (taskId: string) =>
|
||||
appStore.fetchTraceStats(taskId);
|
||||
export const getTraceRawUrl = (taskId: string) =>
|
||||
appStore.getTraceRawUrl(taskId);
|
||||
export const deleteTraces = (taskIds: string[]) =>
|
||||
appStore.deleteTraces(taskIds);
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Toast notification store - Global notification system for the EXO dashboard.
|
||||
*
|
||||
* Usage:
|
||||
* import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte";
|
||||
* addToast({ type: "success", message: "Model launched" });
|
||||
* addToast({ type: "error", message: "Connection lost", persistent: true });
|
||||
*/
|
||||
|
||||
type ToastType = "success" | "error" | "warning" | "info";
|
||||
|
||||
export interface Toast {
|
||||
id: string;
|
||||
type: ToastType;
|
||||
message: string;
|
||||
/** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */
|
||||
duration: number;
|
||||
createdAt: number;
|
||||
}
|
||||
|
||||
interface ToastInput {
|
||||
type: ToastType;
|
||||
message: string;
|
||||
/** If true, toast stays until manually dismissed. Default: false. */
|
||||
persistent?: boolean;
|
||||
/** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */
|
||||
duration?: number;
|
||||
}
|
||||
|
||||
const DEFAULT_DURATIONS: Record<ToastType, number> = {
|
||||
success: 4000,
|
||||
info: 4000,
|
||||
warning: 6000,
|
||||
error: 6000,
|
||||
};
|
||||
|
||||
let toastList = $state<Toast[]>([]);
|
||||
const timers = new Map<string, ReturnType<typeof setTimeout>>();
|
||||
|
||||
function generateId(): string {
|
||||
return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
|
||||
}
|
||||
|
||||
export function addToast(input: ToastInput): string {
|
||||
const id = generateId();
|
||||
const duration = input.persistent
|
||||
? 0
|
||||
: (input.duration ?? DEFAULT_DURATIONS[input.type]);
|
||||
|
||||
const toast: Toast = {
|
||||
id,
|
||||
type: input.type,
|
||||
message: input.message,
|
||||
duration,
|
||||
createdAt: Date.now(),
|
||||
};
|
||||
|
||||
toastList = [...toastList, toast];
|
||||
|
||||
if (duration > 0) {
|
||||
const timer = setTimeout(() => dismissToast(id), duration);
|
||||
timers.set(id, timer);
|
||||
}
|
||||
|
||||
return id;
|
||||
}
|
||||
|
||||
export function dismissToast(id: string): void {
|
||||
const timer = timers.get(id);
|
||||
if (timer) {
|
||||
clearTimeout(timer);
|
||||
timers.delete(id);
|
||||
}
|
||||
toastList = toastList.filter((t) => t.id !== id);
|
||||
}
|
||||
|
||||
/** Dismiss all toasts matching a message (useful for dedup). */
|
||||
export function dismissByMessage(message: string): void {
|
||||
const matching = toastList.filter((t) => t.message === message);
|
||||
for (const t of matching) {
|
||||
dismissToast(t.id);
|
||||
}
|
||||
}
|
||||
|
||||
export function toasts(): Toast[] {
|
||||
return toastList;
|
||||
}
|
||||
@@ -1,5 +1,7 @@
|
||||
<script lang="ts">
|
||||
import "../app.css";
|
||||
import ToastContainer from "$lib/components/ToastContainer.svelte";
|
||||
import ConnectionBanner from "$lib/components/ConnectionBanner.svelte";
|
||||
|
||||
let { children } = $props();
|
||||
</script>
|
||||
@@ -10,5 +12,7 @@
|
||||
</svelte:head>
|
||||
|
||||
<div class="min-h-screen bg-background text-foreground">
|
||||
<ConnectionBanner />
|
||||
{@render children?.()}
|
||||
<ToastContainer />
|
||||
</div>
|
||||
|
||||
+3006
-310
File diff suppressed because it is too large
Load Diff
@@ -412,7 +412,7 @@
|
||||
<div>{col.label}</div>
|
||||
{#if col.diskAvailable != null}
|
||||
<div
|
||||
class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
|
||||
class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
|
||||
>
|
||||
{formatBytes(col.diskAvailable)} free
|
||||
</div>
|
||||
@@ -436,7 +436,7 @@
|
||||
</div>
|
||||
{#if row.prettyName}
|
||||
<div
|
||||
class="text-[10px] text-exo-light-gray/60"
|
||||
class="text-[10px] text-white/60"
|
||||
title={row.modelId}
|
||||
>
|
||||
{row.modelId}
|
||||
@@ -450,7 +450,7 @@
|
||||
title="View model details"
|
||||
>
|
||||
<svg
|
||||
class="w-4 h-4 text-white/30 hover:text-white/60"
|
||||
class="w-4 h-4 text-white/60 hover:text-white/80"
|
||||
viewBox="0 0 24 24"
|
||||
fill="currentColor"
|
||||
>
|
||||
@@ -469,11 +469,11 @@
|
||||
<td class="px-4 py-3 text-center align-middle">
|
||||
{#if cell.kind === "completed"}
|
||||
<div
|
||||
class="flex flex-col items-center gap-0.5"
|
||||
class="flex flex-col items-center gap-1"
|
||||
title="Completed ({formatBytes(cell.totalBytes)})"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5 text-green-400"
|
||||
class="w-7 h-7 text-green-400"
|
||||
viewBox="0 0 20 20"
|
||||
fill="currentColor"
|
||||
>
|
||||
@@ -483,18 +483,18 @@
|
||||
clip-rule="evenodd"
|
||||
></path>
|
||||
</svg>
|
||||
<span class="text-[10px] text-exo-light-gray/70"
|
||||
<span class="text-xs text-white/70"
|
||||
>{formatBytes(cell.totalBytes)}</span
|
||||
>
|
||||
<button
|
||||
type="button"
|
||||
class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
|
||||
class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
|
||||
onclick={() =>
|
||||
deleteDownload(col.nodeId, row.modelId)}
|
||||
title="Delete from this node"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
@@ -517,11 +517,11 @@
|
||||
cell.speed,
|
||||
)} - ETA {formatEta(cell.etaMs)}"
|
||||
>
|
||||
<span class="text-exo-yellow text-xs font-medium"
|
||||
<span class="text-exo-yellow text-sm font-medium"
|
||||
>{clampPercent(cell.percentage).toFixed(1)}%</span
|
||||
>
|
||||
<div
|
||||
class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
|
||||
class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
|
||||
>
|
||||
<div
|
||||
class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
|
||||
@@ -530,25 +530,25 @@
|
||||
).toFixed(1)}%"
|
||||
></div>
|
||||
</div>
|
||||
<span class="text-[9px] text-exo-light-gray/60"
|
||||
<span class="text-[10px] text-white/70"
|
||||
>{formatSpeed(cell.speed)}</span
|
||||
>
|
||||
</div>
|
||||
{:else if cell.kind === "pending"}
|
||||
<div
|
||||
class="flex flex-col items-center gap-0.5"
|
||||
class="flex flex-col items-center gap-1"
|
||||
title={cell.downloaded > 0
|
||||
? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded`
|
||||
? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
|
||||
: "Download pending"}
|
||||
>
|
||||
{#if cell.downloaded > 0 && cell.total > 0}
|
||||
<span class="text-exo-light-gray/70 text-[10px]"
|
||||
<span class="text-white/70 text-xs"
|
||||
>{formatBytes(cell.downloaded)} / {formatBytes(
|
||||
cell.total,
|
||||
)}</span
|
||||
>
|
||||
<div
|
||||
class="w-full h-1 bg-white/10 rounded-full overflow-hidden"
|
||||
class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
|
||||
>
|
||||
<div
|
||||
class="h-full bg-exo-light-gray/40 rounded-full"
|
||||
@@ -558,21 +558,65 @@
|
||||
).toFixed(1)}%"
|
||||
></div>
|
||||
</div>
|
||||
<span class="text-exo-light-gray/40 text-[9px]"
|
||||
>paused</span
|
||||
{#if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Resume download on this node"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path
|
||||
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
></path>
|
||||
</svg>
|
||||
</button>
|
||||
{:else}
|
||||
<span class="text-white/50 text-[10px]">paused</span
|
||||
>
|
||||
{/if}
|
||||
{:else if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Start download on this node"
|
||||
>
|
||||
<svg
|
||||
class="w-6 h-6"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
stroke-width="2"
|
||||
>
|
||||
<path
|
||||
d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
|
||||
stroke-linecap="round"
|
||||
stroke-linejoin="round"
|
||||
></path>
|
||||
</svg>
|
||||
</button>
|
||||
{:else}
|
||||
<span class="text-exo-light-gray/50 text-sm">...</span
|
||||
>
|
||||
<span class="text-white/40 text-sm">...</span>
|
||||
{/if}
|
||||
</div>
|
||||
{:else if cell.kind === "failed"}
|
||||
<div
|
||||
class="flex flex-col items-center gap-0.5"
|
||||
class="flex flex-col items-center gap-1"
|
||||
title="Download failed"
|
||||
>
|
||||
<svg
|
||||
class="w-5 h-5 text-red-400"
|
||||
class="w-7 h-7 text-red-400"
|
||||
viewBox="0 0 20 20"
|
||||
fill="currentColor"
|
||||
>
|
||||
@@ -585,13 +629,13 @@
|
||||
{#if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Retry download on this node"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
@@ -617,13 +661,13 @@
|
||||
{#if row.shardMetadata}
|
||||
<button
|
||||
type="button"
|
||||
class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
|
||||
class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
|
||||
onclick={() =>
|
||||
startDownload(col.nodeId, row.shardMetadata!)}
|
||||
title="Download to this node"
|
||||
>
|
||||
<svg
|
||||
class="w-3.5 h-3.5"
|
||||
class="w-5 h-5"
|
||||
viewBox="0 0 20 20"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import {
|
||||
listTraces,
|
||||
getTraceRawUrl,
|
||||
deleteTraces,
|
||||
type TraceListItem,
|
||||
} from "$lib/stores/app.svelte";
|
||||
import HeaderNav from "$lib/components/HeaderNav.svelte";
|
||||
@@ -10,6 +11,51 @@
|
||||
let traces = $state<TraceListItem[]>([]);
|
||||
let loading = $state(true);
|
||||
let error = $state<string | null>(null);
|
||||
let selectedIds = $state<Set<string>>(new Set());
|
||||
let deleting = $state(false);
|
||||
|
||||
let allSelected = $derived(
|
||||
traces.length > 0 && selectedIds.size === traces.length,
|
||||
);
|
||||
|
||||
function toggleSelect(taskId: string) {
|
||||
const next = new Set(selectedIds);
|
||||
if (next.has(taskId)) {
|
||||
next.delete(taskId);
|
||||
} else {
|
||||
next.add(taskId);
|
||||
}
|
||||
selectedIds = next;
|
||||
}
|
||||
|
||||
function toggleSelectAll() {
|
||||
if (allSelected) {
|
||||
selectedIds = new Set();
|
||||
} else {
|
||||
selectedIds = new Set(traces.map((t) => t.taskId));
|
||||
}
|
||||
}
|
||||
|
||||
async function handleDelete() {
|
||||
if (selectedIds.size === 0) return;
|
||||
const count = selectedIds.size;
|
||||
if (
|
||||
!confirm(
|
||||
`Delete ${count} trace${count === 1 ? "" : "s"}? This cannot be undone.`,
|
||||
)
|
||||
)
|
||||
return;
|
||||
deleting = true;
|
||||
try {
|
||||
await deleteTraces([...selectedIds]);
|
||||
selectedIds = new Set();
|
||||
await refresh();
|
||||
} catch (e) {
|
||||
error = e instanceof Error ? e.message : "Failed to delete traces";
|
||||
} finally {
|
||||
deleting = false;
|
||||
}
|
||||
}
|
||||
|
||||
function formatBytes(bytes: number): string {
|
||||
if (!bytes || bytes <= 0) return "0B";
|
||||
@@ -109,6 +155,16 @@
|
||||
</h1>
|
||||
</div>
|
||||
<div class="flex items-center gap-3">
|
||||
{#if selectedIds.size > 0}
|
||||
<button
|
||||
type="button"
|
||||
class="text-xs font-mono text-red-400 hover:text-red-300 transition-colors uppercase border border-red-500/40 px-2 py-1 rounded"
|
||||
onclick={handleDelete}
|
||||
disabled={deleting}
|
||||
>
|
||||
{deleting ? "Deleting..." : `Delete (${selectedIds.size})`}
|
||||
</button>
|
||||
{/if}
|
||||
<button
|
||||
type="button"
|
||||
class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
|
||||
@@ -143,14 +199,41 @@
|
||||
</div>
|
||||
{:else}
|
||||
<div class="space-y-3">
|
||||
<div class="flex items-center gap-2 px-1">
|
||||
<button
|
||||
type="button"
|
||||
class="text-xs font-mono uppercase transition-colors {allSelected
|
||||
? 'text-exo-yellow'
|
||||
: 'text-exo-light-gray hover:text-exo-yellow'}"
|
||||
onclick={toggleSelectAll}
|
||||
>
|
||||
{allSelected ? "Deselect all" : "Select all"}
|
||||
</button>
|
||||
</div>
|
||||
{#each traces as trace}
|
||||
{@const isSelected = selectedIds.has(trace.taskId)}
|
||||
<!-- svelte-ignore a11y_no_static_element_interactions -->
|
||||
<div
|
||||
class="rounded border border-exo-medium-gray/30 bg-exo-black/30 p-4 flex items-center justify-between gap-4"
|
||||
role="button"
|
||||
tabindex="0"
|
||||
class="w-full text-left rounded border-l-2 border-r border-t border-b transition-all p-4 flex items-center justify-between gap-4 cursor-pointer {isSelected
|
||||
? 'bg-exo-yellow/10 border-l-exo-yellow border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30'
|
||||
: 'bg-exo-black/30 border-l-transparent border-r-exo-medium-gray/30 border-t-exo-medium-gray/30 border-b-exo-medium-gray/30 hover:bg-white/[0.03]'}"
|
||||
onclick={() => toggleSelect(trace.taskId)}
|
||||
onkeydown={(e) => {
|
||||
if (e.key === "Enter" || e.key === " ") {
|
||||
e.preventDefault();
|
||||
toggleSelect(trace.taskId);
|
||||
}
|
||||
}}
|
||||
>
|
||||
<div class="min-w-0 flex-1">
|
||||
<a
|
||||
href="#/traces/{trace.taskId}"
|
||||
class="text-sm font-mono text-white hover:text-exo-yellow transition-colors truncate block"
|
||||
class="text-sm font-mono transition-colors truncate block {isSelected
|
||||
? 'text-exo-yellow'
|
||||
: 'text-white hover:text-exo-yellow'}"
|
||||
onclick={(e) => e.stopPropagation()}
|
||||
>
|
||||
{trace.taskId}
|
||||
</a>
|
||||
@@ -160,7 +243,11 @@
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex items-center gap-2 shrink-0">
|
||||
<!-- svelte-ignore a11y_click_events_have_key_events -->
|
||||
<div
|
||||
class="flex items-center gap-2 shrink-0"
|
||||
onclick={(e) => e.stopPropagation()}
|
||||
>
|
||||
<a
|
||||
href="#/traces/{trace.taskId}"
|
||||
class="text-xs font-mono text-exo-light-gray hover:text-exo-yellow transition-colors uppercase border border-exo-medium-gray/40 px-2 py-1 rounded"
|
||||
|
||||
+475
-23
@@ -4,7 +4,7 @@ This document describes the REST API exposed by the **EXO** service, as implemen
|
||||
|
||||
`src/exo/master/api.py`
|
||||
|
||||
The API is used to manage model instances in the cluster, inspect cluster state, and perform inference using an OpenAI-compatible interface.
|
||||
The API is used to manage model instances in the cluster, inspect cluster state, and perform inference using multiple API-compatible interfaces.
|
||||
|
||||
Base URL example:
|
||||
|
||||
@@ -144,8 +144,59 @@ Placement result.
|
||||
|
||||
Returns the list of available models and their metadata.
|
||||
|
||||
**Query parameters:**
|
||||
|
||||
* `status`: string (optional) - Filter by `downloaded` to show only downloaded models
|
||||
|
||||
**Response:**
|
||||
Array of model descriptors.
|
||||
Array of model descriptors including `is_custom` field for custom HuggingFace models.
|
||||
|
||||
### Add Custom Model
|
||||
|
||||
**POST** `/models/add`
|
||||
|
||||
Add a custom model from HuggingFace hub.
|
||||
|
||||
**Request body (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"model_id": "mlx-community/my-custom-model"
|
||||
}
|
||||
```
|
||||
|
||||
**Response:**
|
||||
Model descriptor for the added model.
|
||||
|
||||
**Security note:**
|
||||
Models with `trust_remote_code` enabled in their configuration require explicit opt-in (default is false) for security.
|
||||
|
||||
### Delete Custom Model
|
||||
|
||||
**DELETE** `/models/custom/{model_id}`
|
||||
|
||||
Delete a user-added custom model card.
|
||||
|
||||
**Path parameters:**
|
||||
|
||||
* `model_id`: string, ID of the custom model to delete
|
||||
|
||||
**Response:**
|
||||
Confirmation JSON with deleted model ID.
|
||||
|
||||
### Search Models
|
||||
|
||||
**GET** `/models/search`
|
||||
|
||||
Search HuggingFace Hub for mlx-community models.
|
||||
|
||||
**Query parameters:**
|
||||
|
||||
* `query`: string (optional) - Search query
|
||||
* `limit`: integer (default: 20) - Maximum number of results
|
||||
|
||||
**Response:**
|
||||
Array of HuggingFace model search results.
|
||||
|
||||
## 4. Inference / Chat Completions
|
||||
|
||||
@@ -168,9 +219,123 @@ Executes a chat completion request using an OpenAI-compatible schema. Supports s
|
||||
}
|
||||
```
|
||||
|
||||
**Request parameters:**
|
||||
|
||||
* `model`: string, required - Model ID to use
|
||||
* `messages`: array, required - Conversation messages
|
||||
* `stream`: boolean (default: false) - Enable streaming responses
|
||||
* `max_tokens`: integer (optional) - Maximum tokens to generate
|
||||
* `temperature`: float (optional) - Sampling temperature
|
||||
* `top_p`: float (optional) - Nucleus sampling parameter
|
||||
* `top_k`: integer (optional) - Top-k sampling parameter
|
||||
* `stop`: string or array (optional) - Stop sequences
|
||||
* `seed`: integer (optional) - Random seed for reproducibility
|
||||
* `enable_thinking`: boolean (optional) - Enable thinking mode for capable models (DeepSeek V3.1, Qwen3, GLM-4.7)
|
||||
* `tools`: array (optional) - Tool definitions for function calling
|
||||
* `logprobs`: boolean (optional) - Return log probabilities
|
||||
* `top_logprobs`: integer (optional) - Number of top log probabilities to return
|
||||
|
||||
**Response:**
|
||||
OpenAI-compatible chat completion response.
|
||||
|
||||
**Streaming response format:**
|
||||
When `stream=true`, returns Server-Sent Events (SSE) with format:
|
||||
|
||||
```
|
||||
data: {"id":"...","object":"chat.completion","created":...,"model":"...","choices":[...]}
|
||||
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
**Non-streaming response includes usage statistics:**
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "...",
|
||||
"object": "chat.completion",
|
||||
"created": 1234567890,
|
||||
"model": "llama-3.2-1b",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I help you?"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 15,
|
||||
"completion_tokens": 8,
|
||||
"total_tokens": 23
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Cancellation:**
|
||||
You can cancel an active generation by closing the HTTP connection. The server detects the disconnection and stops processing.
|
||||
|
||||
### Claude Messages API
|
||||
|
||||
**POST** `/v1/messages`
|
||||
|
||||
Executes a chat completion request using the Claude Messages API format. Supports streaming and non-streaming modes.
|
||||
|
||||
**Request body (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "llama-3.2-1b",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "Hello" }
|
||||
],
|
||||
"max_tokens": 1024,
|
||||
"stream": false
|
||||
}
|
||||
```
|
||||
|
||||
**Streaming response format:**
|
||||
When `stream=true`, returns Server-Sent Events with Claude-specific event types:
|
||||
|
||||
* `message_start` - Message generation started
|
||||
* `content_block_start` - Content block started
|
||||
* `content_block_delta` - Incremental content chunk
|
||||
* `content_block_stop` - Content block completed
|
||||
* `message_delta` - Message metadata updates
|
||||
* `message_stop` - Message generation completed
|
||||
|
||||
**Response:**
|
||||
Claude-compatible messages response.
|
||||
|
||||
### OpenAI Responses API
|
||||
|
||||
**POST** `/v1/responses`
|
||||
|
||||
Executes a chat completion request using the OpenAI Responses API format. Supports streaming and non-streaming modes.
|
||||
|
||||
**Request body (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "llama-3.2-1b",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "Hello" }
|
||||
],
|
||||
"stream": false
|
||||
}
|
||||
```
|
||||
|
||||
**Streaming response format:**
|
||||
When `stream=true`, returns Server-Sent Events with response-specific event types:
|
||||
|
||||
* `response.created` - Response generation started
|
||||
* `response.in_progress` - Response is being generated
|
||||
* `response.output_item.added` - New output item added
|
||||
* `response.output_item.done` - Output item completed
|
||||
* `response.done` - Response generation completed
|
||||
|
||||
**Response:**
|
||||
OpenAI Responses API-compatible response.
|
||||
|
||||
### Benchmarked Chat Completions
|
||||
|
||||
**POST** `/bench/chat/completions`
|
||||
@@ -181,26 +346,177 @@ Same as `/v1/chat/completions`, but also returns performance and generation stat
|
||||
Same schema as `/v1/chat/completions`.
|
||||
|
||||
**Response:**
|
||||
Chat completion plus benchmarking metrics.
|
||||
Chat completion plus benchmarking metrics including:
|
||||
|
||||
## 5. Image Generation & Editing
|
||||
* `prompt_tps` - Tokens per second during prompt processing
|
||||
* `generation_tps` - Tokens per second during generation
|
||||
* `prompt_tokens` - Number of prompt tokens
|
||||
* `generation_tokens` - Number of generated tokens
|
||||
* `peak_memory_usage` - Peak memory used during generation
|
||||
|
||||
### Cancel Command
|
||||
|
||||
**POST** `/v1/cancel/{command_id}`
|
||||
|
||||
Cancels an active generation command (text or image). Notifies workers and closes the stream.
|
||||
|
||||
**Path parameters:**
|
||||
|
||||
* `command_id`: string, ID of the command to cancel
|
||||
|
||||
**Response (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"message": "Command cancelled.",
|
||||
"command_id": "cmd-abc-123"
|
||||
}
|
||||
```
|
||||
|
||||
Returns 404 if the command is not found or already completed.
|
||||
|
||||
## 5. Ollama API Compatibility
|
||||
|
||||
EXO provides Ollama API compatibility for tools like OpenWebUI.
|
||||
|
||||
### Ollama Chat
|
||||
|
||||
**POST** `/ollama/api/chat`
|
||||
**POST** `/ollama/api/api/chat` (alias)
|
||||
**POST** `/ollama/api/v1/chat` (alias)
|
||||
|
||||
Execute a chat request using Ollama API format.
|
||||
|
||||
**Request body (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "llama-3.2-1b",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "Hello" }
|
||||
],
|
||||
"stream": false
|
||||
}
|
||||
```
|
||||
|
||||
**Response:**
|
||||
Ollama-compatible chat response.
|
||||
|
||||
### Ollama Generate
|
||||
|
||||
**POST** `/ollama/api/generate`
|
||||
|
||||
Execute a text generation request using Ollama API format.
|
||||
|
||||
**Request body (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "llama-3.2-1b",
|
||||
"prompt": "Hello",
|
||||
"stream": false
|
||||
}
|
||||
```
|
||||
|
||||
**Response:**
|
||||
Ollama-compatible generation response.
|
||||
|
||||
### Ollama Tags
|
||||
|
||||
**GET** `/ollama/api/tags`
|
||||
**GET** `/ollama/api/api/tags` (alias)
|
||||
**GET** `/ollama/api/v1/tags` (alias)
|
||||
|
||||
Returns list of downloaded models in Ollama tags format.
|
||||
|
||||
**Response:**
|
||||
Array of model tags with metadata.
|
||||
|
||||
### Ollama Show
|
||||
|
||||
**POST** `/ollama/api/show`
|
||||
|
||||
Returns model information in Ollama show format.
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "llama-3.2-1b"
|
||||
}
|
||||
```
|
||||
|
||||
**Response:**
|
||||
Model details including modelfile and family.
|
||||
|
||||
### Ollama PS
|
||||
|
||||
**GET** `/ollama/api/ps`
|
||||
|
||||
Returns list of running models (active instances).
|
||||
|
||||
**Response:**
|
||||
Array of active model instances.
|
||||
|
||||
### Ollama Version
|
||||
|
||||
**GET** `/ollama/api/version`
|
||||
**HEAD** `/ollama/` (alias)
|
||||
**HEAD** `/ollama/api/version` (alias)
|
||||
|
||||
Returns version information for Ollama API compatibility.
|
||||
|
||||
**Response:**
|
||||
|
||||
```json
|
||||
{
|
||||
"version": "exo v1.0"
|
||||
}
|
||||
```
|
||||
|
||||
## 6. Image Generation & Editing
|
||||
|
||||
### Image Generation
|
||||
|
||||
**POST** `/v1/images/generations`
|
||||
|
||||
Executes an image generation request using an OpenAI-compatible schema with additional advanced_params.
|
||||
Executes an image generation request using an OpenAI-compatible schema with additional advanced_params. Supports both streaming and non-streaming modes.
|
||||
|
||||
**Request body (example):**
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": "a robot playing chess",
|
||||
"model": "flux-dev",
|
||||
"model": "exolabs/FLUX.1-dev",
|
||||
"n": 1,
|
||||
"size": "1024x1024",
|
||||
"stream": false,
|
||||
"response_format": "b64_json"
|
||||
}
|
||||
```
|
||||
|
||||
**Request parameters:**
|
||||
|
||||
* `prompt`: string, required - Text description of the image
|
||||
* `model`: string, required - Image model ID
|
||||
* `n`: integer (default: 1) - Number of images to generate
|
||||
* `size`: string (default: "auto") - Image dimensions. Supported sizes:
|
||||
- `512x512`
|
||||
- `768x768`
|
||||
- `1024x768`
|
||||
- `768x1024`
|
||||
- `1024x1024`
|
||||
- `1024x1536`
|
||||
- `1536x1024`
|
||||
- `1024x1365`
|
||||
- `1365x1024`
|
||||
* `stream`: boolean (default: false) - Enable streaming for partial images
|
||||
* `partial_images`: integer (default: 0) - Number of partial images to stream during generation
|
||||
* `response_format`: string (default: "b64_json") - Either `url` or `b64_json`
|
||||
* `quality`: string (default: "medium") - Either `high`, `medium`, or `low`
|
||||
* `output_format`: string (default: "png") - Either `png`, `jpeg`, or `webp`
|
||||
* `advanced_params`: object (optional) - Advanced generation parameters
|
||||
|
||||
**Advanced Parameters (`advanced_params`):**
|
||||
|
||||
| Parameter | Type | Constraints | Description |
|
||||
@@ -210,8 +526,54 @@ Executes an image generation request using an OpenAI-compatible schema with addi
|
||||
| `guidance` | float | 1.0-20.0 | Classifier-free guidance scale |
|
||||
| `negative_prompt` | string | - | Text describing what to avoid in the image |
|
||||
|
||||
**Non-streaming response:**
|
||||
|
||||
```json
|
||||
{
|
||||
"created": 1234567890,
|
||||
"data": [
|
||||
{
|
||||
"b64_json": "iVBORw0KGgoAAAANSUhEUgAA...",
|
||||
"url": null
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Streaming response format:**
|
||||
When `stream=true` and `partial_images > 0`, returns Server-Sent Events:
|
||||
|
||||
```
|
||||
data: {"type":"partial","image_index":0,"partial_index":1,"total_partials":5,"format":"png","data":{"b64_json":"..."}}
|
||||
|
||||
data: {"type":"final","image_index":0,"format":"png","data":{"b64_json":"..."}}
|
||||
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
### Image Editing
|
||||
|
||||
**POST** `/v1/images/edits`
|
||||
|
||||
Executes an image editing request (img2img) using FLUX.1-Kontext-dev or similar models.
|
||||
|
||||
**Request (multipart/form-data):**
|
||||
|
||||
* `image`: file, required - Input image to edit
|
||||
* `prompt`: string, required - Text description of desired changes
|
||||
* `model`: string, required - Image editing model ID (e.g., `exolabs/FLUX.1-Kontext-dev`)
|
||||
* `n`: integer (default: 1) - Number of edited images to generate
|
||||
* `size`: string (optional) - Output image dimensions
|
||||
* `response_format`: string (default: "b64_json") - Either `url` or `b64_json`
|
||||
* `input_fidelity`: string (default: "low") - Either `low` or `high` - Controls how closely the output follows the input image
|
||||
* `stream`: string (default: "false") - Enable streaming
|
||||
* `partial_images`: string (default: "0") - Number of partial images to stream
|
||||
* `quality`: string (default: "medium") - Either `high`, `medium`, or `low`
|
||||
* `output_format`: string (default: "png") - Either `png`, `jpeg`, or `webp`
|
||||
* `advanced_params`: string (optional) - JSON-encoded advanced parameters
|
||||
|
||||
**Response:**
|
||||
OpenAI-compatible image generation response.
|
||||
Same format as `/v1/images/generations`.
|
||||
|
||||
### Benchmarked Image Generation
|
||||
|
||||
@@ -223,16 +585,15 @@ Same as `/v1/images/generations`, but also returns generation statistics.
|
||||
Same schema as `/v1/images/generations`.
|
||||
|
||||
**Response:**
|
||||
Image generation plus benchmarking metrics.
|
||||
Image generation plus benchmarking metrics including:
|
||||
|
||||
### Image Editing
|
||||
|
||||
**POST** `/v1/images/edits`
|
||||
|
||||
Executes an image editing request using an OpenAI-compatible schema with additional advanced_params (same as `/v1/images/generations`).
|
||||
|
||||
**Response:**
|
||||
Same format as `/v1/images/generations`.
|
||||
* `seconds_per_step` - Average time per denoising step
|
||||
* `total_generation_time` - Total generation time
|
||||
* `num_inference_steps` - Number of inference steps used
|
||||
* `num_images` - Number of images generated
|
||||
* `image_width` - Output image width
|
||||
* `image_height` - Output image height
|
||||
* `peak_memory_usage` - Peak memory used during generation
|
||||
|
||||
### Benchmarked Image Editing
|
||||
|
||||
@@ -246,36 +607,127 @@ Same schema as `/v1/images/edits`.
|
||||
**Response:**
|
||||
Same format as `/bench/images/generations`, including `generation_stats`.
|
||||
|
||||
## 6. Complete Endpoint Summary
|
||||
### List Images
|
||||
|
||||
**GET** `/images`
|
||||
|
||||
List all stored images.
|
||||
|
||||
**Response:**
|
||||
Array of image metadata including URLs and expiration times.
|
||||
|
||||
### Get Image
|
||||
|
||||
**GET** `/images/{image_id}`
|
||||
|
||||
Retrieve a stored image by ID.
|
||||
|
||||
**Path parameters:**
|
||||
|
||||
* `image_id`: string, ID of the image
|
||||
|
||||
**Response:**
|
||||
Image file with appropriate content type.
|
||||
|
||||
## 7. Complete Endpoint Summary
|
||||
|
||||
```
|
||||
# General
|
||||
GET /node_id
|
||||
GET /state
|
||||
GET /events
|
||||
|
||||
# Instance Management
|
||||
POST /instance
|
||||
GET /instance/{instance_id}
|
||||
DELETE /instance/{instance_id}
|
||||
|
||||
GET /instance/previews
|
||||
GET /instance/placement
|
||||
POST /place_instance
|
||||
|
||||
# Models
|
||||
GET /models
|
||||
GET /v1/models
|
||||
POST /models/add
|
||||
DELETE /models/custom/{model_id}
|
||||
GET /models/search
|
||||
|
||||
# Text Generation (OpenAI Chat Completions)
|
||||
POST /v1/chat/completions
|
||||
POST /bench/chat/completions
|
||||
|
||||
# Text Generation (Claude Messages API)
|
||||
POST /v1/messages
|
||||
|
||||
# Text Generation (OpenAI Responses API)
|
||||
POST /v1/responses
|
||||
|
||||
# Text Generation (Ollama API)
|
||||
POST /ollama/api/chat
|
||||
POST /ollama/api/api/chat
|
||||
POST /ollama/api/v1/chat
|
||||
POST /ollama/api/generate
|
||||
GET /ollama/api/tags
|
||||
GET /ollama/api/api/tags
|
||||
GET /ollama/api/v1/tags
|
||||
POST /ollama/api/show
|
||||
GET /ollama/api/ps
|
||||
GET /ollama/api/version
|
||||
HEAD /ollama/
|
||||
HEAD /ollama/api/version
|
||||
|
||||
# Command Control
|
||||
POST /v1/cancel/{command_id}
|
||||
|
||||
# Image Generation
|
||||
POST /v1/images/generations
|
||||
POST /bench/images/generations
|
||||
POST /v1/images/edits
|
||||
POST /bench/images/edits
|
||||
GET /images
|
||||
GET /images/{image_id}
|
||||
```
|
||||
|
||||
## 7. Notes
|
||||
## 8. Notes
|
||||
|
||||
* The `/v1/chat/completions` endpoint is compatible with the OpenAI Chat API format, so existing OpenAI clients can be pointed to EXO by changing the base URL.
|
||||
* The `/v1/images/generations` and `/v1/images/edits` endpoints are compatible with the OpenAI Images API format.
|
||||
* The instance placement endpoints allow you to plan and preview cluster allocations before actually creating instances.
|
||||
* The `/events` and `/state` endpoints are primarily intended for operational visibility and debugging.
|
||||
### API Compatibility
|
||||
|
||||
EXO provides multiple API-compatible interfaces:
|
||||
|
||||
* **OpenAI Chat Completions API** - Compatible with OpenAI clients and tools
|
||||
* **Claude Messages API** - Compatible with Anthropic's Claude API format
|
||||
* **OpenAI Responses API** - Compatible with OpenAI's Responses API format
|
||||
* **Ollama API** - Compatible with Ollama and tools like OpenWebUI
|
||||
|
||||
Existing OpenAI, Claude, or Ollama clients can be pointed to EXO by changing the base URL.
|
||||
|
||||
### Custom Models
|
||||
|
||||
You can add custom models from HuggingFace using the `/models/add` endpoint. Custom models are identified by the `is_custom` field in model list responses.
|
||||
|
||||
**Security:** Models requiring `trust_remote_code` must be explicitly enabled (default is false) for security. Only enable this if you trust the model's remote code.
|
||||
|
||||
### Usage Statistics
|
||||
|
||||
Chat completion responses include usage statistics with:
|
||||
|
||||
* `prompt_tokens` - Number of tokens in the prompt
|
||||
* `completion_tokens` - Number of tokens generated
|
||||
* `total_tokens` - Sum of prompt and completion tokens
|
||||
|
||||
### Request Cancellation
|
||||
|
||||
You can cancel active requests by:
|
||||
|
||||
1. Closing the HTTP connection (for streaming requests)
|
||||
2. Calling `/v1/cancel/{command_id}` (for any request)
|
||||
|
||||
The server detects cancellation and stops processing immediately.
|
||||
|
||||
### Instance Placement
|
||||
|
||||
The instance placement endpoints allow you to plan and preview cluster allocations before creating instances. This helps optimize resource usage across nodes.
|
||||
|
||||
### Observability
|
||||
|
||||
The `/events` and `/state` endpoints are primarily intended for operational visibility and debugging.
|
||||
|
||||
@@ -28,6 +28,26 @@ There are currently 5 major systems:
|
||||
|
||||
Implements a distributed algorithm for master election in unstable networking conditions
|
||||
|
||||
## API Layer
|
||||
|
||||
The API system uses multiple adapters to support multiple API formats, converting them to a single request / response type.
|
||||
|
||||
### Adapter Pattern
|
||||
|
||||
Adapters convert between external API formats and EXO's internal types:
|
||||
|
||||
```
|
||||
Chat Completions → [adapter] → TextGenerationTaskParams → Application
|
||||
Claude Messages → [adapter] → TextGenerationTaskParams → Application
|
||||
Responses API → [adapter] → TextGenerationTaskParams → Application
|
||||
Ollama API → [adapter] → TextGenerationTaskParams → Application
|
||||
```
|
||||
|
||||
Each adapter implements two key functions:
|
||||
1. **Request conversion**: Converts API-specific requests to `TextGenerationTaskParams`
|
||||
2. **Response generation**: Converts internal `TokenChunk` streams back to API-specific formats (streaming and non-streaming)
|
||||
|
||||
|
||||
## Topics
|
||||
|
||||
There are currently 5 topics:
|
||||
|
||||
@@ -108,6 +108,7 @@
|
||||
package = pkgsSwift.swiftPackages.swift-format;
|
||||
};
|
||||
shfmt.enable = true;
|
||||
taplo.enable = true;
|
||||
};
|
||||
};
|
||||
|
||||
|
||||
+5
-5
@@ -41,16 +41,16 @@ let
|
||||
|
||||
mlx = stdenv.mkDerivation rec {
|
||||
pname = "mlx";
|
||||
version = let v = "0.30.7.dev20260220+db176851"; in
|
||||
version = let v = "0.30.7.dev20260225+257d5692"; in
|
||||
assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
|
||||
v;
|
||||
pyproject = true;
|
||||
|
||||
src = fetchFromGitHub {
|
||||
owner = "exo-explore";
|
||||
repo = "mlx";
|
||||
rev = "db176851b031631b414890d6bd31d3dfb2e82d25";
|
||||
hash = "sha256-/76lKDMmAjs7Qqh9V0oIBoNPq1kv819JZ+UbVFF8pd4=";
|
||||
owner = "rltakashige";
|
||||
repo = "mlx-jaccl-fix-small-recv";
|
||||
rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
|
||||
hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
|
||||
};
|
||||
|
||||
patches = [
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 6.8 KiB |
Executable
+112
@@ -0,0 +1,112 @@
|
||||
#!/usr/bin/env bash
|
||||
# create-dmg.sh — Build a polished macOS DMG installer for EXO
|
||||
#
|
||||
# Usage:
|
||||
# ./packaging/dmg/create-dmg.sh <app-path> <output-dmg> [volume-name]
|
||||
#
|
||||
# Example:
|
||||
# ./packaging/dmg/create-dmg.sh output/EXO.app EXO-1.0.0.dmg "EXO"
|
||||
#
|
||||
# Creates a DMG with:
|
||||
# - Custom background image with drag-to-Applications arrow
|
||||
# - App icon on left, Applications alias on right
|
||||
# - Proper window size and icon positioning
|
||||
set -euo pipefail
|
||||
|
||||
APP_PATH="${1:?Usage: create-dmg.sh <app-path> <output-dmg> [volume-name]}"
|
||||
OUTPUT_DMG="${2:?Usage: create-dmg.sh <app-path> <output-dmg> [volume-name]}"
|
||||
VOLUME_NAME="${3:-EXO}"
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
BACKGROUND_SCRIPT="${SCRIPT_DIR}/generate-background.py"
|
||||
TEMP_DIR="$(mktemp -d)"
|
||||
DMG_STAGING="${TEMP_DIR}/dmg-root"
|
||||
TEMP_DMG="${TEMP_DIR}/temp.dmg"
|
||||
BACKGROUND_PNG="${TEMP_DIR}/background.png"
|
||||
|
||||
cleanup() { rm -rf "$TEMP_DIR"; }
|
||||
trap cleanup EXIT
|
||||
|
||||
echo "==> Creating DMG installer for ${VOLUME_NAME}"
|
||||
|
||||
# ── Step 1: Generate background image ────────────────────────────────────────
|
||||
if command -v python3 &>/dev/null; then
|
||||
python3 "$BACKGROUND_SCRIPT" "$BACKGROUND_PNG"
|
||||
echo " Background image generated"
|
||||
else
|
||||
echo " Warning: python3 not found, skipping custom background"
|
||||
BACKGROUND_PNG=""
|
||||
fi
|
||||
|
||||
# ── Step 2: Prepare staging directory ─────────────────────────────────────────
|
||||
mkdir -p "$DMG_STAGING"
|
||||
cp -R "$APP_PATH" "$DMG_STAGING/"
|
||||
ln -s /Applications "$DMG_STAGING/Applications"
|
||||
|
||||
# ── Step 3: Create writable DMG ──────────────────────────────────────────────
|
||||
# Calculate required size (app size + 20MB headroom)
|
||||
APP_SIZE_KB=$(du -sk "$APP_PATH" | cut -f1)
|
||||
DMG_SIZE_KB=$((APP_SIZE_KB + 20480))
|
||||
|
||||
hdiutil create \
|
||||
-volname "$VOLUME_NAME" \
|
||||
-size "${DMG_SIZE_KB}k" \
|
||||
-fs HFS+ \
|
||||
-layout SPUD \
|
||||
"$TEMP_DMG"
|
||||
|
||||
# ── Step 4: Mount and configure ──────────────────────────────────────────────
|
||||
MOUNT_DIR=$(hdiutil attach "$TEMP_DMG" -readwrite -noverify | awk -F'\t' '/Apple_HFS/ {gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF); print $NF}')
|
||||
echo " Mounted at: $MOUNT_DIR"
|
||||
|
||||
# Copy contents
|
||||
cp -R "$DMG_STAGING/"* "$MOUNT_DIR/"
|
||||
|
||||
# Add background image
|
||||
if [[ -n $BACKGROUND_PNG && -f $BACKGROUND_PNG ]]; then
|
||||
mkdir -p "$MOUNT_DIR/.background"
|
||||
cp "$BACKGROUND_PNG" "$MOUNT_DIR/.background/background.png"
|
||||
fi
|
||||
|
||||
# ── Step 5: Configure window appearance via AppleScript ──────────────────────
|
||||
# Window: 800×400, app icon on left, Applications on right (matches Ollama layout)
|
||||
# Background image is 1600×740 (2× retina for 800×400 logical window).
|
||||
APP_NAME="$(basename "$APP_PATH")"
|
||||
|
||||
osascript <<APPLESCRIPT
|
||||
tell application "Finder"
|
||||
tell disk "$VOLUME_NAME"
|
||||
open
|
||||
set current view of container window to icon view
|
||||
set toolbar visible of container window to false
|
||||
set statusbar visible of container window to false
|
||||
set bounds of container window to {200, 120, 1000, 520}
|
||||
set opts to icon view options of container window
|
||||
set icon size of opts to 128
|
||||
set text size of opts to 12
|
||||
set arrangement of opts to not arranged
|
||||
if exists file ".background:background.png" then
|
||||
set background picture of opts to file ".background:background.png"
|
||||
end if
|
||||
set position of item "$APP_NAME" of container window to {200, 190}
|
||||
set position of item "Applications" of container window to {600, 190}
|
||||
close
|
||||
open
|
||||
update without registering applications
|
||||
delay 1
|
||||
close
|
||||
end tell
|
||||
end tell
|
||||
APPLESCRIPT
|
||||
|
||||
echo " Window layout configured"
|
||||
|
||||
# Ensure Finder updates are flushed
|
||||
sync
|
||||
|
||||
# ── Step 6: Finalise ─────────────────────────────────────────────────────────
|
||||
hdiutil detach "$MOUNT_DIR" -quiet
|
||||
hdiutil convert "$TEMP_DMG" -format UDZO -imagekey zlib-level=9 -o "$OUTPUT_DMG"
|
||||
|
||||
echo "==> DMG created: $OUTPUT_DMG"
|
||||
echo " Size: $(du -h "$OUTPUT_DMG" | cut -f1)"
|
||||
@@ -0,0 +1,91 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate the DMG background image with a centered drag-to-Applications arrow.
|
||||
|
||||
The output is a 1600×740 retina PNG (2× for 800×400 logical window).
|
||||
Icons are positioned at (200, 190) and (600, 190) in logical coordinates;
|
||||
the arrow is drawn centered between them.
|
||||
|
||||
Usage:
|
||||
python3 generate-background.py [output.png]
|
||||
|
||||
If no output path is given, overwrites the bundled background.png in-place.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
# Retina dimensions (2× logical 800×400)
|
||||
WIDTH = 1600
|
||||
HEIGHT = 740
|
||||
|
||||
# Icon positions in logical coords → retina coords
|
||||
# App icon at (200, 190), Applications at (600, 190)
|
||||
APP_X = 200 * 2 # 400
|
||||
APPS_X = 600 * 2 # 1200
|
||||
ICON_Y = 190 * 2 # 380
|
||||
|
||||
# Arrow drawn between icons, slightly above icon center
|
||||
ARROW_START_X = APP_X + 160 # past the icon
|
||||
ARROW_END_X = APPS_X - 160 # before the Applications icon
|
||||
ARROW_Y = ICON_Y # same height as icons
|
||||
ARROW_RISE = 120 # upward arc height
|
||||
|
||||
|
||||
def draw_arrow(draw: ImageDraw.ImageDraw) -> None:
|
||||
"""Draw a hand-drawn-style curved arrow from app icon toward Applications."""
|
||||
color = (30, 30, 30)
|
||||
line_width = 8
|
||||
|
||||
# Compute bezier curve points for a gentle upward arc
|
||||
points: list[tuple[float, float]] = []
|
||||
steps = 80
|
||||
for i in range(steps + 1):
|
||||
t = i / steps
|
||||
# Quadratic bezier: start → control → end
|
||||
cx = (ARROW_START_X + ARROW_END_X) / 2
|
||||
cy = ARROW_Y - ARROW_RISE
|
||||
x = (1 - t) ** 2 * ARROW_START_X + 2 * (1 - t) * t * cx + t**2 * ARROW_END_X
|
||||
y = (1 - t) ** 2 * ARROW_Y + 2 * (1 - t) * t * cy + t**2 * ARROW_Y
|
||||
points.append((x, y))
|
||||
|
||||
# Draw the curve as connected line segments
|
||||
for i in range(len(points) - 1):
|
||||
draw.line([points[i], points[i + 1]], fill=color, width=line_width)
|
||||
|
||||
# Arrowhead at the end
|
||||
end_x, end_y = points[-1]
|
||||
# Direction from second-to-last to last point
|
||||
prev_x, prev_y = points[-3]
|
||||
angle = math.atan2(end_y - prev_y, end_x - prev_x)
|
||||
head_len = 36
|
||||
head_angle = math.radians(25)
|
||||
|
||||
left_x = end_x - head_len * math.cos(angle - head_angle)
|
||||
left_y = end_y - head_len * math.sin(angle - head_angle)
|
||||
right_x = end_x - head_len * math.cos(angle + head_angle)
|
||||
right_y = end_y - head_len * math.sin(angle + head_angle)
|
||||
|
||||
draw.polygon(
|
||||
[(end_x, end_y), (left_x, left_y), (right_x, right_y)],
|
||||
fill=color,
|
||||
)
|
||||
|
||||
|
||||
def generate_background(output_path: str) -> None:
|
||||
"""Generate a white DMG background with a centered arrow."""
|
||||
img = Image.new("RGBA", (WIDTH, HEIGHT), (255, 255, 255, 255))
|
||||
draw = ImageDraw.Draw(img)
|
||||
draw_arrow(draw)
|
||||
img.save(output_path, "PNG")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
default_output = str(Path(__file__).parent / "background.png")
|
||||
out = sys.argv[1] if len(sys.argv) >= 2 else default_output
|
||||
generate_background(out)
|
||||
print(f"Background image written to {out}")
|
||||
+54
-53
@@ -1,35 +1,35 @@
|
||||
[project]
|
||||
name = "exo"
|
||||
version = "0.3.0"
|
||||
version = "0.3.68"
|
||||
description = "Exo"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.13"
|
||||
dependencies = [
|
||||
"aiofiles>=24.1.0",
|
||||
"aiohttp>=3.12.14",
|
||||
"types-aiofiles>=24.1.0.20250708",
|
||||
"pydantic>=2.11.7",
|
||||
"fastapi>=0.116.1",
|
||||
"filelock>=3.18.0",
|
||||
"rustworkx>=0.17.1",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"psutil>=7.0.0",
|
||||
"loguru>=0.7.3",
|
||||
"exo_pyo3_bindings", # rust bindings
|
||||
"anyio==4.11.0",
|
||||
"mlx; sys_platform == 'darwin'",
|
||||
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
|
||||
"mlx-lm==0.30.7",
|
||||
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
|
||||
"hypercorn>=0.18.0",
|
||||
"openai-harmony>=0.0.8",
|
||||
"httpx>=0.28.1",
|
||||
"tomlkit>=0.14.0",
|
||||
"pillow>=11.0,<12.0", # compatibility with mflux
|
||||
"mflux==0.15.5",
|
||||
"python-multipart>=0.0.21",
|
||||
"msgspec>=0.19.0",
|
||||
"zstandard>=0.23.0",
|
||||
"aiofiles>=24.1.0",
|
||||
"aiohttp>=3.12.14",
|
||||
"types-aiofiles>=24.1.0.20250708",
|
||||
"pydantic>=2.11.7",
|
||||
"fastapi>=0.116.1",
|
||||
"filelock>=3.18.0",
|
||||
"rustworkx>=0.17.1",
|
||||
"huggingface-hub>=0.33.4",
|
||||
"psutil>=7.0.0",
|
||||
"loguru>=0.7.3",
|
||||
"exo_pyo3_bindings", # rust bindings
|
||||
"anyio==4.11.0",
|
||||
"mlx; sys_platform == 'darwin'",
|
||||
"mlx[cpu]==0.30.6; sys_platform == 'linux'",
|
||||
"mlx-lm",
|
||||
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
|
||||
"hypercorn>=0.18.0",
|
||||
"openai-harmony>=0.0.8",
|
||||
"httpx>=0.28.1",
|
||||
"tomlkit>=0.14.0",
|
||||
"pillow>=11.0,<12.0", # compatibility with mflux
|
||||
"mflux==0.15.5",
|
||||
"python-multipart>=0.0.21",
|
||||
"msgspec>=0.19.0",
|
||||
"zstandard>=0.23.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -38,12 +38,12 @@ exo = "exo.main:main"
|
||||
# dependencies only required for development
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"basedpyright>=1.29.0",
|
||||
"pyinstaller>=6.17.0",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
"pytest-env",
|
||||
"ruff>=0.11.13",
|
||||
"basedpyright>=1.29.0",
|
||||
"pyinstaller>=6.17.0",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
"pytest-env",
|
||||
"ruff>=0.11.13",
|
||||
]
|
||||
|
||||
# mlx[cuda] requires a newer version of mlx. the ideal on linux is: default to mlx[cpu] unless[cuda] specified.
|
||||
@@ -57,15 +57,12 @@ dev = [
|
||||
###
|
||||
|
||||
[tool.uv.workspace]
|
||||
members = [
|
||||
"rust/exo_pyo3_bindings",
|
||||
"bench",
|
||||
]
|
||||
members = ["rust/exo_pyo3_bindings", "bench"]
|
||||
|
||||
[tool.uv.sources]
|
||||
exo_pyo3_bindings = { workspace = true }
|
||||
mlx = { git = "https://github.com/exo-explore/mlx.git", branch = "pipe-sidechannel", marker = "sys_platform == 'darwin'" }
|
||||
#mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm", branch = "stable" }
|
||||
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
|
||||
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "fix/float32-logprobs" }
|
||||
# Uncomment to use local mlx/mlx-lm development versions:
|
||||
# mlx = { path = "/Users/Shared/mlx", editable=true }
|
||||
# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
|
||||
@@ -95,7 +92,15 @@ reportUnnecessaryTypeIgnoreComment = "error"
|
||||
pythonVersion = "3.13"
|
||||
pythonPlatform = "Darwin"
|
||||
|
||||
exclude = ["**/.venv", "**/venv", "**/__pycache__", "**/exo_scripts", "**/.direnv", "**/rust", "**/.github"]
|
||||
exclude = [
|
||||
"**/.venv",
|
||||
"**/venv",
|
||||
"**/__pycache__",
|
||||
"**/exo_scripts",
|
||||
"**/.direnv",
|
||||
"**/rust",
|
||||
"**/.github",
|
||||
]
|
||||
stubPath = ".mlx_typings"
|
||||
|
||||
[[tool.basedpyright.executionEnvironments]]
|
||||
@@ -109,17 +114,19 @@ root = "src"
|
||||
[tool.uv]
|
||||
required-version = ">=0.8.6"
|
||||
prerelease = "allow"
|
||||
environments = [
|
||||
"sys_platform == 'darwin'",
|
||||
"sys_platform == 'linux'",
|
||||
]
|
||||
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
|
||||
|
||||
###
|
||||
# ruff configuration
|
||||
###
|
||||
|
||||
[tool.ruff]
|
||||
extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
|
||||
extend-exclude = [
|
||||
"shared/protobufs/**",
|
||||
"*mlx_typings/**",
|
||||
"rust/exo_pyo3_bindings/**",
|
||||
"bench/vendor/**",
|
||||
]
|
||||
|
||||
[tool.ruff.lint]
|
||||
extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
|
||||
@@ -127,13 +134,7 @@ extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
|
||||
[tool.pytest.ini_options]
|
||||
pythonpath = "."
|
||||
asyncio_mode = "auto"
|
||||
markers = [
|
||||
"slow: marks tests as slow (deselected by default)"
|
||||
]
|
||||
env = [
|
||||
"EXO_TESTS=1"
|
||||
]
|
||||
markers = ["slow: marks tests as slow (deselected by default)"]
|
||||
env = ["EXO_TESTS=1"]
|
||||
addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
|
||||
filterwarnings = [
|
||||
"ignore:builtin type Swig:DeprecationWarning",
|
||||
]
|
||||
filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]
|
||||
|
||||
+27
-2
@@ -43,6 +43,27 @@
|
||||
final.setuptools
|
||||
];
|
||||
});
|
||||
# rouge-score and sacrebleu don't declare setuptools as a build dependency
|
||||
rouge-score = prev.rouge-score.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
final.setuptools
|
||||
];
|
||||
});
|
||||
sacrebleu = prev.sacrebleu.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
final.setuptools
|
||||
];
|
||||
});
|
||||
sqlitedict = prev.sqlitedict.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
final.setuptools
|
||||
];
|
||||
});
|
||||
word2number = prev.word2number.overrideAttrs (old: {
|
||||
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
|
||||
final.setuptools
|
||||
];
|
||||
});
|
||||
} // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
|
||||
# Use our pure Nix-built MLX with Metal support (macOS only)
|
||||
mlx = self'.packages.mlx;
|
||||
@@ -96,13 +117,16 @@
|
||||
"lib/python3.13/site-packages/nvidia*"
|
||||
];
|
||||
|
||||
exoVenv = (pythonSet.mkVirtualEnv "exo-env" workspace.deps.default).overrideAttrs {
|
||||
# Exclude bench deps from main env (bench has its own benchVenv)
|
||||
exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
|
||||
|
||||
exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
|
||||
venvIgnoreCollisions = venvCollisionPaths;
|
||||
};
|
||||
|
||||
# Virtual environment with dev dependencies for testing
|
||||
testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
|
||||
workspace.deps.default // {
|
||||
exoDeps // {
|
||||
exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
|
||||
}
|
||||
)).overrideAttrs {
|
||||
@@ -158,6 +182,7 @@
|
||||
exo-test-env = testVenv;
|
||||
} // {
|
||||
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
|
||||
exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
|
||||
exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
|
||||
exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
|
||||
};
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-4bit"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 69593314272
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-6bit"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "6bit"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 100120675296
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-8bit"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 130648036320
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-122B-A10B-bf16"
|
||||
n_layers = 48
|
||||
hidden_size = 3072
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "bf16"
|
||||
base_model = "Qwen3.5 122B A10B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 245125640160
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-27B-4bit"
|
||||
n_layers = 64
|
||||
hidden_size = 5120
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 27B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 16054266848
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-27B-8bit"
|
||||
n_layers = 64
|
||||
hidden_size = 5120
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 27B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 29500943328
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-2B-MLX-8bit"
|
||||
n_layers = 24
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 2B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 2662787264
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-35B-A3B-4bit"
|
||||
n_layers = 40
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 35B A3B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 20391405152
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-35B-A3B-8bit"
|
||||
n_layers = 40
|
||||
hidden_size = 2048
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 35B A3B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 37721130592
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-397B-A17B-4bit"
|
||||
n_layers = 60
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 397B A17B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 223860768352
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-397B-A17B-6bit"
|
||||
n_layers = 60
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "6bit"
|
||||
base_model = "Qwen3.5 397B A17B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 322946674272
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-397B-A17B-8bit"
|
||||
n_layers = 60
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 397B A17B"
|
||||
capabilities = ["text", "thinking", "thinking_toggle"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 422032580192
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-9B-4bit"
|
||||
n_layers = 32
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "4bit"
|
||||
base_model = "Qwen3.5 9B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 5950062560
|
||||
@@ -0,0 +1,12 @@
|
||||
model_id = "mlx-community/Qwen3.5-9B-8bit"
|
||||
n_layers = 32
|
||||
hidden_size = 4096
|
||||
supports_tensor = true
|
||||
tasks = ["TextGeneration"]
|
||||
family = "qwen"
|
||||
quantization = "8bit"
|
||||
base_model = "Qwen3.5 9B"
|
||||
capabilities = ["text", "thinking"]
|
||||
|
||||
[storage_size]
|
||||
in_bytes = 10426433504
|
||||
@@ -26,20 +26,24 @@ networking = { workspace = true }
|
||||
|
||||
# interop
|
||||
pyo3 = { version = "0.27.2", features = [
|
||||
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
|
||||
# "nightly", # enables better-supported GIL integration
|
||||
"experimental-async", # async support in #[pyfunction] & #[pymethods]
|
||||
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
|
||||
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
|
||||
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
|
||||
# "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13
|
||||
# "nightly", # enables better-supported GIL integration
|
||||
"experimental-async", # async support in #[pyfunction] & #[pymethods]
|
||||
#"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation
|
||||
#"py-clone", # adding Clone-ing of `Py<T>` without GIL (may cause panics - remove if panics happen)
|
||||
# "multiple-pymethods", # allows multiple #[pymethods] sections per class
|
||||
|
||||
# integrations with other libraries
|
||||
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
|
||||
# "ordered-float", "rust_decimal", "smallvec",
|
||||
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
|
||||
# integrations with other libraries
|
||||
# "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational",
|
||||
# "ordered-float", "rust_decimal", "smallvec",
|
||||
# "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde",
|
||||
] }
|
||||
pyo3-stub-gen = { version = "0.17.2" }
|
||||
pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime", "testing"] }
|
||||
pyo3-async-runtimes = { version = "0.27.0", features = [
|
||||
"attributes",
|
||||
"tokio-runtime",
|
||||
"testing",
|
||||
] }
|
||||
pyo3-log = "0.13.2"
|
||||
|
||||
# macro dependencies
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
# ruff: noqa: E501, F401
|
||||
|
||||
import builtins
|
||||
import enum
|
||||
import typing
|
||||
|
||||
@typing.final
|
||||
@@ -11,29 +10,6 @@ class AllQueuesFullError(builtins.Exception):
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class ConnectionUpdate:
|
||||
@property
|
||||
def update_type(self) -> ConnectionUpdateType:
|
||||
r"""
|
||||
Whether this is a connection or disconnection event
|
||||
"""
|
||||
@property
|
||||
def peer_id(self) -> builtins.str:
|
||||
r"""
|
||||
Identity of the peer that we have connected to or disconnected from.
|
||||
"""
|
||||
@property
|
||||
def remote_ipv4(self) -> builtins.str:
|
||||
r"""
|
||||
Remote connection's IPv4 address.
|
||||
"""
|
||||
@property
|
||||
def remote_tcp_port(self) -> builtins.int:
|
||||
r"""
|
||||
Remote connection's TCP port.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class Keypair:
|
||||
r"""
|
||||
@@ -58,21 +34,15 @@ class Keypair:
|
||||
Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`.
|
||||
"""
|
||||
|
||||
@typing.final
|
||||
class MessageTooLargeError(builtins.Exception):
|
||||
def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class NetworkingHandle:
|
||||
def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
|
||||
async def connection_update_recv(self) -> ConnectionUpdate:
|
||||
r"""
|
||||
Receives the next `ConnectionUpdate` from networking.
|
||||
"""
|
||||
async def connection_update_recv_many(self, limit: builtins.int) -> builtins.list[ConnectionUpdate]:
|
||||
r"""
|
||||
Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
|
||||
|
||||
For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
|
||||
For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
|
||||
will sleep until a `ConnectionUpdate`s is sent.
|
||||
"""
|
||||
async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
|
||||
r"""
|
||||
Subscribe to a `GossipSub` topic.
|
||||
@@ -91,18 +61,7 @@ class NetworkingHandle:
|
||||
|
||||
If no peers are found that subscribe to this topic, throws `NoPeersSubscribedToTopicError` exception.
|
||||
"""
|
||||
async def gossipsub_recv(self) -> tuple[builtins.str, bytes]:
|
||||
r"""
|
||||
Receives the next message from the `GossipSub` network.
|
||||
"""
|
||||
async def gossipsub_recv_many(self, limit: builtins.int) -> builtins.list[tuple[builtins.str, bytes]]:
|
||||
r"""
|
||||
Receives at most `limit` messages from the `GossipSub` network and returns them.
|
||||
|
||||
For `limit = 0`, an empty collection of messages will be returned immediately.
|
||||
For `limit > 0`, if there are no messages in the channel's queue this method
|
||||
will sleep until a message is sent.
|
||||
"""
|
||||
async def recv(self) -> PyFromSwarm: ...
|
||||
|
||||
@typing.final
|
||||
class NoPeersSubscribedToTopicError(builtins.Exception):
|
||||
@@ -110,11 +69,26 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
|
||||
def __repr__(self) -> builtins.str: ...
|
||||
def __str__(self) -> builtins.str: ...
|
||||
|
||||
@typing.final
|
||||
class ConnectionUpdateType(enum.Enum):
|
||||
r"""
|
||||
Connection or disconnection event discriminant type.
|
||||
"""
|
||||
Connected = ...
|
||||
Disconnected = ...
|
||||
class PyFromSwarm:
|
||||
@typing.final
|
||||
class Connection(PyFromSwarm):
|
||||
__match_args__ = ("peer_id", "connected",)
|
||||
@property
|
||||
def peer_id(self) -> builtins.str: ...
|
||||
@property
|
||||
def connected(self) -> builtins.bool: ...
|
||||
def __new__(cls, peer_id: builtins.str, connected: builtins.bool) -> PyFromSwarm.Connection: ...
|
||||
|
||||
@typing.final
|
||||
class Message(PyFromSwarm):
|
||||
__match_args__ = ("origin", "topic", "data",)
|
||||
@property
|
||||
def origin(self) -> builtins.str: ...
|
||||
@property
|
||||
def topic(self) -> builtins.str: ...
|
||||
@property
|
||||
def data(self) -> bytes: ...
|
||||
def __new__(cls, origin: builtins.str, topic: builtins.str, data: bytes) -> PyFromSwarm.Message: ...
|
||||
|
||||
...
|
||||
|
||||
|
||||
@@ -4,21 +4,18 @@ build-backend = "maturin"
|
||||
|
||||
[project]
|
||||
name = "exo_pyo3_bindings"
|
||||
version = "0.1.0"
|
||||
version = "0.2.1"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
authors = [
|
||||
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" }
|
||||
{ name = "Andrei Cravtov", email = "the.andrei.cravtov@gmail.com" },
|
||||
{ name = "Evan Quiney", email = "evanev7@gmail.com" },
|
||||
]
|
||||
requires-python = ">=3.13"
|
||||
dependencies = []
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"exo_pyo3_bindings",
|
||||
"pytest>=8.4.0",
|
||||
"pytest-asyncio>=1.0.0",
|
||||
]
|
||||
dev = ["exo_pyo3_bindings", "pytest>=8.4.0", "pytest-asyncio>=1.0.0"]
|
||||
|
||||
[tool.maturin]
|
||||
#purelib = true
|
||||
|
||||
@@ -155,6 +155,9 @@ pub(crate) mod ext {
|
||||
fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
// install logger
|
||||
pyo3_log::init();
|
||||
let mut builder = tokio::runtime::Builder::new_multi_thread();
|
||||
builder.enable_all();
|
||||
pyo3_async_runtimes::tokio::init(builder);
|
||||
|
||||
// TODO: for now this is all NOT a submodule, but figure out how to make the submodule system
|
||||
// work with maturin, where the types generate correctly, in the right folder, without
|
||||
|
||||
@@ -1,26 +1,24 @@
|
||||
#![allow(
|
||||
clippy::multiple_inherent_impl,
|
||||
clippy::unnecessary_wraps,
|
||||
clippy::unused_self,
|
||||
clippy::needless_pass_by_value
|
||||
)]
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::r#const::MPSC_CHANNEL_SIZE;
|
||||
use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _};
|
||||
use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _};
|
||||
use crate::ext::{ResultExt as _, TokioMpscSenderExt as _};
|
||||
use crate::ident::PyKeypair;
|
||||
use crate::networking::exception::{
|
||||
PyAllQueuesFullError, PyMessageTooLargeError, PyNoPeersSubscribedToTopicError,
|
||||
};
|
||||
use crate::pyclass;
|
||||
use libp2p::futures::StreamExt as _;
|
||||
use libp2p::gossipsub;
|
||||
use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError};
|
||||
use libp2p::swarm::SwarmEvent;
|
||||
use networking::discovery;
|
||||
use networking::swarm::create_swarm;
|
||||
use futures_lite::{Stream, StreamExt as _};
|
||||
use libp2p::gossipsub::PublishError;
|
||||
use networking::swarm::{FromSwarm, ToSwarm, create_swarm};
|
||||
use pyo3::exceptions::PyRuntimeError;
|
||||
use pyo3::prelude::{PyModule, PyModuleMethods as _};
|
||||
use pyo3::types::PyBytes;
|
||||
use pyo3::{Bound, Py, PyErr, PyResult, PyTraverseError, PyVisit, Python, pymethods};
|
||||
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyclass_enum, gen_stub_pymethods};
|
||||
use std::net::IpAddr;
|
||||
use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, pymethods};
|
||||
use pyo3_stub_gen::derive::{
|
||||
gen_methods_from_python, gen_stub_pyclass, gen_stub_pyclass_complex_enum, gen_stub_pymethods,
|
||||
};
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
|
||||
mod exception {
|
||||
@@ -98,237 +96,78 @@ mod exception {
|
||||
Self::MSG.to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Connection or disconnection event discriminant type.
|
||||
#[gen_stub_pyclass_enum]
|
||||
#[pyclass(eq, eq_int, name = "ConnectionUpdateType")]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
enum PyConnectionUpdateType {
|
||||
Connected = 0,
|
||||
Disconnected,
|
||||
}
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
|
||||
pub struct PyMessageTooLargeError {}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(frozen, name = "ConnectionUpdate")]
|
||||
#[derive(Debug, Clone)]
|
||||
struct PyConnectionUpdate {
|
||||
/// Whether this is a connection or disconnection event
|
||||
#[pyo3(get)]
|
||||
update_type: PyConnectionUpdateType,
|
||||
impl PyMessageTooLargeError {
|
||||
const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
|
||||
|
||||
/// Identity of the peer that we have connected to or disconnected from.
|
||||
#[pyo3(get)]
|
||||
peer_id: String,
|
||||
|
||||
/// Remote connection's IPv4 address.
|
||||
#[pyo3(get)]
|
||||
remote_ipv4: String,
|
||||
|
||||
/// Remote connection's TCP port.
|
||||
#[pyo3(get)]
|
||||
remote_tcp_port: u16,
|
||||
}
|
||||
|
||||
enum ToTask {
|
||||
GossipsubSubscribe {
|
||||
topic: String,
|
||||
result_tx: oneshot::Sender<PyResult<bool>>,
|
||||
},
|
||||
GossipsubUnsubscribe {
|
||||
topic: String,
|
||||
result_tx: oneshot::Sender<bool>,
|
||||
},
|
||||
GossipsubPublish {
|
||||
topic: String,
|
||||
data: Vec<u8>,
|
||||
result_tx: oneshot::Sender<PyResult<MessageId>>,
|
||||
},
|
||||
}
|
||||
|
||||
#[allow(clippy::enum_glob_use)]
|
||||
async fn networking_task(
|
||||
mut swarm: networking::swarm::Swarm,
|
||||
mut to_task_rx: mpsc::Receiver<ToTask>,
|
||||
connection_update_tx: mpsc::Sender<PyConnectionUpdate>,
|
||||
gossipsub_message_tx: mpsc::Sender<(String, Vec<u8>)>,
|
||||
) {
|
||||
use SwarmEvent::*;
|
||||
use ToTask::*;
|
||||
use networking::swarm::BehaviourEvent::*;
|
||||
|
||||
log::info!("RUST: networking task started");
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
message = to_task_rx.recv() => {
|
||||
// handle closed channel
|
||||
let Some(message) = message else {
|
||||
log::info!("RUST: channel closed");
|
||||
break;
|
||||
};
|
||||
|
||||
// dispatch incoming messages
|
||||
match message {
|
||||
GossipsubSubscribe { topic, result_tx } => {
|
||||
// try to subscribe
|
||||
let result = swarm.behaviour_mut()
|
||||
.gossipsub.subscribe(&IdentTopic::new(topic));
|
||||
|
||||
// send response oneshot
|
||||
if let Err(e) = result_tx.send(result.pyerr()) {
|
||||
log::error!("RUST: could not subscribe to gossipsub topic since channel already closed: {e:?}");
|
||||
continue;
|
||||
}
|
||||
}
|
||||
GossipsubUnsubscribe { topic, result_tx } => {
|
||||
// try to unsubscribe from the topic
|
||||
let result = swarm.behaviour_mut()
|
||||
.gossipsub.unsubscribe(&IdentTopic::new(topic));
|
||||
|
||||
// send response oneshot (or exit if connection closed)
|
||||
if let Err(e) = result_tx.send(result) {
|
||||
log::error!("RUST: could not unsubscribe from gossipsub topic since channel already closed: {e:?}");
|
||||
continue;
|
||||
}
|
||||
}
|
||||
GossipsubPublish { topic, data, result_tx } => {
|
||||
// try to publish the data -> catch NoPeersSubscribedToTopic error & convert to correct exception
|
||||
let result = swarm.behaviour_mut().gossipsub.publish(
|
||||
IdentTopic::new(topic), data);
|
||||
let pyresult: PyResult<MessageId> = if let Err(PublishError::NoPeersSubscribedToTopic) = result {
|
||||
Err(exception::PyNoPeersSubscribedToTopicError::new_err())
|
||||
} else if let Err(PublishError::AllQueuesFull(_)) = result {
|
||||
Err(exception::PyAllQueuesFullError::new_err())
|
||||
} else {
|
||||
result.pyerr()
|
||||
};
|
||||
|
||||
// send response oneshot (or exit if connection closed)
|
||||
if let Err(e) = result_tx.send(pyresult) {
|
||||
log::error!("RUST: could not publish gossipsub message since channel already closed: {e:?}");
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// architectural solution to this problem:
|
||||
// create keep_alive behavior who's job it is to dial peers discovered by mDNS (and drop when expired)
|
||||
// -> it will emmit TRUE connected/disconnected events consumable elsewhere
|
||||
//
|
||||
// gossipsub will feed off-of dial attempts created by networking, and that will bootstrap its' peers list
|
||||
// then for actual communication it will dial those peers if need-be
|
||||
swarm_event = swarm.select_next_some() => {
|
||||
match swarm_event {
|
||||
Behaviour(Gossipsub(gossipsub::Event::Message {
|
||||
message: Message {
|
||||
topic,
|
||||
data,
|
||||
..
|
||||
},
|
||||
..
|
||||
})) => {
|
||||
// topic-ID is just the topic hash!!! (since we used identity hasher)
|
||||
let message = (topic.into_string(), data);
|
||||
|
||||
// send incoming message to channel (or exit if connection closed)
|
||||
if let Err(e) = gossipsub_message_tx.send(message).await {
|
||||
log::error!("RUST: could not send incoming gossipsub message since channel already closed: {e}");
|
||||
continue;
|
||||
}
|
||||
},
|
||||
Behaviour(Discovery(discovery::Event::ConnectionEstablished { peer_id, remote_ip, remote_tcp_port, .. })) => {
|
||||
// grab IPv4 string
|
||||
let remote_ipv4 = match remote_ip {
|
||||
IpAddr::V4(ip) => ip.to_string(),
|
||||
IpAddr::V6(ip) => {
|
||||
log::warn!("RUST: ignoring connection to IPv6 address: {ip}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// send connection event to channel (or exit if connection closed)
|
||||
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
|
||||
update_type: PyConnectionUpdateType::Connected,
|
||||
peer_id: peer_id.to_base58(),
|
||||
remote_ipv4,
|
||||
remote_tcp_port,
|
||||
}).await {
|
||||
log::error!("RUST: could not send connection update since channel already closed: {e}");
|
||||
continue;
|
||||
}
|
||||
},
|
||||
Behaviour(Discovery(discovery::Event::ConnectionClosed { peer_id, remote_ip, remote_tcp_port, .. })) => {
|
||||
// grab IPv4 string
|
||||
let remote_ipv4 = match remote_ip {
|
||||
IpAddr::V4(ip) => ip.to_string(),
|
||||
IpAddr::V6(ip) => {
|
||||
log::warn!("RUST: ignoring disconnection from IPv6 address: {ip}");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// send disconnection event to channel (or exit if connection closed)
|
||||
if let Err(e) = connection_update_tx.send(PyConnectionUpdate {
|
||||
update_type: PyConnectionUpdateType::Disconnected,
|
||||
peer_id: peer_id.to_base58(),
|
||||
remote_ipv4,
|
||||
remote_tcp_port,
|
||||
}).await {
|
||||
log::error!("RUST: could not send connection update since channel already closed: {e}");
|
||||
continue;
|
||||
}
|
||||
},
|
||||
e => {
|
||||
log::info!("RUST: other event {e:?}");
|
||||
}
|
||||
}
|
||||
}
|
||||
pub(crate) fn new_err() -> PyErr {
|
||||
PyErr::new::<Self, _>(())
|
||||
}
|
||||
}
|
||||
|
||||
log::info!("RUST: networking task stopped");
|
||||
#[gen_stub_pymethods]
|
||||
#[pymethods]
|
||||
impl PyMessageTooLargeError {
|
||||
#[new]
|
||||
#[pyo3(signature = (*args))]
|
||||
#[allow(unused_variables)]
|
||||
pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
|
||||
Self {}
|
||||
}
|
||||
|
||||
fn __repr__(&self) -> String {
|
||||
format!("MessageTooLargeError(\"{}\")", Self::MSG)
|
||||
}
|
||||
|
||||
fn __str__(&self) -> String {
|
||||
Self::MSG.to_string()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pyclass]
|
||||
#[pyclass(name = "NetworkingHandle")]
|
||||
#[derive(Debug)]
|
||||
struct PyNetworkingHandle {
|
||||
// channels
|
||||
to_task_tx: Option<mpsc::Sender<ToTask>>,
|
||||
connection_update_rx: Mutex<mpsc::Receiver<PyConnectionUpdate>>,
|
||||
gossipsub_message_rx: Mutex<mpsc::Receiver<(String, Vec<u8>)>>,
|
||||
pub to_swarm: mpsc::Sender<ToSwarm>,
|
||||
pub swarm: Arc<Mutex<Pin<Box<dyn Stream<Item = FromSwarm> + Send>>>>,
|
||||
}
|
||||
|
||||
impl Drop for PyNetworkingHandle {
|
||||
fn drop(&mut self) {
|
||||
// TODO: may or may not need to await a "kill-signal" oneshot channel message,
|
||||
// to ensure that the networking task is done BEFORE exiting the clear function...
|
||||
// but this may require GIL?? and it may not be safe to call GIL here??
|
||||
self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
|
||||
}
|
||||
#[gen_stub_pyclass_complex_enum]
|
||||
#[pyclass]
|
||||
enum PyFromSwarm {
|
||||
Connection {
|
||||
peer_id: String,
|
||||
connected: bool,
|
||||
},
|
||||
Message {
|
||||
origin: String,
|
||||
topic: String,
|
||||
data: Py<PyBytes>,
|
||||
},
|
||||
}
|
||||
|
||||
#[allow(clippy::expect_used)]
|
||||
impl PyNetworkingHandle {
|
||||
fn new(
|
||||
to_task_tx: mpsc::Sender<ToTask>,
|
||||
connection_update_rx: mpsc::Receiver<PyConnectionUpdate>,
|
||||
gossipsub_message_rx: mpsc::Receiver<(String, Vec<u8>)>,
|
||||
) -> Self {
|
||||
Self {
|
||||
to_task_tx: Some(to_task_tx),
|
||||
connection_update_rx: Mutex::new(connection_update_rx),
|
||||
gossipsub_message_rx: Mutex::new(gossipsub_message_rx),
|
||||
impl From<FromSwarm> for PyFromSwarm {
|
||||
fn from(value: FromSwarm) -> Self {
|
||||
match value {
|
||||
FromSwarm::Discovered { peer_id } => Self::Connection {
|
||||
peer_id: peer_id.to_base58(),
|
||||
connected: true,
|
||||
},
|
||||
FromSwarm::Expired { peer_id } => Self::Connection {
|
||||
peer_id: peer_id.to_base58(),
|
||||
connected: false,
|
||||
},
|
||||
FromSwarm::Message { from, topic, data } => Self::Message {
|
||||
origin: from.to_base58(),
|
||||
topic: topic,
|
||||
data: data.pybytes(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
const fn to_task_tx(&self) -> &mpsc::Sender<ToTask> {
|
||||
self.to_task_tx
|
||||
.as_ref()
|
||||
.expect("The sender should only be None after de-initialization.")
|
||||
}
|
||||
}
|
||||
|
||||
#[gen_stub_pymethods]
|
||||
@@ -342,97 +181,36 @@ impl PyNetworkingHandle {
|
||||
|
||||
#[new]
|
||||
fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
|
||||
use pyo3_async_runtimes::tokio::get_runtime;
|
||||
|
||||
// create communication channels
|
||||
let (to_task_tx, to_task_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
let (connection_update_tx, connection_update_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
let (gossipsub_message_tx, gossipsub_message_rx) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
|
||||
|
||||
// get identity
|
||||
let identity = identity.borrow().0.clone();
|
||||
|
||||
// create networking swarm (within tokio context!! or it crashes)
|
||||
let swarm = get_runtime()
|
||||
.block_on(async { create_swarm(identity) })
|
||||
.pyerr()?;
|
||||
let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
|
||||
let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
|
||||
|
||||
// spawn tokio task running the networking logic
|
||||
get_runtime().spawn(async move {
|
||||
networking_task(
|
||||
swarm,
|
||||
to_task_rx,
|
||||
connection_update_tx,
|
||||
gossipsub_message_tx,
|
||||
)
|
||||
.await;
|
||||
});
|
||||
Ok(Self::new(
|
||||
to_task_tx,
|
||||
connection_update_rx,
|
||||
gossipsub_message_rx,
|
||||
))
|
||||
Ok(Self {
|
||||
swarm: Arc::new(Mutex::new(swarm)),
|
||||
to_swarm,
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
|
||||
Ok(()) // This is needed purely so `__clear__` can work
|
||||
fn recv<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
let swarm = Arc::clone(&self.swarm);
|
||||
pyo3_async_runtimes::tokio::future_into_py(py, async move {
|
||||
swarm
|
||||
.try_lock()
|
||||
.map_err(|_| PyRuntimeError::new_err("called recv twice concurrently"))?
|
||||
.next()
|
||||
.await
|
||||
.ok_or(PyErr::receiver_channel_closed())
|
||||
.map(PyFromSwarm::from)
|
||||
})
|
||||
}
|
||||
|
||||
#[gen_stub(skip)]
|
||||
fn __clear__(&mut self) {
|
||||
// TODO: may or may not need to await a "kill-signal" oneshot channel message,
|
||||
// to ensure that the networking task is done BEFORE exiting the clear function...
|
||||
// but this may require GIL?? and it may not be safe to call GIL here??
|
||||
self.to_task_tx = None; // Using Option<T> as a trick to force channel to be dropped
|
||||
}
|
||||
|
||||
// ---- Connection update receiver methods ----
|
||||
|
||||
/// Receives the next `ConnectionUpdate` from networking.
|
||||
async fn connection_update_recv(&self) -> PyResult<PyConnectionUpdate> {
|
||||
self.connection_update_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_py()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
}
|
||||
|
||||
/// Receives at most `limit` `ConnectionUpdate`s from networking and returns them.
|
||||
///
|
||||
/// For `limit = 0`, an empty collection of `ConnectionUpdate`s will be returned immediately.
|
||||
/// For `limit > 0`, if there are no `ConnectionUpdate`s in the channel's queue this method
|
||||
/// will sleep until a `ConnectionUpdate`s is sent.
|
||||
async fn connection_update_recv_many(&self, limit: usize) -> PyResult<Vec<PyConnectionUpdate>> {
|
||||
self.connection_update_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_many_py(limit)
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
}
|
||||
|
||||
// TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
|
||||
// so its too dangerous to expose just yet. figure out a better semantics for handling this,
|
||||
// so things don't randomly block
|
||||
// /// Tries to receive the next `ConnectionUpdate` from networking.
|
||||
// fn connection_update_try_recv(&self) -> PyResult<Option<PyConnectionUpdate>> {
|
||||
// self.connection_update_rx.blocking_lock().try_recv_py()
|
||||
// }
|
||||
//
|
||||
// /// Checks if the `ConnectionUpdate` channel is empty.
|
||||
// fn connection_update_is_empty(&self) -> bool {
|
||||
// self.connection_update_rx.blocking_lock().is_empty()
|
||||
// }
|
||||
//
|
||||
// /// Returns the number of `ConnectionUpdate`s in the channel.
|
||||
// fn connection_update_len(&self) -> usize {
|
||||
// self.connection_update_rx.blocking_lock().len()
|
||||
// }
|
||||
|
||||
// ---- Gossipsub management methods ----
|
||||
|
||||
/// Subscribe to a `GossipSub` topic.
|
||||
@@ -442,10 +220,10 @@ impl PyNetworkingHandle {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to subscribe
|
||||
self.to_task_tx()
|
||||
.send_py(ToTask::GossipsubSubscribe {
|
||||
self.to_swarm
|
||||
.send_py(ToSwarm::Subscribe {
|
||||
topic,
|
||||
result_tx: tx,
|
||||
result_sender: tx,
|
||||
})
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?;
|
||||
@@ -454,6 +232,7 @@ impl PyNetworkingHandle {
|
||||
rx.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map_err(|_| PyErr::receiver_channel_closed())?
|
||||
.pyerr()
|
||||
}
|
||||
|
||||
/// Unsubscribes from a `GossipSub` topic.
|
||||
@@ -463,10 +242,10 @@ impl PyNetworkingHandle {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
|
||||
// send off request to unsubscribe
|
||||
self.to_task_tx()
|
||||
.send_py(ToTask::GossipsubUnsubscribe {
|
||||
self.to_swarm
|
||||
.send_py(ToSwarm::Unsubscribe {
|
||||
topic,
|
||||
result_tx: tx,
|
||||
result_sender: tx,
|
||||
})
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?;
|
||||
@@ -485,11 +264,11 @@ impl PyNetworkingHandle {
|
||||
|
||||
// send off request to subscribe
|
||||
let data = Python::attach(|py| Vec::from(data.as_bytes(py)));
|
||||
self.to_task_tx()
|
||||
.send_py(ToTask::GossipsubPublish {
|
||||
self.to_swarm
|
||||
.send_py(ToSwarm::Publish {
|
||||
topic,
|
||||
data,
|
||||
result_tx: tx,
|
||||
result_sender: tx,
|
||||
})
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?;
|
||||
@@ -498,74 +277,35 @@ impl PyNetworkingHandle {
|
||||
let _ = rx
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map_err(|_| PyErr::receiver_channel_closed())??;
|
||||
.map_err(|_| PyErr::receiver_channel_closed())?
|
||||
.map_err(|e| match e {
|
||||
PublishError::AllQueuesFull(_) => PyAllQueuesFullError::new_err(),
|
||||
PublishError::MessageTooLarge => PyMessageTooLargeError::new_err(),
|
||||
PublishError::NoPeersSubscribedToTopic => {
|
||||
PyNoPeersSubscribedToTopicError::new_err()
|
||||
}
|
||||
e => PyRuntimeError::new_err(e.to_string()),
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ---- Gossipsub message receiver methods ----
|
||||
|
||||
/// Receives the next message from the `GossipSub` network.
|
||||
async fn gossipsub_recv(&self) -> PyResult<(String, Py<PyBytes>)> {
|
||||
self.gossipsub_message_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_py()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.map(|(t, d)| (t, d.pybytes()))
|
||||
pyo3_stub_gen::inventory::submit! {
|
||||
gen_methods_from_python! {
|
||||
r#"
|
||||
class PyNetworkingHandle:
|
||||
async def recv() -> PyFromSwarm: ...
|
||||
"#
|
||||
}
|
||||
|
||||
/// Receives at most `limit` messages from the `GossipSub` network and returns them.
|
||||
///
|
||||
/// For `limit = 0`, an empty collection of messages will be returned immediately.
|
||||
/// For `limit > 0`, if there are no messages in the channel's queue this method
|
||||
/// will sleep until a message is sent.
|
||||
async fn gossipsub_recv_many(&self, limit: usize) -> PyResult<Vec<(String, Py<PyBytes>)>> {
|
||||
Ok(self
|
||||
.gossipsub_message_rx
|
||||
.lock()
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await
|
||||
.recv_many_py(limit)
|
||||
.allow_threads_py() // allow-threads-aware async call
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|(t, d)| (t, d.pybytes()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
// TODO: rn this blocks main thread if anything else is awaiting the channel (bc its a mutex)
|
||||
// so its too dangerous to expose just yet. figure out a better semantics for handling this,
|
||||
// so things don't randomly block
|
||||
// /// Tries to receive the next message from the `GossipSub` network.
|
||||
// fn gossipsub_try_recv(&self) -> PyResult<Option<(String, Py<PyBytes>)>> {
|
||||
// Ok(self
|
||||
// .gossipsub_message_rx
|
||||
// .blocking_lock()
|
||||
// .try_recv_py()?
|
||||
// .map(|(t, d)| (t, d.pybytes())))
|
||||
// }
|
||||
//
|
||||
// /// Checks if the `GossipSub` message channel is empty.
|
||||
// fn gossipsub_is_empty(&self) -> bool {
|
||||
// self.gossipsub_message_rx.blocking_lock().is_empty()
|
||||
// }
|
||||
//
|
||||
// /// Returns the number of `GossipSub` messages in the channel.
|
||||
// fn gossipsub_len(&self) -> usize {
|
||||
// self.gossipsub_message_rx.blocking_lock().len()
|
||||
// }
|
||||
}
|
||||
|
||||
pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
|
||||
m.add_class::<exception::PyAllQueuesFullError>()?;
|
||||
m.add_class::<exception::PyMessageTooLargeError>()?;
|
||||
|
||||
m.add_class::<PyConnectionUpdateType>()?;
|
||||
m.add_class::<PyConnectionUpdate>()?;
|
||||
m.add_class::<PyConnectionUpdateType>()?;
|
||||
m.add_class::<PyNetworkingHandle>()?;
|
||||
m.add_class::<PyFromSwarm>()?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -1,16 +1,20 @@
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from exo_pyo3_bindings import Keypair, NetworkingHandle, NoPeersSubscribedToTopicError
|
||||
from exo_pyo3_bindings import (
|
||||
Keypair,
|
||||
NetworkingHandle,
|
||||
NoPeersSubscribedToTopicError,
|
||||
PyFromSwarm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sleep_on_multiple_items() -> None:
|
||||
print("PYTHON: starting handle")
|
||||
h = NetworkingHandle(Keypair.generate_ed25519())
|
||||
h = NetworkingHandle(Keypair.generate())
|
||||
|
||||
ct = asyncio.create_task(_await_cons(h))
|
||||
mt = asyncio.create_task(_await_msg(h))
|
||||
rt = asyncio.create_task(_await_recv(h))
|
||||
|
||||
# sleep for 4 ticks
|
||||
for i in range(4):
|
||||
@@ -22,13 +26,11 @@ async def test_sleep_on_multiple_items() -> None:
|
||||
print("caught it", e)
|
||||
|
||||
|
||||
async def _await_cons(h: NetworkingHandle):
|
||||
async def _await_recv(h: NetworkingHandle):
|
||||
while True:
|
||||
c = await h.connection_update_recv()
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
|
||||
|
||||
async def _await_msg(h: NetworkingHandle):
|
||||
while True:
|
||||
m = await h.gossipsub_recv()
|
||||
print(f"PYTHON: message: {m}")
|
||||
event = await h.recv()
|
||||
match event:
|
||||
case PyFromSwarm.Connection() as c:
|
||||
print(f"PYTHON: connection update: {c}")
|
||||
case PyFromSwarm.Message() as m:
|
||||
print(f"PYTHON: message: {m}")
|
||||
|
||||
@@ -21,13 +21,17 @@ extend = { workspace = true }
|
||||
delegate = { workspace = true }
|
||||
|
||||
# async
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
async-stream = { workspace = true }
|
||||
futures-lite = { workspace = true }
|
||||
futures-timer = { workspace = true }
|
||||
tokio = { workspace = true, features = ["full"] }
|
||||
|
||||
# utility dependencies
|
||||
util = { workspace = true }
|
||||
tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
|
||||
tracing-subscriber = { version = "0.3.19", features = [
|
||||
"default",
|
||||
"env-filter",
|
||||
] }
|
||||
keccak-const = { workspace = true }
|
||||
|
||||
# tracing/logging
|
||||
@@ -35,3 +39,4 @@ log = { workspace = true }
|
||||
|
||||
# networking
|
||||
libp2p = { workspace = true, features = ["full"] }
|
||||
pin-project = "1.1.10"
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
use futures_lite::StreamExt;
|
||||
use libp2p::{gossipsub, identity, swarm::SwarmEvent};
|
||||
use networking::{discovery, swarm};
|
||||
use tokio::{io, io::AsyncBufReadExt as _, select};
|
||||
use libp2p::identity;
|
||||
use networking::swarm;
|
||||
use networking::swarm::{FromSwarm, ToSwarm};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio::{io, io::AsyncBufReadExt as _};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use tracing_subscriber::filter::LevelFilter;
|
||||
|
||||
@@ -11,64 +13,69 @@ async fn main() {
|
||||
.with_env_filter(EnvFilter::from_default_env().add_directive(LevelFilter::INFO.into()))
|
||||
.try_init();
|
||||
|
||||
let (to_swarm, from_client) = mpsc::channel(20);
|
||||
|
||||
// Configure swarm
|
||||
let mut swarm =
|
||||
swarm::create_swarm(identity::Keypair::generate_ed25519()).expect("Swarm creation failed");
|
||||
let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
|
||||
.expect("Swarm creation failed")
|
||||
.into_stream();
|
||||
|
||||
// Create a Gossipsub topic & subscribe
|
||||
let topic = gossipsub::IdentTopic::new("test-net");
|
||||
swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.subscribe(&topic)
|
||||
.expect("Subscribing to topic failed");
|
||||
let (tx, rx) = oneshot::channel();
|
||||
_ = to_swarm
|
||||
.send(ToSwarm::Subscribe {
|
||||
topic: "test-net".to_string(),
|
||||
result_sender: tx,
|
||||
})
|
||||
.await
|
||||
.expect("should send");
|
||||
|
||||
// Read full lines from stdin
|
||||
let mut stdin = io::BufReader::new(io::stdin()).lines();
|
||||
println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub");
|
||||
|
||||
tokio::task::spawn(async move {
|
||||
rx.await
|
||||
.expect("tx not dropped")
|
||||
.expect("subscribe shouldn't fail");
|
||||
loop {
|
||||
if let Ok(Some(line)) = stdin.next_line().await {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
if let Err(e) = to_swarm
|
||||
.send(swarm::ToSwarm::Publish {
|
||||
topic: "test-net".to_string(),
|
||||
data: line.as_bytes().to_vec(),
|
||||
result_sender: tx,
|
||||
})
|
||||
.await
|
||||
{
|
||||
println!("Send error: {e:?}");
|
||||
return;
|
||||
};
|
||||
match rx.await {
|
||||
Ok(Err(e)) => println!("Publish error: {e:?}"),
|
||||
Err(e) => println!("Publish error: {e:?}"),
|
||||
Ok(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Kick it off
|
||||
loop {
|
||||
select! {
|
||||
// on gossipsub outgoing
|
||||
Ok(Some(line)) = stdin.next_line() => {
|
||||
if let Err(e) = swarm
|
||||
.behaviour_mut().gossipsub
|
||||
.publish(topic.clone(), line.as_bytes()) {
|
||||
println!("Publish error: {e:?}");
|
||||
}
|
||||
// on gossipsub outgoing
|
||||
match swarm.next().await {
|
||||
// on gossipsub incoming
|
||||
Some(FromSwarm::Discovered { peer_id }) => {
|
||||
println!("\n\nconnected to {peer_id}\n\n")
|
||||
}
|
||||
event = swarm.next() => match event {
|
||||
// on gossipsub incoming
|
||||
Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message {
|
||||
propagation_source: peer_id,
|
||||
message_id: id,
|
||||
message,
|
||||
}))) => println!(
|
||||
"\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n",
|
||||
String::from_utf8_lossy(&message.data),
|
||||
),
|
||||
|
||||
// on discovery
|
||||
Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e {
|
||||
discovery::Event::ConnectionEstablished {
|
||||
peer_id, connection_id, remote_ip, remote_tcp_port
|
||||
} => {
|
||||
println!("\n\nConnected to: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
|
||||
}
|
||||
discovery::Event::ConnectionClosed {
|
||||
peer_id, connection_id, remote_ip, remote_tcp_port
|
||||
} => {
|
||||
eprintln!("\n\nDisconnected from: {peer_id}; connection ID: {connection_id}; remote IP: {remote_ip}; remote TCP port: {remote_tcp_port}\n\n");
|
||||
}
|
||||
}
|
||||
|
||||
// ignore outgoing errors: those are normal
|
||||
e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); }
|
||||
|
||||
// otherwise log any other event
|
||||
e => { log::info!("Other event {e:?}"); }
|
||||
Some(FromSwarm::Expired { peer_id }) => {
|
||||
println!("\n\ndisconnected from {peer_id}\n\n")
|
||||
}
|
||||
Some(FromSwarm::Message { from, topic, data }) => {
|
||||
println!("{topic}/{from}:\n{}", String::from_utf8_lossy(&data))
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
use crate::alias;
|
||||
use crate::swarm::transport::tcp_transport;
|
||||
pub use behaviour::{Behaviour, BehaviourEvent};
|
||||
use libp2p::{SwarmBuilder, identity};
|
||||
use std::pin::Pin;
|
||||
|
||||
pub type Swarm = libp2p::Swarm<Behaviour>;
|
||||
use crate::swarm::transport::tcp_transport;
|
||||
use crate::{alias, discovery};
|
||||
pub use behaviour::{Behaviour, BehaviourEvent};
|
||||
use futures_lite::{Stream, StreamExt};
|
||||
use libp2p::{PeerId, SwarmBuilder, gossipsub, identity, swarm::SwarmEvent};
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
|
||||
/// The current version of the network: this prevents devices running different versions of the
|
||||
/// software from interacting with each other.
|
||||
@@ -15,8 +17,136 @@ pub type Swarm = libp2p::Swarm<Behaviour>;
|
||||
pub const NETWORK_VERSION: &[u8] = b"v0.0.1";
|
||||
pub const OVERRIDE_VERSION_ENV_VAR: &str = "EXO_LIBP2P_NAMESPACE";
|
||||
|
||||
// Uses oneshot senders to emulate function calling apis while avoiding requiring unique ownership
|
||||
// of the Swarm.
|
||||
pub enum ToSwarm {
|
||||
Unsubscribe {
|
||||
topic: String,
|
||||
result_sender: oneshot::Sender<bool>,
|
||||
},
|
||||
Subscribe {
|
||||
topic: String,
|
||||
result_sender: oneshot::Sender<Result<bool, gossipsub::SubscriptionError>>,
|
||||
},
|
||||
Publish {
|
||||
topic: String,
|
||||
data: Vec<u8>,
|
||||
result_sender: oneshot::Sender<Result<gossipsub::MessageId, gossipsub::PublishError>>,
|
||||
},
|
||||
}
|
||||
pub enum FromSwarm {
|
||||
Message {
|
||||
from: PeerId,
|
||||
topic: String,
|
||||
data: Vec<u8>,
|
||||
},
|
||||
Discovered {
|
||||
peer_id: PeerId,
|
||||
},
|
||||
Expired {
|
||||
peer_id: PeerId,
|
||||
},
|
||||
}
|
||||
|
||||
pub struct Swarm {
|
||||
swarm: libp2p::Swarm<Behaviour>,
|
||||
from_client: mpsc::Receiver<ToSwarm>,
|
||||
}
|
||||
|
||||
impl Swarm {
|
||||
pub fn into_stream(self) -> Pin<Box<dyn Stream<Item = FromSwarm> + Send>> {
|
||||
let Swarm {
|
||||
mut swarm,
|
||||
mut from_client,
|
||||
} = self;
|
||||
let stream = async_stream::stream! {
|
||||
loop {
|
||||
tokio::select! {
|
||||
msg = from_client.recv() => {
|
||||
let Some(msg) = msg else { break };
|
||||
on_message(&mut swarm, msg);
|
||||
}
|
||||
event = swarm.next() => {
|
||||
let Some(event) = event else { break };
|
||||
if let Some(item) = filter_swarm_event(event) {
|
||||
yield item;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
Box::pin(stream)
|
||||
}
|
||||
}
|
||||
|
||||
fn on_message(swarm: &mut libp2p::Swarm<Behaviour>, message: ToSwarm) {
|
||||
match message {
|
||||
ToSwarm::Subscribe {
|
||||
topic,
|
||||
result_sender,
|
||||
} => {
|
||||
let result = swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.subscribe(&gossipsub::IdentTopic::new(topic));
|
||||
_ = result_sender.send(result);
|
||||
}
|
||||
ToSwarm::Unsubscribe {
|
||||
topic,
|
||||
result_sender,
|
||||
} => {
|
||||
let result = swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.unsubscribe(&gossipsub::IdentTopic::new(topic));
|
||||
_ = result_sender.send(result);
|
||||
}
|
||||
ToSwarm::Publish {
|
||||
topic,
|
||||
data,
|
||||
result_sender,
|
||||
} => {
|
||||
let result = swarm
|
||||
.behaviour_mut()
|
||||
.gossipsub
|
||||
.publish(gossipsub::IdentTopic::new(topic), data);
|
||||
_ = result_sender.send(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
|
||||
match event {
|
||||
SwarmEvent::Behaviour(BehaviourEvent::Gossipsub(gossipsub::Event::Message {
|
||||
message:
|
||||
gossipsub::Message {
|
||||
source: Some(peer_id),
|
||||
topic,
|
||||
data,
|
||||
..
|
||||
},
|
||||
..
|
||||
})) => Some(FromSwarm::Message {
|
||||
from: peer_id,
|
||||
topic: topic.into_string(),
|
||||
data,
|
||||
}),
|
||||
SwarmEvent::Behaviour(BehaviourEvent::Discovery(
|
||||
discovery::Event::ConnectionEstablished { peer_id, .. },
|
||||
)) => Some(FromSwarm::Discovered { peer_id }),
|
||||
SwarmEvent::Behaviour(BehaviourEvent::Discovery(discovery::Event::ConnectionClosed {
|
||||
peer_id,
|
||||
..
|
||||
})) => Some(FromSwarm::Expired { peer_id }),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create and configure a swarm which listens to all ports on OS
|
||||
pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
|
||||
pub fn create_swarm(
|
||||
keypair: identity::Keypair,
|
||||
from_client: mpsc::Receiver<ToSwarm>,
|
||||
) -> alias::AnyResult<Swarm> {
|
||||
let mut swarm = SwarmBuilder::with_existing_identity(keypair)
|
||||
.with_tokio()
|
||||
.with_other_transport(tcp_transport)?
|
||||
@@ -25,7 +155,7 @@ pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult<Swarm> {
|
||||
|
||||
// Listen on all interfaces and whatever port the OS assigns
|
||||
swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
|
||||
Ok(swarm)
|
||||
Ok(Swarm { swarm, from_client })
|
||||
}
|
||||
|
||||
mod transport {
|
||||
@@ -133,7 +263,7 @@ mod behaviour {
|
||||
gossipsub::Behaviour::new(
|
||||
MessageAuthenticity::Signed(keypair.clone()),
|
||||
ConfigBuilder::default()
|
||||
.max_transmit_size(1024 * 1024)
|
||||
.max_transmit_size(8 * 1024 * 1024)
|
||||
.validation_mode(ValidationMode::Strict)
|
||||
.build()
|
||||
.expect("the configuration should always be valid"),
|
||||
|
||||
@@ -1,11 +1,8 @@
|
||||
import asyncio
|
||||
import socket
|
||||
from dataclasses import dataclass, field
|
||||
from random import random
|
||||
|
||||
import anyio
|
||||
from anyio import current_time
|
||||
from anyio.abc import TaskGroup
|
||||
from loguru import logger
|
||||
|
||||
from exo.download.download_utils import (
|
||||
@@ -23,13 +20,9 @@ from exo.shared.types.commands import (
|
||||
ForwarderDownloadCommand,
|
||||
StartDownload,
|
||||
)
|
||||
from exo.shared.types.common import NodeId, SessionId, SystemId
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
EventId,
|
||||
# TODO(evan): just for acks, should delete this ASAP
|
||||
GlobalForwarderEvent,
|
||||
LocalForwarderEvent,
|
||||
NodeDownloadProgress,
|
||||
)
|
||||
from exo.shared.types.worker.downloads import (
|
||||
@@ -40,41 +33,28 @@ from exo.shared.types.worker.downloads import (
|
||||
DownloadProgress,
|
||||
)
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Receiver, Sender
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
|
||||
@dataclass
|
||||
class DownloadCoordinator:
|
||||
node_id: NodeId
|
||||
session_id: SessionId
|
||||
shard_downloader: ShardDownloader
|
||||
download_command_receiver: Receiver[ForwarderDownloadCommand]
|
||||
local_event_sender: Sender[LocalForwarderEvent]
|
||||
|
||||
# ack stuff
|
||||
_global_event_receiver: Receiver[GlobalForwarderEvent]
|
||||
_out_for_delivery: dict[EventId, LocalForwarderEvent] = field(default_factory=dict)
|
||||
|
||||
event_sender: Sender[Event]
|
||||
offline: bool = False
|
||||
|
||||
_system_id: SystemId = field(default_factory=SystemId)
|
||||
|
||||
# Local state
|
||||
download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
|
||||
active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
|
||||
|
||||
# Internal event channel for forwarding (initialized in __post_init__)
|
||||
event_sender: Sender[Event] = field(init=False)
|
||||
event_receiver: Receiver[Event] = field(init=False)
|
||||
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
# Per-model throttle for download progress events
|
||||
_last_progress_time: dict[ModelId, float] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.event_sender, self.event_receiver = channel[Event]()
|
||||
if self.offline:
|
||||
self.shard_downloader.set_internet_connection(False)
|
||||
self.shard_downloader.on_progress(self._download_progress_callback)
|
||||
|
||||
def _model_dir(self, model_id: ModelId) -> str:
|
||||
@@ -123,65 +103,16 @@ class DownloadCoordinator:
|
||||
logger.info(
|
||||
f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
|
||||
)
|
||||
if not self.offline:
|
||||
self._test_internet_connection()
|
||||
try:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._command_processor)
|
||||
tg.start_soon(self._forward_events)
|
||||
tg.start_soon(self._emit_existing_download_progress)
|
||||
tg.start_soon(self._resend_out_for_delivery)
|
||||
tg.start_soon(self._clear_ofd)
|
||||
if not self.offline:
|
||||
tg.start_soon(self._check_internet_connection)
|
||||
finally:
|
||||
for task in self.active_downloads.values():
|
||||
task.cancel()
|
||||
|
||||
def _test_internet_connection(self) -> None:
|
||||
# Try multiple endpoints since some ISPs/networks block specific IPs
|
||||
for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
|
||||
try:
|
||||
socket.create_connection((host, 443), timeout=3).close()
|
||||
self.shard_downloader.set_internet_connection(True)
|
||||
logger.debug(f"Internet connectivity: True (via {host})")
|
||||
return
|
||||
except OSError:
|
||||
continue
|
||||
self.shard_downloader.set_internet_connection(False)
|
||||
logger.debug("Internet connectivity: False")
|
||||
|
||||
async def _check_internet_connection(self) -> None:
|
||||
first_connection = True
|
||||
while True:
|
||||
await asyncio.sleep(10)
|
||||
|
||||
# Assume that internet connection is set to False on 443 errors.
|
||||
if self.shard_downloader.internet_connection:
|
||||
continue
|
||||
|
||||
self._test_internet_connection()
|
||||
|
||||
if first_connection and self.shard_downloader.internet_connection:
|
||||
first_connection = False
|
||||
self._tg.start_soon(self._emit_existing_download_progress)
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._tg.cancel_scope.cancel()
|
||||
|
||||
# directly copied from worker
|
||||
async def _resend_out_for_delivery(self) -> None:
|
||||
# This can also be massively tightened, we should check events are at least a certain age before resending.
|
||||
# Exponential backoff would also certainly help here.
|
||||
while True:
|
||||
await anyio.sleep(1 + random())
|
||||
for event in self._out_for_delivery.copy().values():
|
||||
await self.local_event_sender.send(event)
|
||||
|
||||
async def _clear_ofd(self) -> None:
|
||||
with self._global_event_receiver as events:
|
||||
async for event in events:
|
||||
self._out_for_delivery.pop(event.event.event_id, None)
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _command_processor(self) -> None:
|
||||
with self.download_command_receiver as commands:
|
||||
@@ -202,6 +133,16 @@ class DownloadCoordinator:
|
||||
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()
|
||||
current_status = self.download_status[model_id]
|
||||
pending = DownloadPending(
|
||||
shard_metadata=current_status.shard_metadata,
|
||||
node_id=self.node_id,
|
||||
model_directory=self._model_dir(model_id),
|
||||
)
|
||||
self.download_status[model_id] = pending
|
||||
await self.event_sender.send(
|
||||
NodeDownloadProgress(download_progress=pending)
|
||||
)
|
||||
|
||||
async def _start_download(self, shard: ShardMetadata) -> None:
|
||||
model_id = shard.model_card.model_id
|
||||
@@ -355,26 +296,9 @@ class DownloadCoordinator:
|
||||
)
|
||||
del self.download_status[model_id]
|
||||
|
||||
async def _forward_events(self) -> None:
|
||||
idx = 0
|
||||
with self.event_receiver as events:
|
||||
async for event in events:
|
||||
fe = LocalForwarderEvent(
|
||||
origin_idx=idx,
|
||||
origin=self._system_id,
|
||||
session=self.session_id,
|
||||
event=event,
|
||||
)
|
||||
idx += 1
|
||||
logger.debug(
|
||||
f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
|
||||
)
|
||||
await self.local_event_sender.send(fe)
|
||||
self._out_for_delivery[event.event_id] = fe
|
||||
|
||||
async def _emit_existing_download_progress(self) -> None:
|
||||
try:
|
||||
while True:
|
||||
while True:
|
||||
try:
|
||||
logger.debug(
|
||||
"DownloadCoordinator: Fetching and emitting existing download progress..."
|
||||
)
|
||||
@@ -462,8 +386,8 @@ class DownloadCoordinator:
|
||||
logger.debug(
|
||||
"DownloadCoordinator: Done emitting existing download progress."
|
||||
)
|
||||
await anyio.sleep(60)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"DownloadCoordinator: Error emitting existing download progress: {e}"
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"DownloadCoordinator: Error emitting existing download progress: {e}"
|
||||
)
|
||||
await anyio.sleep(60)
|
||||
|
||||
@@ -314,9 +314,13 @@ async def fetch_file_list_with_cache(
|
||||
_fetched_file_lists_this_session.add(cache_key)
|
||||
return file_list
|
||||
except Exception as e:
|
||||
logger.opt(exception=e).warning(
|
||||
"Ran into exception when fetching file list from HF."
|
||||
)
|
||||
|
||||
if await aios.path.exists(cache_file):
|
||||
logger.warning(
|
||||
f"No internet and no cached file list for {model_id} - using local file list"
|
||||
f"No cached file list for {model_id} - using local file list"
|
||||
)
|
||||
async with aiofiles.open(cache_file, "r") as f:
|
||||
return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
|
||||
@@ -823,6 +827,7 @@ async def download_shard(
|
||||
|
||||
for file in filtered_file_list:
|
||||
downloaded_bytes = await get_downloaded_size(target_dir / file.path)
|
||||
final_file_exists = await aios.path.exists(target_dir / file.path)
|
||||
file_progress[file.path] = RepoFileDownloadProgress(
|
||||
repo_id=shard.model_card.model_id,
|
||||
repo_revision=revision,
|
||||
@@ -832,7 +837,9 @@ async def download_shard(
|
||||
total=Memory.from_bytes(file.size or 0),
|
||||
speed=0,
|
||||
eta=timedelta(0),
|
||||
status="complete" if downloaded_bytes == file.size else "not_started",
|
||||
status="complete"
|
||||
if final_file_exists and downloaded_bytes == file.size
|
||||
else "not_started",
|
||||
start_time=time.time(),
|
||||
)
|
||||
|
||||
|
||||
@@ -15,9 +15,11 @@ from exo.shared.types.worker.shards import (
|
||||
)
|
||||
|
||||
|
||||
def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
|
||||
def exo_shard_downloader(
|
||||
max_parallel_downloads: int = 8, offline: bool = False
|
||||
) -> ShardDownloader:
|
||||
return SingletonShardDownloader(
|
||||
CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
|
||||
ResumableShardDownloader(max_parallel_downloads, offline=offline)
|
||||
)
|
||||
|
||||
|
||||
@@ -50,10 +52,6 @@ class SingletonShardDownloader(ShardDownloader):
|
||||
self.shard_downloader = shard_downloader
|
||||
self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
|
||||
|
||||
def set_internet_connection(self, value: bool) -> None:
|
||||
self.internet_connection = value
|
||||
self.shard_downloader.set_internet_connection(value)
|
||||
|
||||
def on_progress(
|
||||
self,
|
||||
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
|
||||
@@ -85,46 +83,10 @@ class SingletonShardDownloader(ShardDownloader):
|
||||
return await self.shard_downloader.get_shard_download_status_for_shard(shard)
|
||||
|
||||
|
||||
class CachedShardDownloader(ShardDownloader):
|
||||
def __init__(self, shard_downloader: ShardDownloader):
|
||||
self.shard_downloader = shard_downloader
|
||||
self.cache: dict[tuple[str, ShardMetadata], Path] = {}
|
||||
|
||||
def set_internet_connection(self, value: bool) -> None:
|
||||
self.internet_connection = value
|
||||
self.shard_downloader.set_internet_connection(value)
|
||||
|
||||
def on_progress(
|
||||
self,
|
||||
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
|
||||
) -> None:
|
||||
self.shard_downloader.on_progress(callback)
|
||||
|
||||
async def ensure_shard(
|
||||
self, shard: ShardMetadata, config_only: bool = False
|
||||
) -> Path:
|
||||
if (shard.model_card.model_id, shard) in self.cache:
|
||||
return self.cache[(shard.model_card.model_id, shard)]
|
||||
|
||||
target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
|
||||
self.cache[(shard.model_card.model_id, shard)] = target_dir
|
||||
return target_dir
|
||||
|
||||
async def get_shard_download_status(
|
||||
self,
|
||||
) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
|
||||
async for path, status in self.shard_downloader.get_shard_download_status():
|
||||
yield path, status
|
||||
|
||||
async def get_shard_download_status_for_shard(
|
||||
self, shard: ShardMetadata
|
||||
) -> RepoDownloadProgress:
|
||||
return await self.shard_downloader.get_shard_download_status_for_shard(shard)
|
||||
|
||||
|
||||
class ResumableShardDownloader(ShardDownloader):
|
||||
def __init__(self, max_parallel_downloads: int = 8):
|
||||
def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
|
||||
self.max_parallel_downloads = max_parallel_downloads
|
||||
self.offline = offline
|
||||
self.on_progress_callbacks: list[
|
||||
Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
|
||||
] = []
|
||||
@@ -151,8 +113,7 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
self.on_progress_wrapper,
|
||||
max_parallel_downloads=self.max_parallel_downloads,
|
||||
allow_patterns=allow_patterns,
|
||||
skip_internet=not self.internet_connection,
|
||||
on_connection_lost=lambda: self.set_internet_connection(False),
|
||||
skip_internet=self.offline,
|
||||
)
|
||||
return target_dir
|
||||
|
||||
@@ -168,8 +129,7 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
shard,
|
||||
self.on_progress_wrapper,
|
||||
skip_download=True,
|
||||
skip_internet=not self.internet_connection,
|
||||
on_connection_lost=lambda: self.set_internet_connection(False),
|
||||
skip_internet=self.offline,
|
||||
)
|
||||
|
||||
semaphore = asyncio.Semaphore(self.max_parallel_downloads)
|
||||
@@ -198,7 +158,6 @@ class ResumableShardDownloader(ShardDownloader):
|
||||
shard,
|
||||
self.on_progress_wrapper,
|
||||
skip_download=True,
|
||||
skip_internet=not self.internet_connection,
|
||||
on_connection_lost=lambda: self.set_internet_connection(False),
|
||||
skip_internet=self.offline,
|
||||
)
|
||||
return progress
|
||||
|
||||
@@ -16,11 +16,6 @@ from exo.shared.types.worker.shards import (
|
||||
|
||||
# TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
|
||||
class ShardDownloader(ABC):
|
||||
internet_connection: bool = False
|
||||
|
||||
def set_internet_connection(self, value: bool) -> None:
|
||||
self.internet_connection = value
|
||||
|
||||
@abstractmethod
|
||||
async def ensure_shard(
|
||||
self, shard: ShardMetadata, config_only: bool = False
|
||||
|
||||
@@ -1,98 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import pytest
|
||||
|
||||
from exo.download.coordinator import DownloadCoordinator
|
||||
from exo.download.shard_downloader import NoopShardDownloader
|
||||
from exo.shared.models.model_cards import ModelCard, ModelTask
|
||||
from exo.shared.types.common import ModelId, NodeId, SessionId
|
||||
from exo.shared.types.events import (
|
||||
GlobalForwarderEvent,
|
||||
LocalForwarderEvent,
|
||||
NodeDownloadProgress,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.downloads import (
|
||||
DownloadPending,
|
||||
)
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata
|
||||
from exo.utils.channels import channel
|
||||
|
||||
# Use the built‑in NoopShardDownloader directly – it already implements the required abstract interface.
|
||||
# No additional subclass is needed for this test.
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_ack_behaviour():
|
||||
# Create channels (type Any for simplicity)
|
||||
_, command_receiver = channel[Any]()
|
||||
local_sender, _ = channel[Any]()
|
||||
global_sender, global_receiver = channel[Any]()
|
||||
|
||||
# Minimal identifiers
|
||||
node_id = NodeId()
|
||||
session_id = SessionId(master_node_id=node_id, election_clock=0)
|
||||
|
||||
# Create a dummy model card and shard metadata
|
||||
model_id = ModelId("test/model")
|
||||
model_card = ModelCard(
|
||||
model_id=model_id,
|
||||
storage_size=Memory.from_bytes(0),
|
||||
n_layers=1,
|
||||
hidden_size=1,
|
||||
supports_tensor=True,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
)
|
||||
shard = PipelineShardMetadata(
|
||||
model_card=model_card,
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
start_layer=0,
|
||||
end_layer=1,
|
||||
n_layers=1,
|
||||
)
|
||||
|
||||
# Instantiate the coordinator with the dummy downloader
|
||||
coord = DownloadCoordinator(
|
||||
node_id=node_id,
|
||||
session_id=session_id,
|
||||
shard_downloader=NoopShardDownloader(),
|
||||
download_command_receiver=command_receiver,
|
||||
local_event_sender=local_sender,
|
||||
_global_event_receiver=global_receiver,
|
||||
)
|
||||
|
||||
async with anyio.create_task_group() as tg:
|
||||
# Start the forwarding and ack‑clearing loops
|
||||
tg.start_soon(coord._forward_events) # pyright: ignore[reportPrivateUsage]
|
||||
tg.start_soon(coord._clear_ofd) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
# Send a pending download progress event via the internal event sender
|
||||
pending = DownloadPending(
|
||||
node_id=node_id,
|
||||
shard_metadata=shard,
|
||||
model_directory="/tmp/model",
|
||||
)
|
||||
await coord.event_sender.send(NodeDownloadProgress(download_progress=pending))
|
||||
# Allow the forwarder to process the event
|
||||
await anyio.sleep(0.1)
|
||||
|
||||
# There should be exactly one entry awaiting ACK
|
||||
assert len(coord._out_for_delivery) == 1 # pyright: ignore[reportPrivateUsage]
|
||||
# Retrieve the stored LocalForwarderEvent
|
||||
stored_fe: LocalForwarderEvent = next(iter(coord._out_for_delivery.values())) # pyright: ignore[reportPrivateUsage]
|
||||
# Simulate receiving a global ack for this event
|
||||
ack = GlobalForwarderEvent(
|
||||
origin_idx=0,
|
||||
origin=node_id,
|
||||
session=session_id,
|
||||
event=stored_fe.event,
|
||||
)
|
||||
await global_sender.send(ack)
|
||||
# Give the clear‑ofd task a moment to process the ack
|
||||
await anyio.sleep(0.1)
|
||||
# The out‑for‑delivery map should now be empty
|
||||
assert len(coord._out_for_delivery) == 0 # pyright: ignore[reportPrivateUsage]
|
||||
# Cancel background tasks
|
||||
tg.cancel_scope.cancel()
|
||||
@@ -0,0 +1,211 @@
|
||||
"""Tests that re-downloading a previously deleted model completes successfully."""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from collections.abc import AsyncIterator, Awaitable
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
from exo.download.coordinator import DownloadCoordinator
|
||||
from exo.download.download_utils import RepoDownloadProgress
|
||||
from exo.download.impl_shard_downloader import SingletonShardDownloader
|
||||
from exo.download.shard_downloader import ShardDownloader
|
||||
from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
|
||||
from exo.shared.types.commands import (
|
||||
DeleteDownload,
|
||||
ForwarderDownloadCommand,
|
||||
StartDownload,
|
||||
)
|
||||
from exo.shared.types.common import NodeId, SystemId
|
||||
from exo.shared.types.events import Event, NodeDownloadProgress
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.worker.downloads import DownloadCompleted
|
||||
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
|
||||
NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
|
||||
MODEL_ID = ModelId("test-org/test-model")
|
||||
|
||||
|
||||
def _make_shard(model_id: ModelId = MODEL_ID) -> ShardMetadata:
|
||||
return PipelineShardMetadata(
|
||||
model_card=ModelCard(
|
||||
model_id=model_id,
|
||||
storage_size=Memory.from_mb(100),
|
||||
n_layers=28,
|
||||
hidden_size=1024,
|
||||
supports_tensor=False,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
),
|
||||
device_rank=0,
|
||||
world_size=1,
|
||||
start_layer=0,
|
||||
end_layer=28,
|
||||
n_layers=28,
|
||||
)
|
||||
|
||||
|
||||
class FakeShardDownloader(ShardDownloader):
|
||||
"""Fake downloader that simulates a successful download by firing the
|
||||
progress callback with status='complete' when ensure_shard is called."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._progress_callbacks: list[
|
||||
Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
|
||||
] = []
|
||||
|
||||
def on_progress(
|
||||
self,
|
||||
callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
|
||||
) -> None:
|
||||
self._progress_callbacks.append(callback)
|
||||
|
||||
async def ensure_shard(
|
||||
self,
|
||||
shard: ShardMetadata,
|
||||
config_only: bool = False, # noqa: ARG002
|
||||
) -> Path:
|
||||
# Simulate a completed download by firing the progress callback
|
||||
progress = RepoDownloadProgress(
|
||||
repo_id=str(shard.model_card.model_id),
|
||||
repo_revision="main",
|
||||
shard=shard,
|
||||
completed_files=1,
|
||||
total_files=1,
|
||||
downloaded=Memory.from_mb(100),
|
||||
downloaded_this_session=Memory.from_mb(100),
|
||||
total=Memory.from_mb(100),
|
||||
overall_speed=0,
|
||||
overall_eta=timedelta(seconds=0),
|
||||
status="complete",
|
||||
)
|
||||
for cb in self._progress_callbacks:
|
||||
await cb(shard, progress)
|
||||
return Path("/fake/models") / shard.model_card.model_id.normalize()
|
||||
|
||||
async def get_shard_download_status(
|
||||
self,
|
||||
) -> AsyncIterator[tuple[Path, RepoDownloadProgress]]:
|
||||
if False: # noqa: SIM108 # empty async generator
|
||||
yield (
|
||||
Path(),
|
||||
RepoDownloadProgress( # pyright: ignore[reportUnreachable]
|
||||
repo_id="",
|
||||
repo_revision="",
|
||||
shard=_make_shard(),
|
||||
completed_files=0,
|
||||
total_files=0,
|
||||
downloaded=Memory.from_bytes(0),
|
||||
downloaded_this_session=Memory.from_bytes(0),
|
||||
total=Memory.from_bytes(0),
|
||||
overall_speed=0,
|
||||
overall_eta=timedelta(seconds=0),
|
||||
status="not_started",
|
||||
),
|
||||
)
|
||||
|
||||
async def get_shard_download_status_for_shard(
|
||||
self,
|
||||
shard: ShardMetadata,
|
||||
) -> RepoDownloadProgress:
|
||||
return RepoDownloadProgress(
|
||||
repo_id=str(shard.model_card.model_id),
|
||||
repo_revision="main",
|
||||
shard=shard,
|
||||
completed_files=0,
|
||||
total_files=1,
|
||||
downloaded=Memory.from_bytes(0),
|
||||
downloaded_this_session=Memory.from_bytes(0),
|
||||
total=Memory.from_mb(100),
|
||||
overall_speed=0,
|
||||
overall_eta=timedelta(seconds=0),
|
||||
status="not_started",
|
||||
)
|
||||
|
||||
|
||||
async def test_re_download_after_delete_completes() -> None:
|
||||
"""A model that was downloaded, deleted, and then re-downloaded should
|
||||
reach DownloadCompleted status. This is an end-to-end test through
|
||||
the DownloadCoordinator."""
|
||||
cmd_send: Sender[ForwarderDownloadCommand]
|
||||
cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
|
||||
event_send, event_recv = channel[Event]()
|
||||
|
||||
fake_downloader = FakeShardDownloader()
|
||||
wrapped_downloader = SingletonShardDownloader(fake_downloader)
|
||||
coordinator = DownloadCoordinator(
|
||||
node_id=NODE_ID,
|
||||
shard_downloader=wrapped_downloader,
|
||||
download_command_receiver=cmd_recv,
|
||||
event_sender=event_send,
|
||||
)
|
||||
|
||||
shard = _make_shard()
|
||||
origin = SystemId("test")
|
||||
|
||||
with patch("exo.download.coordinator.delete_model", new_callable=AsyncMock):
|
||||
# Run the coordinator in the background
|
||||
coordinator_task = asyncio.create_task(coordinator.run())
|
||||
|
||||
try:
|
||||
# 1. Start first download
|
||||
await cmd_send.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=origin,
|
||||
command=StartDownload(target_node_id=NODE_ID, shard_metadata=shard),
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for DownloadCompleted
|
||||
first_completed = await _wait_for_download_completed(event_recv, MODEL_ID)
|
||||
assert first_completed is not None, "First download should complete"
|
||||
|
||||
# 2. Delete the model
|
||||
await cmd_send.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=origin,
|
||||
command=DeleteDownload(target_node_id=NODE_ID, model_id=MODEL_ID),
|
||||
)
|
||||
)
|
||||
# Give the coordinator time to process the delete
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
# 3. Re-download the same model
|
||||
await cmd_send.send(
|
||||
ForwarderDownloadCommand(
|
||||
origin=origin,
|
||||
command=StartDownload(target_node_id=NODE_ID, shard_metadata=shard),
|
||||
)
|
||||
)
|
||||
|
||||
# Wait for second DownloadCompleted — this is the bug: it never arrives
|
||||
second_completed = await _wait_for_download_completed(event_recv, MODEL_ID)
|
||||
assert second_completed is not None, (
|
||||
"Re-download after deletion should complete"
|
||||
)
|
||||
finally:
|
||||
coordinator.shutdown()
|
||||
coordinator_task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await coordinator_task
|
||||
|
||||
|
||||
async def _wait_for_download_completed(
|
||||
event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
|
||||
) -> DownloadCompleted | None:
|
||||
"""Drain events until we see a DownloadCompleted for the given model, or timeout."""
|
||||
try:
|
||||
async with asyncio.timeout(timeout):
|
||||
while True:
|
||||
event = await event_recv.receive()
|
||||
if (
|
||||
isinstance(event, NodeDownloadProgress)
|
||||
and isinstance(event.download_progress, DownloadCompleted)
|
||||
and event.download_progress.shard_metadata.model_card.model_id
|
||||
== model_id
|
||||
):
|
||||
return event.download_progress
|
||||
except TimeoutError:
|
||||
return None
|
||||
+60
-33
@@ -7,7 +7,6 @@ from dataclasses import dataclass, field
|
||||
from typing import Self
|
||||
|
||||
import anyio
|
||||
from anyio.abc import TaskGroup
|
||||
from loguru import logger
|
||||
from pydantic import PositiveInt
|
||||
|
||||
@@ -16,6 +15,7 @@ from exo.download.coordinator import DownloadCoordinator
|
||||
from exo.download.impl_shard_downloader import exo_shard_downloader
|
||||
from exo.master.api import API # TODO: should API be in master?
|
||||
from exo.master.main import Master
|
||||
from exo.routing.event_router import EventRouter
|
||||
from exo.routing.router import Router, get_node_id_keypair
|
||||
from exo.shared.constants import EXO_LOG
|
||||
from exo.shared.election import Election, ElectionResult
|
||||
@@ -23,12 +23,14 @@ from exo.shared.logging import logger_cleanup, logger_setup
|
||||
from exo.shared.types.common import NodeId, SessionId
|
||||
from exo.utils.channels import Receiver, channel
|
||||
from exo.utils.pydantic_ext import CamelCaseModel
|
||||
from exo.utils.task_group import TaskGroup
|
||||
from exo.worker.main import Worker
|
||||
|
||||
|
||||
@dataclass
|
||||
class Node:
|
||||
router: Router
|
||||
event_router: EventRouter
|
||||
download_coordinator: DownloadCoordinator | None
|
||||
worker: Worker | None
|
||||
election: Election # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
|
||||
@@ -38,7 +40,7 @@ class Node:
|
||||
|
||||
node_id: NodeId
|
||||
offline: bool
|
||||
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
@classmethod
|
||||
async def create(cls, args: "Args") -> Self:
|
||||
@@ -52,6 +54,12 @@ class Node:
|
||||
await router.register_topic(topics.ELECTION_MESSAGES)
|
||||
await router.register_topic(topics.CONNECTION_MESSAGES)
|
||||
await router.register_topic(topics.DOWNLOAD_COMMANDS)
|
||||
event_router = EventRouter(
|
||||
session_id,
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
external_outbound=router.sender(topics.LOCAL_EVENTS),
|
||||
external_inbound=router.receiver(topics.GLOBAL_EVENTS),
|
||||
)
|
||||
|
||||
logger.info(f"Starting node {node_id}")
|
||||
|
||||
@@ -59,13 +67,10 @@ class Node:
|
||||
if not args.no_downloads:
|
||||
download_coordinator = DownloadCoordinator(
|
||||
node_id,
|
||||
session_id,
|
||||
exo_shard_downloader(),
|
||||
exo_shard_downloader(offline=args.offline),
|
||||
event_sender=event_router.sender(),
|
||||
download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
|
||||
local_event_sender=router.sender(topics.LOCAL_EVENTS),
|
||||
offline=args.offline,
|
||||
# TODO(evan): remove
|
||||
_global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
|
||||
)
|
||||
else:
|
||||
download_coordinator = None
|
||||
@@ -73,9 +78,8 @@ class Node:
|
||||
if args.spawn_api:
|
||||
api = API(
|
||||
node_id,
|
||||
session_id,
|
||||
port=args.api_port,
|
||||
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
|
||||
event_receiver=event_router.receiver(),
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
election_receiver=router.receiver(topics.ELECTION_MESSAGES),
|
||||
@@ -86,9 +90,8 @@ class Node:
|
||||
if not args.no_worker:
|
||||
worker = Worker(
|
||||
node_id,
|
||||
session_id,
|
||||
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
|
||||
local_event_sender=router.sender(topics.LOCAL_EVENTS),
|
||||
event_receiver=event_router.receiver(),
|
||||
event_sender=event_router.sender(),
|
||||
command_sender=router.sender(topics.COMMANDS),
|
||||
download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
|
||||
)
|
||||
@@ -99,6 +102,7 @@ class Node:
|
||||
master = Master(
|
||||
node_id,
|
||||
session_id,
|
||||
event_sender=event_router.sender(),
|
||||
global_event_sender=router.sender(topics.GLOBAL_EVENTS),
|
||||
local_event_receiver=router.receiver(topics.LOCAL_EVENTS),
|
||||
command_receiver=router.receiver(topics.COMMANDS),
|
||||
@@ -121,6 +125,7 @@ class Node:
|
||||
|
||||
return cls(
|
||||
router,
|
||||
event_router,
|
||||
download_coordinator,
|
||||
worker,
|
||||
election,
|
||||
@@ -136,6 +141,7 @@ class Node:
|
||||
signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
|
||||
signal.signal(signal.SIGTERM, lambda _, __: self.shutdown())
|
||||
tg.start_soon(self.router.run)
|
||||
tg.start_soon(self.event_router.run)
|
||||
tg.start_soon(self.election.run)
|
||||
if self.download_coordinator:
|
||||
tg.start_soon(self.download_coordinator.run)
|
||||
@@ -149,11 +155,11 @@ class Node:
|
||||
|
||||
def shutdown(self):
|
||||
# if this is our second call to shutdown, just sys.exit
|
||||
if self._tg.cancel_scope.cancel_called:
|
||||
if self._tg.cancel_called():
|
||||
import sys
|
||||
|
||||
sys.exit(1)
|
||||
self._tg.cancel_scope.cancel()
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _elect_loop(self):
|
||||
with self.election_result_receiver as results:
|
||||
@@ -170,6 +176,17 @@ class Node:
|
||||
# - Shutdown and re-create the worker
|
||||
# - Shut down and re-create the API
|
||||
|
||||
if result.is_new_master:
|
||||
await anyio.sleep(0)
|
||||
self.event_router.shutdown()
|
||||
self.event_router = EventRouter(
|
||||
result.session_id,
|
||||
self.router.sender(topics.COMMANDS),
|
||||
self.router.receiver(topics.GLOBAL_EVENTS),
|
||||
self.router.sender(topics.LOCAL_EVENTS),
|
||||
)
|
||||
self._tg.start_soon(self.event_router.run)
|
||||
|
||||
if (
|
||||
result.session_id.master_node_id == self.node_id
|
||||
and self.master is not None
|
||||
@@ -183,6 +200,7 @@ class Node:
|
||||
self.master = Master(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
event_sender=self.event_router.sender(),
|
||||
global_event_sender=self.router.sender(topics.GLOBAL_EVENTS),
|
||||
local_event_receiver=self.router.receiver(topics.LOCAL_EVENTS),
|
||||
command_receiver=self.router.receiver(topics.COMMANDS),
|
||||
@@ -205,22 +223,16 @@ class Node:
|
||||
f"Node {result.session_id.master_node_id} elected master"
|
||||
)
|
||||
if result.is_new_master:
|
||||
await anyio.sleep(0)
|
||||
if self.download_coordinator:
|
||||
self.download_coordinator.shutdown()
|
||||
self.download_coordinator = DownloadCoordinator(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
exo_shard_downloader(),
|
||||
exo_shard_downloader(offline=self.offline),
|
||||
event_sender=self.event_router.sender(),
|
||||
download_command_receiver=self.router.receiver(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
),
|
||||
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
|
||||
offline=self.offline,
|
||||
# TODO(evan): remove
|
||||
_global_event_receiver=self.router.receiver(
|
||||
topics.GLOBAL_EVENTS
|
||||
),
|
||||
)
|
||||
self._tg.start_soon(self.download_coordinator.run)
|
||||
if self.worker:
|
||||
@@ -228,11 +240,8 @@ class Node:
|
||||
# TODO: add profiling etc to resource monitor
|
||||
self.worker = Worker(
|
||||
self.node_id,
|
||||
result.session_id,
|
||||
global_event_receiver=self.router.receiver(
|
||||
topics.GLOBAL_EVENTS
|
||||
),
|
||||
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
|
||||
event_receiver=self.event_router.receiver(),
|
||||
event_sender=self.event_router.sender(),
|
||||
command_sender=self.router.sender(topics.COMMANDS),
|
||||
download_command_sender=self.router.sender(
|
||||
topics.DOWNLOAD_COMMANDS
|
||||
@@ -240,7 +249,7 @@ class Node:
|
||||
)
|
||||
self._tg.start_soon(self.worker.run)
|
||||
if self.api:
|
||||
self.api.reset(result.session_id, result.won_clock)
|
||||
self.api.reset(result.won_clock, self.event_router.receiver())
|
||||
else:
|
||||
if self.api:
|
||||
self.api.unpause(result.won_clock)
|
||||
@@ -252,7 +261,7 @@ def main():
|
||||
target = min(max(soft, 65535), hard)
|
||||
resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
|
||||
|
||||
mp.set_start_method("spawn")
|
||||
mp.set_start_method("spawn", force=True)
|
||||
# TODO: Refactor the current verbosity system
|
||||
logger_setup(EXO_LOG, args.verbosity)
|
||||
logger.info("Starting EXO")
|
||||
@@ -261,6 +270,10 @@ def main():
|
||||
if args.offline:
|
||||
logger.info("Running in OFFLINE mode — no internet checks, local models only")
|
||||
|
||||
if args.no_batch:
|
||||
os.environ["EXO_NO_BATCH"] = "1"
|
||||
logger.info("Continuous batching disabled (--no-batch)")
|
||||
|
||||
# Set FAST_SYNCH override env var for runner subprocesses
|
||||
if args.fast_synch is True:
|
||||
os.environ["EXO_FAST_SYNCH"] = "on"
|
||||
@@ -270,9 +283,16 @@ def main():
|
||||
logger.info("FAST_SYNCH forced OFF")
|
||||
|
||||
node = anyio.run(Node.create, args)
|
||||
anyio.run(node.run)
|
||||
logger.info("EXO Shutdown complete")
|
||||
logger_cleanup()
|
||||
try:
|
||||
anyio.run(node.run)
|
||||
except BaseException as exception:
|
||||
logger.opt(exception=exception).critical(
|
||||
"EXO terminated due to unhandled exception"
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
logger.info("EXO Shutdown complete")
|
||||
logger_cleanup()
|
||||
|
||||
|
||||
class Args(CamelCaseModel):
|
||||
@@ -283,7 +303,8 @@ class Args(CamelCaseModel):
|
||||
tb_only: bool = False
|
||||
no_worker: bool = False
|
||||
no_downloads: bool = False
|
||||
offline: bool = False
|
||||
offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
|
||||
no_batch: bool = False
|
||||
fast_synch: bool | None = None # None = auto, True = force on, False = force off
|
||||
|
||||
@classmethod
|
||||
@@ -334,8 +355,14 @@ class Args(CamelCaseModel):
|
||||
parser.add_argument(
|
||||
"--offline",
|
||||
action="store_true",
|
||||
default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
|
||||
help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-batch",
|
||||
action="store_true",
|
||||
help="Disable continuous batching, use sequential generation",
|
||||
)
|
||||
fast_synch_group = parser.add_mutually_exclusive_group()
|
||||
fast_synch_group.add_argument(
|
||||
"--fast-synch",
|
||||
|
||||
@@ -26,7 +26,11 @@ from exo.shared.types.chunks import (
|
||||
ToolCallChunk,
|
||||
)
|
||||
from exo.shared.types.common import CommandId
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
TextGenerationTaskParams,
|
||||
resolve_reasoning_params,
|
||||
)
|
||||
|
||||
|
||||
def chat_request_to_text_generation(
|
||||
@@ -75,6 +79,10 @@ def chat_request_to_text_generation(
|
||||
dumped: dict[str, Any] = msg_copy.model_dump(exclude_none=True)
|
||||
chat_template_messages.append(dumped)
|
||||
|
||||
resolved_effort, resolved_thinking = resolve_reasoning_params(
|
||||
request.reasoning_effort, request.enable_thinking
|
||||
)
|
||||
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=input_messages
|
||||
@@ -89,12 +97,16 @@ def chat_request_to_text_generation(
|
||||
seed=request.seed,
|
||||
stream=request.stream,
|
||||
tools=request.tools,
|
||||
enable_thinking=request.enable_thinking,
|
||||
reasoning_effort=resolved_effort,
|
||||
enable_thinking=resolved_thinking,
|
||||
chat_template_messages=chat_template_messages
|
||||
if chat_template_messages
|
||||
else None,
|
||||
logprobs=request.logprobs or False,
|
||||
top_logprobs=request.top_logprobs,
|
||||
min_p=request.min_p,
|
||||
repetition_penalty=request.repetition_penalty,
|
||||
repetition_context_size=request.repetition_context_size,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -42,7 +42,11 @@ from exo.shared.types.openai_responses import (
|
||||
ResponseTextDoneEvent,
|
||||
ResponseUsage,
|
||||
)
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
from exo.shared.types.text_generation import (
|
||||
InputMessage,
|
||||
TextGenerationTaskParams,
|
||||
resolve_reasoning_params,
|
||||
)
|
||||
|
||||
|
||||
def _format_sse(event: ResponsesStreamEvent) -> str:
|
||||
@@ -119,6 +123,11 @@ def responses_request_to_text_generation(
|
||||
)
|
||||
built_chat_template = chat_template_messages if chat_template_messages else None
|
||||
|
||||
effort_from_reasoning = request.reasoning.effort if request.reasoning else None
|
||||
resolved_effort, resolved_thinking = resolve_reasoning_params(
|
||||
effort_from_reasoning, request.enable_thinking
|
||||
)
|
||||
|
||||
return TextGenerationTaskParams(
|
||||
model=request.model,
|
||||
input=input_value,
|
||||
@@ -132,6 +141,8 @@ def responses_request_to_text_generation(
|
||||
stop=request.stop,
|
||||
seed=request.seed,
|
||||
chat_template_messages=built_chat_template or request.chat_template_messages,
|
||||
reasoning_effort=resolved_effort,
|
||||
enable_thinking=resolved_thinking,
|
||||
)
|
||||
|
||||
|
||||
|
||||
+126
-70
@@ -3,7 +3,7 @@ import contextlib
|
||||
import json
|
||||
import random
|
||||
import time
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterator
|
||||
from collections.abc import AsyncGenerator, Awaitable, Callable, Iterable
|
||||
from datetime import datetime, timezone
|
||||
from http import HTTPStatus
|
||||
from pathlib import Path
|
||||
@@ -11,8 +11,7 @@ from typing import Annotated, Literal, cast
|
||||
from uuid import uuid4
|
||||
|
||||
import anyio
|
||||
from anyio import BrokenResourceError, create_task_group
|
||||
from anyio.abc import TaskGroup
|
||||
from anyio import BrokenResourceError
|
||||
from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
|
||||
@@ -51,6 +50,7 @@ from exo.master.placement import place_instance as get_instance_placements
|
||||
from exo.shared.apply import apply
|
||||
from exo.shared.constants import (
|
||||
DASHBOARD_DIR,
|
||||
EXO_CACHE_HOME,
|
||||
EXO_EVENT_LOG_DIR,
|
||||
EXO_IMAGE_CACHE_DIR,
|
||||
EXO_MAX_CHUNK_SIZE,
|
||||
@@ -73,6 +73,7 @@ from exo.shared.types.api import (
|
||||
BenchChatCompletionResponse,
|
||||
BenchImageGenerationResponse,
|
||||
BenchImageGenerationTaskParams,
|
||||
CancelCommandResponse,
|
||||
ChatCompletionChoice,
|
||||
ChatCompletionMessage,
|
||||
ChatCompletionRequest,
|
||||
@@ -81,6 +82,8 @@ from exo.shared.types.api import (
|
||||
CreateInstanceResponse,
|
||||
DeleteDownloadResponse,
|
||||
DeleteInstanceResponse,
|
||||
DeleteTracesRequest,
|
||||
DeleteTracesResponse,
|
||||
ErrorInfo,
|
||||
ErrorResponse,
|
||||
FinishReason,
|
||||
@@ -140,11 +143,10 @@ from exo.shared.types.commands import (
|
||||
TaskFinished,
|
||||
TextGeneration,
|
||||
)
|
||||
from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.common import CommandId, Id, NodeId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
ChunkGenerated,
|
||||
Event,
|
||||
GlobalForwarderEvent,
|
||||
IndexedEvent,
|
||||
TracesMerged,
|
||||
)
|
||||
@@ -172,9 +174,10 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
|
||||
from exo.shared.types.worker.shards import Sharding
|
||||
from exo.utils.banner import print_startup_banner
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.event_buffer import OrderedBuffer
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
|
||||
ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
|
||||
|
||||
|
||||
def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
|
||||
@@ -194,10 +197,9 @@ class API:
|
||||
def __init__(
|
||||
self,
|
||||
node_id: NodeId,
|
||||
session_id: SessionId,
|
||||
*,
|
||||
port: int,
|
||||
global_event_receiver: Receiver[GlobalForwarderEvent],
|
||||
event_receiver: Receiver[IndexedEvent],
|
||||
command_sender: Sender[ForwarderCommand],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
# This lets us pause the API if an election is running
|
||||
@@ -208,11 +210,9 @@ class API:
|
||||
self._system_id = SystemId()
|
||||
self.command_sender = command_sender
|
||||
self.download_command_sender = download_command_sender
|
||||
self.global_event_receiver = global_event_receiver
|
||||
self.event_receiver = event_receiver
|
||||
self.election_receiver = election_receiver
|
||||
self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
|
||||
self.node_id: NodeId = node_id
|
||||
self.session_id: SessionId = session_id
|
||||
self.last_completed_election: int = 0
|
||||
self.port = port
|
||||
|
||||
@@ -250,19 +250,20 @@ class API:
|
||||
CommandId, Sender[ImageChunk | ErrorChunk]
|
||||
] = {}
|
||||
self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
|
||||
self._tg: TaskGroup | None = None
|
||||
self._tg: TaskGroup = TaskGroup()
|
||||
|
||||
def reset(self, new_session_id: SessionId, result_clock: int):
|
||||
def reset(self, result_clock: int, event_receiver: Receiver[IndexedEvent]):
|
||||
logger.info("Resetting API State")
|
||||
self._event_log.close()
|
||||
self._event_log = DiskEventLog(_API_EVENT_LOG_DIR)
|
||||
self.state = State()
|
||||
self._system_id = SystemId()
|
||||
self.session_id = new_session_id
|
||||
self.event_buffer = OrderedBuffer[Event]()
|
||||
self._text_generation_queues = {}
|
||||
self._image_generation_queues = {}
|
||||
self.unpause(result_clock)
|
||||
self.event_receiver.close()
|
||||
self.event_receiver = event_receiver
|
||||
self._tg.start_soon(self._apply_state)
|
||||
|
||||
def unpause(self, result_clock: int):
|
||||
logger.info("Unpausing API")
|
||||
@@ -322,6 +323,7 @@ class API:
|
||||
self.app.get("/images/{image_id}")(self.get_image)
|
||||
self.app.post("/v1/messages", response_model=None)(self.claude_messages)
|
||||
self.app.post("/v1/responses", response_model=None)(self.openai_responses)
|
||||
self.app.post("/v1/cancel/{command_id}")(self.cancel_command)
|
||||
|
||||
# Ollama API
|
||||
self.app.head("/ollama/")(self.ollama_version)
|
||||
@@ -342,9 +344,12 @@ class API:
|
||||
self.app.post("/download/start")(self.start_download)
|
||||
self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
|
||||
self.app.get("/v1/traces")(self.list_traces)
|
||||
self.app.post("/v1/traces/delete")(self.delete_traces)
|
||||
self.app.get("/v1/traces/{task_id}")(self.get_trace)
|
||||
self.app.get("/v1/traces/{task_id}/stats")(self.get_trace_stats)
|
||||
self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
|
||||
self.app.get("/onboarding")(self.get_onboarding)
|
||||
self.app.post("/onboarding")(self.complete_onboarding)
|
||||
|
||||
async def place_instance(self, payload: PlaceInstanceParams):
|
||||
command = PlaceInstance(
|
||||
@@ -564,6 +569,25 @@ class API:
|
||||
instance_id=instance_id,
|
||||
)
|
||||
|
||||
async def cancel_command(self, command_id: CommandId) -> CancelCommandResponse:
|
||||
"""Cancel an active command by closing its stream and notifying workers."""
|
||||
sender = self._text_generation_queues.get(
|
||||
command_id
|
||||
) or self._image_generation_queues.get(command_id)
|
||||
if sender is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Command not found or already completed",
|
||||
)
|
||||
|
||||
await self._send(TaskCancelled(cancelled_command_id=command_id))
|
||||
sender.close()
|
||||
|
||||
return CancelCommandResponse(
|
||||
message="Command cancelled.",
|
||||
command_id=command_id,
|
||||
)
|
||||
|
||||
async def _token_chunk_stream(
|
||||
self, command_id: CommandId
|
||||
) -> AsyncGenerator[
|
||||
@@ -751,7 +775,7 @@ class API:
|
||||
return resolved_model
|
||||
|
||||
def stream_events(self) -> StreamingResponse:
|
||||
def _generate_json_array(events: Iterator[Event]) -> Iterator[str]:
|
||||
def _generate_json_array(events: Iterable[Event]) -> Iterable[str]:
|
||||
yield "["
|
||||
first = True
|
||||
for event in events:
|
||||
@@ -1562,33 +1586,48 @@ class API:
|
||||
async def search_models(
|
||||
self, query: str = "", limit: int = 20
|
||||
) -> list[HuggingFaceSearchResult]:
|
||||
"""Search HuggingFace Hub for mlx-community models."""
|
||||
from huggingface_hub import list_models
|
||||
"""Search HuggingFace Hub — tries mlx-community first, falls back to all of HuggingFace."""
|
||||
from huggingface_hub import ModelInfo, list_models
|
||||
|
||||
results = list_models(
|
||||
search=query or None,
|
||||
author="mlx-community",
|
||||
sort="downloads",
|
||||
limit=limit,
|
||||
)
|
||||
return [
|
||||
HuggingFaceSearchResult(
|
||||
id=m.id,
|
||||
author=m.author or "",
|
||||
downloads=m.downloads or 0,
|
||||
likes=m.likes or 0,
|
||||
last_modified=str(m.last_modified or ""),
|
||||
tags=list(m.tags or []),
|
||||
def _to_results(models: Iterable[ModelInfo]) -> list[HuggingFaceSearchResult]:
|
||||
return [
|
||||
HuggingFaceSearchResult(
|
||||
id=m.id,
|
||||
author=m.author or "",
|
||||
downloads=m.downloads or 0,
|
||||
likes=m.likes or 0,
|
||||
last_modified=str(m.last_modified or ""),
|
||||
tags=list(m.tags or []),
|
||||
)
|
||||
for m in models
|
||||
]
|
||||
|
||||
# Search mlx-community first
|
||||
mlx_results = _to_results(
|
||||
list_models(
|
||||
search=query or None,
|
||||
author="mlx-community",
|
||||
sort="downloads",
|
||||
limit=limit,
|
||||
)
|
||||
for m in results
|
||||
]
|
||||
)
|
||||
if mlx_results:
|
||||
return mlx_results
|
||||
|
||||
# Fall back to searching all of HuggingFace
|
||||
return _to_results(
|
||||
list_models(
|
||||
search=query or None,
|
||||
sort="downloads",
|
||||
limit=limit,
|
||||
)
|
||||
)
|
||||
|
||||
async def run(self):
|
||||
shutdown_ev = anyio.Event()
|
||||
|
||||
try:
|
||||
async with create_task_group() as tg:
|
||||
self._tg = tg
|
||||
async with self._tg as tg:
|
||||
logger.info("Starting API")
|
||||
tg.start_soon(self._apply_state)
|
||||
tg.start_soon(self._pause_on_new_election)
|
||||
@@ -1603,7 +1642,7 @@ class API:
|
||||
finally:
|
||||
self._event_log.close()
|
||||
self.command_sender.close()
|
||||
self.global_event_receiver.close()
|
||||
self.event_receiver.close()
|
||||
|
||||
async def run_api(self, ev: anyio.Event):
|
||||
cfg = Config()
|
||||
@@ -1620,38 +1659,31 @@ class API:
|
||||
)
|
||||
|
||||
async def _apply_state(self):
|
||||
with self.global_event_receiver as events:
|
||||
async for f_event in events:
|
||||
if f_event.session != self.session_id:
|
||||
continue
|
||||
if f_event.origin != self.session_id.master_node_id:
|
||||
continue
|
||||
self.event_buffer.ingest(f_event.origin_idx, f_event.event)
|
||||
for idx, event in self.event_buffer.drain_indexed():
|
||||
self._event_log.append(event)
|
||||
self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
|
||||
with self.event_receiver as events:
|
||||
async for i_event in events:
|
||||
self._event_log.append(i_event.event)
|
||||
self.state = apply(self.state, i_event)
|
||||
event = i_event.event
|
||||
|
||||
if isinstance(event, ChunkGenerated):
|
||||
if queue := self._image_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._image_generation_queues.pop(
|
||||
event.command_id, None
|
||||
)
|
||||
if queue := self._text_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert not isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._text_generation_queues.pop(event.command_id, None)
|
||||
if isinstance(event, TracesMerged):
|
||||
self._save_merged_trace(event)
|
||||
if isinstance(event, ChunkGenerated):
|
||||
if queue := self._image_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._image_generation_queues.pop(event.command_id, None)
|
||||
if queue := self._text_generation_queues.get(
|
||||
event.command_id, None
|
||||
):
|
||||
assert not isinstance(event.chunk, ImageChunk)
|
||||
try:
|
||||
await queue.send(event.chunk)
|
||||
except BrokenResourceError:
|
||||
self._text_generation_queues.pop(event.command_id, None)
|
||||
if isinstance(event, TracesMerged):
|
||||
self._save_merged_trace(event)
|
||||
|
||||
def _save_merged_trace(self, event: TracesMerged) -> None:
|
||||
traces = [
|
||||
@@ -1715,8 +1747,12 @@ class API:
|
||||
await self._send_download(command)
|
||||
return DeleteDownloadResponse(command_id=command.command_id)
|
||||
|
||||
def _get_trace_path(self, task_id: str) -> Path:
|
||||
return EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
|
||||
@staticmethod
|
||||
def _get_trace_path(task_id: str) -> Path:
|
||||
trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"
|
||||
if not trace_path.resolve().is_relative_to(EXO_TRACING_CACHE_DIR.resolve()):
|
||||
raise HTTPException(status_code=400, detail=f"Invalid task ID: {task_id}")
|
||||
return trace_path
|
||||
|
||||
async def list_traces(self) -> TraceListResponse:
|
||||
traces: list[TraceListItem] = []
|
||||
@@ -1814,3 +1850,23 @@ class API:
|
||||
media_type="application/json",
|
||||
filename=f"trace_{task_id}.json",
|
||||
)
|
||||
|
||||
async def delete_traces(self, request: DeleteTracesRequest) -> DeleteTracesResponse:
|
||||
deleted: list[str] = []
|
||||
not_found: list[str] = []
|
||||
for task_id in request.task_ids:
|
||||
trace_path = self._get_trace_path(task_id)
|
||||
if trace_path.exists():
|
||||
trace_path.unlink()
|
||||
deleted.append(task_id)
|
||||
else:
|
||||
not_found.append(task_id)
|
||||
return DeleteTracesResponse(deleted=deleted, not_found=not_found)
|
||||
|
||||
async def get_onboarding(self) -> JSONResponse:
|
||||
return JSONResponse({"completed": ONBOARDING_COMPLETE_FILE.exists()})
|
||||
|
||||
async def complete_onboarding(self) -> JSONResponse:
|
||||
ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
ONBOARDING_COMPLETE_FILE.write_text("true")
|
||||
return JSONResponse({"completed": True})
|
||||
|
||||
+22
-88
@@ -1,7 +1,6 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import anyio
|
||||
from anyio.abc import TaskGroup
|
||||
from loguru import logger
|
||||
|
||||
from exo.master.event_log import DiskEventLog
|
||||
@@ -36,8 +35,6 @@ from exo.shared.types.events import (
|
||||
IndexedEvent,
|
||||
InputChunkReceived,
|
||||
InstanceDeleted,
|
||||
JacclSideChannelData,
|
||||
JacclSideChannelGathered,
|
||||
LocalForwarderEvent,
|
||||
NodeGatheredInfo,
|
||||
NodeTimedOut,
|
||||
@@ -63,9 +60,9 @@ from exo.shared.types.tasks import (
|
||||
TextGeneration as TextGenerationTask,
|
||||
)
|
||||
from exo.shared.types.worker.instances import InstanceId
|
||||
from exo.shared.types.worker.runners import RunnerId
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.channels import Receiver, Sender
|
||||
from exo.utils.event_buffer import MultiSourceBuffer
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
|
||||
class Master:
|
||||
@@ -75,31 +72,26 @@ class Master:
|
||||
session_id: SessionId,
|
||||
*,
|
||||
command_receiver: Receiver[ForwarderCommand],
|
||||
event_sender: Sender[Event],
|
||||
local_event_receiver: Receiver[LocalForwarderEvent],
|
||||
global_event_sender: Sender[GlobalForwarderEvent],
|
||||
download_command_sender: Sender[ForwarderDownloadCommand],
|
||||
):
|
||||
self.state = State()
|
||||
self._tg: TaskGroup = anyio.create_task_group()
|
||||
self.node_id = node_id
|
||||
self.session_id = session_id
|
||||
self.state = State()
|
||||
self._tg: TaskGroup = TaskGroup()
|
||||
self.command_task_mapping: dict[CommandId, TaskId] = {}
|
||||
self.command_receiver = command_receiver
|
||||
self.local_event_receiver = local_event_receiver
|
||||
self.global_event_sender = global_event_sender
|
||||
self.download_command_sender = download_command_sender
|
||||
send, recv = channel[Event]()
|
||||
self.event_sender: Sender[Event] = send
|
||||
self._loopback_event_receiver: Receiver[Event] = recv
|
||||
self._loopback_event_sender: Sender[LocalForwarderEvent] = (
|
||||
local_event_receiver.clone_sender()
|
||||
)
|
||||
self.event_sender = event_sender
|
||||
self._system_id = SystemId()
|
||||
self._multi_buffer = MultiSourceBuffer[SystemId, Event]()
|
||||
self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master")
|
||||
self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {}
|
||||
self._expected_ranks: dict[TaskId, set[int]] = {}
|
||||
self._jaccl_pending: dict[InstanceId, dict[int, dict[RunnerId, bytes]]] = {}
|
||||
|
||||
async def run(self):
|
||||
logger.info("Starting Master")
|
||||
@@ -108,19 +100,16 @@ class Master:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._event_processor)
|
||||
tg.start_soon(self._command_processor)
|
||||
tg.start_soon(self._loopback_processor)
|
||||
tg.start_soon(self._plan)
|
||||
finally:
|
||||
self._event_log.close()
|
||||
self.global_event_sender.close()
|
||||
self.local_event_receiver.close()
|
||||
self.command_receiver.close()
|
||||
self._loopback_event_sender.close()
|
||||
self._loopback_event_receiver.close()
|
||||
|
||||
async def shutdown(self):
|
||||
logger.info("Stopping Master")
|
||||
self._tg.cancel_scope.cancel()
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _command_processor(self) -> None:
|
||||
with self.command_receiver as commands:
|
||||
@@ -339,17 +328,22 @@ class Master:
|
||||
task_id=task_id,
|
||||
)
|
||||
)
|
||||
case TaskFinished():
|
||||
generated_events.append(
|
||||
TaskDeleted(
|
||||
task_id=self.command_task_mapping[
|
||||
command.finished_command_id
|
||||
]
|
||||
else:
|
||||
logger.warning(
|
||||
f"Nonexistent command {command.cancelled_command_id} cancelled"
|
||||
)
|
||||
)
|
||||
self.command_task_mapping.pop(
|
||||
command.finished_command_id, None
|
||||
)
|
||||
case TaskFinished():
|
||||
if (
|
||||
task_id := self.command_task_mapping.pop(
|
||||
command.finished_command_id, None
|
||||
)
|
||||
) is not None:
|
||||
generated_events.append(TaskDeleted(task_id=task_id))
|
||||
else:
|
||||
logger.warning(
|
||||
f"Finished command {command.finished_command_id} finished"
|
||||
)
|
||||
|
||||
case RequestEventLog():
|
||||
# We should just be able to send everything, since other buffers will ignore old messages
|
||||
# rate limit to 1000 at a time
|
||||
@@ -413,27 +407,6 @@ class Master:
|
||||
self._event_log.append(event)
|
||||
await self._send_event(indexed)
|
||||
|
||||
# After broadcasting JacclSideChannelData, accumulate and
|
||||
# emit gathered result when all runners have contributed.
|
||||
if isinstance(event, JacclSideChannelData):
|
||||
await self._handle_jaccl_side_channel(event)
|
||||
|
||||
async def _loopback_processor(self) -> None:
|
||||
# this would ideally not be necessary.
|
||||
# this is WAY less hacky than how I was working around this before
|
||||
local_index = 0
|
||||
with self._loopback_event_receiver as events:
|
||||
async for event in events:
|
||||
await self._loopback_event_sender.send(
|
||||
LocalForwarderEvent(
|
||||
origin=self._system_id,
|
||||
origin_idx=local_index,
|
||||
session=self.session_id,
|
||||
event=event,
|
||||
)
|
||||
)
|
||||
local_index += 1
|
||||
|
||||
# This function is re-entrant, take care!
|
||||
async def _send_event(self, event: IndexedEvent):
|
||||
# Convenience method since this line is ugly
|
||||
@@ -471,42 +444,3 @@ class Master:
|
||||
del self._pending_traces[task_id]
|
||||
if task_id in self._expected_ranks:
|
||||
del self._expected_ranks[task_id]
|
||||
|
||||
async def _handle_jaccl_side_channel(self, event: JacclSideChannelData) -> None:
|
||||
"""Accumulate SideChannel contributions; when all runners for an instance
|
||||
have submitted for the same sequence, emit JacclSideChannelGathered."""
|
||||
iid = event.instance_id
|
||||
seq = event.sequence
|
||||
|
||||
if iid not in self._jaccl_pending:
|
||||
self._jaccl_pending[iid] = {}
|
||||
if seq not in self._jaccl_pending[iid]:
|
||||
self._jaccl_pending[iid][seq] = {}
|
||||
self._jaccl_pending[iid][seq][event.runner_id] = event.data
|
||||
|
||||
instance = self.state.instances.get(iid)
|
||||
if instance is None:
|
||||
logger.warning(f"JacclSideChannelData for unknown instance {iid}")
|
||||
return
|
||||
|
||||
expected_runners = set(instance.shard_assignments.runner_to_shard.keys())
|
||||
submitted = set(self._jaccl_pending[iid][seq].keys())
|
||||
|
||||
logger.info(
|
||||
f"JACCL side channel: instance={iid} seq={seq} "
|
||||
f"submitted={len(submitted)}/{len(expected_runners)}"
|
||||
)
|
||||
|
||||
if submitted >= expected_runners:
|
||||
gathered = dict(self._jaccl_pending[iid][seq])
|
||||
del self._jaccl_pending[iid][seq]
|
||||
if not self._jaccl_pending[iid]:
|
||||
del self._jaccl_pending[iid]
|
||||
|
||||
await self.event_sender.send(
|
||||
JacclSideChannelGathered(
|
||||
instance_id=iid,
|
||||
sequence=seq,
|
||||
gathered_data=gathered,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -112,7 +112,11 @@ def place_instance(
|
||||
cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
|
||||
]
|
||||
|
||||
if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
|
||||
if command.instance_meta == InstanceMeta.MlxJaccl:
|
||||
if not smallest_rdma_cycles:
|
||||
raise ValueError(
|
||||
"Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
|
||||
)
|
||||
smallest_cycles = smallest_rdma_cycles
|
||||
|
||||
cycles_with_leaf_nodes: list[Cycle] = [
|
||||
@@ -129,6 +133,11 @@ def place_instance(
|
||||
),
|
||||
)
|
||||
|
||||
# Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
|
||||
if len(selected_cycle) == 1:
|
||||
command.instance_meta = InstanceMeta.MlxRing
|
||||
command.sharding = Sharding.Pipeline
|
||||
|
||||
shard_assignments = get_shard_assignments(
|
||||
command.model_card, selected_cycle, command.sharding, node_memory
|
||||
)
|
||||
@@ -138,9 +147,6 @@ def place_instance(
|
||||
instance_id = InstanceId()
|
||||
target_instances = dict(deepcopy(current_instances))
|
||||
|
||||
if len(selected_cycle) == 1:
|
||||
command.instance_meta = InstanceMeta.MlxRing
|
||||
|
||||
match command.instance_meta:
|
||||
case InstanceMeta.MlxJaccl:
|
||||
# TODO(evan): shard assignments should contain information about ranks, this is ugly
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
# pyright: reportUnusedFunction=false, reportAny=false
|
||||
from typing import Any
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from exo.shared.types.common import CommandId
|
||||
|
||||
|
||||
def _make_api() -> Any:
|
||||
"""Create a minimal API instance with cancel route and error handler."""
|
||||
from exo.master.api import API
|
||||
|
||||
app = FastAPI()
|
||||
api = object.__new__(API)
|
||||
api.app = app
|
||||
api._text_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._image_generation_queues = {} # pyright: ignore[reportPrivateUsage]
|
||||
api._send = AsyncMock() # pyright: ignore[reportPrivateUsage]
|
||||
api._setup_exception_handlers() # pyright: ignore[reportPrivateUsage]
|
||||
app.post("/v1/cancel/{command_id}")(api.cancel_command)
|
||||
return api
|
||||
|
||||
|
||||
def test_cancel_nonexistent_command_returns_404() -> None:
|
||||
"""Cancel for an unknown command_id returns 404 in OpenAI error format."""
|
||||
api = _make_api()
|
||||
client = TestClient(api.app)
|
||||
|
||||
response = client.post("/v1/cancel/nonexistent-id")
|
||||
assert response.status_code == 404
|
||||
data: dict[str, Any] = response.json()
|
||||
assert "error" in data
|
||||
assert data["error"]["message"] == "Command not found or already completed"
|
||||
assert data["error"]["type"] == "Not Found"
|
||||
assert data["error"]["code"] == 404
|
||||
|
||||
|
||||
def test_cancel_active_text_generation() -> None:
|
||||
"""Cancel an active text generation command: returns 200, sender.close() called."""
|
||||
api = _make_api()
|
||||
client = TestClient(api.app)
|
||||
|
||||
cid = CommandId("text-cmd-123")
|
||||
sender = MagicMock()
|
||||
api._text_generation_queues[cid] = sender
|
||||
|
||||
response = client.post(f"/v1/cancel/{cid}")
|
||||
assert response.status_code == 200
|
||||
data: dict[str, Any] = response.json()
|
||||
assert data["message"] == "Command cancelled."
|
||||
assert data["command_id"] == str(cid)
|
||||
sender.close.assert_called_once()
|
||||
api._send.assert_called_once()
|
||||
task_cancelled = api._send.call_args[0][0]
|
||||
assert task_cancelled.cancelled_command_id == cid
|
||||
|
||||
|
||||
def test_cancel_active_image_generation() -> None:
|
||||
"""Cancel an active image generation command: returns 200, sender.close() called."""
|
||||
api = _make_api()
|
||||
client = TestClient(api.app)
|
||||
|
||||
cid = CommandId("img-cmd-456")
|
||||
sender = MagicMock()
|
||||
api._image_generation_queues[cid] = sender
|
||||
|
||||
response = client.post(f"/v1/cancel/{cid}")
|
||||
assert response.status_code == 200
|
||||
data: dict[str, Any] = response.json()
|
||||
assert data["message"] == "Command cancelled."
|
||||
assert data["command_id"] == str(cid)
|
||||
sender.close.assert_called_once()
|
||||
api._send.assert_called_once()
|
||||
task_cancelled = api._send.call_args[0][0]
|
||||
assert task_cancelled.cancelled_command_id == cid
|
||||
@@ -17,6 +17,7 @@ from exo.shared.types.commands import (
|
||||
)
|
||||
from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
GlobalForwarderEvent,
|
||||
IndexedEvent,
|
||||
InstanceCreated,
|
||||
@@ -50,6 +51,22 @@ async def test_master():
|
||||
command_sender, co_receiver = channel[ForwarderCommand]()
|
||||
local_event_sender, le_receiver = channel[LocalForwarderEvent]()
|
||||
fcds, _fcdr = channel[ForwarderDownloadCommand]()
|
||||
ev_send, ev_recv = channel[Event]()
|
||||
|
||||
async def mock_event_router():
|
||||
idx = 0
|
||||
sid = SystemId()
|
||||
with ev_recv as master_events:
|
||||
async for event in master_events:
|
||||
await local_event_sender.send(
|
||||
LocalForwarderEvent(
|
||||
origin=sid,
|
||||
origin_idx=idx,
|
||||
session=session_id,
|
||||
event=event,
|
||||
)
|
||||
)
|
||||
idx += 1
|
||||
|
||||
all_events: list[IndexedEvent] = []
|
||||
|
||||
@@ -67,6 +84,7 @@ async def test_master():
|
||||
master = Master(
|
||||
node_id,
|
||||
session_id,
|
||||
event_sender=ev_send,
|
||||
global_event_sender=ge_sender,
|
||||
local_event_receiver=le_receiver,
|
||||
command_receiver=co_receiver,
|
||||
@@ -75,6 +93,7 @@ async def test_master():
|
||||
logger.info("run the master")
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(master.run)
|
||||
tg.start_soon(mock_event_router)
|
||||
|
||||
# inject a NodeGatheredInfo event
|
||||
logger.info("inject a NodeGatheredInfo event")
|
||||
@@ -197,4 +216,5 @@ async def test_master():
|
||||
input=[InputMessage(role="user", content="Hello, how are you?")],
|
||||
)
|
||||
|
||||
ev_send.close()
|
||||
await master.shutdown()
|
||||
|
||||
@@ -14,10 +14,16 @@ from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
|
||||
from exo.shared.topology import Topology
|
||||
from exo.shared.types.commands import PlaceInstance
|
||||
from exo.shared.types.common import CommandId, NodeId
|
||||
from exo.shared.types.events import InstanceCreated, InstanceDeleted
|
||||
from exo.shared.types.events import (
|
||||
InstanceCreated,
|
||||
InstanceDeleted,
|
||||
TaskStatusUpdated,
|
||||
)
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.shared.types.multiaddr import Multiaddr
|
||||
from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
|
||||
from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
|
||||
from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
|
||||
from exo.shared.types.topology import Connection, SocketConnection
|
||||
from exo.shared.types.worker.instances import (
|
||||
Instance,
|
||||
@@ -456,3 +462,117 @@ def test_tensor_rdma_backend_connectivity_matrix(
|
||||
else:
|
||||
ip_part = coordinator.split(":")[0]
|
||||
assert len(ip_part.split(".")) == 4
|
||||
|
||||
|
||||
def _make_task(
|
||||
instance_id: InstanceId,
|
||||
status: TaskStatus = TaskStatus.Running,
|
||||
) -> TextGeneration:
|
||||
return TextGeneration(
|
||||
task_id=TaskId(),
|
||||
task_status=status,
|
||||
instance_id=instance_id,
|
||||
command_id=CommandId(),
|
||||
task_params=TextGenerationTaskParams(
|
||||
model=ModelId("test-model"),
|
||||
input=[InputMessage(role="user", content="hello")],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_get_transition_events_delete_instance_cancels_running_tasks(
|
||||
instance: Instance,
|
||||
):
|
||||
# arrange
|
||||
instance_id = InstanceId()
|
||||
current_instances: dict[InstanceId, Instance] = {instance_id: instance}
|
||||
target_instances: dict[InstanceId, Instance] = {}
|
||||
task = _make_task(instance_id, TaskStatus.Running)
|
||||
tasks = {task.task_id: task}
|
||||
|
||||
# act
|
||||
events = get_transition_events(current_instances, target_instances, tasks)
|
||||
|
||||
# assert – cancellation event should come before the deletion event
|
||||
assert len(events) == 2
|
||||
assert isinstance(events[0], TaskStatusUpdated)
|
||||
assert events[0].task_id == task.task_id
|
||||
assert events[0].task_status == TaskStatus.Cancelled
|
||||
assert isinstance(events[1], InstanceDeleted)
|
||||
assert events[1].instance_id == instance_id
|
||||
|
||||
|
||||
def test_get_transition_events_delete_instance_cancels_pending_tasks(
|
||||
instance: Instance,
|
||||
):
|
||||
# arrange
|
||||
instance_id = InstanceId()
|
||||
current_instances: dict[InstanceId, Instance] = {instance_id: instance}
|
||||
target_instances: dict[InstanceId, Instance] = {}
|
||||
task = _make_task(instance_id, TaskStatus.Pending)
|
||||
tasks = {task.task_id: task}
|
||||
|
||||
# act
|
||||
events = get_transition_events(current_instances, target_instances, tasks)
|
||||
|
||||
# assert
|
||||
assert len(events) == 2
|
||||
assert isinstance(events[0], TaskStatusUpdated)
|
||||
assert events[0].task_id == task.task_id
|
||||
assert events[0].task_status == TaskStatus.Cancelled
|
||||
assert isinstance(events[1], InstanceDeleted)
|
||||
|
||||
|
||||
def test_get_transition_events_delete_instance_ignores_completed_tasks(
|
||||
instance: Instance,
|
||||
):
|
||||
# arrange
|
||||
instance_id = InstanceId()
|
||||
current_instances: dict[InstanceId, Instance] = {instance_id: instance}
|
||||
target_instances: dict[InstanceId, Instance] = {}
|
||||
tasks = {
|
||||
t.task_id: t
|
||||
for t in [
|
||||
_make_task(instance_id, TaskStatus.Complete),
|
||||
_make_task(instance_id, TaskStatus.Failed),
|
||||
_make_task(instance_id, TaskStatus.TimedOut),
|
||||
_make_task(instance_id, TaskStatus.Cancelled),
|
||||
]
|
||||
}
|
||||
|
||||
# act
|
||||
events = get_transition_events(current_instances, target_instances, tasks)
|
||||
|
||||
# assert – only the InstanceDeleted event, no cancellations
|
||||
assert len(events) == 1
|
||||
assert isinstance(events[0], InstanceDeleted)
|
||||
|
||||
|
||||
def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
|
||||
instance: Instance,
|
||||
):
|
||||
# arrange
|
||||
instance_id_a = InstanceId()
|
||||
instance_id_b = InstanceId()
|
||||
current_instances: dict[InstanceId, Instance] = {
|
||||
instance_id_a: instance,
|
||||
instance_id_b: instance,
|
||||
}
|
||||
# only delete instance A, keep instance B
|
||||
target_instances: dict[InstanceId, Instance] = {instance_id_b: instance}
|
||||
|
||||
task_a = _make_task(instance_id_a, TaskStatus.Running)
|
||||
task_b = _make_task(instance_id_b, TaskStatus.Running)
|
||||
tasks = {task_a.task_id: task_a, task_b.task_id: task_b}
|
||||
|
||||
# act
|
||||
events = get_transition_events(current_instances, target_instances, tasks)
|
||||
|
||||
# assert – only task_a should be cancelled
|
||||
cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)]
|
||||
delete_events = [e for e in events if isinstance(e, InstanceDeleted)]
|
||||
assert len(cancel_events) == 1
|
||||
assert cancel_events[0].task_id == task_a.task_id
|
||||
assert cancel_events[0].task_status == TaskStatus.Cancelled
|
||||
assert len(delete_events) == 1
|
||||
assert delete_events[0].instance_id == instance_id_a
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
from enum import Enum
|
||||
|
||||
from exo_pyo3_bindings import ConnectionUpdate, ConnectionUpdateType
|
||||
from exo_pyo3_bindings import PyFromSwarm
|
||||
|
||||
from exo.shared.types.common import NodeId
|
||||
from exo.utils.pydantic_ext import CamelCaseModel
|
||||
@@ -8,30 +6,10 @@ from exo.utils.pydantic_ext import CamelCaseModel
|
||||
"""Serialisable types for Connection Updates/Messages"""
|
||||
|
||||
|
||||
class ConnectionMessageType(Enum):
|
||||
Connected = 0
|
||||
Disconnected = 1
|
||||
|
||||
@staticmethod
|
||||
def from_update_type(update_type: ConnectionUpdateType):
|
||||
match update_type:
|
||||
case ConnectionUpdateType.Connected:
|
||||
return ConnectionMessageType.Connected
|
||||
case ConnectionUpdateType.Disconnected:
|
||||
return ConnectionMessageType.Disconnected
|
||||
|
||||
|
||||
class ConnectionMessage(CamelCaseModel):
|
||||
node_id: NodeId
|
||||
connection_type: ConnectionMessageType
|
||||
remote_ipv4: str
|
||||
remote_tcp_port: int
|
||||
connected: bool
|
||||
|
||||
@classmethod
|
||||
def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage":
|
||||
return cls(
|
||||
node_id=NodeId(update.peer_id),
|
||||
connection_type=ConnectionMessageType.from_update_type(update.update_type),
|
||||
remote_ipv4=update.remote_ipv4,
|
||||
remote_tcp_port=update.remote_tcp_port,
|
||||
)
|
||||
def from_update(cls, update: PyFromSwarm.Connection) -> "ConnectionMessage":
|
||||
return cls(node_id=NodeId(update.peer_id), connected=update.connected)
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
from dataclasses import dataclass, field
|
||||
from random import random
|
||||
|
||||
import anyio
|
||||
from anyio import BrokenResourceError, ClosedResourceError
|
||||
from anyio.abc import CancelScope
|
||||
from loguru import logger
|
||||
|
||||
from exo.shared.types.commands import ForwarderCommand, RequestEventLog
|
||||
from exo.shared.types.common import SessionId, SystemId
|
||||
from exo.shared.types.events import (
|
||||
Event,
|
||||
EventId,
|
||||
GlobalForwarderEvent,
|
||||
IndexedEvent,
|
||||
LocalForwarderEvent,
|
||||
)
|
||||
from exo.utils.channels import Receiver, Sender, channel
|
||||
from exo.utils.event_buffer import OrderedBuffer
|
||||
from exo.utils.task_group import TaskGroup
|
||||
|
||||
|
||||
@dataclass
|
||||
class EventRouter:
|
||||
session_id: SessionId
|
||||
command_sender: Sender[ForwarderCommand]
|
||||
external_inbound: Receiver[GlobalForwarderEvent]
|
||||
external_outbound: Sender[LocalForwarderEvent]
|
||||
_system_id: SystemId = field(init=False, default_factory=SystemId)
|
||||
internal_outbound: list[Sender[IndexedEvent]] = field(
|
||||
init=False, default_factory=list
|
||||
)
|
||||
event_buffer: OrderedBuffer[Event] = field(
|
||||
init=False, default_factory=OrderedBuffer
|
||||
)
|
||||
out_for_delivery: dict[EventId, tuple[float, LocalForwarderEvent]] = field(
|
||||
init=False, default_factory=dict
|
||||
)
|
||||
_tg: TaskGroup = field(init=False, default_factory=TaskGroup)
|
||||
|
||||
_nack_cancel_scope: CancelScope | None = field(init=False, default=None)
|
||||
_nack_attempts: int = field(init=False, default=0)
|
||||
_nack_base_seconds: float = field(init=False, default=0.5)
|
||||
_nack_cap_seconds: float = field(init=False, default=10.0)
|
||||
|
||||
async def run(self):
|
||||
try:
|
||||
async with self._tg as tg:
|
||||
tg.start_soon(self._run_ext_in)
|
||||
tg.start_soon(self._simple_retry)
|
||||
finally:
|
||||
self.external_outbound.close()
|
||||
for send in self.internal_outbound:
|
||||
send.close()
|
||||
|
||||
# can make this better in future
|
||||
async def _simple_retry(self):
|
||||
while True:
|
||||
await anyio.sleep(1 + random())
|
||||
# list here is a shallow clone for shared mutation
|
||||
for e_id, (time, event) in list(self.out_for_delivery.items()):
|
||||
if anyio.current_time() > time + 5:
|
||||
self.out_for_delivery[e_id] = (anyio.current_time(), event)
|
||||
await self.external_outbound.send(event)
|
||||
|
||||
def sender(self) -> Sender[Event]:
|
||||
send, recv = channel[Event]()
|
||||
if self._tg.is_running():
|
||||
self._tg.start_soon(self._ingest, SystemId(), recv)
|
||||
else:
|
||||
self._tg.queue(self._ingest, SystemId(), recv)
|
||||
return send
|
||||
|
||||
def receiver(self) -> Receiver[IndexedEvent]:
|
||||
send, recv = channel[IndexedEvent]()
|
||||
self.internal_outbound.append(send)
|
||||
return recv
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._tg.cancel_tasks()
|
||||
|
||||
async def _ingest(self, system_id: SystemId, recv: Receiver[Event]):
|
||||
idx = 0
|
||||
with recv as events:
|
||||
async for event in events:
|
||||
f_ev = LocalForwarderEvent(
|
||||
origin_idx=idx,
|
||||
origin=system_id,
|
||||
session=self.session_id,
|
||||
event=event,
|
||||
)
|
||||
idx += 1
|
||||
await self.external_outbound.send(f_ev)
|
||||
self.out_for_delivery[event.event_id] = (anyio.current_time(), f_ev)
|
||||
|
||||
async def _run_ext_in(self):
|
||||
buf = OrderedBuffer[Event]()
|
||||
with self.external_inbound as events:
|
||||
async for event in events:
|
||||
if event.session != self.session_id:
|
||||
continue
|
||||
if event.origin != self.session_id.master_node_id:
|
||||
continue
|
||||
|
||||
buf.ingest(event.origin_idx, event.event)
|
||||
event_id = event.event.event_id
|
||||
if event_id in self.out_for_delivery:
|
||||
self.out_for_delivery.pop(event_id)
|
||||
|
||||
drained = buf.drain_indexed()
|
||||
if drained:
|
||||
self._nack_attempts = 0
|
||||
if self._nack_cancel_scope:
|
||||
self._nack_cancel_scope.cancel()
|
||||
|
||||
if not drained and (
|
||||
self._nack_cancel_scope is None
|
||||
or self._nack_cancel_scope.cancel_called
|
||||
):
|
||||
# Request the next index.
|
||||
self._tg.start_soon(self._nack_request, buf.next_idx_to_release)
|
||||
continue
|
||||
|
||||
for idx, event in drained:
|
||||
to_clear = set[int]()
|
||||
for i, sender in enumerate(self.internal_outbound):
|
||||
try:
|
||||
await sender.send(IndexedEvent(idx=idx, event=event))
|
||||
except (ClosedResourceError, BrokenResourceError):
|
||||
to_clear.add(i)
|
||||
for i in sorted(to_clear, reverse=True):
|
||||
self.internal_outbound.pop(i)
|
||||
|
||||
async def _nack_request(self, since_idx: int) -> None:
|
||||
# We request all events after (and including) the missing index.
|
||||
# This function is started whenever we receive an event that is out of sequence.
|
||||
# It is cancelled as soon as we receiver an event that is in sequence.
|
||||
|
||||
if since_idx < 0:
|
||||
logger.warning(f"Negative value encountered for nack request {since_idx=}")
|
||||
since_idx = 0
|
||||
|
||||
with CancelScope() as scope:
|
||||
self._nack_cancel_scope = scope
|
||||
delay: float = self._nack_base_seconds * (2.0**self._nack_attempts)
|
||||
delay = min(self._nack_cap_seconds, delay)
|
||||
self._nack_attempts += 1
|
||||
try:
|
||||
await anyio.sleep(delay)
|
||||
logger.info(
|
||||
f"Nack attempt {self._nack_attempts}: Requesting Event Log from {since_idx}"
|
||||
)
|
||||
await self.command_sender.send(
|
||||
ForwarderCommand(
|
||||
origin=self._system_id,
|
||||
command=RequestEventLog(since_idx=since_idx),
|
||||
)
|
||||
)
|
||||
finally:
|
||||
if self._nack_cancel_scope is scope:
|
||||
self._nack_cancel_scope = None
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user