Compare commits

..

3 Commits

Author SHA1 Message Date
Sami Khan 5b81d55d66 lite node 2026-02-19 22:07:34 +05:00
Sami Khan 3e9eb93f82 exo theme 2026-02-19 04:21:58 +05:00
Sami Khan ab622f79c3 EXO iOS app 2026-02-18 06:40:07 +05:00
380 changed files with 10129 additions and 37359 deletions
-20
View File
@@ -1,20 +0,0 @@
from enum import Enum
class HarmonyEncodingName(Enum):
HARMONY_GPT_OSS = ...
class HarmonyEncoding: ...
class HarmonyError(Exception): ...
class Role(Enum):
ASSISTANT = ...
class StreamableParser:
last_content_delta: str
current_channel: str | None
current_recipient: str | None
def __init__(self, encoding: HarmonyEncoding, role: Role = ...) -> None: ...
def process(self, token_id: int) -> None: ...
def load_harmony_encoding(name: HarmonyEncodingName) -> HarmonyEncoding: ...
-17
View File
@@ -1,17 +0,0 @@
class NvmlMemoryInfo:
used: int
total: int
free: int
class NvmlUtilizationRates:
gpu: int
memory: int
def nvmlInit() -> None: ...
def nvmlShutdown() -> None: ...
def nvmlDeviceGetCount() -> int: ...
def nvmlDeviceGetHandleByIndex(index: int) -> object: ...
def nvmlDeviceGetUtilizationRates(handle: object) -> NvmlUtilizationRates: ...
def nvmlDeviceGetTemperature(handle: object, sensor_type: int) -> int: ...
def nvmlDeviceGetPowerUsage(handle: object) -> int: ...
def nvmlDeviceGetMemoryInfo(handle: object) -> NvmlMemoryInfo: ...
-61
View File
@@ -1,61 +0,0 @@
from typing import Any, Sequence
from torch import backends as backends
from torch import cuda as cuda
from torch import distributed as distributed
__version__: str
class version:
cuda: str
class dtype: ...
bfloat16: dtype
float16: dtype
float32: dtype
int8: dtype
int32: dtype
int64: dtype
long: dtype
float8_e4m3fn: dtype
class Tensor:
shape: Sequence[int]
dtype: dtype
def __getitem__(self, key: Any) -> Tensor: ...
def __setitem__(self, key: Any, value: Any) -> None: ...
def to(self, *args: Any, **kwargs: Any) -> Tensor: ...
def cpu(self) -> Tensor: ...
def detach(self) -> Tensor: ...
def clone(self) -> Tensor: ...
def flatten(self, start_dim: int = 0, end_dim: int = -1) -> Tensor: ...
def view(self, *shape: Any) -> Tensor: ...
def squeeze(self, dim: int = ...) -> Tensor: ...
def unsqueeze(self, dim: int) -> Tensor: ...
def permute(self, *dims: int) -> Tensor: ...
def float(self) -> Tensor: ...
def numpy(self) -> Any: ...
def numel(self) -> int: ...
def nelement(self) -> int: ...
@property
def is_cuda(self) -> bool: ...
@property
def device(self) -> device: ...
def __len__(self) -> int: ...
def data_ptr(self) -> int: ...
def tolist(self) -> Any: ...
def abs(self) -> Tensor: ...
def max(self) -> Tensor: ...
def mean(self) -> Tensor: ...
def sum(self, dim: int = ...) -> Tensor: ...
def item(self) -> float: ...
def tensor(data: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def zeros(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def empty(*size: Any, dtype: dtype | None = None, device: Any = None) -> Tensor: ...
def from_numpy(ndarray: Any) -> Tensor: ...
def inference_mode() -> Any: ...
class device:
def __init__(self, type: str, index: int = ...) -> None: ...
@@ -1 +0,0 @@
from torch.backends import cuda as cuda
@@ -1 +0,0 @@
def is_built() -> bool: ...
-10
View File
@@ -1,10 +0,0 @@
class _DeviceProperties:
total_memory: int
def is_available() -> bool: ...
def get_device_name(device: int) -> str: ...
def get_device_properties(device: int) -> _DeviceProperties: ...
def empty_cache() -> None: ...
def mem_get_info() -> tuple[int, int]: ...
def synchronize() -> None: ...
def max_memory_allocated() -> int: ...
@@ -1,2 +0,0 @@
def is_initialized() -> bool: ...
def destroy_process_group() -> None: ...
-1
View File
@@ -1 +0,0 @@
__version__: str
-2
View File
@@ -1,2 +0,0 @@
class ModelConfig:
max_model_len: int
-18
View File
@@ -1,18 +0,0 @@
from dataclasses import dataclass
@dataclass
class EngineArgs:
model: str = ...
served_model_name: str | list[str] | None = ...
tokenizer: str | None = ...
trust_remote_code: bool = ...
dtype: str = ...
seed: int = ...
max_model_len: int | None = ...
gpu_memory_utilization: float = ...
enforce_eager: bool = ...
tensor_parallel_size: int = ...
pipeline_parallel_size: int = ...
quantization: str | None = ...
load_format: str = ...
enable_sleep_mode: bool = ...
-17
View File
@@ -1,17 +0,0 @@
class CompletionOutput:
index: int
text: str
token_ids: list[int]
cumulative_logprob: float | None
logprobs: object | None
finish_reason: str | None
stop_reason: int | str | None
def finished(self) -> bool: ...
class RequestOutput:
request_id: str
prompt: str | None
prompt_token_ids: list[int] | None
outputs: list[CompletionOutput]
finished: bool
-11
View File
@@ -1,11 +0,0 @@
class SamplingParams:
n: int
temperature: float
top_p: float
top_k: int
min_p: float
seed: int | None
stop: str | list[str] | None
max_tokens: int | None
logprobs: int | None
repetition_penalty: float
@@ -1,3 +0,0 @@
from vllm.tokenizers.protocol import TokenizerLike
__all__ = ["TokenizerLike"]
@@ -1,15 +0,0 @@
from typing import Protocol
class TokenizerLike(Protocol):
@property
def eos_token_id(self) -> int: ...
@property
def vocab_size(self) -> int: ...
def encode(self, text: str, add_special_tokens: bool = ...) -> list[int]: ...
def decode(self, ids: list[int] | int, skip_special_tokens: bool = ...) -> str: ...
def apply_chat_template(
self,
messages: list[dict[str, str]],
tools: list[dict[str, object]] | None = ...,
**kwargs: object,
) -> str | list[int]: ...
-1
View File
@@ -1 +0,0 @@
-1
View File
@@ -1 +0,0 @@
@@ -1,24 +0,0 @@
from collections.abc import Sequence
from vllm.v1.core.kv_cache_utils import BlockPool, KVCacheBlock
from vllm.v1.kv_cache_interface import KVCacheConfig
class KVCacheBlocks:
blocks: tuple[Sequence[KVCacheBlock], ...]
def __init__(self, blocks: tuple[Sequence[KVCacheBlock], ...]) -> None: ...
def get_block_ids(self) -> tuple[list[int], ...]: ...
class KVCacheManager:
block_pool: BlockPool
kv_cache_config: KVCacheConfig
enable_caching: bool
num_kv_cache_groups: int
coordinator: object
def __init__(self, *args: object, **kwargs: object) -> None: ...
def allocate_slots(
self, request: object, num_new_tokens: int, *args: object, **kwargs: object
) -> KVCacheBlocks | None: ...
def get_computed_blocks(self, request: object) -> tuple[KVCacheBlocks, int]: ...
def create_kv_cache_blocks(
self, blocks: tuple[list[KVCacheBlock], ...]
) -> KVCacheBlocks: ...
@@ -1,16 +0,0 @@
class KVCacheBlock:
block_id: int
ref_cnt: int
def __init__(self, block_id: int) -> None: ...
class FreeKVCacheBlockQueue:
def append_n(self, blocks: list[KVCacheBlock]) -> None: ...
def popleft_n(self, n: int) -> list[KVCacheBlock]: ...
class BlockPool:
blocks: list[KVCacheBlock]
free_block_queue: FreeKVCacheBlockQueue
num_gpu_blocks: int
enable_caching: bool
def get_num_free_blocks(self) -> int: ...
def get_new_blocks(self, num_blocks: int) -> list[KVCacheBlock]: ...
@@ -1,22 +0,0 @@
from vllm.config import ModelConfig
from vllm.engine.arg_utils import EngineArgs
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.tokenizers import TokenizerLike
class LLMEngine:
tokenizer: TokenizerLike | None
model_config: ModelConfig
@classmethod
def from_engine_args(cls, engine_args: EngineArgs) -> LLMEngine: ...
def add_request(
self,
request_id: str,
prompt: str,
params: SamplingParams,
arrival_time: float | None = ...,
) -> None: ...
def step(self) -> list[RequestOutput]: ...
def has_unfinished_requests(self) -> bool: ...
def get_tokenizer(self) -> TokenizerLike: ...
@@ -1,23 +0,0 @@
from dataclasses import dataclass
@dataclass
class KVCacheSpec:
block_size: int
num_kv_heads: int
head_size: int
@dataclass
class KVCacheGroupSpec:
layer_names: list[str]
kv_cache_spec: KVCacheSpec
@dataclass
class KVCacheTensorSpec:
shared_by: list[str]
size: int
@dataclass
class KVCacheConfig:
num_blocks: int
kv_cache_groups: list[KVCacheGroupSpec]
kv_cache_tensors: list[KVCacheTensorSpec]
-6
View File
@@ -1,6 +0,0 @@
class Request:
request_id: str
prompt_token_ids: list[int] | None
num_prompt_tokens: int
num_computed_tokens: int
num_tokens: int
@@ -1 +0,0 @@
@@ -1,24 +0,0 @@
import torch
class _CompilationConfig:
static_forward_context: dict[str, object]
class _ModelConfig:
hf_config: object
class GPUModelRunner:
kv_caches: list[torch.Tensor]
compilation_config: _CompilationConfig
model_config: _ModelConfig | None
def _allocate_kv_cache_tensors(
self, kv_cache_config: object
) -> dict[str, torch.Tensor]: ...
def initialize_kv_cache_tensors(
self, kv_cache_config: object, kernel_block_sizes: list[int]
) -> dict[str, torch.Tensor]: ...
def _reshape_kv_cache_tensors(
self,
kv_cache_config: object,
raw_tensors: dict[str, torch.Tensor],
kernel_block_sizes: list[int],
) -> dict[str, torch.Tensor]: ...
@@ -1,6 +0,0 @@
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
class Worker:
model_runner: GPUModelRunner
def determine_available_memory(self) -> int: ...
def initialize_from_config(self, kv_cache_config: object) -> None: ...
-1
View File
@@ -1 +0,0 @@
def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...
-16
View File
@@ -215,22 +215,6 @@ 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.
+1 -1
View File
@@ -200,7 +200,7 @@ class Module(dict):
) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]:
"""Return the submodules that do not contain other modules."""
def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module:
def update(self, parameters: dict, strict: bool = ...) -> Module:
"""Replace the parameters of this Module with the provided ones in the
dict of dicts and lists.
@@ -32,7 +32,6 @@ class Conv1d(Module):
"""
weight: mx.array
bias: mx.array | None
groups: int
def __init__(
self,
-4
View File
@@ -40,10 +40,6 @@ class Linear(Module):
bias (bool, optional): If set to ``False`` then the layer will
not use a bias. Default is ``True``.
"""
weight: mx.array
bias: mx.array | None
def __init__(self, input_dims: int, output_dims: int, bias: bool = ...) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
def to_quantized(
@@ -88,9 +88,6 @@ class RMSNorm(Module):
dims (int): The feature dimension of the input to normalize over
eps (float): A small additive constant for numerical stability
"""
weight: mx.array
def __init__(self, dims: int, eps: float = ...) -> None: ...
def __call__(self, x) -> mx.array: ...
+8 -11
View File
@@ -7,10 +7,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
from mlx.core import MX_ARRAY_TREE
def tree_map(
fn: Callable[..., Any],
tree: Any,
*rest: Any,
is_leaf: Callable[..., bool] | None = ...,
fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ...
) -> Any:
"""Applies ``fn`` to the leaves of the Python tree ``tree`` and
returns a new collection with the results.
@@ -47,11 +44,11 @@ def tree_map(
"""
def tree_map_with_path(
fn: Callable[..., Any],
fn: Callable,
tree: Any,
*rest: Any,
is_leaf: Callable[..., bool] | None = ...,
path: str | None = ...,
is_leaf: Optional[Callable] = ...,
path: Optional[Any] = ...,
) -> Any:
"""Applies ``fn`` to the path and leaves of the Python tree ``tree`` and
returns a new collection with the results.
@@ -83,9 +80,9 @@ def tree_map_with_path(
def tree_flatten(
tree: Any,
prefix: str = ...,
is_leaf: Callable[..., bool] | None = ...,
destination: list[tuple[str, Any]] | dict[str, Any] | None = ...,
) -> list[tuple[str, Any]] | dict[str, Any]:
is_leaf: Optional[Callable] = ...,
destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = ...,
) -> Union[List[Tuple[str, Any]], Dict[str, Any]]:
"""Flattens a Python tree to a list of key, value tuples.
The keys are using the dot notation to define trees of arbitrary depth and
@@ -121,7 +118,7 @@ def tree_flatten(
the Python tree.
"""
def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any:
def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
"""Recreate a Python tree from its flat representation.
.. code-block:: python
+18 -59
View File
@@ -73,11 +73,9 @@ class GenerationResponse:
finish_reason: Optional[str] = ...
def maybe_quantize_kv_cache(
prompt_cache: Any,
quantized_kv_start: int | None,
kv_group_size: int | None,
kv_bits: int | None,
) -> None: ...
prompt_cache, quantized_kv_start, kv_group_size, kv_bits
): # -> None:
...
def generate_step(
prompt: mx.array,
model: nn.Module,
@@ -254,14 +252,7 @@ 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]
@@ -270,71 +261,39 @@ class Batch:
max_tokens: List[int]
num_tokens: List[int]
cache: List[Any]
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]: ...
def __len__(self): # -> int:
...
def filter(self, keep_idx: List[int]): # -> None:
...
def extend(self, other): # -> None:
...
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: Any,
model,
max_tokens: int = ...,
stop_tokens: Optional[set[int]] = ...,
stop_tokens: Optional[set] = ...,
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: 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]]: ...
self, prompts, max_tokens: Union[List[int], int, None] = ...
): # -> list[Any]:
...
def stats(self): # -> BatchStats:
...
def next(self): # -> list[Any]:
...
def batch_generate(
model,
+30 -27
View File
@@ -16,7 +16,7 @@ class Cache(Protocol):
self, keys: mx.array, values: mx.array
) -> tuple[mx.array, mx.array]: ...
@property
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
def state(self) -> tuple[mx.array, mx.array]: ...
@state.setter
def state(self, v) -> None: ...
@@ -92,14 +92,13 @@ class _BaseCache(Cache):
values: mx.array
offset: int
@property
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
def state(self) -> tuple[mx.array, mx.array]: ...
@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: ...
@@ -115,13 +114,15 @@ class ConcatenateKVCache(_BaseCache):
def update_and_fetch(self, keys, values): # -> tuple[Any | array, Any | array]:
...
@property
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
def state(self): # -> tuple[Any | array | None, Any | array | None]:
...
@state.setter
def state(self, v): # -> None:
...
def is_trimmable(self): # -> Literal[True]:
...
def trim(self, n: int) -> int: ...
def trim(self, n): # -> int:
...
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
...
@@ -131,7 +132,10 @@ class QuantizedKVCache(_BaseCache):
def update_and_fetch(self, keys, values): # -> Any:
...
@property
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
def state(
self,
): # -> tuple[Any | tuple[array, array, array] | None, Any | tuple[array, array, array] | None] | Any:
...
@state.setter
def state(self, v): # -> None:
...
@@ -143,7 +147,8 @@ class QuantizedKVCache(_BaseCache):
...
def is_trimmable(self): # -> Literal[True]:
...
def trim(self, n: int) -> int: ...
def trim(self, n): # -> int:
...
def make_mask(self, *args, **kwargs): # -> array | Literal['causal'] | None:
...
@@ -155,30 +160,22 @@ class KVCache(_BaseCache):
@property
def state(
self,
) -> tuple[mx.array | None, mx.array | None]: ...
) -> tuple[array, array]: ...
@state.setter
def state(self, v) -> None: ...
def is_trimmable(self): # -> Literal[True]:
...
def trim(self, n: int) -> int: ...
def trim(self, n): # -> int:
...
def to_quantized(
self, group_size: int = ..., bits: int = ...
) -> QuantizedKVCache: ...
def make_mask(
self, *args: Any, **kwargs: Any
) -> mx.array | Literal["causal"] | None: ...
def make_mask(self, *args, **kwargs): # -> 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]:
@@ -186,7 +183,8 @@ class RotatingKVCache(_BaseCache):
@property
def state(
self,
) -> tuple[mx.array | None, mx.array | None]: ...
): # -> tuple[Any | array, Any | array] | tuple[Any | array | None, Any | array | None]:
...
@state.setter
def state(self, v): # -> None:
...
@@ -198,7 +196,8 @@ class RotatingKVCache(_BaseCache):
...
def is_trimmable(self): # -> bool:
...
def trim(self, n: int) -> int: ...
def trim(self, n): # -> int:
...
def to_quantized(
self, group_size: int = ..., bits: int = ...
) -> QuantizedKVCache: ...
@@ -213,7 +212,8 @@ class ArraysCache(_BaseCache):
...
def __getitem__(self, idx): ...
@property
def state(self) -> tuple[mx.array | None, mx.array | None]: ...
def state(self): # -> list[Any | array] | list[array]:
...
@state.setter
def state(self, v): # -> None:
...
@@ -227,7 +227,8 @@ class ArraysCache(_BaseCache):
In-place extend this cache with the other cache.
"""
def make_mask(self, N: int) -> mx.array | None: ...
def make_mask(self, N: int): # -> array | None:
...
class MambaCache(ArraysCache):
def __init__(self, left_padding: Optional[List[int]] = ...) -> None: ...
@@ -238,7 +239,8 @@ class ChunkedKVCache(KVCache):
...
def update_and_fetch(self, keys, values): # -> tuple[array, array]:
...
def trim(self, n: int) -> int: ...
def trim(self, n): # -> int:
...
@property
def meta_state(self): # -> tuple[str, ...]:
...
@@ -251,9 +253,10 @@ class CacheList(_BaseCache):
def __getitem__(self, idx): ...
def is_trimmable(self): # -> bool:
...
def trim(self, n: int) -> int: ...
def trim(self, n): ...
@property
def state(self) -> list[tuple[mx.array | None, mx.array | None]]: ...
def state(self): # -> list[Any]:
...
@state.setter
def state(self, v): # -> None:
...
@@ -1,46 +0,0 @@
"""Type stubs for mlx_lm.models.glm_moe_dsa"""
from dataclasses import dataclass
from typing import Any, Dict, Optional
from .base import BaseModelArgs
from .deepseek_v32 import Model as DSV32Model
@dataclass
class ModelArgs(BaseModelArgs):
model_type: str
vocab_size: int
hidden_size: int
index_head_dim: int
index_n_heads: int
index_topk: int
intermediate_size: int
moe_intermediate_size: int
num_hidden_layers: int
num_attention_heads: int
num_key_value_heads: int
n_shared_experts: Optional[int]
n_routed_experts: Optional[int]
routed_scaling_factor: float
kv_lora_rank: int
q_lora_rank: int
qk_rope_head_dim: int
v_head_dim: int
qk_nope_head_dim: int
topk_method: str
scoring_func: str
norm_topk_prob: bool
n_group: int
topk_group: int
num_experts_per_tok: int
moe_layer_freq: int
first_k_dense_replace: int
max_position_embeddings: int
rms_norm_eps: float
rope_parameters: Dict[str, Any]
attention_bias: bool
rope_scaling: Dict[str, Any] | None
rope_theta: float | None
class Model(DSV32Model):
def __init__(self, config: ModelArgs) -> None: ...
-154
View File
@@ -1,154 +0,0 @@
from dataclasses import dataclass
from typing import Any, List, Optional, Tuple
import mlx.core as mx
import mlx.nn as nn
from .cache import ArraysCache, KVCache
from .switch_layers import SwitchMLP
@dataclass
class ModelArgs:
model_type: str
vocab_size: int
hidden_size: int
intermediate_size: int
num_hidden_layers: int
max_position_embeddings: int
num_attention_heads: int
num_key_value_heads: int
attention_bias: bool
mamba_num_heads: int
mamba_head_dim: int
mamba_proj_bias: bool
ssm_state_size: int
conv_kernel: int
n_groups: int
mlp_bias: bool
layer_norm_epsilon: float
use_bias: bool
use_conv_bias: bool
hybrid_override_pattern: List[str]
head_dim: Optional[int]
moe_intermediate_size: Optional[int]
moe_shared_expert_intermediate_size: Optional[int]
n_group: Optional[int]
n_routed_experts: Optional[int]
n_shared_experts: Optional[int]
topk_group: Optional[int]
num_experts_per_tok: Optional[int]
norm_topk_prob: Optional[bool]
routed_scaling_factor: Optional[float]
time_step_limit: Optional[Tuple[float, float]]
time_step_min: Optional[float]
time_step_max: Optional[float]
@classmethod
def from_dict(cls, params: dict[str, Any]) -> ModelArgs: ...
def __post_init__(self) -> None: ...
class NemotronHMamba2Mixer(nn.Module):
num_heads: int
hidden_size: int
ssm_state_size: int
conv_kernel_size: int
intermediate_size: int
n_groups: int
head_dim: int
conv_dim: int
conv1d: nn.Conv1d
in_proj: nn.Linear
dt_bias: mx.array
A_log: mx.array
D: mx.array
norm: nn.RMSNorm
heads_per_group: int
out_proj: nn.Linear
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
hidden_states: mx.array,
mask: Optional[mx.array],
cache: Optional[ArraysCache] = None,
) -> mx.array: ...
class NemotronHAttention(nn.Module):
hidden_size: int
num_heads: int
head_dim: int
num_key_value_heads: int
scale: float
q_proj: nn.Linear
k_proj: nn.Linear
v_proj: nn.Linear
o_proj: nn.Linear
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[KVCache] = None,
) -> mx.array: ...
class NemotronHMLP(nn.Module):
up_proj: nn.Linear
down_proj: nn.Linear
def __init__(
self, args: ModelArgs, intermediate_size: Optional[int] = None
) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class NemotronHMoE(nn.Module):
num_experts_per_tok: int
switch_mlp: SwitchMLP
shared_experts: NemotronHMLP
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class NemotronHBlock(nn.Module):
block_type: str
norm: nn.RMSNorm
mixer: NemotronHMamba2Mixer | NemotronHAttention | NemotronHMLP | NemotronHMoE
def __init__(self, args: ModelArgs, block_type: str) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class NemotronHModel(nn.Module):
embeddings: nn.Embedding
layers: list[NemotronHBlock]
norm_f: nn.RMSNorm
fa_idx: int
ssm_idx: int
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
class Model(nn.Module):
args: ModelArgs
backbone: NemotronHModel
lm_head: nn.Linear
model_type: str
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
@property
def layers(self) -> list[NemotronHBlock]: ...
def make_cache(self) -> list[ArraysCache | KVCache]: ...
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
-153
View File
@@ -1,153 +0,0 @@
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]: ...
@@ -1,19 +0,0 @@
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]: ...
-13
View File
@@ -5,18 +5,8 @@ from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .cache import ArraysCache, KVCache
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
@@ -100,8 +90,6 @@ class Qwen3NextModel(nn.Module):
embed_tokens: nn.Embedding
layers: list[Qwen3NextDecoderLayer]
norm: nn.RMSNorm
ssm_idx: int
fa_idx: int
def __init__(self, args: Any) -> None: ...
def __call__(
@@ -124,4 +112,3 @@ class Model(nn.Module):
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
@property
def layers(self) -> list[Qwen3NextDecoderLayer]: ...
def make_cache(self) -> list[ArraysCache | KVCache]: ...
@@ -73,9 +73,6 @@ class SwitchGLU(nn.Module):
def __call__(self, x, indices) -> mx.array: ...
class SwitchMLP(nn.Module):
fc1: SwitchLinear
fc2: SwitchLinear
def __init__(
self,
input_dims: int,
+1 -1
View File
@@ -48,7 +48,7 @@ def make_logits_processors(
logit_bias: Optional[Dict[int, float]] = ...,
repetition_penalty: Optional[float] = ...,
repetition_context_size: Optional[int] = ...,
) -> list[Callable[[mx.array, mx.array], mx.array]]:
): # -> list[Any]:
"""
Make logits processors for use with ``generate_step``.
+4 -4
View File
@@ -39,11 +39,11 @@ class StreamingDetokenizer:
"""
__slots__ = ...
def reset(self) -> None: ...
def add_token(self, token: int) -> None: ...
def finalize(self) -> None: ...
def reset(self): ...
def add_token(self, token): ...
def finalize(self): ...
@property
def last_segment(self) -> str:
def last_segment(self):
"""Return the last segment of readable text since last time this property was accessed."""
class NaiveStreamingDetokenizer(StreamingDetokenizer):
-113
View File
@@ -39,119 +39,6 @@ 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
+125 -35
View File
@@ -141,6 +141,12 @@ version = "0.3.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb"
[[package]]
name = "arrayvec"
version = "0.7.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50"
[[package]]
name = "asn1-rs"
version = "0.7.1"
@@ -216,28 +222,6 @@ 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"
@@ -320,6 +304,19 @@ version = "1.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba"
[[package]]
name = "bigdecimal"
version = "0.4.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "560f42649de9fa436b73517378a147ec21f6c997a546581df4b4b31677828934"
dependencies = [
"autocfg",
"libm",
"num-bigint",
"num-integer",
"num-traits",
]
[[package]]
name = "bimap"
version = "0.6.3"
@@ -519,6 +516,15 @@ version = "0.4.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2f421161cb492475f1661ddc9815a745a1c894592070661180fdec3d4872e9c3"
[[package]]
name = "convert_case"
version = "0.10.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9"
dependencies = [
"unicode-segmentation",
]
[[package]]
name = "core-foundation"
version = "0.9.4"
@@ -740,6 +746,29 @@ dependencies = [
"powerfmt",
]
[[package]]
name = "derive_more"
version = "2.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "10b768e943bed7bf2cab53df09f4bc34bfd217cdb57d971e769874c9a6710618"
dependencies = [
"derive_more-impl",
]
[[package]]
name = "derive_more-impl"
version = "2.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6d286bfdaf75e988b4a78e013ecd79c581e06399ab53fbacd2d916c2f904f30b"
dependencies = [
"convert_case",
"proc-macro2",
"quote",
"rustc_version",
"syn 2.0.111",
"unicode-xid",
]
[[package]]
name = "digest"
version = "0.10.7"
@@ -910,17 +939,22 @@ name = "exo_pyo3_bindings"
version = "0.0.1"
dependencies = [
"delegate",
"derive_more",
"env_logger",
"extend",
"futures-lite",
"futures",
"impl-trait-for-tuples",
"libp2p",
"log",
"networking",
"once_cell",
"pin-project",
"pyo3",
"pyo3-async-runtimes",
"pyo3-log",
"pyo3-stub-gen",
"thiserror 2.0.17",
"thread_local",
"tokio",
"util",
]
@@ -936,12 +970,6 @@ dependencies = [
"syn 2.0.111",
]
[[package]]
name = "fastrand"
version = "2.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
[[package]]
name = "ff"
version = "0.13.1"
@@ -1050,10 +1078,7 @@ version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
dependencies = [
"fastrand",
"futures-core",
"futures-io",
"parking",
"pin-project-lite",
]
@@ -1615,6 +1640,17 @@ dependencies = [
"xmltree",
]
[[package]]
name = "impl-trait-for-tuples"
version = "0.2.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a0eb5a3343abf848c0984fe4604b2b105da9539376e24fc0a3b0007411ae4fd9"
dependencies = [
"proc-macro2",
"quote",
"syn 2.0.111",
]
[[package]]
name = "indexmap"
version = "2.12.1"
@@ -1793,6 +1829,12 @@ version = "0.2.178"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091"
[[package]]
name = "libm"
version = "0.2.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
[[package]]
name = "libp2p"
version = "0.56.0"
@@ -2781,16 +2823,17 @@ dependencies = [
name = "networking"
version = "0.0.1"
dependencies = [
"async-stream",
"delegate",
"derive_more",
"either",
"extend",
"futures-lite",
"futures",
"futures-timer",
"impl-trait-for-tuples",
"keccak-const",
"libp2p",
"log",
"pin-project",
"thiserror 2.0.17",
"tokio",
"tracing-subscriber",
"util",
@@ -2875,6 +2918,17 @@ dependencies = [
"num-traits",
]
[[package]]
name = "num-rational"
version = "0.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
dependencies = [
"num-bigint",
"num-integer",
"num-traits",
]
[[package]]
name = "num-traits"
version = "0.2.19"
@@ -3225,14 +3279,28 @@ version = "0.27.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d"
dependencies = [
"bigdecimal",
"either",
"hashbrown 0.16.1",
"indexmap",
"indoc",
"inventory",
"libc",
"lock_api",
"memoffset",
"num-bigint",
"num-complex",
"num-rational",
"num-traits",
"once_cell",
"ordered-float",
"parking_lot",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
"rust_decimal",
"smallvec",
"unindent",
]
@@ -3673,6 +3741,16 @@ dependencies = [
"tokio",
]
[[package]]
name = "rust_decimal"
version = "1.39.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "35affe401787a9bd846712274d97654355d21b2a2c092a3139aabe31e9022282"
dependencies = [
"arrayvec",
"num-traits",
]
[[package]]
name = "rustc-hash"
version = "1.1.0"
@@ -4537,12 +4615,24 @@ version = "1.0.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5"
[[package]]
name = "unicode-segmentation"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493"
[[package]]
name = "unicode-width"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
[[package]]
name = "unicode-xid"
version = "0.2.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853"
[[package]]
name = "unicode_names2"
version = "1.3.0"
+35 -3
View File
@@ -1,6 +1,10 @@
[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"
@@ -22,21 +26,49 @@ opt-level = 3
networking = { path = "rust/networking" }
util = { path = "rust/util" }
# Proc-macro authoring tools
syn = "2.0"
quote = "1.0"
proc-macro2 = "1.0"
darling = "0.20"
# Macro dependecies
extend = "1.2"
delegate = "0.13"
impl-trait-for-tuples = "0.2"
clap = "4.5"
derive_more = { version = "2.0.1", features = ["display"] }
pin-project = "1"
# Utility dependencies
itertools = "0.14"
thiserror = "2"
internment = "0.8"
recursion = "0.5"
regex = "1.11"
once_cell = "1.21"
thread_local = "1.1"
bon = "3.4"
generativity = "1.1"
anyhow = "1.0"
keccak-const = "0.2"
# Functional generics/lenses frameworks
frunk_core = "0.4"
frunk = "0.4"
frunk_utils = "0.2"
frunk-enum-core = "0.3"
# Async dependencies
async-stream = "0.3"
tokio = "1.46"
futures-lite = "2.6.1"
futures = "0.3"
futures-util = "0.3"
futures-timer = "3.0"
# Data structures
either = "1.15"
ordered-float = "5.0"
ahash = "0.8"
# Tracing/logging
log = "0.4"
+4 -132
View File
@@ -26,8 +26,6 @@ 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
@@ -74,30 +72,16 @@ There are two ways to run exo:
### Run from Source (macOS)
If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly:
```bash
nix run .#exo
```
**Note:** To accept the Cachix binary cache (and avoid the Xcode Metal ToolChain), add to `/etc/nix/nix.conf`:
```
trusted-users = root (or your username)
experimental-features = nix-command flakes
```
Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-daemon`
**Prerequisites:**
- [Xcode](https://developer.apple.com/xcode/) (provides the Metal ToolChain required for MLX compilation)
- [brew](https://github.com/Homebrew/brew) (for simple package management on macOS)
```bash
/bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
```
- [uv](https://github.com/astral-sh/uv) (for Python dependency management)
- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
- [node](https://github.com/nodejs/node) (for building the dashboard)
```bash
brew install uv macmon node
```
@@ -198,8 +182,6 @@ 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.
@@ -279,47 +261,8 @@ 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.
---
@@ -409,85 +352,14 @@ 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 documentation in [docs/api.md](docs/api.md).
- API basic documentation in [docs/api.md](docs/api.md).
- API types and endpoints in [src/exo/master/api.py](src/exo/master/api.py).
---
@@ -546,4 +418,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.
@@ -0,0 +1,628 @@
// !$*UTF8*$!
{
archiveVersion = 1;
classes = {
};
objectVersion = 77;
objects = {
/* Begin PBXBuildFile section */
E09D17522F44F359009C51A3 /* MLXLLM in Frameworks */ = {isa = PBXBuildFile; productRef = E09D17512F44F359009C51A3 /* MLXLLM */; };
E09D17542F44F359009C51A3 /* MLXLMCommon in Frameworks */ = {isa = PBXBuildFile; productRef = E09D17532F44F359009C51A3 /* MLXLMCommon */; };
/* End PBXBuildFile section */
/* Begin PBXContainerItemProxy section */
E09D167D2F44CA20009C51A3 /* PBXContainerItemProxy */ = {
isa = PBXContainerItemProxy;
containerPortal = E09D16672F44CA1E009C51A3 /* Project object */;
proxyType = 1;
remoteGlobalIDString = E09D166E2F44CA1E009C51A3;
remoteInfo = "EXO-iOS";
};
E09D16872F44CA20009C51A3 /* PBXContainerItemProxy */ = {
isa = PBXContainerItemProxy;
containerPortal = E09D16672F44CA1E009C51A3 /* Project object */;
proxyType = 1;
remoteGlobalIDString = E09D166E2F44CA1E009C51A3;
remoteInfo = "EXO-iOS";
};
/* End PBXContainerItemProxy section */
/* Begin PBXFileReference section */
E09D166F2F44CA1E009C51A3 /* EXO-iOS.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = "EXO-iOS.app"; sourceTree = BUILT_PRODUCTS_DIR; };
E09D167C2F44CA20009C51A3 /* EXO-iOSTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "EXO-iOSTests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
E09D16862F44CA20009C51A3 /* EXO-iOSUITests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = "EXO-iOSUITests.xctest"; sourceTree = BUILT_PRODUCTS_DIR; };
/* End PBXFileReference section */
/* Begin PBXFileSystemSynchronizedBuildFileExceptionSet section */
E09D169A2F44CA20009C51A3 /* Exceptions for "EXO-iOS" folder in "EXO-iOS" target */ = {
isa = PBXFileSystemSynchronizedBuildFileExceptionSet;
membershipExceptions = (
Info.plist,
);
target = E09D166E2F44CA1E009C51A3 /* EXO-iOS */;
};
/* End PBXFileSystemSynchronizedBuildFileExceptionSet section */
/* Begin PBXFileSystemSynchronizedRootGroup section */
E09D16712F44CA1E009C51A3 /* EXO-iOS */ = {
isa = PBXFileSystemSynchronizedRootGroup;
exceptions = (
E09D169A2F44CA20009C51A3 /* Exceptions for "EXO-iOS" folder in "EXO-iOS" target */,
);
path = "EXO-iOS";
sourceTree = "<group>";
};
E09D167F2F44CA20009C51A3 /* EXO-iOSTests */ = {
isa = PBXFileSystemSynchronizedRootGroup;
path = "EXO-iOSTests";
sourceTree = "<group>";
};
E09D16892F44CA20009C51A3 /* EXO-iOSUITests */ = {
isa = PBXFileSystemSynchronizedRootGroup;
path = "EXO-iOSUITests";
sourceTree = "<group>";
};
/* End PBXFileSystemSynchronizedRootGroup section */
/* Begin PBXFrameworksBuildPhase section */
E09D166C2F44CA1E009C51A3 /* Frameworks */ = {
isa = PBXFrameworksBuildPhase;
buildActionMask = 2147483647;
files = (
E09D17542F44F359009C51A3 /* MLXLMCommon in Frameworks */,
E09D17522F44F359009C51A3 /* MLXLLM in Frameworks */,
);
runOnlyForDeploymentPostprocessing = 0;
};
E09D16792F44CA20009C51A3 /* Frameworks */ = {
isa = PBXFrameworksBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
E09D16832F44CA20009C51A3 /* Frameworks */ = {
isa = PBXFrameworksBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
/* End PBXFrameworksBuildPhase section */
/* Begin PBXGroup section */
E09D16662F44CA1E009C51A3 = {
isa = PBXGroup;
children = (
E09D16712F44CA1E009C51A3 /* EXO-iOS */,
E09D167F2F44CA20009C51A3 /* EXO-iOSTests */,
E09D16892F44CA20009C51A3 /* EXO-iOSUITests */,
E09D16702F44CA1E009C51A3 /* Products */,
);
sourceTree = "<group>";
};
E09D16702F44CA1E009C51A3 /* Products */ = {
isa = PBXGroup;
children = (
E09D166F2F44CA1E009C51A3 /* EXO-iOS.app */,
E09D167C2F44CA20009C51A3 /* EXO-iOSTests.xctest */,
E09D16862F44CA20009C51A3 /* EXO-iOSUITests.xctest */,
);
name = Products;
sourceTree = "<group>";
};
/* End PBXGroup section */
/* Begin PBXNativeTarget section */
E09D166E2F44CA1E009C51A3 /* EXO-iOS */ = {
isa = PBXNativeTarget;
buildConfigurationList = E09D16902F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOS" */;
buildPhases = (
E09D166B2F44CA1E009C51A3 /* Sources */,
E09D166C2F44CA1E009C51A3 /* Frameworks */,
E09D166D2F44CA1E009C51A3 /* Resources */,
);
buildRules = (
);
dependencies = (
);
fileSystemSynchronizedGroups = (
E09D16712F44CA1E009C51A3 /* EXO-iOS */,
);
name = "EXO-iOS";
packageProductDependencies = (
E09D17512F44F359009C51A3 /* MLXLLM */,
E09D17532F44F359009C51A3 /* MLXLMCommon */,
);
productName = "EXO-iOS";
productReference = E09D166F2F44CA1E009C51A3 /* EXO-iOS.app */;
productType = "com.apple.product-type.application";
};
E09D167B2F44CA20009C51A3 /* EXO-iOSTests */ = {
isa = PBXNativeTarget;
buildConfigurationList = E09D16932F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSTests" */;
buildPhases = (
E09D16782F44CA20009C51A3 /* Sources */,
E09D16792F44CA20009C51A3 /* Frameworks */,
E09D167A2F44CA20009C51A3 /* Resources */,
);
buildRules = (
);
dependencies = (
E09D167E2F44CA20009C51A3 /* PBXTargetDependency */,
);
fileSystemSynchronizedGroups = (
E09D167F2F44CA20009C51A3 /* EXO-iOSTests */,
);
name = "EXO-iOSTests";
packageProductDependencies = (
);
productName = "EXO-iOSTests";
productReference = E09D167C2F44CA20009C51A3 /* EXO-iOSTests.xctest */;
productType = "com.apple.product-type.bundle.unit-test";
};
E09D16852F44CA20009C51A3 /* EXO-iOSUITests */ = {
isa = PBXNativeTarget;
buildConfigurationList = E09D16962F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSUITests" */;
buildPhases = (
E09D16822F44CA20009C51A3 /* Sources */,
E09D16832F44CA20009C51A3 /* Frameworks */,
E09D16842F44CA20009C51A3 /* Resources */,
);
buildRules = (
);
dependencies = (
E09D16882F44CA20009C51A3 /* PBXTargetDependency */,
);
fileSystemSynchronizedGroups = (
E09D16892F44CA20009C51A3 /* EXO-iOSUITests */,
);
name = "EXO-iOSUITests";
packageProductDependencies = (
);
productName = "EXO-iOSUITests";
productReference = E09D16862F44CA20009C51A3 /* EXO-iOSUITests.xctest */;
productType = "com.apple.product-type.bundle.ui-testing";
};
/* End PBXNativeTarget section */
/* Begin PBXProject section */
E09D16672F44CA1E009C51A3 /* Project object */ = {
isa = PBXProject;
attributes = {
BuildIndependentTargetsInParallel = 1;
LastSwiftUpdateCheck = 2620;
LastUpgradeCheck = 2620;
TargetAttributes = {
E09D166E2F44CA1E009C51A3 = {
CreatedOnToolsVersion = 26.2;
};
E09D167B2F44CA20009C51A3 = {
CreatedOnToolsVersion = 26.2;
TestTargetID = E09D166E2F44CA1E009C51A3;
};
E09D16852F44CA20009C51A3 = {
CreatedOnToolsVersion = 26.2;
TestTargetID = E09D166E2F44CA1E009C51A3;
};
};
};
buildConfigurationList = E09D166A2F44CA1E009C51A3 /* Build configuration list for PBXProject "EXO-iOS" */;
developmentRegion = en;
hasScannedForEncodings = 0;
knownRegions = (
en,
Base,
);
mainGroup = E09D16662F44CA1E009C51A3;
minimizedProjectReferenceProxies = 1;
packageReferences = (
E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */,
);
preferredProjectObjectVersion = 77;
productRefGroup = E09D16702F44CA1E009C51A3 /* Products */;
projectDirPath = "";
projectRoot = "";
targets = (
E09D166E2F44CA1E009C51A3 /* EXO-iOS */,
E09D167B2F44CA20009C51A3 /* EXO-iOSTests */,
E09D16852F44CA20009C51A3 /* EXO-iOSUITests */,
);
};
/* End PBXProject section */
/* Begin PBXResourcesBuildPhase section */
E09D166D2F44CA1E009C51A3 /* Resources */ = {
isa = PBXResourcesBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
E09D167A2F44CA20009C51A3 /* Resources */ = {
isa = PBXResourcesBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
E09D16842F44CA20009C51A3 /* Resources */ = {
isa = PBXResourcesBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
/* End PBXResourcesBuildPhase section */
/* Begin PBXSourcesBuildPhase section */
E09D166B2F44CA1E009C51A3 /* Sources */ = {
isa = PBXSourcesBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
E09D16782F44CA20009C51A3 /* Sources */ = {
isa = PBXSourcesBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
E09D16822F44CA20009C51A3 /* Sources */ = {
isa = PBXSourcesBuildPhase;
buildActionMask = 2147483647;
files = (
);
runOnlyForDeploymentPostprocessing = 0;
};
/* End PBXSourcesBuildPhase section */
/* Begin PBXTargetDependency section */
E09D167E2F44CA20009C51A3 /* PBXTargetDependency */ = {
isa = PBXTargetDependency;
target = E09D166E2F44CA1E009C51A3 /* EXO-iOS */;
targetProxy = E09D167D2F44CA20009C51A3 /* PBXContainerItemProxy */;
};
E09D16882F44CA20009C51A3 /* PBXTargetDependency */ = {
isa = PBXTargetDependency;
target = E09D166E2F44CA1E009C51A3 /* EXO-iOS */;
targetProxy = E09D16872F44CA20009C51A3 /* PBXContainerItemProxy */;
};
/* End PBXTargetDependency section */
/* Begin XCBuildConfiguration section */
E09D168E2F44CA20009C51A3 /* Debug */ = {
isa = XCBuildConfiguration;
buildSettings = {
ALWAYS_SEARCH_USER_PATHS = NO;
ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;
CLANG_ANALYZER_NONNULL = YES;
CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
CLANG_ENABLE_MODULES = YES;
CLANG_ENABLE_OBJC_ARC = YES;
CLANG_ENABLE_OBJC_WEAK = YES;
CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
CLANG_WARN_BOOL_CONVERSION = YES;
CLANG_WARN_COMMA = YES;
CLANG_WARN_CONSTANT_CONVERSION = YES;
CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
CLANG_WARN_EMPTY_BODY = YES;
CLANG_WARN_ENUM_CONVERSION = YES;
CLANG_WARN_INFINITE_RECURSION = YES;
CLANG_WARN_INT_CONVERSION = YES;
CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
CLANG_WARN_STRICT_PROTOTYPES = YES;
CLANG_WARN_SUSPICIOUS_MOVE = YES;
CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
CLANG_WARN_UNREACHABLE_CODE = YES;
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
COPY_PHASE_STRIP = NO;
DEBUG_INFORMATION_FORMAT = dwarf;
ENABLE_STRICT_OBJC_MSGSEND = YES;
ENABLE_TESTABILITY = YES;
ENABLE_USER_SCRIPT_SANDBOXING = YES;
GCC_C_LANGUAGE_STANDARD = gnu17;
GCC_DYNAMIC_NO_PIC = NO;
GCC_NO_COMMON_BLOCKS = YES;
GCC_OPTIMIZATION_LEVEL = 0;
GCC_PREPROCESSOR_DEFINITIONS = (
"DEBUG=1",
"$(inherited)",
);
GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
GCC_WARN_UNDECLARED_SELECTOR = YES;
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
GCC_WARN_UNUSED_FUNCTION = YES;
GCC_WARN_UNUSED_VARIABLE = YES;
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
MTL_FAST_MATH = YES;
ONLY_ACTIVE_ARCH = YES;
SDKROOT = iphoneos;
SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)";
SWIFT_OPTIMIZATION_LEVEL = "-Onone";
};
name = Debug;
};
E09D168F2F44CA20009C51A3 /* Release */ = {
isa = XCBuildConfiguration;
buildSettings = {
ALWAYS_SEARCH_USER_PATHS = NO;
ASSETCATALOG_COMPILER_GENERATE_SWIFT_ASSET_SYMBOL_EXTENSIONS = YES;
CLANG_ANALYZER_NONNULL = YES;
CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
CLANG_CXX_LANGUAGE_STANDARD = "gnu++20";
CLANG_ENABLE_MODULES = YES;
CLANG_ENABLE_OBJC_ARC = YES;
CLANG_ENABLE_OBJC_WEAK = YES;
CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
CLANG_WARN_BOOL_CONVERSION = YES;
CLANG_WARN_COMMA = YES;
CLANG_WARN_CONSTANT_CONVERSION = YES;
CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
CLANG_WARN_EMPTY_BODY = YES;
CLANG_WARN_ENUM_CONVERSION = YES;
CLANG_WARN_INFINITE_RECURSION = YES;
CLANG_WARN_INT_CONVERSION = YES;
CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
CLANG_WARN_STRICT_PROTOTYPES = YES;
CLANG_WARN_SUSPICIOUS_MOVE = YES;
CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
CLANG_WARN_UNREACHABLE_CODE = YES;
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
COPY_PHASE_STRIP = NO;
DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
ENABLE_NS_ASSERTIONS = NO;
ENABLE_STRICT_OBJC_MSGSEND = YES;
ENABLE_USER_SCRIPT_SANDBOXING = YES;
GCC_C_LANGUAGE_STANDARD = gnu17;
GCC_NO_COMMON_BLOCKS = YES;
GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
GCC_WARN_UNDECLARED_SELECTOR = YES;
GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
GCC_WARN_UNUSED_FUNCTION = YES;
GCC_WARN_UNUSED_VARIABLE = YES;
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
LOCALIZATION_PREFERS_STRING_CATALOGS = YES;
MTL_ENABLE_DEBUG_INFO = NO;
MTL_FAST_MATH = YES;
SDKROOT = iphoneos;
SWIFT_COMPILATION_MODE = wholemodule;
VALIDATE_PRODUCT = YES;
};
name = Release;
};
E09D16912F44CA20009C51A3 /* Debug */ = {
isa = XCBuildConfiguration;
buildSettings = {
ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
DEVELOPMENT_TEAM = 3M3M67U93M;
ENABLE_PREVIEWS = YES;
GENERATE_INFOPLIST_FILE = YES;
INFOPLIST_FILE = "EXO-iOS/Info.plist";
INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES;
INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES;
INFOPLIST_KEY_UILaunchScreen_Generation = YES;
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
LD_RUNPATH_SEARCH_PATHS = (
"$(inherited)",
"@executable_path/Frameworks",
);
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOS";
PRODUCT_NAME = "$(TARGET_NAME)";
STRING_CATALOG_GENERATE_SYMBOLS = YES;
SWIFT_APPROACHABLE_CONCURRENCY = YES;
SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor;
SWIFT_EMIT_LOC_STRINGS = YES;
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
SWIFT_VERSION = 5.0;
TARGETED_DEVICE_FAMILY = "1,2";
};
name = Debug;
};
E09D16922F44CA20009C51A3 /* Release */ = {
isa = XCBuildConfiguration;
buildSettings = {
ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
DEVELOPMENT_TEAM = 3M3M67U93M;
ENABLE_PREVIEWS = YES;
GENERATE_INFOPLIST_FILE = YES;
INFOPLIST_FILE = "EXO-iOS/Info.plist";
INFOPLIST_KEY_UIApplicationSceneManifest_Generation = YES;
INFOPLIST_KEY_UIApplicationSupportsIndirectInputEvents = YES;
INFOPLIST_KEY_UILaunchScreen_Generation = YES;
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPad = "UIInterfaceOrientationPortrait UIInterfaceOrientationPortraitUpsideDown UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
INFOPLIST_KEY_UISupportedInterfaceOrientations_iPhone = "UIInterfaceOrientationPortrait UIInterfaceOrientationLandscapeLeft UIInterfaceOrientationLandscapeRight";
LD_RUNPATH_SEARCH_PATHS = (
"$(inherited)",
"@executable_path/Frameworks",
);
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOS";
PRODUCT_NAME = "$(TARGET_NAME)";
STRING_CATALOG_GENERATE_SYMBOLS = YES;
SWIFT_APPROACHABLE_CONCURRENCY = YES;
SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor;
SWIFT_EMIT_LOC_STRINGS = YES;
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
SWIFT_VERSION = 5.0;
TARGETED_DEVICE_FAMILY = "1,2";
};
name = Release;
};
E09D16942F44CA20009C51A3 /* Debug */ = {
isa = XCBuildConfiguration;
buildSettings = {
BUNDLE_LOADER = "$(TEST_HOST)";
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
GENERATE_INFOPLIST_FILE = YES;
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSTests";
PRODUCT_NAME = "$(TARGET_NAME)";
STRING_CATALOG_GENERATE_SYMBOLS = NO;
SWIFT_APPROACHABLE_CONCURRENCY = YES;
SWIFT_EMIT_LOC_STRINGS = NO;
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
SWIFT_VERSION = 5.0;
TARGETED_DEVICE_FAMILY = "1,2";
TEST_HOST = "$(BUILT_PRODUCTS_DIR)/EXO-iOS.app/$(BUNDLE_EXECUTABLE_FOLDER_PATH)/EXO-iOS";
};
name = Debug;
};
E09D16952F44CA20009C51A3 /* Release */ = {
isa = XCBuildConfiguration;
buildSettings = {
BUNDLE_LOADER = "$(TEST_HOST)";
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
GENERATE_INFOPLIST_FILE = YES;
IPHONEOS_DEPLOYMENT_TARGET = 26.2;
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSTests";
PRODUCT_NAME = "$(TARGET_NAME)";
STRING_CATALOG_GENERATE_SYMBOLS = NO;
SWIFT_APPROACHABLE_CONCURRENCY = YES;
SWIFT_EMIT_LOC_STRINGS = NO;
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
SWIFT_VERSION = 5.0;
TARGETED_DEVICE_FAMILY = "1,2";
TEST_HOST = "$(BUILT_PRODUCTS_DIR)/EXO-iOS.app/$(BUNDLE_EXECUTABLE_FOLDER_PATH)/EXO-iOS";
};
name = Release;
};
E09D16972F44CA20009C51A3 /* Debug */ = {
isa = XCBuildConfiguration;
buildSettings = {
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
GENERATE_INFOPLIST_FILE = YES;
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSUITests";
PRODUCT_NAME = "$(TARGET_NAME)";
STRING_CATALOG_GENERATE_SYMBOLS = NO;
SWIFT_APPROACHABLE_CONCURRENCY = YES;
SWIFT_EMIT_LOC_STRINGS = NO;
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
SWIFT_VERSION = 5.0;
TARGETED_DEVICE_FAMILY = "1,2";
TEST_TARGET_NAME = "EXO-iOS";
};
name = Debug;
};
E09D16982F44CA20009C51A3 /* Release */ = {
isa = XCBuildConfiguration;
buildSettings = {
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
GENERATE_INFOPLIST_FILE = YES;
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = "com.exo.EXO-iOSUITests";
PRODUCT_NAME = "$(TARGET_NAME)";
STRING_CATALOG_GENERATE_SYMBOLS = NO;
SWIFT_APPROACHABLE_CONCURRENCY = YES;
SWIFT_EMIT_LOC_STRINGS = NO;
SWIFT_UPCOMING_FEATURE_MEMBER_IMPORT_VISIBILITY = YES;
SWIFT_VERSION = 5.0;
TARGETED_DEVICE_FAMILY = "1,2";
TEST_TARGET_NAME = "EXO-iOS";
};
name = Release;
};
/* End XCBuildConfiguration section */
/* Begin XCConfigurationList section */
E09D166A2F44CA1E009C51A3 /* Build configuration list for PBXProject "EXO-iOS" */ = {
isa = XCConfigurationList;
buildConfigurations = (
E09D168E2F44CA20009C51A3 /* Debug */,
E09D168F2F44CA20009C51A3 /* Release */,
);
defaultConfigurationIsVisible = 0;
defaultConfigurationName = Release;
};
E09D16902F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOS" */ = {
isa = XCConfigurationList;
buildConfigurations = (
E09D16912F44CA20009C51A3 /* Debug */,
E09D16922F44CA20009C51A3 /* Release */,
);
defaultConfigurationIsVisible = 0;
defaultConfigurationName = Release;
};
E09D16932F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSTests" */ = {
isa = XCConfigurationList;
buildConfigurations = (
E09D16942F44CA20009C51A3 /* Debug */,
E09D16952F44CA20009C51A3 /* Release */,
);
defaultConfigurationIsVisible = 0;
defaultConfigurationName = Release;
};
E09D16962F44CA20009C51A3 /* Build configuration list for PBXNativeTarget "EXO-iOSUITests" */ = {
isa = XCConfigurationList;
buildConfigurations = (
E09D16972F44CA20009C51A3 /* Debug */,
E09D16982F44CA20009C51A3 /* Release */,
);
defaultConfigurationIsVisible = 0;
defaultConfigurationName = Release;
};
/* End XCConfigurationList section */
/* Begin XCRemoteSwiftPackageReference section */
E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */ = {
isa = XCRemoteSwiftPackageReference;
repositoryURL = "https://github.com/ml-explore/mlx-swift-lm";
requirement = {
kind = upToNextMajorVersion;
minimumVersion = 2.30.3;
};
};
/* End XCRemoteSwiftPackageReference section */
/* Begin XCSwiftPackageProductDependency section */
E09D17512F44F359009C51A3 /* MLXLLM */ = {
isa = XCSwiftPackageProductDependency;
package = E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */;
productName = MLXLLM;
};
E09D17532F44F359009C51A3 /* MLXLMCommon */ = {
isa = XCSwiftPackageProductDependency;
package = E09D17502F44F359009C51A3 /* XCRemoteSwiftPackageReference "mlx-swift-lm" */;
productName = MLXLMCommon;
};
/* End XCSwiftPackageProductDependency section */
};
rootObject = E09D16672F44CA1E009C51A3 /* Project object */;
}
@@ -0,0 +1,7 @@
<?xml version="1.0" encoding="UTF-8"?>
<Workspace
version = "1.0">
<FileRef
location = "self:">
</FileRef>
</Workspace>
@@ -0,0 +1,60 @@
{
"originHash" : "facc0ac7c70363ea20f6cd1235de91dea6b06f0d00190946045a6c8ae753abc2",
"pins" : [
{
"identity" : "mlx-swift",
"kind" : "remoteSourceControl",
"location" : "https://github.com/ml-explore/mlx-swift",
"state" : {
"revision" : "6ba4827fb82c97d012eec9ab4b2de21f85c3b33d",
"version" : "0.30.6"
}
},
{
"identity" : "mlx-swift-lm",
"kind" : "remoteSourceControl",
"location" : "https://github.com/ml-explore/mlx-swift-lm",
"state" : {
"revision" : "360c5052b81cc154b04ee0933597a4ad6db4b8ae",
"version" : "2.30.3"
}
},
{
"identity" : "swift-collections",
"kind" : "remoteSourceControl",
"location" : "https://github.com/apple/swift-collections.git",
"state" : {
"revision" : "7b847a3b7008b2dc2f47ca3110d8c782fb2e5c7e",
"version" : "1.3.0"
}
},
{
"identity" : "swift-jinja",
"kind" : "remoteSourceControl",
"location" : "https://github.com/huggingface/swift-jinja.git",
"state" : {
"revision" : "d81197f35f41445bc10e94600795e68c6f5e94b0",
"version" : "2.3.1"
}
},
{
"identity" : "swift-numerics",
"kind" : "remoteSourceControl",
"location" : "https://github.com/apple/swift-numerics",
"state" : {
"revision" : "0c0290ff6b24942dadb83a929ffaaa1481df04a2",
"version" : "1.1.1"
}
},
{
"identity" : "swift-transformers",
"kind" : "remoteSourceControl",
"location" : "https://github.com/huggingface/swift-transformers",
"state" : {
"revision" : "573e5c9036c2f136b3a8a071da8e8907322403d0",
"version" : "1.1.6"
}
}
],
"version" : 3
}
@@ -0,0 +1,20 @@
{
"colors" : [
{
"color" : {
"color-space" : "srgb",
"components" : {
"alpha" : "1.000",
"blue" : "0x00",
"green" : "0xD7",
"red" : "0xFF"
}
},
"idiom" : "universal"
}
],
"info" : {
"author" : "xcode",
"version" : 1
}
}
Binary file not shown.

After

Width:  |  Height:  |  Size: 10 KiB

@@ -0,0 +1,38 @@
{
"images" : [
{
"filename" : "AppIcon.png",
"idiom" : "universal",
"platform" : "ios",
"size" : "1024x1024"
},
{
"appearances" : [
{
"appearance" : "luminosity",
"value" : "dark"
}
],
"filename" : "AppIcon.png",
"idiom" : "universal",
"platform" : "ios",
"size" : "1024x1024"
},
{
"appearances" : [
{
"appearance" : "luminosity",
"value" : "tinted"
}
],
"filename" : "AppIcon.png",
"idiom" : "universal",
"platform" : "ios",
"size" : "1024x1024"
}
],
"info" : {
"author" : "xcode",
"version" : 1
}
}
@@ -0,0 +1,6 @@
{
"info" : {
"author" : "xcode",
"version" : 1
}
}
@@ -0,0 +1,21 @@
{
"images" : [
{
"filename" : "exo-logo.png",
"idiom" : "universal",
"scale" : "1x"
},
{
"idiom" : "universal",
"scale" : "2x"
},
{
"idiom" : "universal",
"scale" : "3x"
}
],
"info" : {
"author" : "xcode",
"version" : 1
}
}
Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 KiB

+8
View File
@@ -0,0 +1,8 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>com.apple.developer.kernel.increased-memory-limit</key>
<true/>
</dict>
</plist>
+67
View File
@@ -0,0 +1,67 @@
import SwiftUI
import UIKit
@main
struct EXO_iOSApp: App {
@State private var clusterService = ClusterService()
@State private var discoveryService = DiscoveryService()
@State private var localInferenceService = LocalInferenceService()
@State private var chatService: ChatService?
init() {
let darkGray = UIColor(red: 0x1F / 255.0, green: 0x1F / 255.0, blue: 0x1F / 255.0, alpha: 1)
let yellow = UIColor(red: 0xFF / 255.0, green: 0xD7 / 255.0, blue: 0x00 / 255.0, alpha: 1)
let navAppearance = UINavigationBarAppearance()
navAppearance.configureWithOpaqueBackground()
navAppearance.backgroundColor = darkGray
navAppearance.titleTextAttributes = [
.foregroundColor: yellow,
.font: UIFont.monospacedSystemFont(ofSize: 17, weight: .semibold),
]
navAppearance.largeTitleTextAttributes = [
.foregroundColor: yellow,
.font: UIFont.monospacedSystemFont(ofSize: 34, weight: .bold),
]
UINavigationBar.appearance().standardAppearance = navAppearance
UINavigationBar.appearance().compactAppearance = navAppearance
UINavigationBar.appearance().scrollEdgeAppearance = navAppearance
UINavigationBar.appearance().tintColor = yellow
}
var body: some Scene {
WindowGroup {
if let chatService {
RootView()
.environment(clusterService)
.environment(discoveryService)
.environment(chatService)
.environment(localInferenceService)
.preferredColorScheme(.dark)
.task {
await clusterService.attemptAutoReconnect()
discoveryService.startBrowsing()
await localInferenceService.prepareModel()
}
.onChange(of: discoveryService.discoveredClusters) { _, clusters in
guard !clusterService.isConnected,
case .disconnected = clusterService.connectionState,
let first = clusters.first
else { return }
Task {
await clusterService.connectToDiscoveredCluster(
first, using: discoveryService)
}
}
} else {
Color.exoBlack.onAppear {
chatService = ChatService(
clusterService: clusterService,
localInferenceService: localInferenceService
)
}
}
}
}
}
+19
View File
@@ -0,0 +1,19 @@
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
<key>UIUserInterfaceStyle</key>
<string>Dark</string>
<key>CFBundleDisplayName</key>
<string>EXO</string>
<key>NSLocalNetworkUsageDescription</key>
<string>EXO needs local network access to connect to your EXO cluster.</string>
<key>NSBonjourServices</key>
<array>
<string>_exo._tcp</string>
<string>_p2p._tcp</string>
<string>_p2p._udp</string>
<string>_libp2p._udp</string>
</array>
</dict>
</plist>
@@ -0,0 +1,129 @@
import Foundation
// MARK: - Request
struct ChatCompletionRequest: Encodable {
let model: String
let messages: [ChatCompletionMessageParam]
let stream: Bool
let maxTokens: Int?
let temperature: Double?
enum CodingKeys: String, CodingKey {
case model, messages, stream, temperature
case maxTokens = "max_tokens"
}
}
struct ChatCompletionMessageParam: Encodable {
let role: String
let content: String
}
// MARK: - Streaming Response
struct ChatCompletionChunk: Decodable {
let id: String
let model: String?
let choices: [StreamingChoice]
let usage: ChunkUsage?
init(id: String, model: String?, choices: [StreamingChoice], usage: ChunkUsage?) {
self.id = id
self.model = model
self.choices = choices
self.usage = usage
}
}
struct StreamingChoice: Decodable {
let index: Int
let delta: Delta
let finishReason: String?
enum CodingKeys: String, CodingKey {
case index, delta
case finishReason = "finish_reason"
}
init(index: Int, delta: Delta, finishReason: String?) {
self.index = index
self.delta = delta
self.finishReason = finishReason
}
}
struct Delta: Decodable {
let role: String?
let content: String?
init(role: String?, content: String?) {
self.role = role
self.content = content
}
}
struct ChunkUsage: Decodable {
let promptTokens: Int?
let completionTokens: Int?
let totalTokens: Int?
enum CodingKeys: String, CodingKey {
case promptTokens = "prompt_tokens"
case completionTokens = "completion_tokens"
case totalTokens = "total_tokens"
}
init(promptTokens: Int?, completionTokens: Int?, totalTokens: Int?) {
self.promptTokens = promptTokens
self.completionTokens = completionTokens
self.totalTokens = totalTokens
}
}
// MARK: - Non-Streaming Response
struct ChatCompletionResponse: Decodable {
let id: String
let model: String?
let choices: [ResponseChoice]
}
struct ResponseChoice: Decodable {
let index: Int
let message: ResponseMessage
let finishReason: String?
enum CodingKeys: String, CodingKey {
case index, message
case finishReason = "finish_reason"
}
}
struct ResponseMessage: Decodable {
let role: String?
let content: String?
}
// MARK: - Models List
struct ModelListResponse: Decodable {
let data: [ModelInfo]
}
struct ModelInfo: Decodable, Identifiable {
let id: String
let name: String?
}
// MARK: - Error
struct APIErrorResponse: Decodable {
let error: APIErrorInfo
}
struct APIErrorInfo: Decodable {
let message: String
let type: String?
let code: Int?
}
@@ -0,0 +1,26 @@
import Foundation
struct ChatMessage: Identifiable, Equatable {
let id: UUID
let role: Role
var content: String
let timestamp: Date
var isStreaming: Bool
enum Role: String, Codable {
case user
case assistant
case system
}
init(
id: UUID = UUID(), role: Role, content: String, timestamp: Date = Date(),
isStreaming: Bool = false
) {
self.id = id
self.role = role
self.content = content
self.timestamp = timestamp
self.isStreaming = isStreaming
}
}
@@ -0,0 +1,11 @@
import Foundation
struct ConnectionInfo: Codable, Equatable {
let host: String
let port: Int
let nodeId: String?
var baseURL: URL { URL(string: "http://\(host):\(port)")! }
static let defaultPort = 52415
}
@@ -0,0 +1,34 @@
import Foundation
struct Conversation: Identifiable, Codable, Equatable {
let id: UUID
var title: String
var messages: [StoredMessage]
var modelId: String?
let createdAt: Date
init(
id: UUID = UUID(), title: String = "New Chat", messages: [StoredMessage] = [],
modelId: String? = nil, createdAt: Date = Date()
) {
self.id = id
self.title = title
self.messages = messages
self.modelId = modelId
self.createdAt = createdAt
}
}
struct StoredMessage: Identifiable, Codable, Equatable {
let id: UUID
let role: String
var content: String
let timestamp: Date
init(id: UUID = UUID(), role: String, content: String, timestamp: Date = Date()) {
self.id = id
self.role = role
self.content = content
self.timestamp = timestamp
}
}
@@ -0,0 +1,227 @@
import Foundation
@Observable
@MainActor
final class ChatService {
var conversations: [Conversation] = []
var activeConversationId: UUID?
private(set) var isGenerating: Bool = false
private var currentGenerationTask: Task<Void, Never>?
private let clusterService: ClusterService
private let localInferenceService: LocalInferenceService
var canSendMessage: Bool {
clusterService.isConnected || localInferenceService.isAvailable
}
var activeConversation: Conversation? {
guard let id = activeConversationId else { return nil }
return conversations.first { $0.id == id }
}
var activeMessages: [ChatMessage] {
guard let conversation = activeConversation else { return [] }
return conversation.messages.map { stored in
ChatMessage(
id: stored.id,
role: ChatMessage.Role(rawValue: stored.role) ?? .user,
content: stored.content,
timestamp: stored.timestamp
)
}
}
init(clusterService: ClusterService, localInferenceService: LocalInferenceService) {
self.clusterService = clusterService
self.localInferenceService = localInferenceService
loadConversations()
}
// MARK: - Conversation Management
func createConversation(modelId: String? = nil) {
let conversation = Conversation(
modelId: modelId ?? clusterService.availableModels.first?.id)
conversations.insert(conversation, at: 0)
activeConversationId = conversation.id
saveConversations()
}
func deleteConversation(id: UUID) {
conversations.removeAll { $0.id == id }
if activeConversationId == id {
activeConversationId = conversations.first?.id
}
saveConversations()
}
func setActiveConversation(id: UUID) {
activeConversationId = id
}
func setModelForActiveConversation(_ modelId: String) {
guard let index = conversations.firstIndex(where: { $0.id == activeConversationId }) else {
return
}
conversations[index].modelId = modelId
saveConversations()
}
// MARK: - Messaging
func sendMessage(_ text: String) {
guard !text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty else { return }
if activeConversation == nil {
createConversation()
}
guard let index = conversations.firstIndex(where: { $0.id == activeConversationId }) else {
return
}
let userMessage = StoredMessage(role: "user", content: text)
conversations[index].messages.append(userMessage)
if conversations[index].title == "New Chat" {
let preview = String(text.prefix(40))
conversations[index].title = preview + (text.count > 40 ? "..." : "")
}
let modelId: String
if clusterService.isConnected {
guard
let clusterId = conversations[index].modelId
?? clusterService.availableModels.first?.id
else {
let errorMessage = StoredMessage(
role: "assistant", content: "No model selected. Please select a model first.")
conversations[index].messages.append(errorMessage)
saveConversations()
return
}
modelId = clusterId
} else if localInferenceService.isAvailable {
modelId = localInferenceService.defaultModelId
} else {
let errorMessage = StoredMessage(
role: "assistant",
content: "Not connected to a cluster and local model is not available.")
conversations[index].messages.append(errorMessage)
saveConversations()
return
}
conversations[index].modelId = modelId
let assistantMessageId = UUID()
let assistantMessage = StoredMessage(
id: assistantMessageId, role: "assistant", content: "", timestamp: Date())
conversations[index].messages.append(assistantMessage)
let messagesForAPI = conversations[index].messages.dropLast().map { stored in
ChatCompletionMessageParam(role: stored.role, content: stored.content)
}
let request = ChatCompletionRequest(
model: modelId,
messages: Array(messagesForAPI),
stream: true,
maxTokens: 4096,
temperature: nil
)
let conversationId = conversations[index].id
isGenerating = true
currentGenerationTask = Task { [weak self] in
guard let self else { return }
await self.performStreaming(
request: request, conversationId: conversationId,
assistantMessageId: assistantMessageId)
}
saveConversations()
}
func cancelGeneration() {
currentGenerationTask?.cancel()
currentGenerationTask = nil
localInferenceService.cancelGeneration()
isGenerating = false
}
// MARK: - Streaming
private func performStreaming(
request: ChatCompletionRequest, conversationId: UUID, assistantMessageId: UUID
) async {
defer {
isGenerating = false
currentGenerationTask = nil
saveConversations()
}
do {
let stream =
clusterService.isConnected
? clusterService.streamChatCompletion(request: request)
: localInferenceService.streamChatCompletion(request: request)
for try await chunk in stream {
guard !Task.isCancelled else { return }
guard let content = chunk.choices.first?.delta.content, !content.isEmpty else {
continue
}
if let convIndex = conversations.firstIndex(where: { $0.id == conversationId }),
let msgIndex = conversations[convIndex].messages.firstIndex(where: {
$0.id == assistantMessageId
})
{
conversations[convIndex].messages[msgIndex].content += content
}
}
} catch {
if !Task.isCancelled {
if let convIndex = conversations.firstIndex(where: { $0.id == conversationId }),
let msgIndex = conversations[convIndex].messages.firstIndex(where: {
$0.id == assistantMessageId
})
{
if conversations[convIndex].messages[msgIndex].content.isEmpty {
conversations[convIndex].messages[msgIndex].content =
"Error: \(error.localizedDescription)"
}
}
}
}
}
// MARK: - Persistence
private static var storageURL: URL {
let documents = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)
.first!
return documents.appendingPathComponent("exo_conversations.json")
}
private func saveConversations() {
do {
let data = try JSONEncoder().encode(conversations)
try data.write(to: Self.storageURL, options: .atomic)
} catch {
// Save failed silently
}
}
private func loadConversations() {
do {
let data = try Data(contentsOf: Self.storageURL)
conversations = try JSONDecoder().decode([Conversation].self, from: data)
activeConversationId = conversations.first?.id
} catch {
conversations = []
}
}
}
@@ -0,0 +1,246 @@
import Foundation
enum ConnectionState: Equatable {
case disconnected
case connecting
case connected(ConnectionInfo)
}
struct ModelOption: Identifiable, Equatable {
let id: String
let displayName: String
}
@Observable
@MainActor
final class ClusterService {
private(set) var connectionState: ConnectionState = .disconnected
private(set) var availableModels: [ModelOption] = []
private(set) var lastError: String?
private let session: URLSession
private let decoder: JSONDecoder
private var pollingTask: Task<Void, Never>?
private var heartbeatTask: Task<Void, Never>?
private static let connectionInfoKey = "exo_last_connection_info"
var isConnected: Bool {
if case .connected = connectionState { return true }
return false
}
var currentConnection: ConnectionInfo? {
if case .connected(let info) = connectionState { return info }
return nil
}
init(session: URLSession = .shared) {
self.session = session
let decoder = JSONDecoder()
self.decoder = decoder
}
// MARK: - Connection
func connect(to info: ConnectionInfo) async {
connectionState = .connecting
lastError = nil
do {
let url = info.baseURL.appendingPathComponent("node_id")
var request = URLRequest(url: url)
request.timeoutInterval = 5
request.cachePolicy = .reloadIgnoringLocalCacheData
let (_, response) = try await session.data(for: request)
guard let httpResponse = response as? HTTPURLResponse,
(200..<300).contains(httpResponse.statusCode)
else {
throw URLError(.badServerResponse)
}
connectionState = .connected(info)
persistConnection(info)
startPolling()
startHeartbeat()
await fetchModels(baseURL: info.baseURL)
} catch {
connectionState = .disconnected
lastError = "Could not connect to \(info.host):\(info.port)"
}
}
func connectToDiscoveredCluster(
_ cluster: DiscoveredCluster, using discoveryService: DiscoveryService
) async {
guard case .disconnected = connectionState else { return }
connectionState = .connecting
lastError = nil
guard let info = await discoveryService.resolve(cluster) else {
connectionState = .disconnected
lastError = "Could not resolve \(cluster.name)"
return
}
connectionState = .disconnected // reset so connect() can proceed
await connect(to: info)
}
func disconnect() {
stopPolling()
stopHeartbeat()
connectionState = .disconnected
availableModels = []
lastError = nil
}
func attemptAutoReconnect() async {
guard case .disconnected = connectionState,
let info = loadPersistedConnection()
else { return }
await connect(to: info)
}
// MARK: - Polling
private func startPolling(interval: TimeInterval = 2.0) {
stopPolling()
pollingTask = Task { [weak self] in
while !Task.isCancelled {
try? await Task.sleep(for: .seconds(interval))
guard let self, !Task.isCancelled else { return }
guard let connection = self.currentConnection else { return }
await self.fetchModels(baseURL: connection.baseURL)
}
}
}
private func stopPolling() {
pollingTask?.cancel()
pollingTask = nil
}
// MARK: - Heartbeat
private func startHeartbeat(interval: TimeInterval = 10.0) {
stopHeartbeat()
heartbeatTask = Task { [weak self] in
// Send immediately, then on interval
if let self, let connection = self.currentConnection {
await self.sendHeartbeat(baseURL: connection.baseURL)
}
while !Task.isCancelled {
try? await Task.sleep(for: .seconds(interval))
guard let self, !Task.isCancelled else { return }
guard let connection = self.currentConnection else { return }
await self.sendHeartbeat(baseURL: connection.baseURL)
}
}
}
private func stopHeartbeat() {
heartbeatTask?.cancel()
heartbeatTask = nil
}
private func sendHeartbeat(baseURL: URL) async {
do {
let deviceInfo = DeviceInfoService.gather()
let url = baseURL.appendingPathComponent("v1/lite_node/heartbeat")
var request = URLRequest(url: url)
request.httpMethod = "POST"
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.timeoutInterval = 5
request.httpBody = try JSONEncoder().encode(deviceInfo)
let (_, response) = try await session.data(for: request)
if let httpResponse = response as? HTTPURLResponse,
!(200..<300).contains(httpResponse.statusCode) {
// Heartbeat failed silently will retry on next interval
}
} catch {
// Heartbeat failed silently will retry on next interval
}
}
// MARK: - API
private func fetchModels(baseURL: URL) async {
do {
let url = baseURL.appendingPathComponent("models")
var request = URLRequest(url: url)
request.cachePolicy = .reloadIgnoringLocalCacheData
let (data, response) = try await session.data(for: request)
guard let httpResponse = response as? HTTPURLResponse,
(200..<300).contains(httpResponse.statusCode)
else { return }
let list = try decoder.decode(ModelListResponse.self, from: data)
availableModels = list.data.map {
ModelOption(id: $0.id, displayName: $0.name ?? $0.id)
}
} catch {
// Models fetch failed silently will retry on next poll
}
}
func streamChatCompletion(request body: ChatCompletionRequest) -> AsyncThrowingStream<
ChatCompletionChunk, Error
> {
AsyncThrowingStream { continuation in
let task = Task { [weak self] in
guard let self, let connection = self.currentConnection else {
continuation.finish(throwing: URLError(.notConnectedToInternet))
return
}
do {
let url = connection.baseURL.appendingPathComponent("v1/chat/completions")
var request = URLRequest(url: url)
request.httpMethod = "POST"
request.setValue("application/json", forHTTPHeaderField: "Content-Type")
request.httpBody = try JSONEncoder().encode(body)
let (bytes, response) = try await self.session.bytes(for: request)
guard let httpResponse = response as? HTTPURLResponse,
(200..<300).contains(httpResponse.statusCode)
else {
continuation.finish(throwing: URLError(.badServerResponse))
return
}
let parser = SSEStreamParser<ChatCompletionChunk>(
bytes: bytes, decoder: self.decoder)
for try await chunk in parser {
continuation.yield(chunk)
}
continuation.finish()
} catch {
continuation.finish(throwing: error)
}
}
continuation.onTermination = { _ in
task.cancel()
}
}
}
// MARK: - Persistence
private func persistConnection(_ info: ConnectionInfo) {
if let data = try? JSONEncoder().encode(info) {
UserDefaults.standard.set(data, forKey: Self.connectionInfoKey)
}
}
private func loadPersistedConnection() -> ConnectionInfo? {
guard let data = UserDefaults.standard.data(forKey: Self.connectionInfoKey) else {
return nil
}
return try? JSONDecoder().decode(ConnectionInfo.self, from: data)
}
}
@@ -0,0 +1,115 @@
import Foundation
import os
import UIKit
struct DeviceInfo: Encodable {
let nodeId: String
let model: String
let chip: String
let osVersion: String
let friendlyName: String
let ramTotal: Int
let ramAvailable: Int
enum CodingKeys: String, CodingKey {
case nodeId = "node_id"
case model
case chip
case osVersion = "os_version"
case friendlyName = "friendly_name"
case ramTotal = "ram_total"
case ramAvailable = "ram_available"
}
}
enum DeviceInfoService {
private static let liteNodeIdKey = "exo_lite_node_id"
static var liteNodeId: String {
if let existing = UserDefaults.standard.string(forKey: liteNodeIdKey) {
return existing
}
let newId = UUID().uuidString.lowercased()
UserDefaults.standard.set(newId, forKey: liteNodeIdKey)
return newId
}
static func gather() -> DeviceInfo {
let model = modelName()
let chip = chipName(for: model)
let totalRam = Int(ProcessInfo.processInfo.physicalMemory)
let availableRam = availableMemory()
return DeviceInfo(
nodeId: liteNodeId,
model: model,
chip: chip,
osVersion: UIDevice.current.systemVersion,
friendlyName: UIDevice.current.name,
ramTotal: totalRam,
ramAvailable: availableRam
)
}
// MARK: - Private
private static func modelName() -> String {
var systemInfo = utsname()
uname(&systemInfo)
let machine = withUnsafePointer(to: &systemInfo.machine) {
$0.withMemoryRebound(to: CChar.self, capacity: 1) {
String(cString: $0)
}
}
return modelMapping[machine] ?? machine
}
private static func chipName(for model: String) -> String {
let lower = model.lowercased()
if lower.contains("iphone 16 pro") || lower.contains("iphone 16 pro max") {
return "Apple A18 Pro"
} else if lower.contains("iphone 16") {
return "Apple A18"
} else if lower.contains("iphone 15 pro") || lower.contains("iphone 15 pro max") {
return "Apple A17 Pro"
} else if lower.contains("iphone 15") {
return "Apple A16 Bionic"
} else if lower.contains("iphone 14 pro") || lower.contains("iphone 14 pro max") {
return "Apple A16 Bionic"
} else if lower.contains("iphone 14") {
return "Apple A15 Bionic"
}
return "Apple Silicon"
}
private static func availableMemory() -> Int {
return Int(os_proc_available_memory())
}
private static let modelMapping: [String: String] = [
// iPhone 16 series
"iPhone17,1": "iPhone 16 Pro",
"iPhone17,2": "iPhone 16 Pro Max",
"iPhone17,3": "iPhone 16",
"iPhone17,4": "iPhone 16 Plus",
// iPhone 15 series
"iPhone16,1": "iPhone 15 Pro",
"iPhone16,2": "iPhone 15 Pro Max",
"iPhone15,4": "iPhone 15",
"iPhone15,5": "iPhone 15 Plus",
// iPhone 14 series
"iPhone15,2": "iPhone 14 Pro",
"iPhone15,3": "iPhone 14 Pro Max",
"iPhone14,7": "iPhone 14",
"iPhone14,8": "iPhone 14 Plus",
// iPhone 13 series
"iPhone14,2": "iPhone 13 Pro",
"iPhone14,3": "iPhone 13 Pro Max",
"iPhone14,5": "iPhone 13",
"iPhone14,4": "iPhone 13 mini",
// Simulator
"arm64": "iPhone (Simulator)",
"x86_64": "iPhone (Simulator)",
]
}
@@ -0,0 +1,123 @@
import Foundation
import Network
import os
struct DiscoveredCluster: Identifiable, Equatable {
let id: String
let name: String
let endpoint: NWEndpoint
static func == (lhs: DiscoveredCluster, rhs: DiscoveredCluster) -> Bool {
lhs.id == rhs.id && lhs.name == rhs.name
}
}
@Observable
@MainActor
final class DiscoveryService {
private(set) var discoveredClusters: [DiscoveredCluster] = []
private(set) var isSearching = false
private var browser: NWBrowser?
func startBrowsing() {
guard browser == nil else { return }
let browser = NWBrowser(for: .bonjour(type: "_exo._tcp", domain: nil), using: .tcp)
browser.stateUpdateHandler = { [weak self] state in
guard let service = self else { return }
Task { @MainActor in
switch state {
case .ready:
service.isSearching = true
case .failed, .cancelled:
service.isSearching = false
default:
break
}
}
}
browser.browseResultsChangedHandler = { [weak self] results, _ in
guard let service = self else { return }
Task { @MainActor in
service.discoveredClusters = results.compactMap { result in
guard case .service(let name, _, _, _) = result.endpoint else {
return nil
}
return DiscoveredCluster(
id: name,
name: name,
endpoint: result.endpoint
)
}
}
}
browser.start(queue: .main)
self.browser = browser
}
func stopBrowsing() {
browser?.cancel()
browser = nil
isSearching = false
discoveredClusters = []
}
/// Resolve a discovered Bonjour endpoint to an IP address and port, then return a ConnectionInfo.
func resolve(_ cluster: DiscoveredCluster) async -> ConnectionInfo? {
await withCheckedContinuation { continuation in
let didResume = OSAllocatedUnfairLock(initialState: false)
let connection = NWConnection(to: cluster.endpoint, using: .tcp)
connection.stateUpdateHandler = { state in
guard
didResume.withLock({
guard !$0 else { return false }
$0 = true
return true
})
else { return }
switch state {
case .ready:
if let innerEndpoint = connection.currentPath?.remoteEndpoint,
case .hostPort(let host, let port) = innerEndpoint
{
var hostString: String
switch host {
case .ipv4(let addr):
hostString = "\(addr)"
case .ipv6(let addr):
hostString = "\(addr)"
case .name(let name, _):
hostString = name
@unknown default:
hostString = "\(host)"
}
// Strip interface scope suffix (e.g. "%en0")
if let pct = hostString.firstIndex(of: "%") {
hostString = String(hostString[..<pct])
}
let info = ConnectionInfo(
host: hostString,
port: Int(port.rawValue),
nodeId: nil
)
connection.cancel()
continuation.resume(returning: info)
} else {
connection.cancel()
continuation.resume(returning: nil)
}
case .failed, .cancelled:
continuation.resume(returning: nil)
default:
// Not a terminal state allow future callbacks
didResume.withLock { $0 = false }
}
}
connection.start(queue: .global(qos: .userInitiated))
}
}
}
@@ -0,0 +1,201 @@
import Foundation
import MLXLLM
import MLXLMCommon
enum LocalModelState: Equatable {
case notDownloaded
case downloading(progress: Double)
case downloaded
case loading
case ready
case generating
case error(String)
}
@Observable
@MainActor
final class LocalInferenceService {
private(set) var modelState: LocalModelState = .notDownloaded
private var modelContainer: ModelContainer?
private var generationTask: Task<Void, Never>?
let defaultModelId = "mlx-community/Qwen3-0.6B-4bit"
private static let modelDownloadedKey = "exo_local_model_downloaded"
var isReady: Bool {
modelState == .ready
}
var isAvailable: Bool {
modelState == .ready || modelState == .generating
}
init() {
if UserDefaults.standard.bool(forKey: Self.modelDownloadedKey) {
modelState = .downloaded
}
}
// MARK: - Model Lifecycle
func prepareModel() async {
guard modelState == .notDownloaded || modelState == .downloaded else { return }
let wasDownloaded = modelState == .downloaded
if !wasDownloaded {
modelState = .downloading(progress: 0)
} else {
modelState = .loading
}
do {
let container = try await loadModelContainer(
id: defaultModelId
) { [weak self] progress in
guard let self else { return }
Task { @MainActor in
if case .downloading = self.modelState {
self.modelState = .downloading(progress: progress.fractionCompleted)
}
}
}
self.modelContainer = container
UserDefaults.standard.set(true, forKey: Self.modelDownloadedKey)
modelState = .ready
} catch {
modelState = .error(error.localizedDescription)
}
}
func unloadModel() {
cancelGeneration()
modelContainer = nil
modelState = .downloaded
}
// MARK: - Generation
func streamChatCompletion(request: ChatCompletionRequest) -> AsyncThrowingStream<
ChatCompletionChunk, Error
> {
AsyncThrowingStream { continuation in
let task = Task { [weak self] in
guard let self else {
continuation.finish(throwing: LocalInferenceError.serviceUnavailable)
return
}
guard let container = self.modelContainer else {
continuation.finish(throwing: LocalInferenceError.modelNotLoaded)
return
}
await MainActor.run {
self.modelState = .generating
}
defer {
Task { @MainActor [weak self] in
if self?.modelState == .generating {
self?.modelState = .ready
}
}
}
let chunkId = "local-\(UUID().uuidString)"
do {
// Build Chat.Message array from the request
var chatMessages: [Chat.Message] = []
for msg in request.messages {
switch msg.role {
case "system":
chatMessages.append(.system(msg.content))
case "assistant":
chatMessages.append(.assistant(msg.content))
default:
chatMessages.append(.user(msg.content))
}
}
// Use ChatSession for streaming generation
let session = ChatSession(
container,
history: chatMessages,
generateParameters: GenerateParameters(
maxTokens: request.maxTokens ?? 4096,
temperature: Float(request.temperature ?? 0.7)
)
)
// Stream with an empty prompt since history already contains the conversation
let stream = session.streamResponse(to: "")
for try await text in stream {
if Task.isCancelled { break }
let chunk = ChatCompletionChunk(
id: chunkId,
model: request.model,
choices: [
StreamingChoice(
index: 0,
delta: Delta(role: nil, content: text),
finishReason: nil
)
],
usage: nil
)
continuation.yield(chunk)
}
// Send final chunk with finish reason
let finalChunk = ChatCompletionChunk(
id: chunkId,
model: request.model,
choices: [
StreamingChoice(
index: 0,
delta: Delta(role: nil, content: nil),
finishReason: "stop"
)
],
usage: nil
)
continuation.yield(finalChunk)
continuation.finish()
} catch {
continuation.finish(throwing: error)
}
}
self.generationTask = task
continuation.onTermination = { _ in
task.cancel()
}
}
}
func cancelGeneration() {
generationTask?.cancel()
generationTask = nil
if modelState == .generating {
modelState = .ready
}
}
}
enum LocalInferenceError: LocalizedError {
case serviceUnavailable
case modelNotLoaded
var errorDescription: String? {
switch self {
case .serviceUnavailable: "Local inference service is unavailable"
case .modelNotLoaded: "Local model is not loaded"
}
}
}
@@ -0,0 +1,50 @@
import Foundation
struct SSEStreamParser<T: Decodable>: AsyncSequence {
typealias Element = T
let bytes: URLSession.AsyncBytes
let decoder: JSONDecoder
init(bytes: URLSession.AsyncBytes, decoder: JSONDecoder = JSONDecoder()) {
self.bytes = bytes
self.decoder = decoder
}
func makeAsyncIterator() -> AsyncIterator {
AsyncIterator(lines: bytes.lines, decoder: decoder)
}
struct AsyncIterator: AsyncIteratorProtocol {
var lines: AsyncLineSequence<URLSession.AsyncBytes>.AsyncIterator
let decoder: JSONDecoder
init(lines: AsyncLineSequence<URLSession.AsyncBytes>, decoder: JSONDecoder) {
self.lines = lines.makeAsyncIterator()
self.decoder = decoder
}
mutating func next() async throws -> T? {
while let line = try await lines.next() {
let trimmed = line.trimmingCharacters(in: .whitespaces)
guard trimmed.hasPrefix("data: ") else { continue }
let payload = String(trimmed.dropFirst(6))
if payload == "[DONE]" {
return nil
}
guard let data = payload.data(using: .utf8) else { continue }
do {
return try decoder.decode(T.self, from: data)
} catch {
continue
}
}
return nil
}
}
}
@@ -0,0 +1,203 @@
import SwiftUI
struct ChatView: View {
@Environment(ClusterService.self) private var clusterService
@Environment(ChatService.self) private var chatService
@Environment(LocalInferenceService.self) private var localInferenceService
@State private var inputText = ""
@State private var showModelSelector = false
var body: some View {
VStack(spacing: 0) {
modelBar
GradientDivider()
messageList
GradientDivider()
inputBar
}
.background(Color.exoBlack)
.sheet(isPresented: $showModelSelector) {
ModelSelectorView(
models: clusterService.availableModels,
selectedModelId: chatService.activeConversation?.modelId
) { modelId in
chatService.setModelForActiveConversation(modelId)
}
.presentationBackground(Color.exoDarkGray)
}
}
// MARK: - Model Bar
private var useLocalModel: Bool {
!clusterService.isConnected && localInferenceService.isAvailable
}
private var modelBar: some View {
Button {
if !useLocalModel {
showModelSelector = true
}
} label: {
HStack {
Image(systemName: useLocalModel ? "iphone" : "cpu")
.font(.exoCaption)
.foregroundStyle(useLocalModel ? Color.exoYellow : Color.exoLightGray)
if useLocalModel {
Text(localInferenceService.defaultModelId)
.font(.exoSubheadline)
.foregroundStyle(Color.exoForeground)
.lineLimit(1)
} else if let modelId = chatService.activeConversation?.modelId {
Text(modelId)
.font(.exoSubheadline)
.foregroundStyle(Color.exoForeground)
.lineLimit(1)
} else {
Text("SELECT MODEL")
.font(.exoSubheadline)
.tracking(1.5)
.foregroundStyle(Color.exoLightGray)
}
Spacer()
if useLocalModel {
Text("ON-DEVICE")
.font(.exoCaption)
.tracking(1)
.foregroundStyle(Color.exoYellow)
.padding(.horizontal, 6)
.padding(.vertical, 2)
.background(Color.exoYellow.opacity(0.15))
.clipShape(Capsule())
} else {
Image(systemName: "chevron.right")
.font(.caption)
.foregroundStyle(Color.exoLightGray)
}
}
.padding(.horizontal)
.padding(.vertical, 10)
.background(Color.exoDarkGray)
}
.tint(.primary)
.disabled(useLocalModel)
}
// MARK: - Messages
private var messageList: some View {
ScrollViewReader { proxy in
ScrollView {
LazyVStack(spacing: 12) {
if chatService.activeMessages.isEmpty {
emptyState
} else {
ForEach(chatService.activeMessages) { message in
MessageBubbleView(message: message)
.id(message.id)
}
}
}
.padding()
}
.background(Color.exoBlack)
.onChange(of: chatService.activeMessages.last?.content) {
if let lastId = chatService.activeMessages.last?.id {
withAnimation(.easeOut(duration: 0.2)) {
proxy.scrollTo(lastId, anchor: .bottom)
}
}
}
}
}
private var emptyState: some View {
VStack(spacing: 16) {
Spacer(minLength: 80)
ZStack {
Circle()
.stroke(Color.exoYellow.opacity(0.15), lineWidth: 1)
.frame(width: 80, height: 80)
Circle()
.stroke(Color.exoYellow.opacity(0.3), lineWidth: 1)
.frame(width: 56, height: 56)
Circle()
.fill(Color.exoYellow.opacity(0.15))
.frame(width: 32, height: 32)
Circle()
.fill(Color.exoYellow)
.frame(width: 8, height: 8)
.shadow(color: Color.exoYellow.opacity(0.6), radius: 6)
}
Text("AWAITING INPUT")
.font(.exoSubheadline)
.tracking(3)
.foregroundStyle(Color.exoLightGray)
Text("Send a message to begin.")
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray.opacity(0.6))
Spacer(minLength: 80)
}
.padding()
}
// MARK: - Input
private var inputBar: some View {
HStack(alignment: .bottom, spacing: 8) {
TextField("Message...", text: $inputText, axis: .vertical)
.font(.exoBody)
.lineLimit(1...6)
.textFieldStyle(.plain)
.padding(10)
.background(Color.exoMediumGray)
.foregroundStyle(Color.exoForeground)
.clipShape(RoundedRectangle(cornerRadius: 8))
if chatService.isGenerating {
Button {
chatService.cancelGeneration()
} label: {
Image(systemName: "stop.circle.fill")
.font(.title2)
.foregroundStyle(Color.exoDestructive)
}
} else {
Button {
let text = inputText
inputText = ""
chatService.sendMessage(text)
} label: {
Text("SEND")
.font(.exoMono(12, weight: .bold))
.tracking(1)
.foregroundStyle(canSend ? Color.exoBlack : Color.exoLightGray)
.padding(.horizontal, 14)
.padding(.vertical, 8)
.background(canSend ? Color.exoYellow : Color.exoMediumGray)
.clipShape(RoundedRectangle(cornerRadius: 8))
}
.disabled(!canSend)
}
}
.padding(.horizontal)
.padding(.vertical, 8)
.background(Color.exoDarkGray)
}
private var canSend: Bool {
!inputText.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
&& (clusterService.isConnected || localInferenceService.isAvailable)
}
}
@@ -0,0 +1,54 @@
import SwiftUI
struct MessageBubbleView: View {
let message: ChatMessage
private var isAssistant: Bool { message.role == .assistant }
var body: some View {
HStack {
if message.role == .user { Spacer(minLength: 48) }
VStack(alignment: isAssistant ? .leading : .trailing, spacing: 6) {
// Header
HStack(spacing: 4) {
if isAssistant {
Circle()
.fill(Color.exoYellow)
.frame(width: 6, height: 6)
.shadow(color: Color.exoYellow.opacity(0.6), radius: 4)
Text("EXO")
.font(.exoMono(10, weight: .bold))
.tracking(1.5)
.foregroundStyle(Color.exoYellow)
} else {
Text("QUERY")
.font(.exoMono(10, weight: .medium))
.tracking(1.5)
.foregroundStyle(Color.exoLightGray)
}
}
// Bubble
HStack(spacing: 0) {
if isAssistant {
RoundedRectangle(cornerRadius: 1)
.fill(Color.exoYellow.opacity(0.5))
.frame(width: 2)
}
Text(message.content + (message.isStreaming ? " \u{258C}" : ""))
.font(.exoBody)
.textSelection(.enabled)
.foregroundStyle(Color.exoForeground)
.padding(.horizontal, 14)
.padding(.vertical, 10)
}
.background(Color.exoDarkGray)
.clipShape(RoundedRectangle(cornerRadius: 8))
}
if isAssistant { Spacer(minLength: 48) }
}
}
}
@@ -0,0 +1,75 @@
import SwiftUI
struct ModelSelectorView: View {
let models: [ModelOption]
let selectedModelId: String?
let onSelect: (String) -> Void
@Environment(\.dismiss) private var dismiss
var body: some View {
NavigationStack {
List {
if models.isEmpty {
emptyContent
} else {
modelsList
}
}
.scrollContentBackground(.hidden)
.background(Color.exoBlack)
.navigationTitle("SELECT MODEL")
.navigationBarTitleDisplayMode(.inline)
.toolbar {
ToolbarItem(placement: .cancellationAction) {
Button("Cancel") { dismiss() }
.font(.exoSubheadline)
.foregroundStyle(Color.exoYellow)
}
}
}
}
private var emptyContent: some View {
ContentUnavailableView(
"No Models Available",
systemImage: "cpu",
description: Text("Connect to an EXO cluster to see available models.")
)
.foregroundStyle(Color.exoLightGray)
.listRowBackground(Color.exoBlack)
}
private var modelsList: some View {
ForEach(models) { model in
Button {
onSelect(model.id)
dismiss()
} label: {
modelRow(model)
}
.tint(.primary)
.listRowBackground(Color.exoDarkGray)
}
}
private func modelRow(_ model: ModelOption) -> some View {
HStack {
VStack(alignment: .leading, spacing: 2) {
Text(model.displayName)
.font(.exoSubheadline)
.foregroundStyle(Color.exoForeground)
Text(model.id)
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
}
Spacer()
if model.id == selectedModelId {
Image(systemName: "checkmark")
.font(.exoSubheadline)
.foregroundStyle(Color.exoYellow)
}
}
}
}
@@ -0,0 +1,68 @@
import SwiftUI
struct ConnectionStatusBadge: View {
let connectionState: ConnectionState
var localModelState: LocalModelState = .notDownloaded
private var isLocalReady: Bool {
if case .disconnected = connectionState {
return localModelState == .ready || localModelState == .generating
}
return false
}
var body: some View {
HStack(spacing: 6) {
Circle()
.fill(dotColor)
.frame(width: 8, height: 8)
.shadow(color: dotColor.opacity(0.6), radius: 4)
Text(label.uppercased())
.font(.exoMono(10, weight: .medium))
.tracking(1)
.foregroundStyle(Color.exoForeground)
}
.padding(.horizontal, 10)
.padding(.vertical, 5)
.background(backgroundColor)
.clipShape(Capsule())
.overlay(
Capsule()
.stroke(dotColor.opacity(0.3), lineWidth: 1)
)
}
private var dotColor: Color {
if isLocalReady {
return .exoYellow
}
switch connectionState {
case .connected: return .green
case .connecting: return .orange
case .disconnected: return .exoLightGray
}
}
private var label: String {
if isLocalReady {
return "Local"
}
switch connectionState {
case .connected: return "Connected"
case .connecting: return "Connecting"
case .disconnected: return "Disconnected"
}
}
private var backgroundColor: Color {
if isLocalReady {
return Color.exoYellow.opacity(0.1)
}
switch connectionState {
case .connected: return .green.opacity(0.1)
case .connecting: return .orange.opacity(0.1)
case .disconnected: return Color.exoMediumGray.opacity(0.5)
}
}
}
+136
View File
@@ -0,0 +1,136 @@
import SwiftUI
struct RootView: View {
@Environment(ClusterService.self) private var clusterService
@Environment(DiscoveryService.self) private var discoveryService
@Environment(ChatService.self) private var chatService
@Environment(LocalInferenceService.self) private var localInferenceService
@State private var showSettings = false
@State private var showConversations = false
var body: some View {
NavigationStack {
ChatView()
.navigationTitle("EXO")
.navigationBarTitleDisplayMode(.inline)
.toolbar {
ToolbarItem(placement: .topBarLeading) {
conversationMenuButton
}
ToolbarItem(placement: .principal) {
ConnectionStatusBadge(
connectionState: clusterService.connectionState,
localModelState: localInferenceService.modelState
)
}
ToolbarItem(placement: .topBarTrailing) {
Button {
showSettings = true
} label: {
Image(systemName: "gear")
.foregroundStyle(Color.exoYellow)
}
}
}
}
.tint(Color.exoYellow)
.sheet(isPresented: $showSettings) {
SettingsView()
.environment(discoveryService)
.presentationBackground(Color.exoDarkGray)
}
.sheet(isPresented: $showConversations) {
conversationList
.presentationBackground(Color.exoDarkGray)
}
}
// MARK: - Conversations
private var conversationMenuButton: some View {
HStack(spacing: 12) {
Button {
showConversations = true
} label: {
Image(systemName: "sidebar.left")
.foregroundStyle(Color.exoYellow)
}
Button {
chatService.createConversation()
} label: {
Image(systemName: "square.and.pencil")
.foregroundStyle(Color.exoYellow)
}
}
}
private var conversationList: some View {
NavigationStack {
List {
if chatService.conversations.isEmpty {
Text("No conversations yet")
.font(.exoBody)
.foregroundStyle(Color.exoLightGray)
.listRowBackground(Color.exoDarkGray)
} else {
ForEach(chatService.conversations) { conversation in
let isActive = conversation.id == chatService.activeConversationId
Button {
chatService.setActiveConversation(id: conversation.id)
showConversations = false
} label: {
VStack(alignment: .leading, spacing: 4) {
Text(conversation.title)
.font(.exoSubheadline)
.fontWeight(isActive ? .semibold : .regular)
.foregroundStyle(
isActive ? Color.exoYellow : Color.exoForeground
)
.lineLimit(1)
if let modelId = conversation.modelId {
Text(modelId)
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
.lineLimit(1)
}
}
}
.listRowBackground(
isActive
? Color.exoYellow.opacity(0.1)
: Color.exoDarkGray
)
}
.onDelete { indexSet in
for index in indexSet {
chatService.deleteConversation(id: chatService.conversations[index].id)
}
}
}
}
.scrollContentBackground(.hidden)
.background(Color.exoBlack)
.navigationTitle("Conversations")
.navigationBarTitleDisplayMode(.inline)
.toolbar {
ToolbarItem(placement: .confirmationAction) {
Button("Done") { showConversations = false }
.font(.exoSubheadline)
.foregroundStyle(Color.exoYellow)
}
ToolbarItem(placement: .topBarLeading) {
Button {
chatService.createConversation()
} label: {
Image(systemName: "plus")
.foregroundStyle(Color.exoYellow)
}
}
}
}
}
}
@@ -0,0 +1,314 @@
import SwiftUI
struct SettingsView: View {
@Environment(ClusterService.self) private var clusterService
@Environment(DiscoveryService.self) private var discoveryService
@Environment(LocalInferenceService.self) private var localInferenceService
@Environment(\.dismiss) private var dismiss
@State private var host: String = ""
@State private var port: String = "52415"
var body: some View {
NavigationStack {
Form {
localModelSection
nearbyClustersSection
connectionSection
statusSection
}
.scrollContentBackground(.hidden)
.background(Color.exoBlack)
.navigationTitle("Settings")
.navigationBarTitleDisplayMode(.inline)
.toolbar {
ToolbarItem(placement: .confirmationAction) {
Button("Done") { dismiss() }
.font(.exoSubheadline)
.foregroundStyle(Color.exoYellow)
}
}
}
}
// MARK: - Section Headers
private func sectionHeader(_ title: String) -> some View {
Text(title.uppercased())
.font(.exoMono(10, weight: .semibold))
.tracking(2)
.foregroundStyle(Color.exoYellow)
}
// MARK: - Local Model
private var localModelSection: some View {
Section {
HStack {
VStack(alignment: .leading, spacing: 4) {
Text(localInferenceService.defaultModelId)
.font(.exoSubheadline)
.foregroundStyle(Color.exoForeground)
Text(localModelStatusText)
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
}
Spacer()
localModelActionButton
}
.listRowBackground(Color.exoDarkGray)
if case .downloading(let progress) = localInferenceService.modelState {
ProgressView(value: progress)
.tint(Color.exoYellow)
.listRowBackground(Color.exoDarkGray)
}
} header: {
sectionHeader("Local Model")
} footer: {
Text(
"When disconnected from a cluster, messages are processed on-device using this model."
)
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray.opacity(0.7))
}
}
private var localModelStatusText: String {
switch localInferenceService.modelState {
case .notDownloaded: "Not downloaded"
case .downloading(let progress): "Downloading \(Int(progress * 100))%..."
case .downloaded: "Downloaded — not loaded"
case .loading: "Loading into memory..."
case .ready: "Ready"
case .generating: "Generating..."
case .error(let message): "Error: \(message)"
}
}
@ViewBuilder
private var localModelActionButton: some View {
switch localInferenceService.modelState {
case .notDownloaded:
exoButton("Download") {
Task { await localInferenceService.prepareModel() }
}
case .downloading:
ProgressView()
.controlSize(.small)
.tint(Color.exoYellow)
case .downloaded:
exoButton("Load") {
Task { await localInferenceService.prepareModel() }
}
case .loading:
ProgressView()
.controlSize(.small)
.tint(Color.exoYellow)
case .ready, .generating:
exoButton("Unload") {
localInferenceService.unloadModel()
}
case .error:
exoButton("Retry", destructive: true) {
Task { await localInferenceService.prepareModel() }
}
}
}
private func exoButton(_ title: String, destructive: Bool = false, action: @escaping () -> Void)
-> some View
{
let borderColor = destructive ? Color.exoDestructive : Color.exoYellow
return Button(action: action) {
Text(title.uppercased())
.font(.exoMono(11, weight: .semibold))
.tracking(1)
.foregroundStyle(borderColor)
.padding(.horizontal, 10)
.padding(.vertical, 5)
.overlay(
RoundedRectangle(cornerRadius: 6)
.stroke(borderColor, lineWidth: 1)
)
}
}
// MARK: - Nearby Clusters
private var nearbyClustersSection: some View {
Section {
if discoveryService.discoveredClusters.isEmpty {
if discoveryService.isSearching {
HStack {
ProgressView()
.tint(Color.exoYellow)
.padding(.trailing, 8)
Text("Searching for clusters...")
.font(.exoBody)
.foregroundStyle(Color.exoLightGray)
}
.listRowBackground(Color.exoDarkGray)
} else {
Text("No clusters found")
.font(.exoBody)
.foregroundStyle(Color.exoLightGray)
.listRowBackground(Color.exoDarkGray)
}
} else {
ForEach(discoveryService.discoveredClusters) { cluster in
HStack {
VStack(alignment: .leading) {
Text(cluster.name)
.font(.exoBody)
.foregroundStyle(Color.exoForeground)
}
Spacer()
exoButton("Connect") {
Task {
await clusterService.connectToDiscoveredCluster(
cluster, using: discoveryService
)
if clusterService.isConnected {
dismiss()
}
}
}
}
.listRowBackground(Color.exoDarkGray)
}
}
} header: {
sectionHeader("Nearby Clusters")
}
}
// MARK: - Manual Connection
private var connectionSection: some View {
Section {
TextField("IP Address (e.g. 192.168.1.42)", text: $host)
.font(.exoBody)
.keyboardType(.decimalPad)
.textContentType(.URL)
.autocorrectionDisabled()
.foregroundStyle(Color.exoForeground)
.listRowBackground(Color.exoDarkGray)
TextField("Port", text: $port)
.font(.exoBody)
.keyboardType(.numberPad)
.foregroundStyle(Color.exoForeground)
.listRowBackground(Color.exoDarkGray)
Button {
Task {
let portNum = Int(port) ?? ConnectionInfo.defaultPort
let info = ConnectionInfo(host: host, port: portNum, nodeId: nil)
await clusterService.connect(to: info)
if clusterService.isConnected {
dismiss()
}
}
} label: {
Text(clusterService.isConnected ? "RECONNECT" : "CONNECT")
.font(.exoMono(13, weight: .semibold))
.tracking(1.5)
.foregroundStyle(
host.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
? Color.exoLightGray : Color.exoYellow
)
.frame(maxWidth: .infinity)
}
.disabled(host.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty)
.listRowBackground(Color.exoDarkGray)
} header: {
sectionHeader("Manual Connection")
}
}
// MARK: - Status
private var statusSection: some View {
Section {
if let connection = clusterService.currentConnection {
LabeledContent {
Text(connection.host)
.font(.exoCaption)
.foregroundStyle(Color.exoForeground)
} label: {
Text("Host")
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
}
.listRowBackground(Color.exoDarkGray)
LabeledContent {
Text("\(connection.port)")
.font(.exoCaption)
.foregroundStyle(Color.exoForeground)
} label: {
Text("Port")
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
}
.listRowBackground(Color.exoDarkGray)
if let nodeId = connection.nodeId {
LabeledContent {
Text(String(nodeId.prefix(12)) + "...")
.font(.exoCaption)
.foregroundStyle(Color.exoForeground)
} label: {
Text("Node ID")
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
}
.listRowBackground(Color.exoDarkGray)
}
LabeledContent {
Text("\(clusterService.availableModels.count)")
.font(.exoCaption)
.foregroundStyle(Color.exoForeground)
} label: {
Text("Models")
.font(.exoCaption)
.foregroundStyle(Color.exoLightGray)
}
.listRowBackground(Color.exoDarkGray)
Button(role: .destructive) {
clusterService.disconnect()
} label: {
Text("DISCONNECT")
.font(.exoMono(13, weight: .semibold))
.tracking(1.5)
.foregroundStyle(Color.exoDestructive)
.frame(maxWidth: .infinity)
}
.listRowBackground(Color.exoDarkGray)
} else {
if let error = clusterService.lastError {
Label {
Text(error)
.font(.exoCaption)
} icon: {
Image(systemName: "exclamationmark.triangle")
}
.foregroundStyle(Color.exoDestructive)
.listRowBackground(Color.exoDarkGray)
} else {
Text("Not connected")
.font(.exoBody)
.foregroundStyle(Color.exoLightGray)
.listRowBackground(Color.exoDarkGray)
}
}
} header: {
sectionHeader("Status")
}
}
}
@@ -0,0 +1,51 @@
import SwiftUI
// MARK: - EXO Color Palette
extension Color {
/// Primary background near-black (#121212)
static let exoBlack = Color(red: 0x12 / 255.0, green: 0x12 / 255.0, blue: 0x12 / 255.0)
/// Card / surface background (#1F1F1F)
static let exoDarkGray = Color(red: 0x1F / 255.0, green: 0x1F / 255.0, blue: 0x1F / 255.0)
/// Input field / elevated surface (#353535)
static let exoMediumGray = Color(red: 0x35 / 255.0, green: 0x35 / 255.0, blue: 0x35 / 255.0)
/// Secondary text (#999999)
static let exoLightGray = Color(red: 0x99 / 255.0, green: 0x99 / 255.0, blue: 0x99 / 255.0)
/// Accent yellow matches dashboard (#FFD700)
static let exoYellow = Color(red: 0xFF / 255.0, green: 0xD7 / 255.0, blue: 0x00 / 255.0)
/// Primary foreground text (#E5E5E5)
static let exoForeground = Color(red: 0xE5 / 255.0, green: 0xE5 / 255.0, blue: 0xE5 / 255.0)
/// Destructive / error (#E74C3C)
static let exoDestructive = Color(red: 0xE7 / 255.0, green: 0x4C / 255.0, blue: 0x3C / 255.0)
}
// MARK: - EXO Typography (SF Mono via .monospaced design)
extension Font {
/// Monospaced font at a given size and weight.
static func exoMono(_ size: CGFloat, weight: Font.Weight = .regular) -> Font {
.system(size: size, weight: weight, design: .monospaced)
}
/// Body text 15pt monospaced
static let exoBody: Font = .system(size: 15, weight: .regular, design: .monospaced)
/// Caption 11pt monospaced
static let exoCaption: Font = .system(size: 11, weight: .regular, design: .monospaced)
/// Subheadline 13pt monospaced medium
static let exoSubheadline: Font = .system(size: 13, weight: .medium, design: .monospaced)
/// Headline 17pt monospaced semibold
static let exoHeadline: Font = .system(size: 17, weight: .semibold, design: .monospaced)
}
// MARK: - Reusable Gradient Divider
struct GradientDivider: View {
var body: some View {
LinearGradient(
colors: [.clear, Color.exoYellow.opacity(0.3), .clear],
startPoint: .leading,
endPoint: .trailing
)
.frame(height: 1)
}
}
@@ -0,0 +1,18 @@
//
// EXO_iOSTests.swift
// EXO-iOSTests
//
// Created by Sami Khan on 2026-02-17.
//
import Testing
@testable import EXO_iOS
struct EXO_iOSTests {
@Test func example() async throws {
// Write your test here and use APIs like `#expect(...)` to check expected conditions.
}
}
@@ -0,0 +1,41 @@
//
// EXO_iOSUITests.swift
// EXO-iOSUITests
//
// Created by Sami Khan on 2026-02-17.
//
import XCTest
final class EXO_iOSUITests: XCTestCase {
override func setUpWithError() throws {
// Put setup code here. This method is called before the invocation of each test method in the class.
// In UI tests it is usually best to stop immediately when a failure occurs.
continueAfterFailure = false
// In UI tests its important to set the initial state - such as interface orientation - required for your tests before they run. The setUp method is a good place to do this.
}
override func tearDownWithError() throws {
// Put teardown code here. This method is called after the invocation of each test method in the class.
}
@MainActor
func testExample() throws {
// UI tests must launch the application that they test.
let app = XCUIApplication()
app.launch()
// Use XCTAssert and related functions to verify your tests produce the correct results.
}
@MainActor
func testLaunchPerformance() throws {
// This measures how long it takes to launch your application.
measure(metrics: [XCTApplicationLaunchMetric()]) {
XCUIApplication().launch()
}
}
}
@@ -0,0 +1,33 @@
//
// EXO_iOSUITestsLaunchTests.swift
// EXO-iOSUITests
//
// Created by Sami Khan on 2026-02-17.
//
import XCTest
final class EXO_iOSUITestsLaunchTests: XCTestCase {
override class var runsForEachTargetApplicationUIConfiguration: Bool {
true
}
override func setUpWithError() throws {
continueAfterFailure = false
}
@MainActor
func testLaunch() throws {
let app = XCUIApplication()
app.launch()
// Insert steps here to perform after app launch but before taking a screenshot,
// such as logging into a test account or navigating somewhere in the app
let attachment = XCTAttachment(screenshot: app.screenshot())
attachment.name = "Launch Screen"
attachment.lifetime = .keepAlways
add(attachment)
}
}
+145 -212
View File
@@ -15,23 +15,14 @@ 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
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 bugReportInFlight = false
@State private var bugReportMessage: String?
@State private var uninstallInProgress = false
@State private var pendingNamespace: String = ""
@State private var pendingHFToken: String = ""
@@ -267,79 +258,139 @@ struct ContentView: View {
VStack(alignment: .leading, spacing: 0) {
if controller.status != .stopped {
dashboardButton
baseURLRow
Divider()
.padding(.vertical, 8)
} else {
Divider()
.padding(.vertical, 4)
}
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) {
advancedSection
.padding(.bottom, 8)
controlButton(title: "Quit", tint: .secondary) {
controller.stop()
NSApplication.shared.terminate(nil)
}
}
}
private var dashboardButton: some View {
HoverButton(
title: "Web Dashboard",
tint: .primary,
trailingSystemImage: "arrow.up.right"
) {
guard let url = URL(string: "http://localhost:52415/") else { return }
NSWorkspace.shared.open(url)
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 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")
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 {
guard let url = URL(string: "http://localhost:52415/") else { return }
NSWorkspace.shared.open(url)
} label: {
HStack {
Image(systemName: "arrow.up.right.square")
.imageScale(.small)
.foregroundColor(baseURLCopied ? .green : .secondary)
.contentTransition(.symbolEffect(.replace))
Text("Dashboard")
.fontWeight(.medium)
Spacer()
}
.buttonStyle(.plain)
.help("Copy API base URL")
.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))
)
}
.padding(.vertical, 4)
.padding(.horizontal, 8)
.buttonStyle(.plain)
.padding(.bottom, 4)
}
private func collapseButton(isExpanded: Binding<Bool>) -> some View {
@@ -560,115 +611,39 @@ struct ContentView: View {
}
private var sendBugReportButton: some View {
VStack(alignment: .leading, spacing: 6) {
switch bugReportPhase {
case .idle:
Button {
bugReportPhase = .prompting
bugReportUserDescription = ""
} label: {
HStack {
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))
)
VStack(alignment: .leading, spacing: 4) {
Button {
Task {
await sendBugReport()
}
.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)
} label: {
HStack {
if bugReportInFlight {
ProgressView()
.scaleEffect(0.6)
}
Text("Send Bug Report")
.font(.caption)
.fontWeight(.semibold)
Spacer()
}
.padding(8)
.padding(.vertical, 6)
.padding(.horizontal, 8)
.background(
RoundedRectangle(cornerRadius: 6)
.fill(Color.accentColor.opacity(0.06))
.fill(Color.accentColor.opacity(0.12))
)
}
.buttonStyle(.plain)
.disabled(bugReportInFlight)
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 = ""
}
if let message = bugReportMessage {
Text(message)
.font(.caption2)
.buttonStyle(.plain)
.foregroundColor(.secondary)
}
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)
}
.fixedSize(horizontal: false, vertical: true)
}
}
.animation(.easeInOut(duration: 0.2), value: bugReportPhase)
}
private var processToggleBinding: Binding<Bool> {
@@ -712,58 +687,16 @@ struct ContentView: View {
}
private func sendBugReport() async {
bugReportPhase = .sending("Collecting logs...")
bugReportInFlight = true
bugReportMessage = "Collecting logs..."
let service = BugReportService()
let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
do {
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)
}
let outcome = try await service.sendReport(isManual: true)
bugReportMessage = outcome.message
} catch {
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)
bugReportMessage = error.localizedDescription
}
bugReportInFlight = false
}
private func showUninstallConfirmationAlert() {
+1 -15
View File
@@ -21,9 +21,7 @@ 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() {
@@ -45,13 +43,12 @@ 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: 5)
controller.scheduleLaunch(after: 15)
service.startPolling()
networkStatus.startPolling()
}
@@ -65,19 +62,8 @@ 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)
}
+3 -59
View File
@@ -5,8 +5,6 @@ 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 {
@@ -61,17 +59,6 @@ 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?
@@ -126,9 +113,6 @@ 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")
@@ -142,37 +126,11 @@ final class ExoProcessController: ObservableObject {
return
}
process.terminationHandler = nil
status = .stopped
guard process.isRunning else {
self.process = nil
return
if process.isRunning {
process.terminate()
}
let proc = process
self.process = nil
Task.detached {
proc.interrupt()
for _ in 0..<50 {
if !proc.isRunning { return }
try? await Task.sleep(nanoseconds: 100_000_000)
}
if proc.isRunning {
proc.terminate()
}
for _ in 0..<30 {
if !proc.isRunning { return }
try? await Task.sleep(nanoseconds: 100_000_000)
}
if proc.isRunning {
kill(proc.processIdentifier, SIGKILL)
}
}
status = .stopped
}
func restart() {
@@ -180,17 +138,6 @@ 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)))
@@ -276,9 +223,6 @@ 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 {
+3 -9
View File
@@ -38,8 +38,7 @@ struct BugReportService {
func sendReport(
baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
now: Date = Date(),
isManual: Bool = false,
userDescription: String? = nil
isManual: Bool = false
) async throws -> BugReportOutcome {
let timestamp = Self.runTimestampString(now)
let dayPrefix = Self.dayPrefixString(now)
@@ -61,8 +60,7 @@ struct BugReportService {
ifconfig: ifconfigText,
debugInfo: debugInfo,
isManual: isManual,
clusterTbBridgeStatus: clusterTbBridgeStatus,
userDescription: userDescription
clusterTbBridgeStatus: clusterTbBridgeStatus
)
let eventLogFiles = readAllEventLogs()
@@ -308,8 +306,7 @@ struct BugReportService {
ifconfig: String,
debugInfo: DebugInfo,
isManual: Bool,
clusterTbBridgeStatus: [[String: Any]]?,
userDescription: String? = nil
clusterTbBridgeStatus: [[String: Any]]?
) -> Data? {
let system = readSystemMetadata()
let exo = readExoMetadata()
@@ -326,9 +323,6 @@ struct BugReportService {
if let tbStatus = clusterTbBridgeStatus {
payload["cluster_thunderbolt_bridge"] = tbStatus
}
if let desc = userDescription, !desc.isEmpty {
payload["user_description"] = desc
}
return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
}
-200
View File
@@ -1,200 +0,0 @@
import AppKit
import SwiftUI
/// A popover callout anchored to the menu bar icon on every launch,
/// pointing the user to the web dashboard with an arrow connecting to the icon.
@MainActor
final class FirstLaunchPopout {
private var popover: NSPopover?
private var countdownTask: Task<Void, Never>?
private static let dashboardURL = "http://localhost:52415/"
/// Called when the user completes onboarding (clicks Open Dashboard or dismisses).
var onComplete: (() -> Void)?
func show() {
guard popover == nil else { return }
// The status bar button may not exist yet on first launch; retry generously.
showWithRetry(attemptsRemaining: 15)
}
private func showWithRetry(attemptsRemaining: Int) {
guard attemptsRemaining > 0 else {
// Exhausted retries fall back to just opening the dashboard directly.
openDashboard()
onComplete?()
return
}
guard let button = Self.findStatusItemButton() else {
DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { [weak self] in
self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1)
}
return
}
let pop = NSPopover()
pop.behavior = .applicationDefined
pop.animates = true
pop.contentSize = NSSize(width: 280, height: 120)
pop.contentViewController = NSHostingController(
rootView: WelcomeCalloutView(
countdownDuration: 10,
onDismiss: { [weak self] in
self?.onComplete?()
self?.dismiss()
},
onOpen: { [weak self] in
self?.openDashboard()
self?.onComplete?()
self?.dismiss()
}
)
)
self.popover = pop
pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY)
// Auto-open dashboard after 10s then dismiss
countdownTask = Task {
try? await Task.sleep(nanoseconds: 10_000_000_000)
if !Task.isCancelled {
openDashboard()
onComplete?()
dismiss()
}
}
}
func dismiss() {
countdownTask?.cancel()
countdownTask = nil
guard let pop = popover else { return }
popover = nil
pop.performClose(nil)
}
private func openDashboard() {
guard let url = URL(string: Self.dashboardURL) else { return }
NSWorkspace.shared.open(url)
}
/// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra.
/// Walks the view hierarchy to find the actual button rather than the content view.
private static func findStatusItemButton() -> NSView? {
for window in NSApp.windows {
let className = NSStringFromClass(type(of: window))
// Match NSStatusBarWindow or any internal SwiftUI status bar window
guard className.contains("StatusBar") || className.contains("MenuBarExtra") else {
continue
}
if let content = window.contentView {
if let button = findButton(in: content) {
return button
}
// Fall back to the content view itself if it has a non-zero frame
if content.frame.width > 0 {
return content
}
}
}
return nil
}
/// Recursively searches the view hierarchy for an NSStatusBarButton.
private static func findButton(in view: NSView) -> NSView? {
let className = NSStringFromClass(type(of: view))
if className.contains("StatusBarButton") || className.contains("StatusItem") {
return view
}
for subview in view.subviews {
if let found = findButton(in: subview) {
return found
}
}
return nil
}
}
/// Minimal welcome callout friendly pointer, not a wall of text.
/// Rendered inside the NSPopover which provides its own chrome and arrow.
private struct WelcomeCalloutView: View {
let countdownDuration: Int
let onDismiss: () -> Void
let onOpen: () -> Void
@State private var countdown: Int
@State private var timerTask: Task<Void, Never>?
init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) {
self.countdownDuration = countdownDuration
self.onDismiss = onDismiss
self.onOpen = onOpen
self._countdown = State(initialValue: countdownDuration)
}
var body: some View {
VStack(alignment: .leading, spacing: 10) {
HStack(alignment: .top) {
Text("EXO is running")
.font(.system(.headline, design: .rounded))
.fontWeight(.semibold)
.foregroundColor(.primary)
Spacer()
Button {
onDismiss()
} label: {
Image(systemName: "xmark.circle.fill")
.font(.system(size: 14))
.foregroundStyle(.tertiary)
}
.buttonStyle(.plain)
}
Text("Run your first model here:")
.font(.system(.subheadline, design: .default))
.foregroundColor(.secondary)
HStack {
Button {
onOpen()
} label: {
Label("Open Dashboard", systemImage: "arrow.up.right.square")
.font(.system(.caption, design: .default))
.fontWeight(.medium)
}
.buttonStyle(.borderedProminent)
.tint(.accentColor)
.controlSize(.small)
Spacer()
if countdown > 0 {
Text("Launching in \(countdown) secs...")
.font(.system(.caption2, design: .default))
.foregroundColor(.secondary.opacity(0.6))
.monospacedDigit()
}
}
}
.padding(14)
.onAppear {
startCountdown()
}
.onDisappear {
timerTask?.cancel()
timerTask = nil
}
}
private func startCountdown() {
timerTask = Task {
while countdown > 0 {
try? await Task.sleep(nanoseconds: 1_000_000_000)
if !Task.isCancelled {
countdown -= 1
}
}
}
}
}
-489
View File
@@ -1,489 +0,0 @@
import AppKit
import SwiftUI
/// Native macOS Settings window following Apple HIG.
/// Organized into General, Model, Advanced, and About sections.
struct SettingsView: View {
@EnvironmentObject private var controller: ExoProcessController
@EnvironmentObject private var updater: SparkleUpdater
@EnvironmentObject private var networkStatusService: NetworkStatusService
@EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
@EnvironmentObject private var stateService: ClusterStateService
@State private var pendingNamespace: String = ""
@State private var pendingHFToken: String = ""
@State private var pendingEnableImageModels = false
@State private var pendingOfflineMode = false
@State private var needsRestart = false
@State private var bugReportInFlight = false
@State private var bugReportMessage: String?
@State private var uninstallInProgress = false
var body: some View {
TabView {
generalTab
.tabItem {
Label("General", systemImage: "gear")
}
modelTab
.tabItem {
Label("Model", systemImage: "cube")
}
advancedTab
.tabItem {
Label("Advanced", systemImage: "wrench.and.screwdriver")
}
aboutTab
.tabItem {
Label("About", systemImage: "info.circle")
}
}
.frame(width: 450, height: 400)
.onAppear {
pendingNamespace = controller.customNamespace
pendingHFToken = controller.hfToken
pendingEnableImageModels = controller.enableImageModels
pendingOfflineMode = controller.offlineMode
needsRestart = false
}
}
// MARK: - General Tab
private var generalTab: some View {
Form {
Section {
LabeledContent("Cluster Namespace") {
TextField("default", text: $pendingNamespace)
.textFieldStyle(.roundedBorder)
.frame(width: 200)
}
Text("Nodes with the same namespace form a cluster. Leave empty for default.")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
LabeledContent("HuggingFace Token") {
SecureField("optional", text: $pendingHFToken)
.textFieldStyle(.roundedBorder)
.frame(width: 200)
}
Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
Toggle("Offline Mode", isOn: $pendingOfflineMode)
Text("Skip internet checks and use only locally available models.")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
HStack {
Spacer()
Button("Save & Restart") {
applyGeneralSettings()
}
.disabled(!hasGeneralChanges)
}
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - Model Tab
private var modelTab: some View {
Form {
Section {
Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels)
Text("Allow text-to-image and image-to-image models in the model picker.")
.font(.caption)
.foregroundColor(.secondary)
}
Section {
HStack {
Spacer()
Button("Save & Restart") {
applyModelSettings()
}
.disabled(!hasModelChanges)
}
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - Advanced Tab
private var advancedTab: some View {
Form {
Section("Onboarding") {
HStack {
VStack(alignment: .leading) {
Text("Reset Onboarding")
Text("Opens the dashboard and resets the onboarding wizard.")
.font(.caption)
.foregroundColor(.secondary)
}
Spacer()
Button("Reset") {
guard let url = URL(string: "http://localhost:52415/?reset-onboarding")
else { return }
NSWorkspace.shared.open(url)
}
}
}
Section("Debug Info") {
LabeledContent("Thunderbolt Bridge") {
Text(thunderboltStatusText)
.foregroundColor(thunderboltStatusColor)
}
VStack(alignment: .leading, spacing: 2) {
clusterThunderboltBridgeView
}
VStack(alignment: .leading, spacing: 2) {
interfaceIpList
}
VStack(alignment: .leading, spacing: 2) {
rdmaStatusView
}
sendBugReportButton
}
Section("Danger Zone") {
Button(role: .destructive) {
showUninstallConfirmationAlert()
} label: {
HStack {
Text("Uninstall EXO")
Spacer()
Image(systemName: "trash")
.imageScale(.small)
}
}
.disabled(uninstallInProgress)
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - About Tab
private var aboutTab: some View {
Form {
Section {
LabeledContent("Version") {
Text(buildTag)
.textSelection(.enabled)
}
LabeledContent("Commit") {
Text(buildCommit)
.font(.system(.body, design: .monospaced))
.textSelection(.enabled)
}
}
Section {
Button("Check for Updates") {
updater.checkForUpdates()
}
}
}
.formStyle(.grouped)
.padding()
}
// MARK: - Debug Info Views (moved from ContentView)
private var thunderboltStatusText: String {
switch networkStatusService.status.thunderboltBridgeState {
case .some(.disabled):
return "Disabled"
case .some(.deleted):
return "Deleted"
case .some(.enabled):
return "Enabled"
case nil:
return "Unknown"
}
}
private var thunderboltStatusColor: Color {
switch networkStatusService.status.thunderboltBridgeState {
case .some(.disabled), .some(.deleted):
return .green
case .some(.enabled):
return .red
case nil:
return .secondary
}
}
private var clusterThunderboltBridgeView: some View {
let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
let localNodeId = stateService.localNodeId
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
return VStack(alignment: .leading, spacing: 1) {
if bridgeStatuses.isEmpty {
Text("Cluster TB Bridge: No data")
.font(.caption2)
.foregroundColor(.secondary)
} else {
Text("Cluster TB Bridge Status:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
if let status = bridgeStatuses[nodeId] {
let nodeName =
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
let isLocal = nodeId == localNodeId
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
let statusText =
!status.exists
? "N/A"
: (status.enabled ? "Enabled" : "Disabled")
let color: Color =
!status.exists
? .secondary
: (status.enabled ? .red : .green)
Text("\(prefix) \(statusText)")
.font(.caption2)
.foregroundColor(color)
}
}
}
}
}
private var interfaceIpList: some View {
let statuses = networkStatusService.status.interfaceStatuses
return VStack(alignment: .leading, spacing: 1) {
Text("Interfaces (en0en7):")
.font(.caption2)
.foregroundColor(.secondary)
if statuses.isEmpty {
Text(" Unknown")
.font(.caption2)
.foregroundColor(.secondary)
} else {
ForEach(statuses, id: \.interfaceName) { status in
let ipText = status.ipAddress ?? "No IP"
Text(" \(status.interfaceName): \(ipText)")
.font(.caption2)
.foregroundColor(status.ipAddress == nil ? .red : .green)
}
}
}
}
private var rdmaStatusView: some View {
let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
let localNodeId = stateService.localNodeId
let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
let localDevices = networkStatusService.status.localRdmaDevices
let localPorts = networkStatusService.status.localRdmaActivePorts
return VStack(alignment: .leading, spacing: 1) {
if rdmaStatuses.isEmpty {
Text("Cluster RDMA: No data")
.font(.caption2)
.foregroundColor(.secondary)
} else {
Text("Cluster RDMA Status:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
if let status = rdmaStatuses[nodeId] {
let nodeName =
nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
let isLocal = nodeId == localNodeId
let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):"
let statusText = status.enabled ? "Enabled" : "Disabled"
let color: Color = status.enabled ? .green : .orange
Text("\(prefix) \(statusText)")
.font(.caption2)
.foregroundColor(color)
}
}
}
if !localDevices.isEmpty {
Text(" Local Devices: \(localDevices.joined(separator: ", "))")
.font(.caption2)
.foregroundColor(.secondary)
}
if !localPorts.isEmpty {
Text(" Local Active Ports:")
.font(.caption2)
.foregroundColor(.secondary)
ForEach(localPorts, id: \.device) { port in
Text(" \(port.device) port \(port.port): \(port.state)")
.font(.caption2)
.foregroundColor(.green)
}
}
}
}
private var sendBugReportButton: some View {
VStack(alignment: .leading, spacing: 4) {
Button {
Task {
await sendBugReport()
}
} label: {
HStack {
if bugReportInFlight {
ProgressView()
.scaleEffect(0.6)
}
Text("Send Bug Report")
.font(.caption)
.fontWeight(.semibold)
Spacer()
}
}
.disabled(bugReportInFlight)
if let message = bugReportMessage {
Text(message)
.font(.caption2)
.foregroundColor(.secondary)
.fixedSize(horizontal: false, vertical: true)
}
}
}
// MARK: - Actions
private func sendBugReport() async {
bugReportInFlight = true
bugReportMessage = "Collecting logs..."
let service = BugReportService()
do {
let outcome = try await service.sendReport(isManual: true)
bugReportMessage = outcome.message
} catch {
bugReportMessage = error.localizedDescription
}
bugReportInFlight = false
}
private func showUninstallConfirmationAlert() {
let alert = NSAlert()
alert.messageText = "Uninstall EXO"
alert.informativeText = """
This will remove EXO and all its system components:
Network configuration daemon
Launch at login registration
EXO network location
The app will be moved to Trash.
"""
alert.alertStyle = .warning
alert.addButton(withTitle: "Uninstall")
alert.addButton(withTitle: "Cancel")
if let uninstallButton = alert.buttons.first {
uninstallButton.hasDestructiveAction = true
}
let response = alert.runModal()
if response == .alertFirstButtonReturn {
performUninstall()
}
}
private func performUninstall() {
uninstallInProgress = true
controller.cancelPendingLaunch()
controller.stop()
stateService.stopPolling()
DispatchQueue.global(qos: .utility).async {
do {
try NetworkSetupHelper.uninstall()
DispatchQueue.main.async {
LaunchAtLoginHelper.disable()
self.moveAppToTrash()
DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
NSApplication.shared.terminate(nil)
}
}
} catch {
DispatchQueue.main.async {
let errorAlert = NSAlert()
errorAlert.messageText = "Uninstall Failed"
errorAlert.informativeText = error.localizedDescription
errorAlert.alertStyle = .critical
errorAlert.addButton(withTitle: "OK")
errorAlert.runModal()
self.uninstallInProgress = false
}
}
}
}
private func moveAppToTrash() {
guard let appURL = Bundle.main.bundleURL as URL? else { return }
do {
try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
} catch {
// If we can't trash the app, that's OK - user can do it manually
}
}
// MARK: - Helpers
private var hasGeneralChanges: Bool {
pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
|| pendingOfflineMode != controller.offlineMode
}
private var hasModelChanges: Bool {
pendingEnableImageModels != controller.enableImageModels
}
private func applyGeneralSettings() {
controller.customNamespace = pendingNamespace
controller.hfToken = pendingHFToken
controller.offlineMode = pendingOfflineMode
restartIfRunning()
}
private func applyModelSettings() {
controller.enableImageModels = pendingEnableImageModels
restartIfRunning()
}
private func restartIfRunning() {
if controller.status == .running || controller.status == .starting {
controller.restart()
}
}
private var buildTag: String {
Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
}
private var buildCommit: String {
Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
}
}
@@ -1,47 +0,0 @@
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
}
}
+3 -1
View File
@@ -2,4 +2,6 @@
#
# 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",
]
-292
View File
@@ -1,292 +0,0 @@
# 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
File diff suppressed because it is too large Load Diff
+499 -182
View File
@@ -1,51 +1,29 @@
# type: ignore
#!/usr/bin/env python3
"""Tool-calling eval for exo's OpenAI-compatible API.
Tests whether models correctly:
- Trigger tool calls when appropriate
- Return valid JSON arguments matching function schemas
- Handle multi-turn tool use (call -> result -> final answer)
- Avoid calling tools when unnecessary
Start exo with a model first, then run:
uv run python tool_call_eval.py --model <model-id>
uv run python tool_call_eval.py --model <model-id> --host 10.0.0.5 --port 52415
uv run python tool_call_eval.py --model <model-id> --repeat 3
uv run python tool_call_eval.py --model <model-id> --scenarios weather_simple calculator_multi_turn
"""
# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
from __future__ import annotations
import argparse
import contextlib
import http.client
import itertools
import json
import os
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
from urllib.parse import urlencode
from harness import (
ExoClient,
ExoHttpError,
add_common_instance_args,
ensure_cuda_available,
instance_id_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
run_planning_phase,
settle_and_fetch_placements,
validate_vllm_args,
wait_for_instance_gone,
wait_for_instance_ready,
)
from loguru import logger
from transformers import AutoTokenizer
# Backoff constants for cluster settling retry
_SETTLE_INITIAL_BACKOFF_S = 1.0
_SETTLE_MAX_BACKOFF_S = 60.0
_SETTLE_BACKOFF_MULTIPLIER = 2.0
# Monkey-patch for transformers 5.x compatibility
# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
# which was moved in transformers 5.0.0rc2
@@ -78,7 +56,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
model_path = Path(
snapshot_download(
model_id,
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
allow_patterns=["*.json", "*.py", "*.tiktoken"],
)
)
@@ -125,6 +103,154 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
class ExoHttpError(RuntimeError):
def __init__(self, status: int, reason: str, body_preview: str):
super().__init__(f"HTTP {status} {reason}: {body_preview}")
self.status = status
class ExoClient:
def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
self.host = host
self.port = port
self.timeout_s = timeout_s
def request_json(
self,
method: str,
path: str,
params: dict[str, Any] | None = None,
body: dict[str, Any] | None = None,
headers: dict[str, str] | None = None,
) -> Any:
if not path.startswith("/"):
path = "/" + path
if params:
path = path + "?" + urlencode(params)
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
try:
payload: bytes | None = None
hdrs: dict[str, str] = {"Accept": "application/json"}
if body is not None:
payload = json.dumps(body).encode("utf-8")
hdrs["Content-Type"] = "application/json"
if headers:
hdrs.update(headers)
conn.request(method.upper(), path, body=payload, headers=hdrs)
resp = conn.getresponse()
raw = resp.read()
text = raw.decode("utf-8", errors="replace") if raw else ""
if resp.status >= 400:
raise ExoHttpError(resp.status, resp.reason, text[:300])
if not text:
return None
return json.loads(text)
finally:
conn.close()
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
return self.request_json("POST", "/bench/chat/completions", body=payload)
def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
if len(instance) != 1:
raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
tag = next(iter(instance))
inner = instance[tag]
if not isinstance(inner, dict):
raise TypeError(f"payload for {tag} must be dict, got {type(inner)}")
return inner
def instance_id_from_instance(instance: dict[str, Any]) -> str:
inner = unwrap_instance(instance)
return str(inner["instanceId"])
def nodes_used_in_instance(instance: dict[str, Any]) -> int:
inner = unwrap_instance(instance)
return len(inner["shardAssignments"]["nodeToRunner"])
def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
inner = unwrap_instance(instance)
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
return list(runner_to_shard.keys())
def runner_ready(runner: dict[str, Any]) -> bool:
return "RunnerReady" in runner
def runner_failed(runner: dict[str, Any]) -> bool:
return "RunnerFailed" in runner
def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
if "RunnerFailed" in runner:
return runner["RunnerFailed"].get("errorMessage")
return None
def wait_for_instance_ready(
client: ExoClient, instance_id: str, timeout: float = 24000.0
) -> None:
start_time = time.time()
instance_existed = False
while time.time() - start_time < timeout:
state = client.request_json("GET", "/state")
instances = state.get("instances", {})
if instance_id not in instances:
if instance_existed:
# Instance was deleted after being created - likely due to runner failure
raise RuntimeError(
f"Instance {instance_id} was deleted (runner may have failed)"
)
time.sleep(0.1)
continue
instance_existed = True
instance = instances[instance_id]
runner_ids = runner_ids_from_instance(instance)
runners = state.get("runners", {})
# Check for failed runners first
for rid in runner_ids:
runner = runners.get(rid, {})
if runner_failed(runner):
error_msg = get_runner_failed_message(runner) or "Unknown error"
raise RuntimeError(f"Runner {rid} failed: {error_msg}")
if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
return
time.sleep(0.1)
raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
def wait_for_instance_gone(
client: ExoClient, instance_id: str, timeout: float = 3.0
) -> None:
start_time = time.time()
while time.time() - start_time < timeout:
try:
client.request_json("GET", f"/instance/{instance_id}")
time.sleep(0.4)
except ExoHttpError as e:
if e.status == 404:
return
raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
def format_peak_memory(b: float) -> str:
for unit in ["B", "KB", "MB", "GB", "TB"]:
if b < 1024.0:
@@ -143,6 +269,184 @@ def parse_int_list(values: list[str]) -> list[int]:
return items
def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]:
models = client.request_json("GET", "/models") or {}
data = models.get("data") or []
for m in data:
if m.get("name").lower() == model_arg.lower():
short_id = str(m["name"])
full_id = str(m.get("hugging_face_id") or m["name"])
return short_id, full_id
for m in data:
if m.get("hugging_face_id") == model_arg:
short_id = str(m["name"])
full_id = str(m["hugging_face_id"])
return short_id, full_id
raise ValueError(f"Model not found in /models: {model_arg}")
def run_planning_phase(
client: ExoClient,
full_model_id: str,
preview: dict[str, Any],
danger_delete: bool,
timeout: float,
settle_deadline: float | None,
) -> None:
"""Check disk space and ensure model is downloaded before benchmarking."""
# Get model size from /models
models = client.request_json("GET", "/models") or {}
model_bytes = 0
for m in models.get("data", []):
if m.get("hugging_face_id") == full_model_id:
model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
break
if not model_bytes:
logger.warning(
f"Could not determine size for {full_model_id}, skipping disk check"
)
return
# Get nodes from preview
inner = unwrap_instance(preview["instance"])
node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
state = client.request_json("GET", "/state")
downloads = state.get("downloads", {})
node_disk = state.get("nodeDisk", {})
for node_id in node_ids:
node_downloads = downloads.get(node_id, [])
# Check if model already downloaded on this node
already_downloaded = any(
"DownloadCompleted" in p
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
]
== full_model_id
for p in node_downloads
)
if already_downloaded:
continue
# Wait for disk info if settle_deadline is set
disk_info = node_disk.get(node_id, {})
backoff = _SETTLE_INITIAL_BACKOFF_S
while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
remaining = settle_deadline - time.monotonic()
logger.info(
f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
state = client.request_json("GET", "/state")
node_disk = state.get("nodeDisk", {})
disk_info = node_disk.get(node_id, {})
if not disk_info:
logger.warning(f"No disk info for {node_id}, skipping space check")
continue
avail = disk_info.get("available", {}).get("inBytes", 0)
if avail >= model_bytes:
continue
if not danger_delete:
raise RuntimeError(
f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
)
# Delete from smallest to largest
completed = [
(
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
],
p["DownloadCompleted"]["totalBytes"]["inBytes"],
)
for p in node_downloads
if "DownloadCompleted" in p
]
for del_model, size in sorted(completed, key=lambda x: x[1]):
logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
client.request_json("DELETE", f"/download/{node_id}/{del_model}")
avail += size
if avail >= model_bytes:
break
if avail < model_bytes:
raise RuntimeError(f"Could not free enough space on {node_id}")
# Start downloads (idempotent)
for node_id in node_ids:
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
shard = runner_to_shard[runner_id]
client.request_json(
"POST",
"/download/start",
body={
"targetNodeId": node_id,
"shardMetadata": shard,
},
)
logger.info(f"Started download on {node_id}")
# Wait for downloads
start = time.time()
while time.time() - start < timeout:
state = client.request_json("GET", "/state")
downloads = state.get("downloads", {})
all_done = True
for node_id in node_ids:
done = any(
"DownloadCompleted" in p
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
"modelCard"
]["modelId"]
== full_model_id
for p in downloads.get(node_id, [])
)
failed = [
p["DownloadFailed"]["errorMessage"]
for p in downloads.get(node_id, [])
if "DownloadFailed" in p
and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
"modelId"
]
== full_model_id
]
if failed:
raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
if not done:
all_done = False
if all_done:
return
time.sleep(1)
raise TimeoutError("Downloads did not complete in time")
def placement_filter(instance_meta: str, wanted: str) -> bool:
s = (instance_meta or "").lower()
if wanted == "both":
return ("ring" in s) or ("jaccl" in s)
return wanted in s
def sharding_filter(sharding: str, wanted: str) -> bool:
s = (sharding or "").lower()
if wanted == "both":
return ("pipeline" in s) or ("tensor" in s)
return wanted in s
def run_one_completion(
client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer
) -> tuple[dict[str, Any], int]:
@@ -234,12 +538,76 @@ class PromptSizer:
return content, tok
def fetch_and_filter_placements(
client: ExoClient, full_model_id: str, args: argparse.Namespace
) -> list[dict[str, Any]]:
previews_resp = client.request_json(
"GET", "/instance/previews", params={"model_id": full_model_id}
)
previews = previews_resp.get("previews") or []
selected: list[dict[str, Any]] = []
for p in previews:
if p.get("error") is not None:
continue
if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
continue
if not sharding_filter(str(p.get("sharding", "")), args.sharding):
continue
instance = p.get("instance")
if not isinstance(instance, dict):
continue
n = nodes_used_in_instance(instance)
# Skip tensor ring single node as it is pointless when pipeline ring
if n == 1 and (
(args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
or (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
):
continue
if (
args.skip_pipeline_jaccl
and (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
and (
args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
)
):
continue
if (
args.skip_tensor_ring
and (
args.instance_meta == "both"
and "ring" in p.get("instance_meta", "").lower()
)
and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
):
continue
if args.min_nodes <= n <= args.max_nodes:
selected.append(p)
return selected
def main() -> int:
ap = argparse.ArgumentParser(
prog="exo-bench",
description="Benchmark exo model throughput across placement previews.",
)
add_common_instance_args(ap)
ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
ap.add_argument(
"--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(
"--pp",
nargs="+",
@@ -253,13 +621,35 @@ def main() -> int:
help="Generation lengths (ints). Accepts commas.",
)
ap.add_argument(
"--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
"--max-nodes",
type=int,
default=4,
help="Only consider placements using <= this many nodes.",
)
ap.add_argument(
"--concurrency",
nargs="+",
default=["1"],
help="Concurrency levels (ints). Accepts commas. E.g. --concurrency 1,2,4,8. Default 1.",
"--min-nodes",
type=int,
default=1,
help="Only consider placements using >= this many nodes.",
)
ap.add_argument(
"--instance-meta", choices=["ring", "jaccl", "both"], default="both"
)
ap.add_argument(
"--sharding", choices=["pipeline", "tensor", "both"], default="both"
)
ap.add_argument(
"--skip-pipeline-jaccl",
action="store_true",
help="Skip pipeline+jaccl placements, as it's often pointless.",
)
ap.add_argument(
"--skip-tensor-ring",
action="store_true",
help="Skip tensor+ring placements, as it's so slow.",
)
ap.add_argument(
"--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
)
ap.add_argument(
"--warmup",
@@ -267,6 +657,9 @@ def main() -> int:
default=0,
help="Warmup runs per placement (uses first pp/tg).",
)
ap.add_argument(
"--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
)
ap.add_argument(
"--json-out",
default="bench/results.json",
@@ -281,8 +674,18 @@ def main() -> int:
action="store_true",
help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
)
ap.add_argument(
"--settle-timeout",
type=float,
default=0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
ap.add_argument(
"--danger-delete-downloads",
action="store_true",
help="Delete existing models from smallest to largest to make room for benchmark model.",
)
args = ap.parse_args()
validate_vllm_args(args)
pp_list = parse_int_list(args.pp)
tg_list = parse_int_list(args.tg)
@@ -292,10 +695,6 @@ 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)
@@ -307,11 +706,7 @@ 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)
if args.ensure_cuda:
ensure_cuda_available(client)
short_id, full_model_id = resolve_model_short_id(
client, args.model, force_download=args.force_download
)
short_id, full_model_id = resolve_model_short_id(client, args.model)
tokenizer = load_tokenizer_for_bench(full_model_id)
if tokenizer is None:
@@ -324,10 +719,24 @@ def main() -> int:
logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
raise
selected = settle_and_fetch_placements(
client, full_model_id, args, settle_timeout=args.settle_timeout
settle_deadline = (
time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
)
selected = fetch_and_filter_placements(client, full_model_id, args)
if not selected and settle_deadline:
backoff = _SETTLE_INITIAL_BACKOFF_S
while not selected and time.monotonic() < settle_deadline:
remaining = settle_deadline - time.monotonic()
logger.warning(
f"No valid placements yet (cluster may still be settling). "
f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
selected = fetch_and_filter_placements(client, full_model_id, args)
if not selected:
logger.error("No valid placements matched your filters.")
return 1
@@ -351,12 +760,8 @@ def main() -> int:
if args.dry_run:
return 0
settle_deadline = (
time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
)
logger.info("Planning phase: checking downloads...")
download_duration_s = run_planning_phase(
run_planning_phase(
client,
full_model_id,
selected[0],
@@ -364,10 +769,6 @@ def main() -> int:
args.timeout,
settle_deadline,
)
if download_duration_s is not None:
logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
else:
logger.info("Download: model already cached")
all_rows: list[dict[str, Any]] = []
@@ -410,131 +811,47 @@ def main() -> int:
pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
for pp, tg in pp_tg_pairs:
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:
valid_gen_tps = [
x["stats"]["generation_tps"]
for x, _ in batch_results
if x["stats"]["generation_tps"] > 0
]
agg_gen_tps = (
mean(valid_gen_tps) if valid_gen_tps else 0.0
)
gen_tps = agg_gen_tps / concurrency
logger.info(
f"[concurrent {concurrency}x] "
f"agg_gen_tps={agg_gen_tps:.2f} "
f"gen_tps={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"] / x["concurrency"]
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
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
)
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,
}
)
runs.append(row)
all_rows.append(row)
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)
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)
finally:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
-1380
View File
File diff suppressed because it is too large Load Diff
-536
View File
@@ -1,536 +0,0 @@
# type: ignore
from __future__ import annotations
import argparse
import http.client
import json
import os
import time
from typing import Any
from urllib.parse import urlencode
from loguru import logger
_SETTLE_INITIAL_BACKOFF_S = 1.0
_SETTLE_MAX_BACKOFF_S = 60.0
_SETTLE_BACKOFF_MULTIPLIER = 2.0
class ExoHttpError(RuntimeError):
def __init__(self, status: int, reason: str, body_preview: str):
super().__init__(f"HTTP {status} {reason}: {body_preview}")
self.status = status
class ExoClient:
def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
self.host = host
self.port = port
self.timeout_s = timeout_s
def request_json(
self,
method: str,
path: str,
params: dict[str, Any] | None = None,
body: dict[str, Any] | None = None,
headers: dict[str, str] | None = None,
) -> Any:
if not path.startswith("/"):
path = "/" + path
if params:
path = path + "?" + urlencode(params)
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
try:
payload: bytes | None = None
hdrs: dict[str, str] = {"Accept": "application/json"}
if body is not None:
payload = json.dumps(body).encode("utf-8")
hdrs["Content-Type"] = "application/json"
if headers:
hdrs.update(headers)
conn.request(method.upper(), path, body=payload, headers=hdrs)
resp = conn.getresponse()
raw = resp.read()
text = raw.decode("utf-8", errors="replace") if raw else ""
if resp.status >= 400:
raise ExoHttpError(resp.status, resp.reason, text[:300])
if not text:
return None
return json.loads(text)
finally:
conn.close()
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
return self.request_json("POST", "/bench/chat/completions", body=payload)
def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
if len(instance) != 1:
raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
tag = next(iter(instance))
inner = instance[tag]
if not isinstance(inner, dict):
raise TypeError(f"payload for {tag} must be dict, got {type(inner)}")
return inner
def instance_id_from_instance(instance: dict[str, Any]) -> str:
inner = unwrap_instance(instance)
return str(inner["instanceId"])
def nodes_used_in_instance(instance: dict[str, Any]) -> int:
inner = unwrap_instance(instance)
return len(inner["shardAssignments"]["nodeToRunner"])
def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
inner = unwrap_instance(instance)
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
return list(runner_to_shard.keys())
def runner_ready(runner: dict[str, Any]) -> bool:
return "RunnerReady" in runner
def runner_failed(runner: dict[str, Any]) -> bool:
return "RunnerFailed" in runner
def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
if "RunnerFailed" in runner:
return runner["RunnerFailed"].get("errorMessage")
return None
def wait_for_instance_ready(
client: ExoClient, instance_id: str, timeout: float = 24000.0
) -> None:
start_time = time.time()
instance_existed = False
while time.time() - start_time < timeout:
state = client.request_json("GET", "/state")
instances = state.get("instances", {})
if instance_id not in instances:
if instance_existed:
# Instance was deleted after being created - likely due to runner failure
raise RuntimeError(
f"Instance {instance_id} was deleted (runner may have failed)"
)
time.sleep(0.1)
continue
instance_existed = True
instance = instances[instance_id]
runner_ids = runner_ids_from_instance(instance)
runners = state.get("runners", {})
# Check for failed runners first
for rid in runner_ids:
runner = runners.get(rid, {})
if runner_failed(runner):
error_msg = get_runner_failed_message(runner) or "Unknown error"
raise RuntimeError(f"Runner {rid} failed: {error_msg}")
if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
return
time.sleep(0.1)
raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}")
def wait_for_instance_gone(
client: ExoClient, instance_id: str, timeout: float = 3.0
) -> None:
start_time = time.time()
while time.time() - start_time < timeout:
try:
client.request_json("GET", f"/instance/{instance_id}")
time.sleep(0.4)
except ExoHttpError as e:
if e.status == 404:
return
raise
raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
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 []
for m in data:
if (m.get("name") or "").lower() == model_arg.lower():
short_id = str(m["name"])
full_id = str(m.get("hugging_face_id") or m["name"])
return short_id, full_id
for m in data:
if m.get("hugging_face_id") == model_arg:
short_id = str(m["name"])
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}")
def validate_vllm_args(args: argparse.Namespace) -> None:
if args.instance_meta != "vllm":
return
if args.sharding == "tensor":
raise SystemExit(
"--instance-meta vllm is incompatible with --sharding tensor (vllm is pipeline-only)"
)
if args.min_nodes > 1:
raise SystemExit(
"--instance-meta vllm is incompatible with --min-nodes > 1 (vllm is single-node)"
)
if args.max_nodes > 1:
raise SystemExit(
"--instance-meta vllm is incompatible with --max-nodes > 1 (vllm is single-node)"
)
def ensure_cuda_available(client: ExoClient) -> None:
capabilities = client.request_json("GET", "/capabilities")
if not capabilities or not capabilities.get("vllm_available"):
raise SystemExit(
"--ensure-cuda: vllm is not available on the exo cluster (no CUDA capability)"
)
def placement_filter(instance_meta: str, wanted: str) -> bool:
s = (instance_meta or "").lower()
if wanted == "both":
return ("ring" in s) or ("jaccl" in s)
return wanted in s
def sharding_filter(sharding: str, wanted: str) -> bool:
s = (sharding or "").lower()
if wanted == "both":
return ("pipeline" in s) or ("tensor" in s)
return wanted in s
def fetch_and_filter_placements(
client: ExoClient, full_model_id: str, args: argparse.Namespace
) -> list[dict[str, Any]]:
previews_resp = client.request_json(
"GET", "/instance/previews", params={"model_id": full_model_id}
)
previews = previews_resp.get("previews") or []
selected: list[dict[str, Any]] = []
for p in previews:
if p.get("error") is not None:
continue
if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta):
continue
if not sharding_filter(str(p.get("sharding", "")), args.sharding):
continue
instance = p.get("instance")
if not isinstance(instance, dict):
continue
n = nodes_used_in_instance(instance)
# Skip tensor ring single node as it is pointless when pipeline ring
if n == 1 and (
(args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
or (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
):
continue
if (
args.skip_pipeline_jaccl
and (
args.instance_meta == "both"
and "jaccl" in p.get("instance_meta", "").lower()
)
and (
args.sharding == "both" and "pipeline" in p.get("sharding", "").lower()
)
):
continue
if (
args.skip_tensor_ring
and (
args.instance_meta == "both"
and "ring" in p.get("instance_meta", "").lower()
)
and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower())
):
continue
if args.min_nodes <= n <= args.max_nodes:
selected.append(p)
return selected
def settle_and_fetch_placements(
client: ExoClient,
full_model_id: str,
args: argparse.Namespace,
settle_timeout: float = 0,
) -> list[dict[str, Any]]:
selected = fetch_and_filter_placements(client, full_model_id, args)
if not selected and settle_timeout > 0:
backoff = _SETTLE_INITIAL_BACKOFF_S
deadline = time.monotonic() + settle_timeout
while not selected and time.monotonic() < deadline:
remaining = deadline - time.monotonic()
logger.warning(
f"No valid placements yet (cluster may still be settling). "
f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
selected = fetch_and_filter_placements(client, full_model_id, args)
return selected
def run_planning_phase(
client: ExoClient,
full_model_id: str,
preview: dict[str, Any],
danger_delete: bool,
timeout: float,
settle_deadline: float | None,
) -> float | None:
"""Check disk space and ensure model is downloaded before benchmarking.
Returns the wall-clock download duration in seconds if a fresh download
was needed, or None if the model was already cached on all nodes.
"""
# Get model size from /models
models = client.request_json("GET", "/models") or {}
model_bytes = 0
for m in models.get("data", []):
if m.get("hugging_face_id") == full_model_id:
model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
break
if not model_bytes:
logger.warning(
f"Could not determine size for {full_model_id}, skipping disk check"
)
return None
# Get nodes from preview
inner = unwrap_instance(preview["instance"])
node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
runner_to_shard = inner["shardAssignments"]["runnerToShard"]
state = client.request_json("GET", "/state")
downloads = state.get("downloads", {})
node_disk = state.get("nodeDisk", {})
needs_download = False
for node_id in node_ids:
node_downloads = downloads.get(node_id, [])
# Check if model already downloaded on this node
already_downloaded = any(
"DownloadCompleted" in p
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
]
== full_model_id
for p in node_downloads
)
if already_downloaded:
continue
needs_download = True
# Wait for disk info if settle_deadline is set
disk_info = node_disk.get(node_id, {})
backoff = _SETTLE_INITIAL_BACKOFF_S
while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
remaining = settle_deadline - time.monotonic()
logger.info(
f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
)
time.sleep(min(backoff, remaining))
backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
state = client.request_json("GET", "/state")
node_disk = state.get("nodeDisk", {})
disk_info = node_disk.get(node_id, {})
if not disk_info:
logger.warning(f"No disk info for {node_id}, skipping space check")
continue
avail = disk_info.get("available", {}).get("inBytes", 0)
if avail >= model_bytes:
continue
if not danger_delete:
raise RuntimeError(
f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
)
# Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
completed = [
(
unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
"modelId"
],
p["DownloadCompleted"]["total"]["inBytes"],
)
for p in node_downloads
if "DownloadCompleted" in p
and not p["DownloadCompleted"].get("readOnly", False)
]
for del_model, size in sorted(completed, key=lambda x: x[1]):
logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
client.request_json("DELETE", f"/download/{node_id}/{del_model}")
avail += size
if avail >= model_bytes:
break
if avail < model_bytes:
raise RuntimeError(f"Could not free enough space on {node_id}")
# Start downloads (idempotent)
download_t0 = time.perf_counter() if needs_download else None
for node_id in node_ids:
runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
shard = runner_to_shard[runner_id]
client.request_json(
"POST",
"/download/start",
body={
"targetNodeId": node_id,
"shardMetadata": shard,
},
)
logger.info(f"Started download on {node_id}")
# Wait for downloads
start = time.time()
while time.time() - start < timeout:
state = client.request_json("GET", "/state")
downloads = state.get("downloads", {})
all_done = True
for node_id in node_ids:
done = any(
"DownloadCompleted" in p
and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
"modelCard"
]["modelId"]
== full_model_id
for p in downloads.get(node_id, [])
)
failed = [
p["DownloadFailed"]["errorMessage"]
for p in downloads.get(node_id, [])
if "DownloadFailed" in p
and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
"modelId"
]
== full_model_id
]
if failed:
raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
if not done:
all_done = False
if all_done:
if download_t0 is not None:
return time.perf_counter() - download_t0
return None
time.sleep(1)
raise TimeoutError("Downloads did not complete in time")
def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
ap.add_argument(
"--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,
default=4,
help="Only consider placements using <= this many nodes.",
)
ap.add_argument(
"--min-nodes",
type=int,
default=1,
help="Only consider placements using >= this many nodes.",
)
ap.add_argument(
"--instance-meta", choices=["ring", "jaccl", "vllm", "both"], default="both"
)
ap.add_argument(
"--sharding", choices=["pipeline", "tensor", "both"], default="both"
)
ap.add_argument(
"--skip-pipeline-jaccl",
action="store_true",
help="Skip pipeline+jaccl placements, as it's often pointless.",
)
ap.add_argument(
"--skip-tensor-ring",
action="store_true",
help="Skip tensor+ring placements, as it's so slow.",
)
ap.add_argument(
"--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)."
)
ap.add_argument(
"--settle-timeout",
type=float,
default=0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
ap.add_argument(
"--danger-delete-downloads",
action="store_true",
help="Delete existing models from smallest to largest to make room for benchmark model.",
)
ap.add_argument(
"--ensure-cuda",
action="store_true",
help="Verify the exo cluster has CUDA/vllm capability; error if not.",
)
-255
View File
@@ -1,255 +0,0 @@
# 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()
+5 -12
View File
@@ -4,18 +4,11 @@ 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",
"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",
"loguru>=0.7.3",
"transformers>=5.0.0",
"huggingface-hub>=0.33.4",
"tiktoken>=0.12.0",
"jinja2>=3.1.0",
]
[build-system]
-306
View File
@@ -1,306 +0,0 @@
# Tool definitions — each becomes an OpenAI function tool.
# All scenarios get all tools unless they specify a `tools` list.
[tools.get_current_weather]
description = "Get the current weather in a given location"
required = ["location"]
[tools.get_current_weather.properties.location]
type = "string"
description = "City and state, e.g. San Francisco, CA"
[tools.get_current_weather.properties.unit]
type = "string"
enum = ["celsius", "fahrenheit"]
description = "Temperature unit"
[tools.calculate]
description = "Evaluate a mathematical expression and return the numeric result"
required = ["expression"]
[tools.calculate.properties.expression]
type = "string"
description = "The math expression to evaluate, e.g. '2 + 3 * 4'"
[tools.search_products]
description = "Search for products in a catalog by query, category, and price"
required = ["query"]
[tools.search_products.properties.query]
type = "string"
description = "Search query string"
[tools.search_products.properties.category]
type = "string"
enum = ["electronics", "clothing", "food", "books"]
description = "Product category to filter by"
[tools.search_products.properties.max_price]
type = "number"
description = "Maximum price in USD"
[tools.create_todos]
description = "Create a structured todo list"
required = ["todos"]
[tools.create_todos.properties.todos]
type = "array"
description = "List of todo items"
[tools.create_todos.properties.todos.items]
type = "object"
required = ["content", "status", "priority"]
[tools.create_todos.properties.todos.items.properties.content]
type = "string"
description = "The todo item text"
[tools.create_todos.properties.todos.items.properties.status]
type = "string"
description = "Status: pending, in_progress, or completed"
[tools.create_todos.properties.todos.items.properties.priority]
type = "string"
description = "Priority: low, normal, or high"
# -- Should call a tool --
[[scenarios]]
name = "weather_simple"
description = "Basic weather query -> get_current_weather"
expect_tool_call = true
expected_function = "get_current_weather"
required_arg_keys = ["location"]
[[scenarios.messages]]
role = "user"
content = "What's the weather like in Tokyo right now?"
[[scenarios]]
name = "calculator_simple"
description = "Math question -> calculate"
expect_tool_call = true
expected_function = "calculate"
required_arg_keys = ["expression"]
[[scenarios.messages]]
role = "user"
content = "Use the calculator to compute 3847 * 926 + 17293"
[[scenarios]]
name = "search_with_filters"
description = "Product search with category and price filter"
expect_tool_call = true
expected_function = "search_products"
required_arg_keys = ["query"]
[[scenarios.messages]]
role = "user"
content = "Find me electronics under $50"
# -- Multi-turn: tool call then follow-up --
[[scenarios]]
name = "weather_multi_turn"
description = "Weather query -> tool result -> natural language summary"
expect_tool_call = true
expected_function = "get_current_weather"
required_arg_keys = ["location"]
[scenarios.tool_result]
temperature = "18C"
condition = "partly cloudy"
humidity = "65%"
wind = "12 km/h NW"
[[scenarios.messages]]
role = "user"
content = "What's the weather in Paris?"
[[scenarios]]
name = "calculator_multi_turn"
description = "Math query -> tool result -> model reports the answer"
expect_tool_call = true
expected_function = "calculate"
required_arg_keys = ["expression"]
[scenarios.tool_result]
result = 491682
[[scenarios.messages]]
role = "user"
content = "Use the calculator to compute 1847 * 263 + 5921"
[[scenarios]]
name = "search_multi_turn"
description = "Search query -> tool result -> model summarizes products"
expect_tool_call = true
expected_function = "search_products"
required_arg_keys = ["query"]
[[scenarios.tool_result.results]]
name = "Hands-On Machine Learning"
price = 45.99
rating = 4.8
[[scenarios.tool_result.results]]
name = "Deep Learning with Python"
price = 39.99
rating = 4.6
[[scenarios.messages]]
role = "user"
content = "Search for books about machine learning"
# -- Sequential tool calls --
[[scenarios]]
name = "chained_tool_calls_same"
description = "Thinking + weather(Tokyo) -> result -> model must call weather(London)"
expect_tool_call = true
expected_function = "get_current_weather"
required_arg_keys = ["location"]
[[scenarios.messages]]
role = "user"
content = "Compare the weather in Tokyo and London."
[[scenarios.messages]]
role = "assistant"
content = "I'll check both cities. Let me start with Tokyo."
[[scenarios.messages.tool_calls]]
id = "call_1"
name = "get_current_weather"
arguments = { location = "Tokyo" }
[[scenarios.messages]]
role = "tool"
tool_call_id = "call_1"
content = '{"temperature": "25C", "condition": "sunny"}'
[[scenarios]]
name = "chained_tool_calls_different"
description = "Thinking + weather(Berlin) -> result -> model must call calculator"
expect_tool_call = true
expected_function = "calculate"
required_arg_keys = ["expression"]
[[scenarios.messages]]
role = "user"
content = "What's the weather in Berlin, and also use the calculator to compute 4819 * 37 + 291."
[[scenarios.messages]]
role = "assistant"
content = "I'll handle both. Let me check Berlin's weather first."
[[scenarios.messages.tool_calls]]
id = "call_2"
name = "get_current_weather"
arguments = { location = "Berlin" }
[[scenarios.messages]]
role = "tool"
tool_call_id = "call_2"
content = '{"temperature": "12C", "condition": "rainy"}'
[[scenarios]]
name = "chained_tool_calls_three"
description = "Two prior thinking+tool calls -> results -> model must make a third"
expect_tool_call = true
expected_function = "get_current_weather"
required_arg_keys = ["location"]
[[scenarios.messages]]
role = "user"
content = "Compare weather in Tokyo, Paris, and London."
[[scenarios.messages]]
role = "assistant"
content = "I'll check all three cities. Starting with Tokyo."
[[scenarios.messages.tool_calls]]
id = "call_3"
name = "get_current_weather"
arguments = { location = "Tokyo" }
[[scenarios.messages]]
role = "tool"
tool_call_id = "call_3"
content = '{"temperature": "25C", "condition": "sunny"}'
[[scenarios.messages]]
role = "assistant"
content = "Got Tokyo. Now checking Paris."
[[scenarios.messages.tool_calls]]
id = "call_4"
name = "get_current_weather"
arguments = { location = "Paris" }
[[scenarios.messages]]
role = "tool"
tool_call_id = "call_4"
content = '{"temperature": "18C", "condition": "cloudy"}'
# -- Nested object schema (regression for lossy chat template rendering) --
[[scenarios]]
name = "nested_schema_tool_call"
description = "Tool call with nested object array schema -> create_todos"
expect_tool_call = true
expected_function = "create_todos"
required_arg_keys = ["todos"]
nested_array_key = "todos"
required_item_keys = ["content", "status", "priority"]
tools = ["create_todos"]
[[scenarios.messages]]
role = "user"
content = "Create a todo list with 3 items to learn Python"
# -- Tool name integrity (regression for harmony token leaking into name) --
[tools.glob]
description = "Search for files matching a glob pattern in the codebase"
required = ["pattern"]
[tools.glob.properties.pattern]
type = "string"
description = "The glob pattern to match files against, e.g. '**/*.py'"
[tools.glob.properties.path]
type = "string"
description = "The directory to search in"
[[scenarios]]
name = "tool_name_integrity"
description = "Tool name must not contain harmony tokens like <|channel|>"
expect_tool_call = true
expected_function = "glob"
required_arg_keys = ["pattern"]
tools = ["glob"]
[[scenarios.messages]]
role = "user"
content = "Find all Python files in the src directory"
# -- Should NOT call a tool --
[[scenarios]]
name = "no_tool_joke"
description = "Joke request should NOT trigger any tool"
expect_tool_call = false
[[scenarios.messages]]
role = "user"
content = "Tell me a funny joke about cats."
[[scenarios]]
name = "no_tool_factual"
description = "Factual question answerable from training data"
expect_tool_call = false
[[scenarios.messages]]
role = "user"
content = "What is the capital of Japan?"
+4 -4
View File
@@ -2,10 +2,10 @@
#
# Shared constraints applied to ALL benchmarks in this file.
constraints = [
"All(MacOsBuild(=25D125))",
"Hosts(=1)",
"All(Chip(m3_ultra))",
"All(GpuCores(=80))",
"All(MacOsBuild(=25D125))",
"Hosts(=1)",
"All(Chip(m3_ultra))",
"All(GpuCores(=80))",
]
[topology]
-377
View File
@@ -1,377 +0,0 @@
# type: ignore
import argparse
import json
import os
import statistics
import sys
import tempfile
import time
import mlx.core as mx
DTYPE_MAP = {
"float32": (mx.float32, 4),
"float16": (mx.float16, 2),
"bfloat16": (mx.bfloat16, 2),
}
SIZES = [
1 * 1024,
4 * 1024,
16 * 1024,
64 * 1024,
256 * 1024,
1 * 1024 * 1024,
4 * 1024 * 1024,
16 * 1024 * 1024,
64 * 1024 * 1024,
256 * 1024 * 1024,
1 * 1024 * 1024 * 1024,
2 * 1024 * 1024 * 1024,
4 * 1024 * 1024 * 1024,
8 * 1024 * 1024 * 1024,
]
def format_bytes(n: int) -> str:
if n >= 1024 * 1024 * 1024:
return f"{n / (1024 * 1024 * 1024):.0f} GB"
if n >= 1024 * 1024:
return f"{n / (1024 * 1024):.0f} MB"
if n >= 1024:
return f"{n / 1024:.0f} KB"
return f"{n} B"
def format_time(seconds: float) -> str:
if seconds >= 1.0:
return f"{seconds:.3f} s"
if seconds >= 0.001:
return f"{seconds * 1000:.2f} ms"
return f"{seconds * 1_000_000:.1f} us"
def format_bandwidth(bytes_per_sec: float) -> str:
if bytes_per_sec >= 1024 * 1024 * 1024:
return f"{bytes_per_sec / (1024 * 1024 * 1024):.2f} GB/s"
if bytes_per_sec >= 1024 * 1024:
return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s"
return f"{bytes_per_sec / 1024:.1f} KB/s"
def barrier(group: mx.distributed.Group) -> None:
mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group))
def init_ring(
rank: int, self_ip: str, peer_ip: str, port: int, tmpdir: str
) -> mx.distributed.Group:
if rank == 0:
hosts = [f"{self_ip}:{port}", f"{peer_ip}:{port}"]
else:
hosts = [f"{peer_ip}:{port}", f"{self_ip}:{port}"]
hostfile = os.path.join(tmpdir, "hosts.json")
with open(hostfile, "w") as f:
json.dump(hosts, f)
for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
os.environ.pop(var, None)
os.environ["MLX_HOSTFILE"] = hostfile
os.environ["MLX_RANK"] = str(rank)
return mx.distributed.init(backend="ring", strict=True)
def init_jaccl(
rank: int, interface: str, coordinator: str, port: int, tmpdir: str
) -> mx.distributed.Group:
devices = [[None, interface], [interface, None]]
devfile = os.path.join(tmpdir, "devices.json")
with open(devfile, "w") as f:
json.dump(devices, f)
for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"):
os.environ.pop(var, None)
os.environ["MLX_IBV_DEVICES"] = devfile
os.environ["MLX_RANK"] = str(rank)
if rank == 0:
os.environ["MLX_JACCL_COORDINATOR"] = f"0.0.0.0:{port}"
else:
os.environ["MLX_JACCL_COORDINATOR"] = coordinator
return mx.distributed.init(backend="jaccl", strict=True)
def bench_unidirectional(
group: mx.distributed.Group,
rank: int,
size_bytes: int,
dtype: mx.Dtype,
element_size: int,
warmup: int,
iterations: int,
) -> list[float]:
n_elements = size_bytes // element_size
tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
mx.eval(tensor)
for _ in range(warmup):
if rank == 0:
sent = mx.distributed.send(tensor, dst=1, group=group)
mx.eval(sent)
else:
received = mx.distributed.recv_like(tensor, src=0, group=group)
mx.eval(received)
barrier(group)
times: list[float] = []
for _ in range(iterations):
barrier(group)
t0 = time.perf_counter()
if rank == 0:
sent = mx.distributed.send(tensor, dst=1, group=group)
mx.eval(sent)
else:
received = mx.distributed.recv_like(tensor, src=0, group=group)
mx.eval(received)
barrier(group)
t1 = time.perf_counter()
times.append(t1 - t0)
return times
def bench_rtt(
group: mx.distributed.Group,
rank: int,
size_bytes: int,
dtype: mx.Dtype,
element_size: int,
warmup: int,
iterations: int,
) -> list[float]:
n_elements = size_bytes // element_size
tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
mx.eval(tensor)
for _ in range(warmup):
if rank == 0:
sent = mx.distributed.send(tensor, dst=1, group=group)
mx.eval(sent)
received = mx.distributed.recv_like(tensor, src=1, group=group)
mx.eval(received)
else:
received = mx.distributed.recv_like(tensor, src=0, group=group)
mx.eval(received)
sent = mx.distributed.send(received, dst=0, group=group)
mx.eval(sent)
barrier(group)
times: list[float] = []
for _ in range(iterations):
barrier(group)
t0 = time.perf_counter()
if rank == 0:
sent = mx.distributed.send(tensor, dst=1, group=group)
mx.eval(sent)
received = mx.distributed.recv_like(tensor, src=1, group=group)
mx.eval(received)
else:
received = mx.distributed.recv_like(tensor, src=0, group=group)
mx.eval(received)
sent = mx.distributed.send(received, dst=0, group=group)
mx.eval(sent)
barrier(group)
t1 = time.perf_counter()
times.append(t1 - t0)
return times
def bench_all_gather(
group: mx.distributed.Group,
rank: int,
size_bytes: int,
dtype: mx.Dtype,
element_size: int,
warmup: int,
iterations: int,
) -> list[float]:
n_elements = (size_bytes // 2) // element_size
tensor = mx.random.normal(shape=(n_elements,)).astype(dtype)
mx.eval(tensor)
for _ in range(warmup):
gathered = mx.distributed.all_gather(tensor, group=group)
mx.eval(gathered)
barrier(group)
times: list[float] = []
for _ in range(iterations):
barrier(group)
t0 = time.perf_counter()
gathered = mx.distributed.all_gather(tensor, group=group)
mx.eval(gathered)
t1 = time.perf_counter()
times.append(t1 - t0)
return times
def print_table(title: str, rows: list[dict[str, str]]) -> None:
print(f"\n=== {title} ===")
headers = ["Size", "Median", "Min", "Max", "Bandwidth"]
widths = [
max(len(h), max((len(r[h]) for r in rows), default=0)) + 2 for h in headers
]
header_line = "".join(h.ljust(w) for h, w in zip(headers, widths, strict=True))
print(header_line)
print("-" * len(header_line))
for row in rows:
print("".join(row[h].ljust(w) for h, w in zip(headers, widths, strict=True)))
def run_bench(
name: str,
bench_fn,
group: mx.distributed.Group,
rank: int,
dtype: mx.Dtype,
element_size: int,
warmup: int,
iterations: int,
bw_multiplier: int = 1,
) -> None:
rows: list[dict[str, str]] = []
for size in SIZES:
if rank == 0:
print(f" {name}: {format_bytes(size)}...", end="", flush=True)
times = bench_fn(group, rank, size, dtype, element_size, warmup, iterations)
if rank == 0:
med = statistics.median(times)
mn = min(times)
mx_ = max(times)
bw = (size * bw_multiplier) / med
rows.append(
{
"Size": format_bytes(size),
"Median": format_time(med),
"Min": format_time(mn),
"Max": format_time(mx_),
"Bandwidth": format_bandwidth(bw),
}
)
print(f" {format_bandwidth(bw)}")
if rank == 0:
print_table(name, rows)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="MLX Distributed Communication Benchmark"
)
subparsers = parser.add_subparsers(dest="backend", required=True)
ring_parser = subparsers.add_parser("ring")
ring_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
ring_parser.add_argument("--self-ip", required=True)
ring_parser.add_argument("--peer-ip", required=True)
ring_parser.add_argument("--port", type=int, default=5555)
jaccl_parser = subparsers.add_parser("jaccl")
jaccl_parser.add_argument("--rank", type=int, required=True, choices=[0, 1])
jaccl_parser.add_argument("--interface", required=True)
jaccl_parser.add_argument(
"--coordinator",
type=str,
default=None,
help="IP:PORT of rank 0 (required for rank 1)",
)
jaccl_parser.add_argument(
"--port", type=int, default=9999, help="Coordinator port (rank 0 only)"
)
for p in [ring_parser, jaccl_parser]:
p.add_argument("--warmup", type=int, default=3)
p.add_argument("--iterations", type=int, default=10)
p.add_argument("--dtype", choices=list(DTYPE_MAP.keys()), default="float32")
args = parser.parse_args()
if args.backend == "jaccl" and args.rank == 1 and args.coordinator is None:
jaccl_parser.error("--coordinator is required for rank 1")
return args
def main() -> int:
args = parse_args()
dtype, element_size = DTYPE_MAP[args.dtype]
with tempfile.TemporaryDirectory() as tmpdir:
if args.backend == "ring":
print(f"Initializing ring backend (rank {args.rank})...")
group = init_ring(args.rank, args.self_ip, args.peer_ip, args.port, tmpdir)
else:
print(f"Initializing jaccl backend (rank {args.rank})...")
group = init_jaccl(
args.rank, args.interface, args.coordinator or "", args.port, tmpdir
)
print(f"Rank {group.rank()} of {group.size()} initialized")
barrier(group)
if args.rank == 0:
print("\nMLX Distributed Communication Benchmark")
print(
f"Backend: {args.backend} | Dtype: {args.dtype} | Warmup: {args.warmup} | Iterations: {args.iterations}"
)
run_bench(
"Unidirectional (rank 0 -> rank 1)",
bench_unidirectional,
group,
args.rank,
dtype,
element_size,
args.warmup,
args.iterations,
)
run_bench(
"Round-Trip (ping-pong)",
bench_rtt,
group,
args.rank,
dtype,
element_size,
args.warmup,
args.iterations,
bw_multiplier=2,
)
run_bench(
"All-Gather",
bench_all_gather,
group,
args.rank,
dtype,
element_size,
args.warmup,
args.iterations,
)
if args.rank == 0:
print("\nDone.")
else:
print("Rank 1 complete.")
return 0
if __name__ == "__main__":
try:
sys.exit(main())
except KeyboardInterrupt:
print("\nInterrupted.")
sys.exit(1)
View File
-593
View File
@@ -1,593 +0,0 @@
# 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
View File

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