Compare commits

..
63 Commits
Author SHA1 Message Date
Evan a618d02277 more fixes 2026-04-14 12:45:27 +01:00
Evan 25eb488d29 more fixes 2026-04-14 11:17:32 +01:00
Evan 5b34e04a6c lint 2026-04-14 10:30:59 +01:00
Evan 4037634662 fix pyproject 2026-04-14 10:30:59 +01:00
Evan 164c081a65 Merge branch 'main' into leo/prefill-decode-really 2026-04-14 10:06:33 +01:00
Ryuichi Leo Takashige 39f0893a90 tmp 2026-04-05 01:39:26 +01:00
Ryuichi Leo Takashige 5d22805a77 Benchmarking 2026-04-01 19:51:03 +01:00
Ryuichi Leo Takashige 973e4db085 Unnecessary further optimizations 3 2026-03-22 16:02:31 +00:00
Ryuichi Leo Takashige 016de1803b Unnecessary further optimizations 2 2026-03-20 22:01:50 +00:00
Ryuichi Leo Takashige 60a6ac1125 Unnecessary further optimizations 2026-03-19 21:56:47 +00:00
Ryuichi Leo Takashige 5422e831ce Extra QOL that needs to be reworked 2026-03-19 20:59:29 +00:00
Ryuichi Leo Takashige 03ea3cf6cd Performance optimizations 2026-03-19 19:07:41 +00:00
Ryuichi Leo Takashige 6fa2cc1265 Non overlapping case 2026-03-18 23:41:16 +00:00
Ryuichi Leo Takashige 04197fe27b Some QOL 2026-03-18 23:12:38 +00:00
Ryuichi Leo Takashige d1490444a1 Implement prefill/decode really 2026-03-18 21:14:15 +00:00
Ryuichi Leo Takashige ba472da84f Implement prefill/decode 2026-03-18 18:20:05 +00:00
Ryuichi Leo Takashige f208586092 Implement prefill/decode 2026-03-18 10:39:22 +00:00
Ryuichi Leo Takashige be731d3a85 Merge main 2026-03-17 19:05:55 +00:00
Ryuichi Leo Takashige 655185cfe7 Address comments 4 - defer to the warmup into the exo batch generator and vllm batch engine and don't store model on the generators. 2026-03-17 19:00:12 +00:00
Ryuichi Leo Takashige 1dd9c28842 Address comments 3, mainly refactors 2026-03-17 18:31:41 +00:00
Ryuichi Leo Takashige cacd26e63c Type error lol 2026-03-17 18:00:39 +00:00
Ryuichi Leo Takashige 6a3eb2f37d close() 2026-03-17 17:31:15 +00:00
Ryuichi Leo Takashige e1df77bc4c No more future annotations 2026-03-17 17:07:32 +00:00
Ryuichi Leo Takashige e78e53df6e Address comments 2 including vllm capability in state 2026-03-17 16:55:08 +00:00
Ryuichi Leo Takashige c70d9006e8 Address comments including task id interface 2026-03-17 15:31:55 +00:00
Ryuichi Leo Takashige 72cd8552ae Merge branch 'main' into leo/dgx-spark-integrations 2026-03-17 13:58:49 +00:00
Ryuichi Leo Takashige 8cd1308336 Tidy pass 1 2026-03-17 00:38:16 +00:00
Ryuichi Leo Takashige 04dcdbd127 Merge main 2026-03-16 23:01:48 +00:00
Ryuichi Leo Takashige ec5d62f935 Strip vllm generator 2026-03-16 22:11:35 +00:00
Ryuichi Leo Takashige e96f084051 Distributed callbacks 2026-03-16 21:17:14 +00:00
Ryuichi Leo Takashige dc68ddbac0 Prompt formatting 2026-03-16 20:37:38 +00:00
Ryuichi Leo Takashige 073f8c1690 add batching 2026-03-16 19:25:45 +00:00
Ryuichi Leo Takashige 3c29d0dd4c test prefix caching 2026-03-16 16:43:51 +00:00
rltakashige 594ed99734 new uv lock for fastsafetensors 2026-03-13 17:08:49 +00:00
Ryuichi Leo Takashige e9e23e556e Have loading progress 2026-03-13 12:58:34 +00:00
Ryuichi Leo Takashige 169ea2a5e8 Fix GPT OSS by not retokenizing prompts 2026-03-12 18:02:43 +00:00
Ryuichi Leo Takashige 4a7901c548 Fix patches 2026-03-12 17:42:02 +00:00
Ryuichi Leo Takashige 7bb5cb4fc7 Allow memory profiling to be unstable 2026-03-12 17:34:53 +00:00
Ryuichi Leo Takashige 493e342f83 Skip impossible shardings 2026-03-12 17:33:12 +00:00
Ryuichi Leo Takashige 283b1809c9 Move VLLM runner into VLLM engine 2026-03-12 17:20:30 +00:00
Ryuichi Leo Takashige 35030119e3 ExoBench and ExoEval for CUDA 2026-03-12 17:09:51 +00:00
rltakashige 585dfe3549 Merge branch 'main' into leo/dgx-spark-integrations 2026-03-12 15:41:29 +00:00
Ryuichi Leo Takashige 5d7a005a13 Destroy process group on Keyboard Interrupt 2026-03-12 15:33:32 +00:00
Ryuichi Leo Takashige 87b7c5ef8b Pass CI 2026-03-12 15:24:49 +00:00
Ryuichi Leo Takashige 957ebbd21f Set max token length as max context length if no max tokens set 2026-03-12 15:15:39 +00:00
rltakashige 4b6dd7588f lockgit status 2026-03-12 14:41:33 +00:00
Ryuichi Leo Takashige 3f4f7c9ba6 Only do for aarch64 linux 2026-03-12 13:42:23 +00:00
Ryuichi Leo Takashige 1331465ba0 Add missing runner features 2026-03-12 10:51:37 +00:00
Ryuichi Leo Takashige 8f94727f14 Ignore missing modules if type stubs exist 2026-03-12 00:20:28 +00:00
Ryuichi Leo Takashige 9ee23ee0d3 Make vllm inference runner closer to the normal inference runner 2026-03-11 23:41:34 +00:00
Ryuichi Leo Takashige f75d36cbe0 Fix cache patch 2026-03-11 22:28:45 +00:00
Ryuichi Leo Takashige 2683ac7b61 Add Torch typings 2026-03-11 21:44:10 +00:00
Ryuichi Leo Takashige 404b9769ac Download models without model.safetensors.index 2026-03-11 21:32:07 +00:00
Ryuichi Leo Takashige 3e097f7243 only download a single copy of the model. 2026-03-11 20:57:29 +00:00
Ryuichi Leo Takashige 0c8615f25c Only import VLLM once.. 2026-03-11 20:49:19 +00:00
Ryuichi Leo Takashige ba35a4ba13 Move VLLM into the runner and add type stubs 2026-03-11 19:15:55 +00:00
Ryuichi Leo Takashige 9a83fa6cdf Patch VLLM to load multiple models dynamically 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 659c1bc737 Progress: Run EXO-CUDA through nix! 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 34df811b92 Vibe coding design baby 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige ca5870a2e8 Some Linux Laptop/Desktop detection and goodbye penguin 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 6be6ea5fd2 Fix placement preview 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige 5e9d27b753 Show Sparks and Linux in topology 2026-03-11 18:17:59 +00:00
Ryuichi Leo Takashige cfc8f09004 Fast direct USB connectivity 2026-03-11 18:17:59 +00:00
347 changed files with 10014 additions and 21634 deletions

No files matched your search

+20
View File
@@ -0,0 +1,20 @@
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
@@ -0,0 +1,17 @@
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
@@ -0,0 +1,61 @@
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: ...
@@ -0,0 +1 @@
from torch.backends import cuda as cuda
@@ -0,0 +1 @@
def is_built() -> bool: ...
+10
View File
@@ -0,0 +1,10 @@
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: ...
@@ -0,0 +1,2 @@
def is_initialized() -> bool: ...
def destroy_process_group() -> None: ...
+1
View File
@@ -0,0 +1 @@
__version__: str
+2
View File
@@ -0,0 +1,2 @@
class ModelConfig:
max_model_len: int
+18
View File
@@ -0,0 +1,18 @@
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
@@ -0,0 +1,17 @@
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
@@ -0,0 +1,11 @@
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
@@ -0,0 +1,3 @@
from vllm.tokenizers.protocol import TokenizerLike
__all__ = ["TokenizerLike"]
@@ -0,0 +1,15 @@
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
@@ -0,0 +1 @@
+1
View File
@@ -0,0 +1 @@
@@ -0,0 +1,24 @@
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: ...
@@ -0,0 +1,16 @@
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]: ...
@@ -0,0 +1,22 @@
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: ...
@@ -0,0 +1,23 @@
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
@@ -0,0 +1,6 @@
class Request:
request_id: str
prompt_token_ids: list[int] | None
num_prompt_tokens: int
num_computed_tokens: int
num_tokens: int
@@ -0,0 +1 @@
@@ -0,0 +1,24 @@
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]: ...
@@ -0,0 +1,6 @@
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
@@ -0,0 +1 @@
def extract_layer_index(layer_name: str, num_attn_module: int) -> int: ...
+7
View File
@@ -1 +1,8 @@
use flake
# creates .venv if doesn't exist and loads its environment
export VIRTUAL_ENV=".venv"
if ! [ -d "./$VIRTUAL_ENV" ]; then
uv venv
fi
layout python
+3 -115
View File
@@ -32,6 +32,7 @@ jobs:
SPARKLE_ED25519_PRIVATE: ${{ secrets.SPARKLE_ED25519_PRIVATE }}
SPARKLE_S3_BUCKET: ${{ secrets.SPARKLE_S3_BUCKET }}
SPARKLE_S3_PREFIX: ${{ secrets.SPARKLE_S3_PREFIX }}
EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT: ${{ secrets.EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT }}
AWS_REGION: ${{ secrets.AWS_REGION }}
EXO_BUILD_NUMBER: ${{ github.run_number }}
EXO_LIBP2P_NAMESPACE: ${{ github.ref_name }}
@@ -238,80 +239,6 @@ jobs:
# Export keychain path for other steps
echo "BUILD_KEYCHAIN_PATH=$KEYCHAIN_PATH" >> $GITHUB_ENV
# ============================================================
# Pre-flight credential / profile validation
# Runs BEFORE the ~16 min build so auth/expiry failures surface in <1 min.
# ============================================================
- name: Validate Apple notarization credentials
env:
APPLE_NOTARIZATION_USERNAME: ${{ secrets.APPLE_NOTARIZATION_USERNAME }}
APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
run: |
# All-or-nothing: either all three creds are set, or none are.
CRED_COUNT=0
for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
[[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
done
if [[ "$CRED_COUNT" -eq 0 ]]; then
echo "No notarization credentials configured — skipping notarization for this build."
exit 0
fi
if [[ "$CRED_COUNT" -ne 3 ]]; then
echo "ERROR: partial notarization credentials set ($CRED_COUNT/3). Aborting before build."
exit 1
fi
# Cheap, ~5s, auth-only call. Fails instantly with a clear message if
# the app-specific password is stale, wrong team-id, etc.
echo "Verifying Apple notarization credentials via notarytool history..."
if ! xcrun notarytool history \
--apple-id "$APPLE_NOTARIZATION_USERNAME" \
--password "$APPLE_NOTARIZATION_PASSWORD" \
--team-id "$APPLE_NOTARIZATION_TEAM" >/dev/null; then
echo "ERROR: notarytool rejected the provided credentials. Fix before rerunning."
echo "Common causes: app-specific password expired/revoked, wrong team-id,"
echo "Apple ID not on the team, or 2FA not configured for this Apple ID."
exit 1
fi
echo "Apple notarization credentials OK."
- name: Validate provisioning profile expiry
run: |
PROFILE="$HOME/Library/Developer/Xcode/UserData/Provisioning Profiles/EXO.provisionprofile"
if [[ ! -f "$PROFILE" ]]; then
echo "ERROR: provisioning profile not found at $PROFILE"
exit 1
fi
EXPIRY=$(security cms -D -i "$PROFILE" | plutil -extract ExpirationDate raw -o - - 2>/dev/null || true)
if [[ -z "$EXPIRY" ]]; then
echo "WARNING: could not read ExpirationDate from provisioning profile; skipping expiry check."
exit 0
fi
# Try a couple of known plutil date formats. If none parse, skip the check rather
# than risk a false-positive "expired" block on a format we didn't anticipate.
EXPIRY_EPOCH=""
for fmt in "%Y-%m-%dT%H:%M:%SZ" "%Y-%m-%d %H:%M:%S %z" "%Y-%m-%d %H:%M:%S +0000"; do
if parsed=$(date -j -f "$fmt" "$EXPIRY" +%s 2>/dev/null); then
EXPIRY_EPOCH="$parsed"
break
fi
done
if [[ -z "$EXPIRY_EPOCH" ]]; then
echo "WARNING: could not parse ExpirationDate '$EXPIRY'; skipping expiry check."
exit 0
fi
NOW_EPOCH=$(date +%s)
if [[ "$EXPIRY_EPOCH" -le "$NOW_EPOCH" ]]; then
echo "ERROR: provisioning profile expired on $EXPIRY. Regenerate it before rerunning."
exit 1
fi
DAYS_LEFT=$(( (EXPIRY_EPOCH - NOW_EPOCH) / 86400 ))
echo "Provisioning profile valid until $EXPIRY ($DAYS_LEFT days remaining)."
if [[ "$DAYS_LEFT" -lt 14 ]]; then
echo "WARNING: profile expires in under 14 days — regenerate soon."
fi
# ============================================================
# Build the bundle
# ============================================================
@@ -346,6 +273,7 @@ jobs:
EXO_BUILD_COMMIT="$GITHUB_SHA" \
SPARKLE_FEED_URL="$SPARKLE_FEED_URL" \
SPARKLE_ED25519_PUBLIC="$SPARKLE_ED25519_PUBLIC" \
EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT="$EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT" \
CODE_SIGNING_IDENTITY="$SIGNING_IDENTITY" \
CODE_SIGN_INJECT_BASE_ENTITLEMENTS=YES
mkdir -p ../../output
@@ -378,41 +306,11 @@ jobs:
APPLE_NOTARIZATION_PASSWORD: ${{ secrets.APPLE_NOTARIZATION_PASSWORD }}
APPLE_NOTARIZATION_TEAM: ${{ secrets.APPLE_NOTARIZATION_TEAM }}
run: |
set -o pipefail
cd output
security unlock-keychain -p "$MACOS_CERTIFICATE_PASSWORD" "$BUILD_KEYCHAIN_PATH"
SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}')
# Fail fast if notarization creds are partial. All-or-nothing.
CRED_COUNT=0
for v in "$APPLE_NOTARIZATION_USERNAME" "$APPLE_NOTARIZATION_PASSWORD" "$APPLE_NOTARIZATION_TEAM"; do
[[ -n "$v" ]] && CRED_COUNT=$((CRED_COUNT + 1))
done
if [[ "$CRED_COUNT" -ne 0 && "$CRED_COUNT" -ne 3 ]]; then
echo "ERROR: partial Apple notarization credentials set ($CRED_COUNT/3). Aborting."
exit 1
fi
/usr/bin/codesign --deep --force --timestamp --options runtime \
--sign "$SIGNING_IDENTITY" EXO.app
# Pre-flight: verify the signed app BEFORE building DMG and submitting to Apple.
# If this fails, notarization will fail too — cheap way to fail in seconds, not 15 minutes.
echo "===== codesign --verify EXO.app ====="
if ! /usr/bin/codesign --verify --deep --strict --verbose=2 EXO.app; then
echo "ERROR: EXO.app failed codesign verification. Dumping signing status of every executable:"
find EXO.app -type f \( -perm -111 -o -name "*.dylib" -o -name "*.so" -o -name "*.framework" \) -print0 |
while IFS= read -r -d '' f; do
printf -- '--- %s\n' "$f"
/usr/bin/codesign -dv --verbose=2 "$f" 2>&1 | sed 's/^/ /' || true
done
exit 1
fi
# Gatekeeper assessment. A failure here strongly predicts notarization rejection.
echo "===== spctl assessment (predicts notarization outcome) ====="
/usr/bin/spctl -a -vvv -t install EXO.app || echo "WARNING: spctl assessment failed — notarization is likely to fail too."
mkdir -p dmg-root
cp -R EXO.app dmg-root/
ln -s /Applications dmg-root/Applications
@@ -420,22 +318,12 @@ jobs:
hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME"
/usr/bin/codesign --force --timestamp --options runtime \
--sign "$SIGNING_IDENTITY" "$DMG_NAME"
echo "===== codesign --verify DMG ====="
if ! /usr/bin/codesign --verify --verbose=2 "$DMG_NAME"; then
echo "ERROR: DMG failed codesign verification."
exit 1
fi
if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then
echo "===== notarytool submit ====="
# `|| true` so set -e doesn't abort before we can echo output / fetch the log.
# We rely on the parsed STATUS below to decide pass/fail.
SUBMISSION_OUTPUT=$(xcrun notarytool submit "$DMG_NAME" \
--apple-id "$APPLE_NOTARIZATION_USERNAME" \
--password "$APPLE_NOTARIZATION_PASSWORD" \
--team-id "$APPLE_NOTARIZATION_TEAM" \
--wait --timeout 15m 2>&1) || true
--wait --timeout 15m 2>&1)
echo "$SUBMISSION_OUTPUT"
SUBMISSION_ID=$(echo "$SUBMISSION_OUTPUT" | awk 'tolower($1)=="id:" && $2 ~ /^[0-9a-fA-F-]+$/ {print $2; exit}')
+3
View File
@@ -91,6 +91,9 @@ jobs:
nix build .#metal-toolchain
fi
# Build mlx (depends on metal-toolchain)
nix build .#mlx
- name: Build all Nix outputs
run: |
nix flake show --json | jq -r '
+5 -6
View File
@@ -1767,12 +1767,12 @@ def clip(
array: The clipped array.
"""
def compile[F: Callable[..., object]](
fun: F,
def compile(
fun: Callable,
inputs: object | None = ...,
outputs: object | None = ...,
shapeless: bool = ...,
) -> F:
) -> Callable:
"""
Returns a compiled function which produces the same output as ``fun``.
@@ -2915,8 +2915,8 @@ def gather_mm(
a: array,
b: array,
/,
lhs_indices: array | None = ...,
rhs_indices: array | None = ...,
lhs_indices: array,
rhs_indices: array,
*,
sorted_indices: bool = ...,
stream: Stream | Device | None = ...,
@@ -4707,7 +4707,6 @@ def softmax(
/,
axis: int | Sequence[int] | None = ...,
*,
precise: bool = ...,
stream: Stream | Device | None = ...,
) -> array:
"""
-4
View File
@@ -57,10 +57,6 @@ class Module(dict):
def __init__(self) -> None:
"""Should be called by the subclasses of ``Module``."""
def __getitem__(self, key: str) -> mx.array | Module: ...
def get(
self, key: str, default: mx.array | Module | None = ...
) -> mx.array | Module | None: ...
@property
def training(self): # -> bool:
"""Boolean indicating if the model is in training mode."""
+4 -5
View File
@@ -383,12 +383,11 @@ class GenerationBatch:
state_machines: List[SequenceStateMachine]
max_tokens: List[int]
_current_tokens: Optional[mx.array]
_current_logprobs: mx.array | List[mx.array]
_next_tokens: Optional[mx.array]
_next_logprobs: mx.array | List[mx.array]
_token_context: List[Any]
_current_logprobs: List[mx.array]
_next_tokens: mx.array
_next_logprobs: List[mx.array]
_token_context: List[mx.array]
_num_tokens: List[int]
_matcher_states: List[Any]
def __init__(
self,
+5 -5
View File
@@ -3,7 +3,7 @@ This type stub file was generated by pyright.
"""
from dataclasses import dataclass
from typing import Any, Optional
from typing import Optional
import mlx.core as mx
@@ -37,10 +37,10 @@ def quantized_scaled_dot_product_attention(
bits: int = ...,
) -> mx.array: ...
def scaled_dot_product_attention(
queries: mx.array,
keys: mx.array,
values: mx.array,
cache: Optional[Any],
queries,
keys,
values,
cache,
scale: float,
mask: Optional[mx.array],
sinks: Optional[mx.array] = ...,
+6 -3
View File
@@ -191,10 +191,13 @@ class RotatingKVCache(_BaseCache):
def state(self, v): # -> None:
...
@property
def meta_state(self) -> tuple[str, ...]: ...
def meta_state(self): # -> tuple[str, ...]:
...
@meta_state.setter
def meta_state(self, v: tuple[str, ...]) -> None: ...
def is_trimmable(self) -> bool: ...
def meta_state(self, v): # -> None:
...
def is_trimmable(self): # -> bool:
...
def trim(self, n: int) -> int: ...
def to_quantized(
self, group_size: int = ..., bits: int = ...
-280
View File
@@ -1,280 +0,0 @@
"""Type stubs for mlx_lm.models.deepseek_v4"""
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs
from .cache import ArraysCache, RotatingKVCache
from .switch_layers import SwitchGLU
@dataclass
class ModelArgs(BaseModelArgs):
model_type: str
vocab_size: int
hidden_size: 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: int
num_experts_per_tok: int
head_dim: int
qk_rope_head_dim: int
q_lora_rank: int
o_lora_rank: int
o_groups: int
sliding_window: int
hc_mult: int
hc_sinkhorn_iters: int
hc_eps: float
compress_ratios: Optional[List[int]]
compress_rope_theta: float
rope_theta: float
rope_scaling: Optional[Dict[str, Any]]
rms_norm_eps: float
swiglu_limit: float
attention_bias: bool
max_position_embeddings: int
class DeepseekV4RoPE(nn.Module):
dims: int
freqs: mx.array
def __init__(
self,
dims: int,
base: float,
scaling_config: Optional[Dict[str, Any]] = None,
) -> None: ...
def __call__(
self,
x: mx.array,
offset: int = 0,
inverse: bool = False,
) -> mx.array: ...
class HyperConnection(nn.Module):
dim: int
hc_mult: int
norm_eps: float
def __init__(
self,
dim: int,
hc_mult: int,
norm_eps: float,
sinkhorn_iters: int,
hc_eps: float,
) -> None: ...
class HyperHead(nn.Module):
dim: int
hc_mult: int
def __init__(
self,
dim: int,
hc_mult: int,
norm_eps: float,
hc_eps: float,
) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class Compressor(nn.Module):
dim: int
head_dim: int
rope_head_dim: int
compress_ratio: int
overlap: bool
wkv_gate: nn.Linear
ape: mx.array
norm: nn.RMSNorm
rope: DeepseekV4RoPE
def __init__(
self,
dim: int,
compress_ratio: int,
head_dim: int,
rope_head_dim: int,
rms_norm_eps: float,
rope: DeepseekV4RoPE,
) -> None: ...
def __call__(
self,
x: mx.array,
cache: "DeepseekV4Cache",
offset: Any,
key: str = ...,
) -> mx.array: ...
class Indexer(nn.Module):
def __init__(
self,
args: ModelArgs,
compress_ratio: int,
rope: DeepseekV4RoPE,
) -> None: ...
class _CompressorBranch:
buffer_kv: Optional[mx.array]
buffer_gate: Optional[mx.array]
prev_kv: Optional[mx.array]
prev_gate: Optional[mx.array]
pool: Optional[mx.array]
buffer_lengths: Optional[List[int]]
pool_lengths: Optional[List[int]]
buffer_count: int
_new_pool_lengths: Optional[List[int]]
def __init__(self) -> None: ...
class DeepseekV4Cache:
local: RotatingKVCache
offset: int
keys: Optional[mx.array]
values: Optional[mx.array]
state: Any
meta_state: Any
nbytes: int
_branches: Dict[str, _CompressorBranch]
_pending_lengths: Optional[List[int]]
def __init__(self, sliding_window: int) -> None: ...
def update_and_fetch(
self, keys: mx.array, values: mx.array
) -> tuple[mx.array, mx.array]: ...
def is_trimmable(self) -> bool: ...
def trim(self, n: int) -> int: ...
def empty(self) -> bool: ...
def size(self) -> int: ...
def prepare(
self,
*,
left_padding: Optional[List[int]] = None,
lengths: Optional[List[int]] = None,
right_padding: Optional[List[int]] = None,
) -> None: ...
def finalize(self) -> None: ...
def filter(self, batch_indices: mx.array) -> None: ...
def extend(self, other: "DeepseekV4Cache") -> None: ...
def extract(self, idx: int) -> "DeepseekV4Cache": ...
@classmethod
def merge(cls, caches: List["DeepseekV4Cache"]) -> "DeepseekV4Cache": ...
class V4Attention(nn.Module):
args: ModelArgs
layer_id: int
dim: int
n_heads: int
head_dim: int
rope_head_dim: int
nope_head_dim: int
n_groups: int
q_lora_rank: int
o_lora_rank: int
window: int
eps: float
scale: float
compress_ratio: int
wqkv_a: nn.Linear
q_norm: nn.RMSNorm
wq_b: nn.Linear
kv_norm: nn.RMSNorm
attn_sink: mx.array
wo_a: nn.Linear
wo_b: nn.Linear
rope: DeepseekV4RoPE
compressor: Compressor
indexer: Indexer
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class DeepseekV4MLP(nn.Module):
gate_proj: nn.Linear
up_proj: nn.Linear
down_proj: nn.Linear
def __init__(
self,
hidden_size: int,
intermediate_size: int,
swiglu_limit: float = 0.0,
) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class MoEGate(nn.Module):
weight: mx.array
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
def __call__(
self, x: mx.array, input_ids: mx.array
) -> tuple[mx.array, mx.array]: ...
class DeepseekV4MoE(nn.Module):
num_experts_per_tok: int
switch_mlp: SwitchGLU
gate: MoEGate
shared_experts: DeepseekV4MLP
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
def __call__(self, x: mx.array, input_ids: mx.array) -> mx.array: ...
class DeepseekV4Block(nn.Module):
attn_norm: nn.RMSNorm
attn: V4Attention
hc_attn: HyperConnection
ffn_norm: nn.RMSNorm
ffn: DeepseekV4MoE
hc_ffn: HyperConnection
def __init__(self, args: ModelArgs, layer_id: int) -> None: ...
def __call__(
self,
h: mx.array,
cache: Optional[Any],
input_ids: mx.array,
) -> mx.array: ...
class DeepseekV4Model(nn.Module):
args: ModelArgs
vocab_size: int
embed_tokens: nn.Embedding
layers: list[DeepseekV4Block]
norm: nn.RMSNorm
hc_head: HyperHead
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[List[Any]] = None,
) -> mx.array: ...
class Model(nn.Module):
args: ModelArgs
model_type: str
model: DeepseekV4Model
lm_head: nn.Linear
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[List[Any]] = None,
) -> mx.array: ...
def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
def make_cache(self) -> list[RotatingKVCache | DeepseekV4Cache]: ...
@property
def layers(self) -> list[DeepseekV4Block]: ...
-103
View File
@@ -1,103 +0,0 @@
"""Type stubs for mlx_lm.models.gpt_oss"""
from dataclasses import dataclass
from typing import Any, List, Optional
import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs
from .cache import KVCache
from .switch_layers import SwitchGLU
@dataclass
class ModelArgs(BaseModelArgs):
model_type: str
hidden_size: int
intermediate_size: int
num_hidden_layers: int
num_attention_heads: int
num_key_value_heads: int
num_local_experts: int
num_experts_per_tok: int
vocab_size: int
rms_norm_eps: float
sliding_window: int
layer_types: Optional[List[str]]
def mlx_topk(a: mx.array, k: int, axis: int = -1) -> tuple[mx.array, mx.array]: ...
class AttentionBlock(nn.Module):
head_dim: int
num_attention_heads: int
num_key_value_heads: int
num_key_value_groups: int
sinks: mx.array
q_proj: nn.Linear
k_proj: nn.Linear
v_proj: nn.Linear
o_proj: nn.Linear
sm_scale: float
rope: nn.Module
def __init__(self, config: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class TransformerBlock(nn.Module):
self_attn: AttentionBlock
mlp: MLPBlock
def __init__(self, config: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class MLPBlock(nn.Module):
hidden_size: int
num_local_experts: int
num_experts_per_tok: int
experts: SwitchGLU
router: nn.Linear
sharding_group: Optional[mx.distributed.Group]
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class GptOssMoeModel(nn.Module):
embed_tokens: nn.Embedding
norm: nn.RMSNorm
layer_types: List[str]
layers: list[TransformerBlock]
window_size: int
swa_idx: int
ga_idx: int
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
class Model(nn.Module):
model_type: str
model: GptOssMoeModel
lm_head: nn.Linear
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
@property
def layers(self) -> list[nn.Module]: ...
def make_cache(self) -> list[KVCache]: ...
-94
View File
@@ -1,94 +0,0 @@
"""Type stubs for mlx_lm.models.minimax"""
from dataclasses import dataclass
from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .base import BaseModelArgs
from .switch_layers import SwitchGLU
@dataclass
class ModelArgs(BaseModelArgs):
model_type: str
hidden_size: int
intermediate_size: int
num_hidden_layers: int
num_attention_heads: int
num_key_value_heads: int
num_local_experts: int
num_experts_per_tok: int
max_position_embeddings: int
class MiniMaxAttention(nn.Module):
num_heads: int
num_attention_heads: int
num_key_value_heads: int
head_dim: int
scale: float
q_proj: nn.Linear
k_proj: nn.Linear
v_proj: nn.Linear
o_proj: nn.Linear
q_norm: nn.Module
k_norm: nn.Module
rope: nn.Module
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class MiniMaxSparseMoeBlock(nn.Module):
num_experts_per_tok: int
gate: nn.Linear
switch_mlp: SwitchGLU
e_score_correction_bias: mx.array
sharding_group: Optional[mx.distributed.Group]
def __init__(self, args: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class MiniMaxDecoderLayer(nn.Module):
self_attn: MiniMaxAttention
block_sparse_moe: MiniMaxSparseMoeBlock
input_layernorm: nn.RMSNorm
post_attention_layernorm: nn.RMSNorm
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array: ...
class MiniMaxModel(nn.Module):
embed_tokens: nn.Embedding
layers: list[MiniMaxDecoderLayer]
norm: nn.RMSNorm
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
class Model(nn.Module):
model_type: str
model: MiniMaxModel
lm_head: nn.Linear
def __init__(self, args: ModelArgs) -> None: ...
def __call__(
self,
inputs: mx.array,
cache: Optional[Any] = None,
) -> mx.array: ...
@property
def layers(self) -> list[MiniMaxDecoderLayer]: ...
-14
View File
@@ -92,15 +92,6 @@ class NemotronHAttention(nn.Module):
cache: Optional[KVCache] = None,
) -> mx.array: ...
class MoEGate(nn.Module):
config: ModelArgs
top_k: int
norm_topk_prob: bool
weight: mx.array
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
class NemotronHMLP(nn.Module):
up_proj: nn.Linear
down_proj: nn.Linear
@@ -111,14 +102,9 @@ class NemotronHMLP(nn.Module):
def __call__(self, x: mx.array) -> mx.array: ...
class NemotronHMoE(nn.Module):
config: ModelArgs
num_experts_per_tok: int
moe_latent_size: Optional[int]
switch_mlp: SwitchMLP
gate: MoEGate
shared_experts: NemotronHMLP
fc1_latent_proj: nn.Linear
fc2_latent_proj: nn.Linear
def __init__(self, config: ModelArgs) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
@@ -71,7 +71,6 @@ class Qwen3NextAttention(nn.Module):
class Qwen3NextSparseMoeBlock(nn.Module):
norm_topk_prob: bool
num_experts: int
num_experts_per_tok: int
top_k: int
gate: nn.Linear
switch_mlp: SwitchGLU
-4
View File
@@ -48,10 +48,6 @@ def make_logits_processors(
logit_bias: Optional[Dict[int, float]] = ...,
repetition_penalty: Optional[float] = ...,
repetition_context_size: Optional[int] = ...,
presence_penalty: Optional[float] = ...,
presence_context_size: Optional[int] = ...,
frequency_penalty: Optional[float] = ...,
frequency_context_size: Optional[int] = ...,
) -> list[Callable[[mx.array, mx.array], mx.array]]:
"""
Make logits processors for use with ``generate_step``.
File diff suppressed because it is too large. Load diff
+1 -10
View File
@@ -584,18 +584,9 @@ struct ContentView: View {
case .prompting:
VStack(alignment: .leading, spacing: 6) {
VStack(alignment: .leading, spacing: 2) {
Text("Tell us what went wrong (optional)")
.font(.caption2)
.foregroundColor(.secondary)
Text(
"A quick description of what you were doing and what happened helps us track down the bug for you."
)
Text("What's the issue? (optional)")
.font(.caption2)
.foregroundColor(.secondary)
.opacity(0.8)
.fixedSize(horizontal: false, vertical: true)
}
TextEditor(text: $bugReportUserDescription)
.font(.caption2)
.frame(height: 60)
+3 -43
View File
@@ -9,14 +9,11 @@ private let enableImageModelsKey = "EXOEnableImageModels"
private let offlineModeKey = "EXOOfflineMode"
private let fastSynchEnabledKey = "EXOFastSynchEnabled"
private let onboardingCompletedKey = "EXOOnboardingCompleted"
private let defaultModelsDirKey = "EXODefaultModelsDir"
private let additionalModelsDirsKey = "EXOAdditionalModelsDirs"
private let readOnlyModelsDirsKey = "EXOReadOnlyModelsDirs"
private let customEnvironmentVariablesKey = "EXOCustomEnvironmentVariables"
/// A user-defined environment variable that is injected into the exo child
/// process at launch. Used as an escape hatch for env vars that don't have
/// first-class typed UI in Settings.
/// process at launch. Used to pass arbitrary key/value settings to exo
/// without having to add first-class UI for each one.
struct CustomEnvironmentVariable: Codable, Identifiable, Equatable {
var id: UUID
var key: String
@@ -109,30 +106,6 @@ final class ExoProcessController: ObservableObject {
UserDefaults.standard.set(fastSynchEnabled, forKey: fastSynchEnabledKey)
}
}
@Published var defaultModelsDir: String = {
return UserDefaults.standard.string(forKey: defaultModelsDirKey) ?? ""
}()
{
didSet {
UserDefaults.standard.set(defaultModelsDir, forKey: defaultModelsDirKey)
}
}
@Published var additionalModelsDirs: String = {
return UserDefaults.standard.string(forKey: additionalModelsDirsKey) ?? ""
}()
{
didSet {
UserDefaults.standard.set(additionalModelsDirs, forKey: additionalModelsDirsKey)
}
}
@Published var readOnlyModelsDirs: String = {
return UserDefaults.standard.string(forKey: readOnlyModelsDirsKey) ?? ""
}()
{
didSet {
UserDefaults.standard.set(readOnlyModelsDirs, forKey: readOnlyModelsDirsKey)
}
}
@Published var customEnvironmentVariables: [CustomEnvironmentVariable] = {
guard
let data = UserDefaults.standard.data(forKey: customEnvironmentVariablesKey),
@@ -391,21 +364,8 @@ final class ExoProcessController: ObservableObject {
environment["PATH"] = paths.joined(separator: ":")
let trimmedDefaultModelsDir = defaultModelsDir.trimmingCharacters(in: .whitespaces)
if !trimmedDefaultModelsDir.isEmpty {
environment["EXO_DEFAULT_MODELS_DIR"] = trimmedDefaultModelsDir
}
let trimmedAdditionalModelsDirs = additionalModelsDirs.trimmingCharacters(in: .whitespaces)
if !trimmedAdditionalModelsDirs.isEmpty {
environment["EXO_MODELS_DIRS"] = trimmedAdditionalModelsDirs
}
let trimmedReadOnlyModelsDirs = readOnlyModelsDirs.trimmingCharacters(in: .whitespaces)
if !trimmedReadOnlyModelsDirs.isEmpty {
environment["EXO_MODELS_READ_ONLY_DIRS"] = trimmedReadOnlyModelsDirs
}
// Apply user-defined arbitrary environment variables last so that
// power users can override any of the typed fields above when
// power users can override any of the built-in keys above when
// necessary. Empty keys are ignored.
for variable in customEnvironmentVariables {
let trimmedKey = variable.key.trimmingCharacters(in: .whitespaces)
+1 -1
View File
@@ -9,7 +9,7 @@
<key>EXOBuildCommit</key>
<string>$(EXO_BUILD_COMMIT)</string>
<key>EXOBugReportPresignedUrlEndpoint</key>
<string>https://reports.exolabs.net/presigned-urls</string>
<string>$(EXO_BUG_REPORT_PRESIGNED_URL_ENDPOINT)</string>
<key>NSLocalNetworkUsageDescription</key>
<string>EXO needs local network access to discover and connect to other devices in your cluster for distributed AI inference.</string>
<key>NSBonjourServices</key>
+14 -109
View File
@@ -16,9 +16,6 @@ struct SettingsView: View {
@State private var pendingEnableImageModels = false
@State private var pendingOfflineMode = false
@State private var pendingFastSynchEnabled = false
@State private var pendingDefaultModelsDir: String = ""
@State private var pendingAdditionalModelsDirs: String = ""
@State private var pendingReadOnlyModelsDirs: String = ""
@State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
@State private var needsRestart = false
@State private var bugReportInFlight = false
@@ -48,7 +45,7 @@ struct SettingsView: View {
Label("About", systemImage: "info.circle")
}
}
.frame(width: 640, height: 560)
.frame(width: 450, height: 400)
.onAppear {
pendingNamespace = controller.customNamespace
pendingHFToken = controller.hfToken
@@ -56,9 +53,6 @@ struct SettingsView: View {
pendingEnableImageModels = controller.enableImageModels
pendingOfflineMode = controller.offlineMode
pendingFastSynchEnabled = controller.fastSynchEnabled
pendingDefaultModelsDir = controller.defaultModelsDir
pendingAdditionalModelsDirs = controller.additionalModelsDirs
pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
pendingCustomEnvironmentVariables = controller.customEnvironmentVariables
needsRestart = false
}
@@ -70,9 +64,9 @@ struct SettingsView: View {
Form {
Section {
LabeledContent("Cluster Namespace") {
TextField("", text: $pendingNamespace, prompt: Text("default"))
TextField("default", text: $pendingNamespace)
.textFieldStyle(.roundedBorder)
.frame(width: 260)
.frame(width: 200)
}
Text("Nodes with the same namespace form a cluster. Leave empty for default.")
.font(.caption)
@@ -81,9 +75,9 @@ struct SettingsView: View {
Section {
LabeledContent("HuggingFace Token") {
SecureField("", text: $pendingHFToken, prompt: Text("optional"))
SecureField("optional", text: $pendingHFToken)
.textFieldStyle(.roundedBorder)
.frame(width: 260)
.frame(width: 200)
}
Text("Required for gated models. Get yours at huggingface.co/settings/tokens")
.font(.caption)
@@ -92,9 +86,9 @@ struct SettingsView: View {
Section {
LabeledContent("HuggingFace Endpoint") {
TextField("", text: $pendingHFEndpoint, prompt: Text("default"))
TextField("default", text: $pendingHFEndpoint)
.textFieldStyle(.roundedBorder)
.frame(width: 260)
.frame(width: 200)
}
Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
.font(.caption)
@@ -228,57 +222,10 @@ struct SettingsView: View {
private var environmentTab: some View {
Form {
Section("Models Directories") {
LabeledContent("Default Models Directory") {
TextField(
"",
text: $pendingDefaultModelsDir,
prompt: Text("~/.exo/models")
)
.textFieldStyle(.roundedBorder)
.font(.system(.body, design: .monospaced))
.frame(width: 260)
}
Text("Sets EXO_DEFAULT_MODELS_DIR. Where models are downloaded.")
.font(.caption)
.foregroundColor(.secondary)
LabeledContent("Additional Directories") {
TextField(
"",
text: $pendingAdditionalModelsDirs,
prompt: Text("optional, colon-separated")
)
.textFieldStyle(.roundedBorder)
.font(.system(.body, design: .monospaced))
.frame(width: 260)
}
Text("Sets EXO_MODELS_DIRS. Extra writable model directories.")
.font(.caption)
.foregroundColor(.secondary)
LabeledContent("Read-Only Directories") {
TextField(
"",
text: $pendingReadOnlyModelsDirs,
prompt: Text("optional, colon-separated")
)
.textFieldStyle(.roundedBorder)
.font(.system(.body, design: .monospaced))
.frame(width: 260)
}
Text("Sets EXO_MODELS_READ_ONLY_DIRS. Never written to.")
.font(.caption)
.foregroundColor(.secondary)
}
Section("Custom Environment Variables") {
Text(
"Escape hatch for env vars that don't have typed fields above. "
+ "Values here override the typed fields on conflict."
)
.font(.caption)
.foregroundColor(.secondary)
Text("Passed to the exo process at launch. Override built-in defaults here.")
.font(.caption)
.foregroundColor(.secondary)
if pendingCustomEnvironmentVariables.isEmpty {
Text("No custom variables.")
@@ -552,24 +499,15 @@ struct SettingsView: View {
let alert = NSAlert()
alert.messageText = "Uninstall EXO"
alert.informativeText = """
This will remove EXO and all its components:
This will remove EXO and all its system components:
• Network configuration daemon
• Launch at login registration
• EXO network location
• EXO data directory (~/.exo)
The app will be moved to Trash.
"""
alert.alertStyle = .warning
let checkbox = NSButton(
checkboxWithTitle: "Keep downloaded models (~/.exo/models)",
target: nil, action: nil)
checkbox.state = .off
checkbox.sizeToFit()
alert.accessoryView = checkbox
alert.addButton(withTitle: "Uninstall")
alert.addButton(withTitle: "Cancel")
@@ -579,11 +517,11 @@ struct SettingsView: View {
let response = alert.runModal()
if response == .alertFirstButtonReturn {
performUninstall(keepModels: checkbox.state == .on)
performUninstall()
}
}
private func performUninstall(keepModels: Bool) {
private func performUninstall() {
uninstallInProgress = true
controller.cancelPendingLaunch()
@@ -593,7 +531,6 @@ struct SettingsView: View {
DispatchQueue.global(qos: .utility).async {
do {
try NetworkSetupHelper.uninstall()
try Self.removeExoDirectory(keepModels: keepModels)
DispatchQueue.main.async {
LaunchAtLoginHelper.disable()
@@ -617,23 +554,6 @@ struct SettingsView: View {
}
}
private static func removeExoDirectory(keepModels: Bool) throws {
let fm = FileManager.default
let exoDir = ExoProcessController.exoDirectoryURL
guard fm.fileExists(atPath: exoDir.path) else { return }
if !keepModels {
try fm.removeItem(at: exoDir)
return
}
let contents = try fm.contentsOfDirectory(
at: exoDir, includingPropertiesForKeys: nil, options: [])
for entry in contents where entry.lastPathComponent != "models" {
try? fm.removeItem(at: entry)
}
}
private func moveAppToTrash() {
guard let appURL = Bundle.main.bundleURL as URL? else { return }
do {
@@ -660,10 +580,7 @@ struct SettingsView: View {
}
private var hasEnvironmentChanges: Bool {
pendingDefaultModelsDir != controller.defaultModelsDir
|| pendingAdditionalModelsDirs != controller.additionalModelsDirs
|| pendingReadOnlyModelsDirs != controller.readOnlyModelsDirs
|| pendingCustomEnvironmentVariables != controller.customEnvironmentVariables
pendingCustomEnvironmentVariables != controller.customEnvironmentVariables
}
private func applyGeneralSettings() {
@@ -685,17 +602,6 @@ struct SettingsView: View {
}
private func applyEnvironmentSettings() {
controller.defaultModelsDir = pendingDefaultModelsDir.trimmingCharacters(
in: .whitespaces)
controller.additionalModelsDirs = pendingAdditionalModelsDirs.trimmingCharacters(
in: .whitespaces)
controller.readOnlyModelsDirs = pendingReadOnlyModelsDirs.trimmingCharacters(
in: .whitespaces)
pendingDefaultModelsDir = controller.defaultModelsDir
pendingAdditionalModelsDirs = controller.additionalModelsDirs
pendingReadOnlyModelsDirs = controller.readOnlyModelsDirs
// Trim whitespace from keys and drop empty ones so that the stored
// form matches what is actually injected into the child process and
// hasEnvironmentChanges doesn't show a stale diff after save.
@@ -723,7 +629,6 @@ struct SettingsView: View {
pendingCustomEnvironmentVariables = sanitized
controller.customEnvironmentVariables = sanitized
restartIfRunning()
}
@@ -30,7 +30,7 @@ final class SettingsWindowController: ObservableObject {
let hostingView = NSHostingView(rootView: settingsView)
let newWindow = NSWindow(
contentRect: NSRect(x: 0, y: 0, width: 640, height: 560),
contentRect: NSRect(x: 0, y: 0, width: 450, height: 400),
styleMask: [.titled, .closable],
backing: .buffered,
defer: false
+7 -63
View File
@@ -3,55 +3,25 @@
# EXO Uninstaller Script
#
# This script removes all EXO system components that persist after deleting the app.
# Run with: sudo ./uninstall-exo.sh [--keep-models]
#
# Options:
# --keep-models Preserve ~/.exo/models when removing the EXO data directory.
# Run with: sudo ./uninstall-exo.sh
#
# Components removed:
# - LaunchDaemon: /Library/LaunchDaemons/io.exo.networksetup.plist
# - Network script: /Library/Application Support/EXO/
# - Log files: /var/log/io.exo.networksetup.*
# - Network location: "exo"
# - EXO data directory: ~/.exo (or all of ~/.exo except models/ when --keep-models is set)
# - Launch at login registration
#
set -euo pipefail
KEEP_MODELS=0
for arg in "$@"; do
case "$arg" in
--keep-models)
KEEP_MODELS=1
;;
-h | --help)
echo "Usage: sudo ./uninstall-exo.sh [--keep-models]"
echo " --keep-models Preserve ~/.exo/models when removing the EXO data directory."
exit 0
;;
*)
echo "Unknown argument: $arg" >&2
echo "Usage: sudo ./uninstall-exo.sh [--keep-models]" >&2
exit 2
;;
esac
done
LABEL="io.exo.networksetup"
# Current script path. Older installs used a different filename; keep the
# legacy path here so a fresh uninstall still cleans up upgraded machines.
CURRENT_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge.sh"
LEGACY_SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
SCRIPT_DEST="/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
PLIST_DEST="/Library/LaunchDaemons/io.exo.networksetup.plist"
LOG_OUT="/var/log/${LABEL}.log"
LOG_ERR="/var/log/${LABEL}.err.log"
APP_BUNDLE_ID="io.exo.EXO"
# Resolve the invoking user's home, even when run via sudo.
USER_HOME="$(eval echo "~${SUDO_USER:-$USER}")"
EXO_DIR="$USER_HOME/.exo"
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
@@ -99,17 +69,11 @@ else
echo_warn "LaunchDaemon plist not found (already removed?)"
fi
# Remove the script (current and legacy filenames) — backwards-compatible:
# tolerate either, both, or neither being present.
removed_any_script=0
for script in "$CURRENT_SCRIPT_DEST" "$LEGACY_SCRIPT_DEST"; do
if [[ -f $script ]]; then
rm -f "$script"
echo_info "Removed network setup script: $script"
removed_any_script=1
fi
done
if [[ $removed_any_script -eq 0 ]]; then
# Remove the script and parent directory
if [[ -f $SCRIPT_DEST ]]; then
rm -f "$SCRIPT_DEST"
echo_info "Removed network setup script"
else
echo_warn "Network setup script not found (already removed?)"
fi
@@ -151,22 +115,6 @@ if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge";
echo_info "Re-enabled Thunderbolt Bridge"
fi
# Remove EXO data directory (~/.exo)
EXO_DIR_REMOVED=""
if [[ -d $EXO_DIR ]]; then
if [[ $KEEP_MODELS == "1" && -d "$EXO_DIR/models" ]]; then
find "$EXO_DIR" -mindepth 1 -maxdepth 1 ! -name models -exec rm -rf {} +
EXO_DIR_REMOVED="kept_models"
echo_info "Removed ~/.exo (preserved models/)"
else
rm -rf "$EXO_DIR"
EXO_DIR_REMOVED="full"
echo_info "Removed ~/.exo"
fi
else
echo_warn "~/.exo not found (already removed?)"
fi
# Note about launch at login registration
# SMAppService-based login items cannot be removed from a shell script.
# They can only be unregistered from within the app itself or manually via System Settings.
@@ -196,10 +144,6 @@ echo " • Network setup LaunchDaemon"
echo " • Network configuration script"
echo " • Log files"
echo " • 'exo' network location"
case "$EXO_DIR_REMOVED" in
full) echo " • EXO data directory (~/.exo)" ;;
kept_models) echo " • EXO data directory (~/.exo, models preserved)" ;;
esac
echo ""
echo "Your network has been restored to use the 'Automatic' location."
echo "Thunderbolt Bridge has been re-enabled (if present)."
+1 -17
View File
@@ -28,7 +28,7 @@ Chat template formatting means that it may be impossible to attain very small pp
When a request reaches the server via the `/bench/chat/completions` endpoint, three things change compared to a normal chat completion:
- **KV prefix cache is disabled by default**. Every request starts from a cold cache, ensuring prefill timing is not affected by prior requests. See [Prefix Cache Mode](#prefix-cache-mode) for the `--use-prefix-cache` option.
- **KV prefix cache is disabled**. Every request starts from a cold cache, ensuring prefill timing is not affected by prior requests.
- **EOS tokens are banned**. A logits processor suppresses all end-of-sequence tokens, forcing the model to generate exactly `max_tokens` tokens. This guarantees consistent generation length for fair TPS comparison — the model cannot short-circuit a run by stopping early.
- **No model output parsing**. The bench collection path concatenates raw token text without any model-specific post-processing (thinking tag extraction, structured output handling, etc.). This is to avoid model outputs such as tool parsing or any structural mistakes from breaking the benchmark - we are testing for speed; see Exo-Eval for performance metrics.
@@ -94,22 +94,6 @@ agg_gen_tps = per_req_tps * concurrency
---
## Prefix Cache Mode
When `--use-prefix-cache` is passed, the KV prefix cache remains active during benchmarking. This speeds up repeated runs by skipping redundant prefill work, which is useful when prompt processing is not the focus of the benchmark (e.g. when measuring generation throughput or power consumption across many configurations).
Each response includes a `prefix_cache_hit` field (`"none"`, `"partial"`, or `"exact"`):
- **none**: Cold prefill — no cached KV state was available. The reported `prompt_tps` is the real prefill throughput.
- **partial**: A prefix of the prompt was found in cache. Only the remaining tokens were prefilled. The reported `prompt_tps` reflects the real throughput on the uncached portion. This occurs when multiple ascending `--pp` values share a common prefix (e.g. `--pp 1000,5000` — the 5000-token prompt reuses the 1000-token cache entry and prefills the remaining 4000 tokens).
- **exact**: The entire prompt was found in cache (e.g. same `--pp` value on a `--repeat`). No prefill work was done. The reported `prompt_tps` is the TPS from when the cache entry was originally created, not a new measurement.
**Prompt TPS is approximate in this mode.** Exact-hit runs report the stored TPS from the original cold/partial prefill rather than a freshly measured value. For accurate cold prefill numbers, run without `--use-prefix-cache`.
Ascending `--pp` order (e.g. `--pp 1000,5000,10000`) gives the most useful data: each size gets a meaningful partial hit except the first which is cold. Descending order produces exact hits with approximate TPS from a longer prompt's original run.
---
## Warmup
Before measurement begins, `--warmup N` (default: 0) discarded requests are sent using the first pp/tg pair. Warmup results are not included in the output.
+22 -77
View File
@@ -7,7 +7,7 @@
# name, patterns, reasoning
#
# Optional per-model overrides (CLI flags take priority over these):
# temperature, top_p, max_tokens, reasoning_effort, enable_thinking
# temperature, top_p, max_tokens, reasoning_effort
#
# Fallback defaults (when no per-model config):
# reasoning: temperature=1.0, max_tokens=131072, reasoning_effort="high"
@@ -18,9 +18,10 @@
# ─── Qwen3.5 (Feb 2026) ─────────────────────────────────────────────
# Source: HuggingFace model cards (Qwen/Qwen3.5-*)
# Model card recommends: temp=0.6, top_p=0.95, top_k=20
# We omit top_k to match vllm eval (which doesn't set it).
# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin).
# 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"
@@ -28,8 +29,7 @@ patterns = ["Qwen3.5-2B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 9B"
@@ -37,8 +37,7 @@ patterns = ["Qwen3.5-9B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 27B"
@@ -46,17 +45,15 @@ patterns = ["Qwen3.5-27B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 35B A3B"
patterns = ["Qwen3.5-35B-A3B"]
reasoning = true
temperature = 0.6
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 122B A10B"
@@ -64,8 +61,7 @@ patterns = ["Qwen3.5-122B-A10B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
[[model]]
name = "Qwen3.5 397B A17B"
@@ -73,14 +69,12 @@ patterns = ["Qwen3.5-397B-A17B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 81920
# ─── Qwen3 (Apr 2025) ───────────────────────────────────────────────
# Source: HuggingFace model cards (Qwen/Qwen3-*)
# Model card recommends: temp=0.6, top_p=0.95, top_k=20
# We omit top_k to match vllm eval (which doesn't set it).
# Non-thinking: temp=0.7, top_p=0.8
# 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]]
@@ -89,7 +83,6 @@ patterns = ["Qwen3-0.6B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -98,7 +91,6 @@ patterns = ["Qwen3-30B-A3B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -107,7 +99,6 @@ patterns = ["Qwen3-235B-A22B"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -116,7 +107,6 @@ patterns = ["Qwen3-Next-80B-A3B-Thinking"]
reasoning = true
temperature = 0.6
top_p = 0.95
enable_thinking = true
max_tokens = 38912
[[model]]
@@ -139,9 +129,9 @@ max_tokens = 16384
name = "Qwen3 Coder Next"
patterns = ["Qwen3-Coder-Next"]
reasoning = false
temperature = 1.0
top_p = 0.95
max_tokens = 121072
temperature = 0.7
top_p = 0.8
max_tokens = 16384
# ─── GPT-OSS (OpenAI) ───────────────────────────────────────────────
# Source: OpenAI GitHub README + HuggingFace discussion #21
@@ -175,38 +165,10 @@ patterns = ["DeepSeek-V3.1"]
reasoning = true
temperature = 0.0
[[model]]
name = "DeepSeek V3.2"
patterns = ["DeepSeek-V3.2"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
# ─── NVIDIA Nemotron ───────────────────────────────────────────────────
# Source: HuggingFace model cards
# All variants: temp=1.0, top_p=0.95, enable_thinking=true
[[model]]
name = "Nemotron Cascade 2 30B A3B"
patterns = ["Nemotron-Cascade-2-30B-A3B"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
[[model]]
name = "Nemotron 3 Super 120B A12B"
patterns = ["Nemotron-3-Super-120B-A12B", "NVIDIA-Nemotron-3-Super-120B-A12B"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
# ─── GLM (ZhipuAI / THUDM) ──────────────────────────────────────────
# Source: HuggingFace model cards + generation_config.json + docs.z.ai
# GLM 4.5+: temp=1.0, top_p=0.95
# max_tokens=121072 to match vllm eval (131072 context - 10000 safety margin)
# Reasoning tasks: 131072 max_tokens; coding/SWE tasks: temp=0.7
[[model]]
name = "GLM-5"
@@ -214,8 +176,7 @@ patterns = ["GLM-5"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 131072
[[model]]
name = "GLM 4.5 Air"
@@ -230,8 +191,7 @@ patterns = ["GLM-4.7-"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 131072
# Note: matches both GLM-4.7 and GLM-4.7-Flash
# ─── Kimi (Moonshot AI) ─────────────────────────────────────────────
@@ -253,8 +213,7 @@ patterns = ["Kimi-K2.5"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
max_tokens = 131072
[[model]]
name = "Kimi K2 Instruct"
@@ -264,17 +223,7 @@ temperature = 0.6
# ─── MiniMax ─────────────────────────────────────────────────────────
# Source: HuggingFace model cards + generation_config.json
# All models: temp=1.0, top_p=0.95
# max_tokens=90000 to match vllm eval (100000 context - 10000 safety margin)
[[model]]
name = "MiniMax M2.7"
patterns = ["MiniMax-M2.7"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 90000
# All models: temp=1.0, top_p=0.95, top_k=40
[[model]]
name = "MiniMax M2.5"
@@ -282,8 +231,6 @@ patterns = ["MiniMax-M2.5"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 90000
[[model]]
name = "MiniMax M2.1"
@@ -304,8 +251,6 @@ patterns = ["Step-3.5-Flash"]
reasoning = true
temperature = 1.0
top_p = 0.95
enable_thinking = true
max_tokens = 121072
# ─── Llama (Meta) ───────────────────────────────────────────────────
# Source: generation_config.json + meta-llama/llama-models generation.py
+42 -82
View File
@@ -3,13 +3,11 @@ from __future__ import annotations
import argparse
import contextlib
import io
import json
import os
import sys
import time
import tomllib
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Literal
@@ -20,11 +18,13 @@ from harness import (
ExoHttpError,
add_common_instance_args,
capture_cluster_snapshot,
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,
)
@@ -211,7 +211,7 @@ def _openai_build_request(
"model": model,
"messages": messages,
"tools": tools,
"max_tokens": 4096,
"max_tokens": 16384,
"temperature": 0.0,
}
return "/v1/chat/completions", body
@@ -278,7 +278,7 @@ def _openai_build_followup(
"model": model,
"messages": followup_messages,
"tools": tools,
"max_tokens": 4096,
"max_tokens": 16384,
"temperature": 0.0,
}
return "/v1/chat/completions", body
@@ -381,7 +381,7 @@ def _claude_build_request(
"model": model,
"messages": claude_messages,
"tools": claude_tools,
"max_tokens": 4096,
"max_tokens": 16384,
"temperature": 0.0,
}
if system_content is not None:
@@ -491,7 +491,7 @@ def _claude_build_followup(
"model": model,
"messages": claude_messages,
"tools": claude_tools,
"max_tokens": 4096,
"max_tokens": 16384,
"temperature": 0.0,
}
if system_content is not None:
@@ -915,12 +915,6 @@ Examples:
default=1,
help="Repeat each scenario N times (default: 1)",
)
parser.add_argument(
"--concurrency",
type=int,
default=1,
help="Run up to N scenarios in parallel against the same instance (default: 1)",
)
parser.add_argument(
"--scenarios",
nargs="*",
@@ -942,13 +936,7 @@ Examples:
help="Write JSON results to stdout instead of file",
)
args = parser.parse_args()
if args.concurrency < 1:
print(
f"--concurrency must be >= 1 (got {args.concurrency})",
file=sys.stderr,
)
sys.exit(2)
validate_vllm_args(args)
all_scenarios = load_scenarios(SCENARIOS_PATH)
if args.scenarios:
@@ -968,6 +956,8 @@ Examples:
log = sys.stderr if args.stdout else sys.stdout
exo = ExoClient(args.host, args.port, timeout_s=args.timeout)
if args.ensure_cuda:
ensure_cuda_available(exo)
_short_id, full_model_id = resolve_model_short_id(exo, args.model)
selected = settle_and_fetch_placements(
@@ -1025,72 +1015,42 @@ Examples:
cluster_snapshot = capture_cluster_snapshot(exo)
all_results: list[ScenarioResult] = []
tasks: list[tuple[int, Scenario, ApiName]] = [
(run_idx, scenario, api_name)
for run_idx in range(args.repeat)
for scenario in scenarios
for api_name in api_names
]
def _run_one(
http_client: httpx.Client,
task: tuple[int, Scenario, ApiName],
) -> tuple[tuple[int, Scenario, ApiName], list[ScenarioResult], str]:
run_idx, scenario, api_name = task
buf = io.StringIO()
run_tag = f"[run {run_idx + 1}/{args.repeat}]" if args.repeat > 1 else ""
print(
f"\n {run_tag}[{api_name:>9}] {scenario.name}: {scenario.description}",
file=buf,
)
scenario_results = run_scenario(
http_client,
args.host,
args.port,
full_model_id,
scenario,
api_name,
args.timeout,
args.verbose,
)
for r in scenario_results:
status = "PASS" if r.passed else "FAIL"
print(
f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)",
file=buf,
)
for check_name, check_ok in r.checks.items():
mark = "+" if check_ok else "-"
print(f" {mark} {check_name}", file=buf)
if r.error:
print(f" ! {r.error}", file=buf)
return task, scenario_results, buf.getvalue()
try:
with httpx.Client() as http_client:
if args.concurrency == 1:
current_run = -1
for task in tasks:
run_idx = task[0]
if args.repeat > 1 and run_idx != current_run:
print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
current_run = run_idx
_, scenario_results, buffered = _run_one(http_client, task)
all_results.extend(scenario_results)
log.write(buffered)
log.flush()
else:
print(
f"Running {len(tasks)} tasks with concurrency={args.concurrency}",
file=log,
)
with ThreadPoolExecutor(max_workers=args.concurrency) as pool:
futures = [pool.submit(_run_one, http_client, t) for t in tasks]
for fut in as_completed(futures):
_, scenario_results, buffered = fut.result()
for run_idx in range(args.repeat):
if args.repeat > 1:
print(f"\n--- Run {run_idx + 1}/{args.repeat} ---", file=log)
for scenario in scenarios:
for api_name in api_names:
print(
f"\n [{api_name:>9}] {scenario.name}: {scenario.description}",
file=log,
)
scenario_results = run_scenario(
http_client,
args.host,
args.port,
full_model_id,
scenario,
api_name,
args.timeout,
args.verbose,
)
all_results.extend(scenario_results)
log.write(buffered)
log.flush()
for r in scenario_results:
status = "PASS" if r.passed else "FAIL"
print(
f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)",
file=log,
)
for check_name, check_ok in r.checks.items():
mark = "+" if check_ok else "-"
print(f" {mark} {check_name}", file=log)
if r.error:
print(f" ! {r.error}", file=log)
finally:
try:
exo.request_json("DELETE", f"/instance/{instance_id}")
+100 -270
View File
@@ -35,13 +35,14 @@ from harness import (
ExoHttpError,
add_common_instance_args,
capture_cluster_snapshot,
find_existing_instance,
ensure_cuda_available,
instance_id_from_instance,
node_ids_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,
)
@@ -80,7 +81,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
model_path = Path(
snapshot_download(
model_id,
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
)
)
@@ -123,48 +124,8 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
return hf_tokenizer
# TODO: Change back to using only transformers
try:
return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
except (AttributeError, ValueError):
from huggingface_hub import snapshot_download
from transformers import PretrainedConfig
model_path = Path(
snapshot_download(
model_id,
allow_patterns=[
"*.json",
"*.py",
"tokenizer.model",
"*.tiktoken",
"tiktoken.model",
"*.txt",
"*.jsonl",
"*.jinja",
],
)
)
stub_kwargs: dict[str, Any] = {}
config_file = model_path / "config.json"
if config_file.exists():
with open(config_file) as f:
raw = json.load(f)
for key in (
"model_type",
"max_position_embeddings",
"vocab_size",
"bos_token_id",
"eos_token_id",
"pad_token_id",
):
if key in raw:
stub_kwargs[key] = raw[key]
return AutoTokenizer.from_pretrained(
str(model_path),
config=PretrainedConfig(**stub_kwargs),
trust_remote_code=True,
)
# Default: use AutoTokenizer
return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
def format_peak_memory(b: float) -> str:
@@ -271,79 +232,28 @@ def parse_int_list(values: list[str]) -> list[int]:
def run_one_completion(
client: ExoClient,
model_id: str,
pp_hint: int,
tg: int,
prompt_sizer: PromptSizer,
*,
use_prefix_cache: bool = False,
stream: bool = False,
client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer
) -> tuple[dict[str, Any], int]:
content, pp_tokens = prompt_sizer.build(pp_hint)
payload: dict[str, Any] = {
"model": model_id,
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
"logprobs": False,
"use_prefix_cache": use_prefix_cache,
}
if not stream:
payload["stream"] = False
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
preview = content[:200] if content else ""
else:
tokens = 0
first_token_time = None
t0 = time.perf_counter()
text_parts: list[str] = []
stats = None
stats = out.get("generation_stats")
for raw_line in client.stream_bench_chat_completions(payload):
line = raw_line.strip()
if line.startswith(": generation_stats "):
with contextlib.suppress(json.JSONDecodeError):
stats = json.loads(line[len(": generation_stats ") :])
continue
if not line.startswith("data: "):
continue
data = line[6:]
if data == "[DONE]":
break
try:
chunk = json.loads(data)
delta = chunk.get("choices", [{}])[0].get("delta", {})
if delta.get("content"):
if first_token_time is None:
first_token_time = time.perf_counter()
tokens += 1
text_parts.append(delta["content"])
except json.JSONDecodeError:
pass
elapsed = time.perf_counter() - t0
preview = "".join(text_parts)[:200]
if not stats:
ttft = (first_token_time - t0) if first_token_time else elapsed
gen_time = elapsed - ttft if tokens > 1 else elapsed
gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
stats = {
"prompt_tokens": pp_tokens,
"generation_tokens": tokens,
"prompt_tps": round(prompt_tps, 2),
"generation_tps": round(gen_tps, 2),
"peak_memory_usage": {"inBytes": 0},
}
# Extract preview, handling None content (common for thinking models)
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
content = message.get("content") or ""
preview = content[:200] if content else ""
return {
"elapsed_s": elapsed,
@@ -363,19 +273,9 @@ class PromptSizer:
def _make_counter(tokenizer: Any) -> Callable[[str], int]:
def count_fn(user_content: str) -> int:
messages = [{"role": "user", "content": user_content}]
try:
ids = tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True
)
except ValueError:
# Models without a Jinja chat template (e.g. DeepSeek V4 which
# ships its own Python encoder). Use the exo-side V4 encoder.
from exo.worker.engines.mlx.deepseek_v4_encoding import (
encode_messages as encode_v4,
)
prompt = encode_v4(messages, thinking_mode="thinking")
ids = tokenizer.encode(prompt, add_special_tokens=False)
ids = tokenizer.apply_chat_template(
messages, tokenize=True, add_generation_prompt=True
)
# Fix for transformers 5.x
if hasattr(ids, "input_ids"):
ids = ids.input_ids
@@ -470,11 +370,6 @@ def main() -> int:
action="store_true",
help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
)
ap.add_argument(
"--stream",
action="store_true",
help="Use /bench/chat/completions with streaming SSE response (bench=True still applies: no EOS detection, no KV cache).",
)
ap.add_argument(
"--no-system-metrics",
action="store_true",
@@ -486,12 +381,8 @@ def main() -> int:
default=1.0,
help="System metrics polling interval in seconds (default: 1.0).",
)
ap.add_argument(
"--use-prefix-cache",
action="store_true",
help="Enable KV prefix cache during bench (default: disabled for cold-cache measurements).",
)
args = ap.parse_args()
validate_vllm_args(args)
pp_list = parse_int_list(args.pp)
tg_list = parse_int_list(args.tg)
@@ -506,15 +397,6 @@ def main() -> int:
logger.error("--concurrency values must be >= 1")
return 2
if args.use_prefix_cache:
logger.warning(
"--use-prefix-cache: prompt TPS will be approximate. See METHODOLOGY.md for details."
)
if pp_list != sorted(pp_list):
logger.warning(
"--pp values are not in ascending order: prompt TPS will be less accurate. Use ascending --pp for best results."
)
# Log pairing mode
use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
if use_combinations:
@@ -525,6 +407,8 @@ 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
)
@@ -540,124 +424,81 @@ def main() -> int:
logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
raise
# Optionally reuse a running instance for this model
reused_instance_id: str | None = None
if args.reuse_instance:
existing = find_existing_instance(client, full_model_id)
if existing:
reused_instance_id = existing
logger.info(f"Reusing existing instance {reused_instance_id}")
else:
logger.warning(
"--reuse-instance: no existing instance found, creating a new one"
)
selected = settle_and_fetch_placements(
client, full_model_id, args, settle_timeout=args.settle_timeout
)
if reused_instance_id is not None:
# Use the existing instance directly — skip placement iteration
selected = []
download_duration_s = None
if not selected:
logger.error("No valid placements matched your filters.")
return 1
selected.sort(
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
-nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
logger.info(f"placements: {len(selected)}")
for p in selected:
logger.info(
f" - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
)
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(
client,
full_model_id,
selected[0],
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
if download_duration_s is not None:
logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)")
else:
selected = settle_and_fetch_placements(
client, full_model_id, args, settle_timeout=args.settle_timeout
)
if not selected:
logger.error("No valid placements matched your filters.")
return 1
selected.sort(
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
logger.debug(f"exo-bench model: short_id={short_id} full_id={full_model_id}")
logger.info(f"placements: {len(selected)}")
for p in selected:
logger.info(
f" - {p['sharding']} / {p['instance_meta']} / nodes={nodes_used_in_instance(p['instance'])}"
)
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(
client,
full_model_id,
selected[0],
args.danger_delete_downloads,
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")
logger.info("Download: model already cached")
cluster_snapshot = capture_cluster_snapshot(client)
all_rows: list[dict[str, Any]] = []
all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
# If reusing an existing instance, run a single benchmark pass against it
if reused_instance_id is not None:
selected = [None]
for preview in selected:
created_instance = False
if preview is not None:
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
sharding = str(preview["sharding"])
instance_meta = str(preview["instance_meta"])
n_nodes = nodes_used_in_instance(instance)
sharding = str(preview["sharding"])
instance_meta = str(preview["instance_meta"])
n_nodes = nodes_used_in_instance(instance)
logger.info("=" * 80)
logger.info(
f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
)
logger.info("=" * 80)
logger.info(
f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
)
# Delete any existing instances to free resources before placing
try:
state = client.request_json("GET", "/state")
for old_id in list(state.get("instances", {}).keys()):
logger.info(f"Deleting stale instance {old_id}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{old_id}")
if state.get("instances"):
time.sleep(2)
except Exception as e:
logger.warning(f"Failed to clean up stale instances: {e}")
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
except (RuntimeError, TimeoutError) as e:
logger.error(f"Failed to initialize placement: {e}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{instance_id}")
continue
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
except (RuntimeError, TimeoutError) as e:
logger.error(f"Failed to initialize placement: {e}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{instance_id}")
continue
time.sleep(1)
created_instance = True
else:
instance_id = reused_instance_id
sharding = "reused"
instance_meta = "reused"
n_nodes = 0
logger.info("=" * 80)
logger.info(f"Using existing instance {instance_id}")
time.sleep(1)
sampler: SystemMetricsSampler | None = None
if not args.no_system_metrics and preview is not None:
if not args.no_system_metrics:
nids = node_ids_from_instance(instance)
sampler = SystemMetricsSampler(
ExoClient(args.host, args.port, timeout_s=30),
@@ -666,20 +507,11 @@ def main() -> int:
)
sampler.start()
def _do_one(c: ExoClient, pp: int, tg: int) -> tuple[dict[str, Any], int]:
return run_one_completion(
c,
full_model_id,
pp,
tg,
prompt_sizer,
use_prefix_cache=args.use_prefix_cache,
stream=args.stream,
)
try:
for i in range(args.warmup):
_do_one(client, pp_list[0], tg_list[0])
run_one_completion(
client, full_model_id, pp_list[0], tg_list[0], prompt_sizer
)
logger.debug(f" warmup {i + 1}/{args.warmup} done")
# If pp and tg lists have same length, run in tandem (zip)
@@ -701,7 +533,9 @@ def main() -> int:
# Sequential: single request
try:
inf_t0 = time.monotonic()
row, actual_pp_tokens = _do_one(client, pp, tg)
row, actual_pp_tokens = run_one_completion(
client, full_model_id, pp, tg, prompt_sizer
)
inference_windows.append((inf_t0, time.monotonic()))
except Exception as e:
logger.error(e)
@@ -737,7 +571,6 @@ def main() -> int:
"stream": False,
"max_tokens": tg,
"logprobs": False,
"use_prefix_cache": args.use_prefix_cache,
}
barrier = threading.Barrier(concurrency)
batch_start = threading.Event()
@@ -850,12 +683,10 @@ def main() -> int:
gen_tps = per_req_tps * concurrency
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
)
def _peak_bytes(s: dict[str, Any]) -> float:
pm = s["peak_memory_usage"]
return pm.get("inBytes") or pm.get("in_bytes", 0)
peak = mean(_peak_bytes(x["stats"]) for x in runs)
summary = (
f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
f"prompt_tokens={ptok} gen_tokens={gtok} "
@@ -880,16 +711,15 @@ def main() -> int:
if placement_metrics:
all_system_metrics.update(placement_metrics)
if created_instance and instance_id is not None:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
wait_for_instance_gone(client, instance_id)
logger.debug(f"Deleted instance {instance_id}")
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
wait_for_instance_gone(client, instance_id)
logger.debug(f"Deleted instance {instance_id}")
time.sleep(5)
time.sleep(5)
output: dict[str, Any] = {"runs": all_rows}
if cluster_snapshot:
+61 -427
View File
@@ -47,12 +47,13 @@ from harness import (
ExoHttpError,
add_common_instance_args,
capture_cluster_snapshot,
find_existing_instance,
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,
)
@@ -63,15 +64,6 @@ from loguru import logger
# ---------------------------------------------------------------------------
MAX_RETRIES = 30
INSTANCE_HEALTH_CHECK_AFTER = (
3 # Check instance health after this many consecutive failures
)
class InstanceFailedError(RuntimeError):
"""Raised when the exo instance is detected as failed/gone."""
DEFAULT_MAX_TOKENS = 16_384
REASONING_MAX_TOKENS = 131_072
TEMPERATURE_NON_REASONING = 0.0
@@ -281,7 +273,7 @@ def run_humaneval_test(
@dataclass
class QuestionResult:
question_id: int | str
question_id: int
prompt: str
response: str
extracted_answer: str | None
@@ -291,11 +283,7 @@ class QuestionResult:
prompt_tokens: int = 0
completion_tokens: int = 0
reasoning_tokens: int = 0
reasoning_content: str = ""
finish_reason: str = ""
elapsed_s: float = 0.0
power_watts: float = 0.0
energy_joules: float = 0.0
@dataclass
@@ -531,10 +519,6 @@ class ApiResult:
prompt_tokens: int
completion_tokens: int
reasoning_tokens: int
reasoning_content: str = ""
finish_reason: str = ""
power_watts: float = 0.0
energy_joules: float = 0.0
async def _call_api(
@@ -548,9 +532,6 @@ async def _call_api(
system_message: str | None = None,
reasoning_effort: str | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float | None = None,
enable_thinking: bool | None = None,
) -> ApiResult:
messages = []
if system_message:
@@ -567,12 +548,6 @@ async def _call_api(
body["reasoning_effort"] = reasoning_effort
if top_p is not None:
body["top_p"] = top_p
if top_k is not None:
body["top_k"] = top_k
if min_p is not None:
body["min_p"] = min_p
if enable_thinking is not None:
body["enable_thinking"] = enable_thinking
resp = await client.post(
f"{base_url}/v1/chat/completions",
@@ -581,40 +556,19 @@ async def _call_api(
)
resp.raise_for_status()
data = resp.json()
choice = data["choices"][0]
message = choice["message"]
content = message.get("content") or ""
reasoning_content = message.get("reasoning_content") or ""
finish_reason = choice.get("finish_reason") or ""
# For thinking models, empty content is expected when finish_reason is "length"
if not content.strip() and finish_reason != "length" and not reasoning_content:
content = data["choices"][0]["message"]["content"]
if not content or not content.strip():
raise ValueError("Empty response from model")
usage = data.get("usage", {})
details = usage.get("completion_tokens_details", {})
power = data.get("power_usage") or {}
return ApiResult(
content=content,
prompt_tokens=usage.get("prompt_tokens", 0),
completion_tokens=usage.get("completion_tokens", 0),
reasoning_tokens=details.get("reasoning_tokens", 0) if details else 0,
reasoning_content=reasoning_content,
finish_reason=finish_reason,
power_watts=power.get("total_avg_sys_power_watts", 0.0),
energy_joules=power.get("total_energy_joules", 0.0),
)
async def _check_instance_health(base_url: str) -> bool:
"""Return True if the exo instance is still reachable."""
try:
async with httpx.AsyncClient() as c:
resp = await c.get(f"{base_url}/models", timeout=5.0)
return resp.status_code == 200
except Exception:
return False
async def call_with_retries(
client: httpx.AsyncClient,
base_url: str,
@@ -626,14 +580,8 @@ async def call_with_retries(
system_message: str | None = None,
reasoning_effort: str | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float | None = None,
enable_thinking: bool | None = None,
instance_failed: asyncio.Event | None = None,
) -> ApiResult | None:
for attempt in range(MAX_RETRIES):
if instance_failed and instance_failed.is_set():
raise InstanceFailedError("Instance already marked as failed")
try:
return await _call_api(
client,
@@ -646,30 +594,8 @@ async def call_with_retries(
system_message,
reasoning_effort,
top_p,
top_k,
min_p,
enable_thinking,
)
except Exception as e:
is_conn_error = isinstance(
e,
(
httpx.ConnectError,
httpx.RemoteProtocolError,
ConnectionRefusedError,
OSError,
),
)
if (
is_conn_error
and attempt >= INSTANCE_HEALTH_CHECK_AFTER
and not await _check_instance_health(base_url)
):
if instance_failed:
instance_failed.set()
raise InstanceFailedError(
f"Instance is down after {attempt + 1} failures: {e}"
) from e
if attempt < MAX_RETRIES - 1:
wait = min(2**attempt, 60)
logger.warning(
@@ -694,16 +620,10 @@ async def evaluate_benchmark(
max_tokens: int,
concurrency: int = 1,
limit: int | None = None,
offset: int = 0,
timeout: float | None = None,
reasoning_effort: str | None = None,
top_p: float | None = None,
top_k: int | None = None,
min_p: float | None = None,
enable_thinking: bool | None = None,
difficulty: str | None = None,
checkpoint_path: Path | None = None,
release_version: str | None = None,
) -> list[QuestionResult]:
"""Run a benchmark. Returns per-question results."""
import datasets
@@ -734,21 +654,7 @@ async def evaluate_benchmark(
ds = ds.filter(lambda x: x["difficulty"] == difficulty)
logger.info(f"Filtered to {len(ds)} {difficulty} problems")
if release_version and "release_version" in ds.column_names:
ds = ds.filter(lambda x: x["release_version"] == release_version)
logger.info(
f"Filtered to {len(ds)} problems with release_version={release_version}"
)
# Sort by question_id to match LCB runner ordering (scenario_router.py:60).
# This ensures [offset:offset+limit] slices select the same problems as vllm.
if "question_id" in ds.column_names:
ds = ds.sort("question_id")
total = len(ds)
if offset > 0:
ds = ds.select(range(min(offset, total), total))
total = len(ds)
if limit and limit < total:
ds = ds.select(range(limit))
total = limit
@@ -756,13 +662,6 @@ async def evaluate_benchmark(
logger.info(
f"Evaluating {benchmark_name}: {total} questions, concurrency={concurrency}, "
f"temperature={temperature}, max_tokens={max_tokens}"
+ (f", top_k={top_k}" if top_k is not None else "")
+ (f", min_p={min_p}" if min_p is not None else "")
+ (
f", enable_thinking={enable_thinking}"
if enable_thinking is not None
else ""
)
)
if config.kind == "code":
@@ -770,64 +669,16 @@ async def evaluate_benchmark(
"Code benchmarks execute model-generated code. Use a sandboxed environment."
)
# Load checkpoint for resume
checkpoint_data: dict[str | int, dict[str, Any]] = {}
if checkpoint_path and checkpoint_path.exists():
with open(checkpoint_path) as f:
for line in f:
entry = json.loads(line)
checkpoint_data[entry["question_id"]] = entry
logger.info(f"Loaded {len(checkpoint_data)} checkpointed results")
semaphore = asyncio.Semaphore(concurrency)
instance_failed = asyncio.Event()
results: list[QuestionResult | None] = [None] * total
completed = 0
lock = asyncio.Lock()
def _get_question_id(idx: int, doc: dict) -> str | int:
"""Get a stable question ID for checkpointing."""
if benchmark_name == "livecodebench":
return doc.get("question_id", idx)
elif benchmark_name == "humaneval":
return doc.get("task_id", idx)
return idx
async def process_question(
idx: int, doc: dict, http_client: httpx.AsyncClient
) -> None:
nonlocal completed
system_msg = None
question_id = _get_question_id(idx, doc)
# Bail out early if instance is already dead
if instance_failed.is_set():
return
# Check checkpoint
if question_id in checkpoint_data:
cached = checkpoint_data[question_id]
results[idx] = QuestionResult(
question_id=question_id,
prompt=cached.get("prompt", ""),
response=cached.get("response", ""),
extracted_answer=cached.get("extracted_answer"),
gold_answer=cached.get("gold_answer", ""),
correct=cached.get("correct", False),
error=cached.get("error"),
prompt_tokens=cached.get("prompt_tokens", 0),
completion_tokens=cached.get("completion_tokens", 0),
reasoning_tokens=cached.get("reasoning_tokens", 0),
reasoning_content=cached.get("reasoning_content", ""),
finish_reason=cached.get("finish_reason", ""),
elapsed_s=cached.get("elapsed_s", 0.0),
power_watts=cached.get("power_watts", 0.0),
energy_joules=cached.get("energy_joules", 0.0),
)
async with lock:
completed += 1
logger.info(f" [{completed}/{total}] {question_id} (cached)")
return
if benchmark_name == "gpqa_diamond":
prompt, gold = format_gpqa_question(doc, idx)
@@ -848,50 +699,24 @@ async def evaluate_benchmark(
raise ValueError(f"Unknown benchmark: {benchmark_name}")
async with semaphore:
if instance_failed.is_set():
return
t0 = time.monotonic()
try:
# Race the API call against the instance_failed event
api_task = asyncio.create_task(
call_with_retries(
http_client,
base_url,
model,
prompt,
temperature,
max_tokens,
timeout,
system_message=system_msg,
reasoning_effort=reasoning_effort,
top_p=top_p,
top_k=top_k,
min_p=min_p,
enable_thinking=enable_thinking,
instance_failed=instance_failed,
)
)
failed_waiter = asyncio.create_task(instance_failed.wait())
done, pending = await asyncio.wait(
[api_task, failed_waiter],
return_when=asyncio.FIRST_COMPLETED,
)
for p in pending:
p.cancel()
with contextlib.suppress(asyncio.CancelledError):
await p
if instance_failed.is_set() and api_task not in done:
logger.error(f"Instance failed, aborting {question_id}")
return
api_result = api_task.result()
except InstanceFailedError:
logger.error(f"Instance failed, skipping {question_id}")
return
api_result = await call_with_retries(
http_client,
base_url,
model,
prompt,
temperature,
max_tokens,
timeout,
system_message=system_msg,
reasoning_effort=reasoning_effort,
top_p=top_p,
)
elapsed = time.monotonic() - t0
if api_result is None:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response="",
extracted_answer=None,
@@ -906,17 +731,13 @@ async def evaluate_benchmark(
"prompt_tokens": api_result.prompt_tokens,
"completion_tokens": api_result.completion_tokens,
"reasoning_tokens": api_result.reasoning_tokens,
"reasoning_content": api_result.reasoning_content,
"finish_reason": api_result.finish_reason,
"elapsed_s": elapsed,
"power_watts": api_result.power_watts,
"energy_joules": api_result.energy_joules,
}
if config.kind == "mc":
extracted = extract_mc_answer(response, valid_letters)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=extracted,
@@ -930,7 +751,7 @@ async def evaluate_benchmark(
check_aime_answer(extracted, int(gold)) if extracted else False
)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=extracted,
@@ -944,7 +765,7 @@ async def evaluate_benchmark(
code = extract_code_block(response, preserve_indent=keep_indent)
if code is None:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -959,7 +780,7 @@ async def evaluate_benchmark(
code,
)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer="pass" if passed else "fail",
@@ -974,7 +795,7 @@ async def evaluate_benchmark(
exec_meta["sample"],
)
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer="pass" if passed else "fail",
@@ -985,7 +806,7 @@ async def evaluate_benchmark(
)
else:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -996,7 +817,7 @@ async def evaluate_benchmark(
)
else:
result = QuestionResult(
question_id=question_id,
question_id=idx,
prompt=prompt,
response=response,
extracted_answer=None,
@@ -1008,82 +829,24 @@ async def evaluate_benchmark(
results[idx] = result
# Write checkpoint (skip infra failures so they get retried on resume,
# but keep wrong answers — they are legitimate results)
if checkpoint_path is not None and result.response:
_write_checkpoint(checkpoint_path, result)
async with lock:
completed += 1
n = completed
# Log progress
thinking_info = ""
if result.reasoning_content:
thinking_info = f", {len(result.reasoning_content)} chars thinking"
logger.info(
f" [{n}/{total}] {question_id}: {len(result.response)} chars{thinking_info}, "
f"tokens: {result.prompt_tokens}+{result.completion_tokens} "
f"[{result.finish_reason}]"
+ (f" {result.extracted_answer}" if result.extracted_answer else "")
)
async def _health_monitor() -> None:
"""Periodically check if the instance is still alive."""
# Wait a bit before first check to let things start
await asyncio.sleep(10)
while not instance_failed.is_set():
if not await _check_instance_health(base_url):
# Double-check to avoid false positives
await asyncio.sleep(2)
if not await _check_instance_health(base_url):
logger.error("Health monitor: instance is down!")
instance_failed.set()
return
await asyncio.sleep(5)
if n % max(1, total // 20) == 0 or n == total:
correct_so_far = sum(1 for r in results if r is not None and r.correct)
answered = sum(1 for r in results if r is not None)
logger.info(
f" [{n}/{total}] {correct_so_far}/{answered} correct "
f"({correct_so_far / max(answered, 1):.1%})"
)
async with httpx.AsyncClient() as http_client:
monitor = asyncio.create_task(_health_monitor())
tasks = [process_question(i, doc, http_client) for i, doc in enumerate(ds)]
await asyncio.gather(*tasks)
monitor.cancel()
with contextlib.suppress(asyncio.CancelledError):
await monitor
if instance_failed.is_set():
completed_count = sum(1 for r in results if r is not None)
logger.error(
f"Instance failed! Completed {completed_count}/{total} problems. "
f"Checkpoint saved — restart to resume remaining problems."
)
raise InstanceFailedError("Instance failed during evaluation")
return [r for r in results if r is not None]
def _write_checkpoint(path: Path, result: QuestionResult) -> None:
"""Append a single result to the JSONL checkpoint file."""
entry = {
"question_id": result.question_id,
"prompt": result.prompt,
"response": result.response,
"extracted_answer": result.extracted_answer,
"gold_answer": result.gold_answer,
"correct": result.correct,
"error": result.error,
"prompt_tokens": result.prompt_tokens,
"completion_tokens": result.completion_tokens,
"reasoning_tokens": result.reasoning_tokens,
"reasoning_content": result.reasoning_content,
"finish_reason": result.finish_reason,
"elapsed_s": round(result.elapsed_s, 2),
"power_watts": round(result.power_watts, 2),
"energy_joules": round(result.energy_joules, 2),
}
with open(path, "a") as f:
f.write(json.dumps(entry) + "\n")
# ---------------------------------------------------------------------------
# Results display
# ---------------------------------------------------------------------------
@@ -1106,8 +869,6 @@ def print_results(
total_elapsed = sum(r.elapsed_s for r in results)
wall_clock = max(r.elapsed_s for r in results) if results else 0.0
avg_gen_tps = total_completion_tokens / total_elapsed if total_elapsed > 0 else 0.0
total_energy = sum(r.energy_joules for r in results)
avg_power = sum(r.power_watts for r in results) / max(total, 1)
label = f"[c={concurrency}] " if concurrency is not None else ""
print(f"\n{label}{benchmark_name}: {correct}/{total} ({accuracy:.1%})")
@@ -1119,10 +880,6 @@ def print_results(
f" | total time: {total_elapsed:.1f}s wall clock: {wall_clock:.1f}s"
)
print(tok_line)
if total_energy > 0:
print(
f" power: avg {avg_power:.1f}W | total energy: {total_energy:.1f}J ({total_energy / 3600:.2f}Wh)"
)
if errors:
print(f" API errors: {errors}")
if no_extract:
@@ -1141,8 +898,6 @@ def print_results(
"total_elapsed_s": total_elapsed,
"wall_clock_s": wall_clock,
"avg_gen_tps": avg_gen_tps,
"avg_power_watts": avg_power,
"total_energy_joules": total_energy,
}
@@ -1300,11 +1055,7 @@ def save_results(
"prompt_tokens": r.prompt_tokens,
"completion_tokens": r.completion_tokens,
"reasoning_tokens": r.reasoning_tokens,
"reasoning_content": r.reasoning_content,
"finish_reason": r.finish_reason,
"elapsed_s": round(r.elapsed_s, 2),
"power_watts": round(r.power_watts, 2),
"energy_joules": round(r.energy_joules, 2),
}
for r in results
],
@@ -1320,15 +1071,6 @@ def save_results(
# ---------------------------------------------------------------------------
def _checkpoint_path(
results_dir: str, benchmark: str, model: str, concurrency: int
) -> Path:
"""Return the JSONL checkpoint path for a benchmark run."""
out_dir = Path(results_dir) / model.replace("/", "_") / benchmark
out_dir.mkdir(parents=True, exist_ok=True)
return out_dir / f"c{concurrency}.checkpoint.jsonl"
def parse_int_list(values: list[str]) -> list[int]:
items: list[int] = []
for v in values:
@@ -1356,12 +1098,6 @@ def main() -> int:
default=None,
help="Max questions per benchmark (for fast iteration).",
)
ap.add_argument(
"--offset",
type=int,
default=0,
help="Skip first N questions (0-based).",
)
reasoning_group = ap.add_mutually_exclusive_group()
reasoning_group.add_argument(
@@ -1381,8 +1117,6 @@ def main() -> int:
"--temperature", type=float, default=None, help="Override temperature."
)
ap.add_argument("--top-p", type=float, default=None, help="Override top_p.")
ap.add_argument("--top-k", type=int, default=None, help="Override top_k.")
ap.add_argument("--min-p", type=float, default=None, help="Override min_p.")
ap.add_argument(
"--max-tokens", type=int, default=None, help="Override max output tokens."
)
@@ -1416,34 +1150,19 @@ def main() -> int:
choices=["easy", "medium", "hard"],
help="Filter by difficulty (livecodebench only). E.g. --difficulty hard",
)
ap.add_argument(
"--release-version",
default=None,
help="LCB dataset release version (livecodebench only). E.g. release_v5",
)
ap.add_argument(
"--results-dir",
default="eval_results",
help="Directory for result JSON files (default: eval_results).",
)
ap.add_argument(
"--enable-thinking",
type=lambda v: v.lower() in ("true", "1", "yes"),
default=None,
help="Enable thinking mode for models that support it.",
)
ap.add_argument(
"--force",
"--skip-instance-setup",
action="store_true",
help="Discard any existing checkpoint and run from scratch.",
)
ap.add_argument(
"--keep-instance",
action="store_true",
help="Skip deleting the instance after eval (for chaining runs).",
help="Skip exo instance management (assumes model is already running).",
)
args, _ = ap.parse_known_args()
validate_vllm_args(args)
# Resolve tasks
if args.tasks:
@@ -1460,27 +1179,16 @@ def main() -> int:
# Instance management
client = ExoClient(args.host, args.port, timeout_s=args.timeout)
if args.ensure_cuda:
ensure_cuda_available(client)
instance_id: str | None = None
created_instance = False
_short_id, full_model_id = resolve_model_short_id(
client,
args.model,
force_download=args.force_download,
)
# Optionally reuse a running instance for this model
if args.reuse_instance:
existing = find_existing_instance(client, full_model_id)
if existing:
instance_id = existing
logger.info(f"Reusing existing instance {instance_id}")
else:
logger.warning(
"--reuse-instance: no existing instance found, creating a new one"
)
if instance_id is None:
if not args.skip_instance_setup:
short_id, full_model_id = resolve_model_short_id(
client,
args.model,
force_download=args.force_download,
)
selected = settle_and_fetch_placements(
client,
full_model_id,
@@ -1495,7 +1203,7 @@ def main() -> int:
key=lambda p: (
str(p.get("instance_meta", "")),
str(p.get("sharding", "")),
nodes_used_in_instance(p["instance"]),
-nodes_used_in_instance(p["instance"]),
),
reverse=True,
)
@@ -1522,18 +1230,6 @@ def main() -> int:
if download_duration is not None:
logger.info(f"Download: {download_duration:.1f}s")
# Delete any existing instances to free resources before placing
try:
state = client.request_json("GET", "/state")
for old_id in list(state.get("instances", {}).keys()):
logger.info(f"Deleting stale instance {old_id}")
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{old_id}")
if state.get("instances"):
time.sleep(2)
except Exception as e:
logger.warning(f"Failed to clean up stale instances: {e}")
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
@@ -1543,9 +1239,10 @@ def main() -> int:
client.request_json("DELETE", f"/instance/{instance_id}")
return 1
time.sleep(1)
created_instance = True
cluster_snapshot = capture_cluster_snapshot(client)
cluster_snapshot = capture_cluster_snapshot(client)
else:
full_model_id = args.model
cluster_snapshot = None
# Auto-detect reasoning from model config
model_config = load_model_config(full_model_id)
@@ -1599,57 +1296,16 @@ def main() -> int:
reasoning_effort = str(cfg["reasoning_effort"])
else:
reasoning_effort = "high" if is_reasoning else None
if args.top_k is not None:
top_k: int | None = args.top_k
elif "top_k" in cfg:
top_k = int(cfg["top_k"])
else:
top_k = None
if args.min_p is not None:
min_p: float | None = args.min_p
elif "min_p" in cfg:
min_p = float(cfg["min_p"])
else:
min_p = None
if args.enable_thinking is not None:
enable_thinking: bool | None = args.enable_thinking
elif "enable_thinking" in cfg:
enable_thinking = bool(cfg["enable_thinking"])
else:
enable_thinking = None
base_url = f"http://{args.host}:{args.port}"
logger.info(f"Model: {full_model_id}")
logger.info(
f"Settings: temperature={temperature}, max_tokens={max_tokens}, "
+ (f"top_p={top_p}, " if top_p is not None else "")
+ (f"top_k={top_k}, " if top_k is not None else "")
+ (f"min_p={min_p}, " if min_p is not None else "")
+ f"reasoning={'yes' if is_reasoning else 'no'}"
+ (f", reasoning_effort={reasoning_effort}" if reasoning_effort else "")
+ (
f", enable_thinking={enable_thinking}"
if enable_thinking is not None
else ""
)
)
# Common kwargs for evaluate_benchmark
eval_kwargs: dict[str, Any] = {
"reasoning_effort": reasoning_effort,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"enable_thinking": enable_thinking,
"difficulty": args.difficulty,
"offset": args.offset,
"release_version": args.release_version,
}
try:
if args.compare_concurrency:
concurrency_levels = parse_int_list(args.compare_concurrency)
@@ -1658,11 +1314,6 @@ def main() -> int:
for c in concurrency_levels:
logger.info(f"\n{'=' * 50}")
logger.info(f"Running {task_name} at concurrency={c}")
checkpoint_path = _checkpoint_path(
args.results_dir, task_name, full_model_id, c
)
if args.force and checkpoint_path.exists():
checkpoint_path.unlink()
results = asyncio.run(
evaluate_benchmark(
task_name,
@@ -1673,8 +1324,9 @@ def main() -> int:
concurrency=c,
limit=args.limit,
timeout=args.request_timeout,
checkpoint_path=checkpoint_path,
**eval_kwargs,
reasoning_effort=reasoning_effort,
top_p=top_p,
difficulty=args.difficulty,
)
)
if results:
@@ -1689,18 +1341,10 @@ def main() -> int:
cluster=cluster_snapshot,
)
results_by_c[c] = results
# Clean up checkpoint on success
if checkpoint_path.exists():
checkpoint_path.unlink()
if len(results_by_c) >= 2:
print_comparison(task_name, results_by_c)
else:
for task_name in task_names:
checkpoint_path = _checkpoint_path(
args.results_dir, task_name, full_model_id, args.num_concurrent
)
if args.force and checkpoint_path.exists():
checkpoint_path.unlink()
results = asyncio.run(
evaluate_benchmark(
task_name,
@@ -1711,8 +1355,9 @@ def main() -> int:
concurrency=args.num_concurrent,
limit=args.limit,
timeout=args.request_timeout,
checkpoint_path=checkpoint_path,
**eval_kwargs,
reasoning_effort=reasoning_effort,
top_p=top_p,
difficulty=args.difficulty,
)
)
if results:
@@ -1726,25 +1371,14 @@ def main() -> int:
scores,
cluster=cluster_snapshot,
)
# Clean up checkpoint on success
if checkpoint_path.exists():
checkpoint_path.unlink()
finally:
if created_instance and instance_id is not None:
if args.keep_instance:
logger.info(f"Keeping instance {instance_id} (--keep-instance)")
else:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
try:
wait_for_instance_gone(client, instance_id)
except TimeoutError:
logger.warning(
f"Timed out waiting for instance {instance_id} to be deleted"
)
if instance_id is not None:
try:
client.request_json("DELETE", f"/instance/{instance_id}")
except ExoHttpError as e:
if e.status != 404:
raise
wait_for_instance_gone(client, instance_id)
return 0
+43 -62
View File
@@ -6,7 +6,6 @@ import http.client
import json
import os
import time
from collections.abc import Iterator
from typing import Any
from urllib.parse import urlencode
@@ -70,29 +69,9 @@ class ExoClient:
def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
return self.request_json("POST", "/bench/chat/completions", body=payload)
def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
"""POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
payload = {**payload, "stream": True}
data = json.dumps(payload).encode("utf-8")
conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
try:
conn.request(
"POST",
"/bench/chat/completions",
body=data,
headers={
"Content-Type": "application/json",
"Accept": "text/event-stream",
},
)
resp = conn.getresponse()
if resp.status >= 400:
raw = resp.read().decode("utf-8", errors="replace")
raise ExoHttpError(resp.status, resp.reason, raw[:300])
for line in resp:
yield line.decode("utf-8", errors="replace")
finally:
conn.close()
def post_bench_disaggregated(self, payload: dict[str, Any]) -> dict[str, Any]:
payload["disaggregated"] = True
return self.request_json("POST", "/bench/chat/completions", body=payload)
def get_state_path(self, path: str) -> Any:
try:
@@ -278,6 +257,31 @@ def resolve_model_short_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":
@@ -293,15 +297,11 @@ def sharding_filter(sharding: str, wanted: str) -> bool:
def fetch_and_filter_placements(
client: ExoClient,
full_model_id: str,
args: argparse.Namespace,
node_id: str | None = None,
client: ExoClient, full_model_id: str, args: argparse.Namespace
) -> list[dict[str, Any]]:
params: dict[str, str] = {"model_id": full_model_id}
if node_id is not None:
params["node_ids"] = node_id
previews_resp = client.request_json("GET", "/instance/previews", params=params)
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]] = []
@@ -361,9 +361,8 @@ def settle_and_fetch_placements(
full_model_id: str,
args: argparse.Namespace,
settle_timeout: float = 0,
node_id: str | None = None,
) -> list[dict[str, Any]]:
selected = fetch_and_filter_placements(client, full_model_id, args, node_id=node_id)
selected = fetch_and_filter_placements(client, full_model_id, args)
if not selected and settle_timeout > 0:
backoff = _SETTLE_INITIAL_BACKOFF_S
@@ -376,9 +375,7 @@ def settle_and_fetch_placements(
)
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, node_id=node_id
)
selected = fetch_and_filter_placements(client, full_model_id, args)
return selected
@@ -494,8 +491,9 @@ def run_planning_phase(
)
logger.info(f"Started download on {node_id}")
# Wait for downloads (no timeout — poll until complete or failed)
while True:
# Wait for downloads
start = time.time()
while time.time() - start < timeout:
all_done = True
for node_id in node_ids:
node_downloads = client.get_node_downloads(node_id) or []
@@ -545,24 +543,9 @@ def run_planning_phase(
if download_t0 is not None:
return time.perf_counter() - download_t0
return None
time.sleep(10)
time.sleep(1)
def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
"""Find an existing running instance for the given model."""
try:
state = client.request_json("GET", "/state")
except Exception:
return None
for inst_id, inst in state.get("instances", {}).items():
# Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
for _inst_type, inner in inst.items():
if not isinstance(inner, dict):
continue
sa = inner.get("shardAssignments", {})
if sa.get("modelId") == model_id:
return inst_id
return None
raise TimeoutError("Downloads did not complete in time")
def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
@@ -589,9 +572,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
help="Only consider placements using >= this many nodes.",
)
ap.add_argument(
"--instance-meta",
choices=["ring", "jaccl", "vllm", "both"],
default="both",
"--instance-meta", choices=["ring", "jaccl", "vllm", "both"], default="both"
)
ap.add_argument(
"--sharding", choices=["pipeline", "tensor", "both"], default="both"
@@ -612,7 +593,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
ap.add_argument(
"--settle-timeout",
type=float,
default=60.0,
default=0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
ap.add_argument(
@@ -621,7 +602,7 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
help="Delete existing models from smallest to largest to make room for benchmark model.",
)
ap.add_argument(
"--reuse-instance",
"--ensure-cuda",
action="store_true",
help="Reuse an existing running instance for this model instead of creating a new one.",
help="Verify the exo cluster has CUDA/vllm capability; error if not.",
)
-36
View File
@@ -1,36 +0,0 @@
# Prefill/Decode disaggregation benchmark config.
#
# Top-level keys are bench-wide. [prefill] and [decode] sections set per-side
# placement filters and (optionally) per-side model.
#
# Example:
# uv run python bench/prefill_decode_bench.py --config bench/prefill-decode.toml
host = "james"
port = 52415
timeout = 7200.0
settle_timeout = 60.0
# Workload
pp = [4096, 8192]
tg = [128]
repeat = 1
warmup = 0
json_out = "bench/prefill_decode_results.json"
[prefill]
model = "sakamakismile/Qwen3.6-27B-NVFP4"
node = "gx10-de89"
instance_meta = "vllm"
sharding = "pipeline"
min_nodes = 1
max_nodes = 1
[decode]
model = "mlx-community/Qwen3.6-27B-4bit"
node = "Ryuichis MacBook Pro"
instance_meta = "ring"
sharding = "pipeline"
min_nodes = 1
max_nodes = 1
-869
View File
@@ -1,869 +0,0 @@
# type: ignore
#!/usr/bin/env python3
"""Disaggregated prefill-decode benchmark for exo (MLX → MLX).
Spins up two MLX instances on the cluster, marks one as Prefill source and
the other as Decode target via /v1/instance-links, then sends chat
completions to the API. The master routes the request to the decode
instance and stamps `prefill_endpoint` pointing at the prefill instance
the worker decides per-request whether to ship prefill remotely
(uncached_count > REMOTE_PREFILL_MIN_TOKENS).
Usage:
uv run python bench/prefill_decode_bench.py --model <id> --pp 2048,8192 --tg 128
uv run python bench/prefill_decode_bench.py --model <id> --pp 4096 --tg 128 --repeat 3
uv run python bench/prefill_decode_bench.py --model <id> --pp 2048 --tg 128 --dry-run
"""
from __future__ import annotations
import argparse
import contextlib
import copy
import itertools
import json
import sys
import time
import tomllib
from pathlib import Path
from statistics import mean
from typing import Any
from exo_bench import (
PromptSizer,
SystemMetricsSampler,
format_peak_memory,
load_tokenizer_for_bench,
parse_int_list,
)
from harness import (
ExoClient,
ExoHttpError,
add_common_instance_args,
instance_id_from_instance,
node_ids_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
run_planning_phase,
settle_and_fetch_placements,
unwrap_instance,
wait_for_instance_gone,
wait_for_instance_ready,
)
from loguru import logger
def _node_id_to_friendly(client: ExoClient) -> dict[str, str]:
identities = client.get_node_identities() or {}
out: dict[str, str] = {}
for node_id, identity in identities.items():
if isinstance(identity, dict):
name = identity.get("friendlyName") or identity.get("friendly_name")
if isinstance(name, str):
out[str(node_id)] = name
return out
def _placement_node_friendly_names(
placement: dict[str, Any], id_to_friendly: dict[str, str]
) -> list[str]:
instance = placement["instance"]
return [id_to_friendly.get(nid, nid) for nid in node_ids_from_instance(instance)]
def _filter_by_node(
placements: list[dict[str, Any]],
friendly_name: str,
id_to_friendly: dict[str, str],
) -> list[dict[str, Any]]:
target = friendly_name.lower()
matched: list[dict[str, Any]] = []
for p in placements:
names = [n.lower() for n in _placement_node_friendly_names(p, id_to_friendly)]
if any(target == n or target in n for n in names):
matched.append(p)
return matched
def _node_id_by_friendly(id_to_friendly: dict[str, str], target: str) -> str | None:
target_lc = target.lower()
for nid, name in id_to_friendly.items():
if target_lc == name.lower() or target_lc in name.lower():
return nid
return None
def _load_toml(path: str) -> dict[str, Any]:
with Path(path).open("rb") as f:
return tomllib.load(f)
_TOP_LEVEL_TOML_KEYS = {
"host",
"port",
"timeout",
"settle_timeout",
"model",
"pp",
"tg",
"repeat",
"warmup",
"json_out",
"instance_meta",
"sharding",
"min_nodes",
"max_nodes",
"force_download",
"danger_delete_downloads",
"all_combinations",
}
def _inject_toml_into_argv() -> None:
"""If --config X is in sys.argv, pre-load it and inject required CLI args
(--model, --pp, --tg) so argparse's required=True checks pass."""
argv = sys.argv
if "--config" not in argv:
return
idx = argv.index("--config")
if idx + 1 >= len(argv):
return
cfg_path = argv[idx + 1]
cfg = _load_toml(cfg_path)
decode = cfg.get("decode", {})
def _has(flag: str) -> bool:
return any(a == flag or a.startswith(flag + "=") for a in argv)
# --model: prefer top-level, then [decode].model
if not _has("--model"):
model = cfg.get("model") or decode.get("model")
if model:
argv += ["--model", str(model)]
if not _has("--pp"):
pp = cfg.get("pp")
if pp:
argv += (
["--pp", *(str(x) for x in pp)]
if isinstance(pp, list)
else [
"--pp",
str(pp),
]
)
if not _has("--tg"):
tg = cfg.get("tg")
if tg:
argv += (
["--tg", *(str(x) for x in tg)]
if isinstance(tg, list)
else [
"--tg",
str(tg),
]
)
def _merge_toml_into_args(args: argparse.Namespace, cfg: dict[str, Any]) -> None:
"""Apply top-level toml keys onto args namespace where args has a default."""
for key, value in cfg.items():
if key in {"prefill", "decode"}:
continue
if key not in _TOP_LEVEL_TOML_KEYS:
continue
attr = key
current = getattr(args, attr, None)
if current in (None, [], False):
setattr(args, attr, value)
def _side_args(
base: argparse.Namespace, overrides: dict[str, Any]
) -> argparse.Namespace:
out = copy.copy(base)
for k in (
"instance_meta",
"sharding",
"min_nodes",
"max_nodes",
"skip_pipeline_jaccl",
"skip_tensor_ring",
):
if k in overrides:
setattr(out, k, overrides[k])
return out
def _pick_two_distinct_placements(
placements: list[dict[str, Any]],
) -> tuple[dict[str, Any], dict[str, Any]] | None:
if len(placements) < 2:
return None
seen_nodes: set[tuple[str, ...]] = set()
chosen: list[dict[str, Any]] = []
for p in placements:
nodes = tuple(sorted(str(n) for n in p.get("nodes", [])))
if nodes in seen_nodes:
continue
seen_nodes.add(nodes)
chosen.append(p)
if len(chosen) == 2:
return chosen[0], chosen[1]
return None
def _create_instance_link(
client: ExoClient,
prefill_instance_id: str,
decode_instance_id: str,
) -> str:
out = client.request_json(
"POST",
"/v1/instance-links",
body={
"prefill_instances": [prefill_instance_id],
"decode_instances": [decode_instance_id],
},
)
return str(out.get("commandId", ""))
def _list_instance_links(client: ExoClient) -> list[dict[str, Any]]:
out = client.request_json("GET", "/v1/instance-links")
return out if isinstance(out, list) else []
def _delete_instance_link(client: ExoClient, link_id: str) -> None:
client.request_json("DELETE", f"/v1/instance-links/{link_id}")
def run_one(
client: ExoClient,
model_id: str,
pp_hint: int,
tg: int,
prompt_sizer: PromptSizer,
) -> tuple[dict[str, Any], int]:
content, pp_tokens = prompt_sizer.build(pp_hint)
payload: dict[str, Any] = {
"model": model_id,
"messages": [{"role": "user", "content": content}],
"stream": False,
"max_tokens": tg,
}
t0 = time.perf_counter()
out = client.post_bench_chat_completions(payload)
elapsed = time.perf_counter() - t0
stats = out.get("generation_stats")
choices = out.get("choices") or [{}]
message = choices[0].get("message", {}) if choices else {}
text = message.get("content") or ""
preview = text[:200] if text else ""
return {
"elapsed_s": elapsed,
"output_text_preview": preview,
"stats": stats,
}, pp_tokens
def _run_phase(
*,
client: ExoClient,
label: str,
pp_tg_pairs: list[tuple[int, int]],
model_id: str,
prompt_sizer: PromptSizer,
warmup: int,
repeat: int,
common_meta: dict[str, Any],
sampler: SystemMetricsSampler | None = None,
) -> list[dict[str, Any]]:
logger.info(f"=== phase: {label} (model={model_id}) ===")
rows: list[dict[str, Any]] = []
for i in range(warmup):
run_one(client, model_id, pp_tg_pairs[0][0], pp_tg_pairs[0][1], prompt_sizer)
logger.debug(f" warmup {i + 1}/{warmup} done")
for pp, tg in pp_tg_pairs:
logger.info(f"--- {label}: pp={pp} tg={tg} ---")
runs: list[dict[str, Any]] = []
inference_windows: list[tuple[float, float]] = []
for r in range(repeat):
time.sleep(2)
try:
inf_t0 = time.monotonic()
row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
inference_windows.append((inf_t0, time.monotonic()))
except Exception as e:
logger.error(e)
continue
row.update(common_meta)
row.update(
{
"phase": label,
"phase_model_id": model_id,
"pp_tokens": actual_pp_tokens,
"tg": tg,
"repeat_index": r,
}
)
runs.append(row)
rows.append(row)
if runs:
prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
gen_tps = mean(x["stats"]["generation_tps"] for x in runs)
ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
gtok = mean(x["stats"]["generation_tokens"] for x in runs)
peak = mean(x["stats"]["peak_memory_usage"]["inBytes"] for x in runs)
avg_elapsed = mean(x["elapsed_s"] for x in runs)
energy_str = ""
if sampler is not None and inference_windows:
joules = sum(
sampler.energy_between(t0, t1) for t0, t1 in inference_windows
)
inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
avg_watts = joules / inf_seconds if inf_seconds > 0 else 0.0
energy_per_run = joules / len(runs) if runs else 0.0
energy_str = (
f" energy={joules:.1f}J ({avg_watts:.1f}W avg over "
f"{inf_seconds:.1f}s inference, {energy_per_run:.1f}J/run)"
)
for run_row, (t0, t1) in zip(runs, inference_windows, strict=False):
run_row["energy_joules"] = sampler.energy_between(t0, t1)
run_row["inference_window_s"] = t1 - t0
logger.info(
f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} "
f"prompt_tokens={ptok} gen_tokens={gtok} "
f"peak_memory={format_peak_memory(peak)} "
f"avg_elapsed={avg_elapsed:.2f}s{energy_str}"
)
time.sleep(2)
return rows
def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, float]]:
grouped: dict[tuple[int, int], list[dict[str, Any]]] = {}
for r in rows:
key = (int(r["pp_tokens"]), int(r["tg"]))
grouped.setdefault(key, []).append(r)
out: dict[tuple[int, int], dict[str, float]] = {}
for key, runs in grouped.items():
energy_runs = [x.get("energy_joules") for x in runs if "energy_joules" in x]
window_runs = [
x.get("inference_window_s") for x in runs if "inference_window_s" in x
]
out[key] = {
"prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
"gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
"elapsed_s": mean(x["elapsed_s"] for x in runs),
"prompt_tokens": mean(x["stats"]["prompt_tokens"] for x in runs),
"gen_tokens": mean(x["stats"]["generation_tokens"] for x in runs),
"energy_j": mean(energy_runs) if energy_runs else 0.0,
"inference_window_s": mean(window_runs) if window_runs else 0.0,
}
return out
def _normalised_seconds(summary: dict[str, float], pp: int, tg: int) -> float | None:
"""Wall-clock time implied by reported tps for the *configured* pp/tg.
elapsed_s is not comparable across phases when models EOS at different
lengths. This formula reconstructs "what would this phase take to do
pp prompt tokens + tg generation tokens" using its own reported rates.
"""
p_tps = summary.get("prompt_tps", 0.0)
g_tps = summary.get("gen_tps", 0.0)
if p_tps <= 0 or g_tps <= 0:
return None
return pp / p_tps + tg / g_tps
def _print_diff(
disagg_rows: list[dict[str, Any]],
decode_alone_rows: list[dict[str, Any]],
prefill_alone_rows: list[dict[str, Any]],
) -> None:
disagg = _summarise(disagg_rows)
decode_alone = _summarise(decode_alone_rows)
prefill_alone = _summarise(prefill_alone_rows)
keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
width = 110
for key in sorted(keys):
pp, tg = key
logger.info("" * width)
logger.info(f" pp={pp} tg={tg}")
logger.info("" * width)
logger.info(
f" {'phase':<16} {'elapsed':>9} {'norm':>9} "
f"{'prompt_tps':>11} {'gen_tps':>8} "
f"{'p_tok':>6} {'g_tok':>6} "
f"{'energy':>9} {'avg_W':>7}"
)
for label, summary in (
("disaggregated", disagg.get(key)),
("decode_alone", decode_alone.get(key)),
("prefill_alone", prefill_alone.get(key)),
):
if summary is None:
logger.info(
f" {label:<16} {'':>9} {'':>9} "
f"{'':>11} {'':>8} {'':>6} {'':>6} "
f"{'':>9} {'':>7}"
)
continue
norm = _normalised_seconds(summary, pp, tg)
norm_str = f"{norm:>8.2f}s" if norm is not None else f"{'':>9}"
energy = summary.get("energy_j", 0.0)
window = summary.get("inference_window_s", 0.0)
energy_str = f"{energy:>8.1f}J" if energy > 0 else f"{'':>9}"
avg_w = energy / window if window > 0 else 0.0
avg_w_str = f"{avg_w:>6.1f}W" if avg_w > 0 else f"{'':>7}"
logger.info(
f" {label:<16} "
f"{summary['elapsed_s']:>8.2f}s "
f"{norm_str} "
f"{summary['prompt_tps']:>11.1f} "
f"{summary['gen_tps']:>8.2f} "
f"{summary['prompt_tokens']:>6.0f} "
f"{summary['gen_tokens']:>6.0f} "
f"{energy_str} "
f"{avg_w_str}"
)
d = disagg.get(key)
da = decode_alone.get(key)
pa = prefill_alone.get(key)
d_norm = _normalised_seconds(d, pp, tg) if d else None
if d_norm and da:
da_norm = _normalised_seconds(da, pp, tg)
if da_norm:
logger.info(
f" norm speedup vs decode_alone: {da_norm / d_norm:.2f}x "
f"(prefill {d['prompt_tps'] / da['prompt_tps']:.2f}x, "
f"decode {d['gen_tps'] / da['gen_tps']:.2f}x)"
)
if d_norm and pa:
pa_norm = _normalised_seconds(pa, pp, tg)
if pa_norm:
logger.info(
f" norm speedup vs prefill_alone: {pa_norm / d_norm:.2f}x "
f"(prefill {d['prompt_tps'] / pa['prompt_tps']:.2f}x, "
f"decode {d['gen_tps'] / pa['gen_tps']:.2f}x)"
)
logger.info("" * width)
def main() -> int:
_inject_toml_into_argv()
ap = argparse.ArgumentParser(
prog="prefill-decode-bench",
description="Benchmark MLX-MLX disaggregated prefill/decode via instance links.",
)
add_common_instance_args(ap)
ap.add_argument(
"--pp",
nargs="+",
required=True,
help="Prompt-size hints (ints, must be >1000). Accepts commas.",
)
ap.add_argument(
"--tg",
nargs="+",
required=True,
help="Generation lengths (ints). Accepts commas.",
)
ap.add_argument(
"--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair."
)
ap.add_argument(
"--warmup",
type=int,
default=0,
help="Warmup runs (uses first pp/tg).",
)
ap.add_argument(
"--json-out",
default="bench/prefill_decode_results.json",
help="Write raw per-run results JSON to this path.",
)
ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
ap.add_argument(
"--dry-run", action="store_true", help="List selected placements and exit."
)
ap.add_argument(
"--all-combinations",
action="store_true",
help="Force all pp×tg combinations even when lists have equal length.",
)
ap.add_argument(
"--prefill-model",
default=None,
help="Model id for the prefill instance. Defaults to --model.",
)
ap.add_argument(
"--prefill-node",
default=None,
help="friendly_name of the node hosting the prefill instance.",
)
ap.add_argument(
"--decode-node",
default=None,
help="friendly_name of the node hosting the decode instance.",
)
ap.add_argument(
"--config",
default=None,
help="TOML config file. CLI flags override toml values.",
)
ap.add_argument(
"--compare-baseline",
action="store_true",
help="Also run each (pp,tg) pair without the prefill/decode link "
"(decode instance does its own prefill) and report the diff.",
)
args = ap.parse_args()
cfg = _load_toml(args.config) if args.config else {}
_merge_toml_into_args(args, cfg)
prefill_overrides = cfg.get("prefill", {}) if cfg else {}
decode_overrides = cfg.get("decode", {}) if cfg else {}
if args.prefill_model is None and "model" in prefill_overrides:
args.prefill_model = prefill_overrides["model"]
if args.prefill_node is None and "node" in prefill_overrides:
args.prefill_node = prefill_overrides["node"]
if args.decode_node is None and "node" in decode_overrides:
args.decode_node = decode_overrides["node"]
if "model" in decode_overrides and not args.model:
args.model = decode_overrides["model"]
pp_list = parse_int_list(args.pp)
tg_list = parse_int_list(args.tg)
if not pp_list or not tg_list:
logger.error("pp and tg lists must be non-empty")
return 2
for pp in pp_list:
if pp <= 1000:
logger.error(
f"pp={pp} must be >1000 (remote prefill triggers when uncached >1000)"
)
return 2
if args.repeat <= 0:
logger.error("--repeat must be >= 1")
return 2
use_combinations = args.all_combinations or len(pp_list) != len(tg_list)
if use_combinations:
logger.info(
f"pp/tg mode: combinations (product) — {len(pp_list) * len(tg_list)} pairs"
)
else:
logger.info(f"pp/tg mode: tandem (zip) — {len(pp_list)} pairs")
client = ExoClient(args.host, args.port, timeout_s=args.timeout)
decode_short_id, decode_full_id = resolve_model_short_id(
client, args.model, force_download=args.force_download
)
if args.prefill_model:
prefill_short_id, prefill_full_id = resolve_model_short_id(
client, args.prefill_model, force_download=args.force_download
)
else:
prefill_short_id, prefill_full_id = decode_short_id, decode_full_id
tokenizer = load_tokenizer_for_bench(decode_full_id)
if tokenizer is None:
raise RuntimeError("[prefill-decode-bench] decode tokenizer load failed")
try:
decode_prompt_sizer = PromptSizer(tokenizer)
except Exception:
logger.error("[prefill-decode-bench] decode prompt sizing failed")
raise
if prefill_full_id == decode_full_id:
prefill_prompt_sizer = decode_prompt_sizer
else:
prefill_tokenizer = load_tokenizer_for_bench(prefill_full_id)
if prefill_tokenizer is None:
raise RuntimeError("[prefill-decode-bench] prefill tokenizer load failed")
prefill_prompt_sizer = PromptSizer(prefill_tokenizer)
id_to_friendly = _node_id_to_friendly(client)
prefill_args = _side_args(args, prefill_overrides)
decode_args = _side_args(args, decode_overrides)
if prefill_full_id == decode_full_id and prefill_overrides == decode_overrides:
placements = settle_and_fetch_placements(
client, decode_full_id, args, settle_timeout=args.settle_timeout
)
prefill_candidates = (
_filter_by_node(placements, args.prefill_node, id_to_friendly)
if args.prefill_node
else placements
)
decode_candidates = (
_filter_by_node(placements, args.decode_node, id_to_friendly)
if args.decode_node
else placements
)
if args.prefill_node and not prefill_candidates:
logger.error(f"No placement on prefill node {args.prefill_node!r}.")
return 1
if args.decode_node and not decode_candidates:
logger.error(f"No placement on decode node {args.decode_node!r}.")
return 1
if args.prefill_node and args.decode_node:
prefill_p = prefill_candidates[0]
decode_p = decode_candidates[0]
else:
pair = _pick_two_distinct_placements(placements)
if pair is None:
logger.error(
"Need at least two distinct-node MLX placements for the same model."
)
return 1
prefill_p, decode_p = pair
if args.prefill_node:
prefill_p = prefill_candidates[0]
if args.decode_node:
decode_p = decode_candidates[0]
else:
prefill_node_id = (
_node_id_by_friendly(id_to_friendly, args.prefill_node)
if args.prefill_node
else None
)
decode_node_id = (
_node_id_by_friendly(id_to_friendly, args.decode_node)
if args.decode_node
else None
)
if args.prefill_node and prefill_node_id is None:
logger.error(f"Unknown node {args.prefill_node!r}.")
return 1
if args.decode_node and decode_node_id is None:
logger.error(f"Unknown node {args.decode_node!r}.")
return 1
prefill_placements = settle_and_fetch_placements(
client,
prefill_full_id,
prefill_args,
settle_timeout=args.settle_timeout,
node_id=prefill_node_id,
)
decode_placements = settle_and_fetch_placements(
client,
decode_full_id,
decode_args,
settle_timeout=args.settle_timeout,
node_id=decode_node_id,
)
if not prefill_placements:
logger.error(
f"No placement found for prefill model {prefill_full_id}"
f"{f' on node {args.prefill_node!r}' if args.prefill_node else ''}."
)
return 1
if not decode_placements:
logger.error(
f"No placement found for decode model {decode_full_id}"
f"{f' on node {args.decode_node!r}' if args.decode_node else ''}."
)
return 1
prefill_p = prefill_placements[0]
decode_p = decode_placements[0]
prefill_node_names = _placement_node_friendly_names(prefill_p, id_to_friendly)
decode_node_names = _placement_node_friendly_names(decode_p, id_to_friendly)
_ = unwrap_instance
prefill_instance = prefill_p["instance"]
decode_instance = decode_p["instance"]
prefill_id = instance_id_from_instance(prefill_instance)
decode_id = instance_id_from_instance(decode_instance)
prefill_meta = str(prefill_p.get("instance_meta", ""))
decode_meta = str(decode_p.get("instance_meta", ""))
prefill_nodes = nodes_used_in_instance(prefill_instance)
decode_nodes = nodes_used_in_instance(decode_instance)
logger.info("=" * 80)
logger.info(
f"PREFILL: {prefill_meta} / nodes={prefill_nodes} ({','.join(prefill_node_names)}) "
f"/ {prefill_short_id} ({prefill_full_id}) / instance_id={prefill_id}"
)
logger.info(
f"DECODE: {decode_meta} / nodes={decode_nodes} ({','.join(decode_node_names)}) "
f"/ {decode_short_id} ({decode_full_id}) / instance_id={decode_id}"
)
if args.dry_run:
return 0
settle_deadline = (
time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
)
logger.info("Planning phase: prefill...")
run_planning_phase(
client,
prefill_full_id,
prefill_p,
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
logger.info("Planning phase: decode...")
run_planning_phase(
client,
decode_full_id,
decode_p,
args.danger_delete_downloads,
args.timeout,
settle_deadline,
)
if use_combinations:
pp_tg_pairs = list(itertools.product(pp_list, tg_list))
else:
pp_tg_pairs = list(zip(pp_list, tg_list, strict=True))
common_meta = {
"decode_model_short_id": decode_short_id,
"decode_model_id": decode_full_id,
"prefill_model_short_id": prefill_short_id,
"prefill_model_id": prefill_full_id,
"prefill_instance_id": prefill_id,
"prefill_instance_meta": prefill_meta,
"prefill_nodes": prefill_nodes,
"decode_instance_id": decode_id,
"decode_instance_meta": decode_meta,
"decode_nodes": decode_nodes,
}
all_rows: list[dict[str, Any]] = []
disagg_rows: list[dict[str, Any]] = []
decode_alone_rows: list[dict[str, Any]] = []
prefill_alone_rows: list[dict[str, Any]] = []
link_id = ""
prefill_alive = False
decode_alive = False
sampler_nodes = sorted(
{
*node_ids_from_instance(prefill_instance),
*node_ids_from_instance(decode_instance),
}
)
sampler = SystemMetricsSampler(
ExoClient(args.host, args.port, timeout_s=30), sampler_nodes
)
sampler.start()
try:
logger.info("Creating prefill instance...")
client.request_json("POST", "/instance", body={"instance": prefill_instance})
wait_for_instance_ready(client, prefill_id)
prefill_alive = True
logger.info("Prefill instance ready")
if args.compare_baseline:
time.sleep(2)
prefill_alone_rows = _run_phase(
client=client,
label="prefill_alone",
pp_tg_pairs=pp_tg_pairs,
model_id=prefill_full_id,
prompt_sizer=prefill_prompt_sizer,
warmup=args.warmup,
repeat=args.repeat,
common_meta=common_meta,
sampler=sampler,
)
all_rows.extend(prefill_alone_rows)
logger.info("Creating decode instance...")
client.request_json("POST", "/instance", body={"instance": decode_instance})
wait_for_instance_ready(client, decode_id)
decode_alive = True
logger.info("Decode instance ready")
logger.info("Linking instances (prefill → decode)...")
_create_instance_link(client, prefill_id, decode_id)
time.sleep(1)
links = _list_instance_links(client)
if not links:
logger.error("Link did not appear in state.")
return 1
link_id = str(links[-1].get("linkId") or links[-1].get("link_id") or "")
logger.info(f"Link created: {link_id}")
time.sleep(2)
disagg_rows = _run_phase(
client=client,
label="disaggregated",
pp_tg_pairs=pp_tg_pairs,
model_id=decode_full_id,
prompt_sizer=decode_prompt_sizer,
warmup=args.warmup,
repeat=args.repeat,
common_meta=common_meta,
sampler=sampler,
)
all_rows.extend(disagg_rows)
if args.compare_baseline:
logger.info("Removing link and prefill instance to isolate decode_alone.")
with contextlib.suppress(ExoHttpError):
if link_id:
_delete_instance_link(client, link_id)
link_id = ""
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{prefill_id}")
wait_for_instance_gone(client, prefill_id)
prefill_alive = False
time.sleep(2)
decode_alone_rows = _run_phase(
client=client,
label="decode_alone",
pp_tg_pairs=pp_tg_pairs,
model_id=decode_full_id,
prompt_sizer=decode_prompt_sizer,
warmup=args.warmup,
repeat=args.repeat,
common_meta=common_meta,
sampler=sampler,
)
all_rows.extend(decode_alone_rows)
_print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
finally:
sampler.stop()
with contextlib.suppress(ExoHttpError):
if link_id:
_delete_instance_link(client, link_id)
if decode_alive:
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{decode_id}")
wait_for_instance_gone(client, decode_id)
if prefill_alive:
with contextlib.suppress(ExoHttpError):
client.request_json("DELETE", f"/instance/{prefill_id}")
wait_for_instance_gone(client, prefill_id)
logger.debug("Deleted both instances")
if args.stdout:
json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
elif args.json_out:
with open(args.json_out, "w", encoding="utf-8") as f:
json.dump(all_rows, f, indent=2, ensure_ascii=False)
logger.debug(f"\nWrote results JSON: {args.json_out}")
return 0
if __name__ == "__main__":
sys.exit(main())
+377
View File
@@ -0,0 +1,377 @@
# 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)
File renamed without changes.
+24
View File
@@ -0,0 +1,24 @@
#!/bin/bash
set -e
export PATH="/opt/homebrew/bin:$PATH"
echo "=== Starting overnight bench runs at $(date) ==="
echo "--- [4/8] Qwen3.5-122B-A10B-GPTQ-Int4 ---"
echo "Skipping because Int 4"
#uv run bench/exo_bench.py --force-download --model "Qwen/Qwen3.5-122B-A10B-GPTQ-Int4" --pp 700 --tg 36000 --repeat 1
echo "--- [5/8] Qwen3.5-27B-FP8 ---"
#uv run bench/exo_bench.py --force-download --model "Qwen/Qwen3.5-27B-FP8" --pp 700 --tg 35133 --repeat 1
echo "--- [6/8] GLM-4.7-Flash-bf16 ---"
uv run bench/exo_bench.py --force-download --model "mlx-community/GLM-4.7-Flash-bf16" --pp 700 --tg 29000 --repeat 1
echo "--- [7/8] NVIDIA-Nemotron-3-Nano-30B-A3B (23000,1200) ---"
uv run bench/exo_bench.py --force-download --model "mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-MLX-BF16" --pp 700 --tg 23000,1200 --repeat 1
echo "--- [8/8] Qwen3.5-27B-bf16 ---"
uv run bench/exo_bench.py --force-download --model "mlx-community/Qwen3.5-27B-bf16" --pp 700 --tg 35400 --repeat 1
echo "=== All bench runs complete at $(date) ==="
@@ -9,6 +9,10 @@
quantization: string;
}
function normalizeBaseModel(s: string): string {
return s.toLowerCase().replace(/[-_]/g, " ").trim();
}
// Auto mode tier list (for when user just starts typing)
export const AUTO_TIERS: string[][] = [
// Tier 1 (frontier)
@@ -43,8 +47,9 @@
/** Return the tier index (0 = best) for a base_model name. */
export function getAutoTierIndex(baseModel: string): number {
const norm = normalizeBaseModel(baseModel);
for (let i = 0; i < AUTO_TIERS.length; i++) {
if (AUTO_TIERS[i].includes(baseModel)) return i;
if (AUTO_TIERS[i].some((t) => normalizeBaseModel(t) === norm)) return i;
}
return AUTO_TIERS.length; // not in any tier → lowest priority
}
@@ -60,7 +65,7 @@
const variants = modelList
.filter(
(m) =>
m.base_model === baseModel &&
normalizeBaseModel(m.base_model) === normalizeBaseModel(baseModel) &&
(m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
(m.storage_size_megabytes || 0) > 0,
)
@@ -162,7 +167,7 @@
/** For a given base_model name, find the biggest quant variant that fits in memory. */
function pickBestVariant(baseModel: string): ChatModelInfo | null {
const variants = models
.filter((m) => m.base_model === baseModel && fitsInMemory(m))
.filter((m) => normalizeBaseModel(m.base_model) === normalizeBaseModel(baseModel) && fitsInMemory(m))
.sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
return variants[0] ?? null;
}
@@ -202,7 +202,6 @@
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
else if (instanceTag === "VllmInstance") instanceType = "vLLM";
let sharding: string | null = null;
const inst = instance as {
@@ -1,8 +1,5 @@
<script lang="ts">
import { browser } from "$app/environment";
import { featureFlags } from "$lib/stores/app.svelte";
const showAdvanced = $derived(featureFlags()["disaggregation"] === true);
interface Props {
showHome?: boolean;
@@ -300,28 +297,5 @@
</svg>
<span class="hidden sm:inline">Integrations</span>
</a>
{#if showAdvanced}
<a
href="/#/advanced"
class="text-xs md:text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-1.5 md:gap-2 cursor-pointer"
title="Advanced cluster settings"
>
<svg
class="w-4 h-4"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2"
stroke-linecap="round"
stroke-linejoin="round"
>
<circle cx="12" cy="12" r="3" />
<path
d="M19.4 15a1.65 1.65 0 0 0 .33 1.82l.06.06a2 2 0 0 1 0 2.83 2 2 0 0 1-2.83 0l-.06-.06a1.65 1.65 0 0 0-1.82-.33 1.65 1.65 0 0 0-1 1.51V21a2 2 0 0 1-4 0v-.09A1.65 1.65 0 0 0 9 19.4a1.65 1.65 0 0 0-1.82.33l-.06.06a2 2 0 0 1-2.83 0 2 2 0 0 1 0-2.83l.06-.06a1.65 1.65 0 0 0 .33-1.82 1.65 1.65 0 0 0-1.51-1H3a2 2 0 0 1 0-4h.09A1.65 1.65 0 0 0 4.6 9a1.65 1.65 0 0 0-.33-1.82l-.06-.06a2 2 0 0 1 0-2.83 2 2 0 0 1 2.83 0l.06.06a1.65 1.65 0 0 0 1.82.33H9a1.65 1.65 0 0 0 1-1.51V3a2 2 0 0 1 4 0v.09a1.65 1.65 0 0 0 1 1.51 1.65 1.65 0 0 0 1.82-.33l.06-.06a2 2 0 0 1 2.83 0 2 2 0 0 1 0 2.83l-.06.06a1.65 1.65 0 0 0-.33 1.82V9a1.65 1.65 0 0 0 1.51 1H21a2 2 0 0 1 0 4h-.09a1.65 1.65 0 0 0-1.51 1z"
/>
</svg>
<span class="hidden sm:inline">Advanced</span>
</a>
{/if}
</nav>
</header>
@@ -1,6 +1,4 @@
<script lang="ts">
import { copyText } from "$lib/utils/clipboard";
interface Props {
title: string;
subtitle: string;
@@ -18,17 +16,11 @@
}: Props = $props();
let copied = $state(false);
let failed = $state(false);
async function copyToClipboard() {
const ok = await copyText(config);
if (ok) {
copied = true;
setTimeout(() => (copied = false), 2000);
} else {
failed = true;
setTimeout(() => (failed = false), 2000);
}
await navigator.clipboard.writeText(config);
copied = true;
setTimeout(() => (copied = false), 2000);
}
</script>
@@ -45,11 +37,9 @@
class="px-3 py-1.5 text-xs rounded border transition-all duration-200 cursor-pointer
{copied
? 'border-green-500/50 text-green-400 bg-green-500/10'
: failed
? 'border-red-500/50 text-red-400 bg-red-500/10'
: 'border-exo-light-gray/30 text-exo-light-gray hover:border-exo-yellow/50 hover:text-exo-yellow'}"
: 'border-exo-light-gray/30 text-exo-light-gray hover:border-exo-yellow/50 hover:text-exo-yellow'}"
>
{copied ? "Copied!" : failed ? "Copy failed" : "Copy"}
{copied ? "Copied!" : "Copy"}
</button>
</div>
{#if description}
+60 -8
View File
@@ -23,7 +23,7 @@
} | null;
nodes?: Record<string, NodeInfo>;
sharding?: "Pipeline" | "Tensor";
runtime?: "MlxRing" | "MlxJaccl" | "Vllm";
runtime?: "MlxRing" | "MlxJaccl";
onLaunch?: () => void;
tags?: string[];
apiPreview?: PlacementPreview | null;
@@ -279,7 +279,7 @@
let placementNodes: Array<{
id: string;
deviceName: string;
deviceType: "macbook" | "studio" | "mini" | "unknown";
deviceType: "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown";
totalGB: number;
currentUsedGB: number;
modelUsageGB: number;
@@ -578,17 +578,13 @@
class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
title={runtime === "MlxRing"
? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
: runtime === "MlxJaccl"
? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
: "vLLM: NVIDIA CUDA inference engine."}
: "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
>
{runtime === "MlxRing"
? "MLX Ring"
: runtime === "MlxJaccl"
? "MLX RDMA"
: runtime === "Vllm"
? "vLLM"
: runtime}
: runtime}
</span>
</div>
@@ -1071,6 +1067,62 @@
/>
{/if}
</g>
{:else if node.deviceType === "linux"}
<!-- Linux Tux penguin icon -->
{@const sz = node.iconSize}
{@const sc = sz / 100}
<g transform="translate({-sz / 2}, {-sz / 2})">
<!-- Body -->
<path
d="M50 8c-8 0-14 6-14 13 0 4 2 8 5 10-8 4-16 14-16 28v12c0 4 2 7 5 9l-6 4c-2 1-3 3-3 5v3c0 2 2 4 4 4h10l4-4h22l4 4h10c2 0 4-2 4-4v-3c0-2-1-4-3-5l-6-4c3-2 5-5 5-9V59c0-14-8-24-16-28 3-2 5-6 5-10 0-7-6-13-14-13z"
transform="scale({sc})"
fill="#1a1a1a"
stroke={node.isUsed ? "#FFD700" : "#4B5563"}
stroke-width={1.5 / sc}
/>
<!-- Belly -->
<path
d="M38 52c0-8 5-15 12-15s12 7 12 15v14c0 4-5 7-12 7s-12-3-12-7V52z"
transform="scale({sc})"
fill="rgba(220,220,220,0.85)"
/>
<!-- Eyes -->
<circle cx={44 * sc} cy={16 * sc} r={2.5 * sc} fill="white" />
<circle cx={56 * sc} cy={16 * sc} r={2.5 * sc} fill="white" />
<circle
cx={44 * sc}
cy={16 * sc}
r={1.2 * sc}
fill="#1a1a1a"
/>
<circle
cx={56 * sc}
cy={16 * sc}
r={1.2 * sc}
fill="#1a1a1a"
/>
<!-- Beak -->
<path
d="M{46 * sc} {22 * sc} L{50 * sc} {27 * sc} L{54 *
sc} {22 * sc} Z"
fill="#E8A317"
/>
<!-- Feet -->
<ellipse
cx={42 * sc}
cy={94 * sc}
rx={6 * sc}
ry={2.5 * sc}
fill="#E8A317"
/>
<ellipse
cx={58 * sc}
cy={94 * sc}
rx={6 * sc}
ry={2.5 * sc}
fill="#E8A317"
/>
</g>
{:else}
<!-- Unknown device - hexagon -->
<g
@@ -9,7 +9,6 @@
capabilities?: string[];
family?: string;
is_custom?: boolean;
requires_vllm?: boolean;
}
interface ModelGroup {
@@ -20,7 +19,6 @@
variants: ModelInfo[];
smallestVariant: ModelInfo;
hasMultipleVariants: boolean;
requiresVllm: boolean;
}
type DownloadAvailability = {
@@ -215,14 +213,6 @@
<span class="font-mono text-sm text-white truncate">
{group.name}
</span>
{#if group.requiresVllm}
<span
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
title="Requires vLLM runtime"
>
vLLM
</span>
{/if}
<!-- Capability icons -->
{#each group.capabilities.filter((c) => c !== "text") as cap}
{#if cap === "thinking"}
@@ -533,15 +523,6 @@
{variant.quantization || "default"}
</span>
{#if variant.requires_vllm}
<span
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
title="Requires vLLM runtime"
>
vLLM
</span>
{/if}
<!-- Size -->
<span
class="text-xs font-mono flex-1 {getSizeClassForFitStatus(
@@ -647,7 +628,6 @@
variants: [variant],
smallestVariant: variant,
hasMultipleVariants: false,
requiresVllm: variant.requires_vllm === true,
});
}}
title="View variant details"
@@ -22,7 +22,6 @@
is_custom?: boolean;
tasks?: string[];
hugging_face_id?: string;
requires_vllm?: boolean;
}
interface ModelGroup {
@@ -33,7 +32,6 @@
variants: ModelInfo[];
smallestVariant: ModelInfo;
hasMultipleVariants: boolean;
requiresVllm: boolean;
}
interface FilterState {
@@ -381,12 +379,16 @@
return hfTrendingModels;
});
function normalizeBaseModel(s: string): string {
return s.toLowerCase().replace(/[-_]/g, " ").trim();
}
// Group models by base_model
const groupedModels = $derived.by((): ModelGroup[] => {
const groups = new Map<string, ModelGroup>();
for (const model of models) {
const groupId = model.base_model || model.id;
const groupId = normalizeBaseModel(model.base_model || model.id);
const groupName = model.base_model || model.name || model.id;
if (!groups.has(groupId)) {
@@ -398,7 +400,6 @@
variants: [],
smallestVariant: model,
hasMultipleVariants: false,
requiresVllm: true,
});
}
@@ -433,7 +434,6 @@
(a.storage_size_megabytes || 0) - (b.storage_size_megabytes || 0),
);
group.hasMultipleVariants = group.variants.length > 1;
group.requiresVllm = group.variants.every((v) => v.requires_vllm);
}
// Convert to array and sort by smallest variant size (biggest first)
@@ -584,14 +584,13 @@
const model = models.find((m) => m.id === id);
if (model) {
result.push({
id: model.base_model || model.id,
id: normalizeBaseModel(model.base_model || model.id),
name: model.name || model.id,
capabilities: model.capabilities || ["text"],
family: model.family || "",
variants: [model],
smallestVariant: model,
hasMultipleVariants: false,
requiresVllm: model.requires_vllm === true,
});
}
}
@@ -1170,17 +1169,6 @@
<span class="text-white/40">Variants:</span>
<span class="text-white/70">{infoGroup.variants.length}</span>
</div>
{#if infoGroup.requiresVllm}
<div class="flex items-center gap-2">
<span class="text-white/40">Runtime:</span>
<span
class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 tracking-wider uppercase"
>
vLLM
</span>
<span class="text-white/40 text-[11px]">required</span>
</div>
{/if}
{#if infoGroup.variants.length > 0}
<div class="mt-3 pt-3 border-t border-exo-yellow/10">
<span class="text-white/40">Available quantizations:</span>
@@ -1,565 +0,0 @@
<script lang="ts">
import { onMount, onDestroy } from "svelte";
import FamilyLogos from "$lib/components/FamilyLogos.svelte";
import {
instances,
instanceLinks,
nodeIdentities,
refreshState,
createInstanceLink,
updateInstanceLink,
deleteInstanceLink,
type Instance,
} from "$lib/stores/app.svelte";
import { deriveBaseModel, deriveFamily } from "$lib/utils/model_family";
type InstanceWrapper = {
MlxRingInstance?: Instance;
MlxJacclInstance?: Instance;
VllmInstance?: Instance;
};
let interval: ReturnType<typeof setInterval> | null = null;
onMount(() => {
refreshState();
interval = setInterval(refreshState, 3000);
});
onDestroy(() => {
if (interval) clearInterval(interval);
});
type InstanceRow = {
id: string;
modelId: string;
family: string;
baseModel: string;
nodeNames: string[];
nodeCount: number;
};
const instanceRows = $derived.by<InstanceRow[]>(() => {
const rows: InstanceRow[] = [];
const ids = nodeIdentities();
for (const [id, raw] of Object.entries(instances())) {
const wrapper = raw as InstanceWrapper;
const inst =
wrapper.MlxRingInstance ??
wrapper.MlxJacclInstance ??
wrapper.VllmInstance;
const modelId = inst?.shardAssignments?.modelId ?? "";
const nodeToRunner = inst?.shardAssignments?.nodeToRunner ?? {};
const nodeIds = Object.keys(nodeToRunner);
const nodeNames = nodeIds
.map((nodeId) => ids[nodeId]?.friendlyName ?? nodeId.slice(0, 6))
.filter((name) => !!name);
rows.push({
id,
modelId,
family: deriveFamily(modelId),
baseModel: deriveBaseModel(modelId),
nodeNames,
nodeCount: nodeIds.length,
});
}
rows.sort((a, b) => a.modelId.localeCompare(b.modelId));
return rows;
});
const instanceById = $derived(
Object.fromEntries(instanceRows.map((r) => [r.id, r])),
);
type LinkRow = {
linkId: string;
prefill: string[];
decode: string[];
families: string[];
multiNode: boolean;
};
const linkRows = $derived.by<LinkRow[]>(() => {
const rows: LinkRow[] = [];
for (const [, link] of Object.entries(instanceLinks())) {
const fams = new Set<string>();
let multiNode = false;
for (const id of [...link.prefillInstances, ...link.decodeInstances]) {
const r = instanceById[id];
if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
if (r && r.nodeCount > 1) multiNode = true;
}
rows.push({
linkId: link.linkId,
prefill: link.prefillInstances,
decode: link.decodeInstances,
families: Array.from(fams),
multiNode,
});
}
return rows;
});
let editingLinkId = $state<string | null>(null);
let editingPrefill = $state<Set<string>>(new Set());
let editingDecode = $state<Set<string>>(new Set());
let saving = $state(false);
let errorMessage = $state<string | null>(null);
function startCreate() {
editingLinkId = "new";
editingPrefill = new Set();
editingDecode = new Set();
errorMessage = null;
}
function startEdit(row: LinkRow) {
editingLinkId = row.linkId;
editingPrefill = new Set(row.prefill);
editingDecode = new Set(row.decode);
errorMessage = null;
}
function cancelEdit() {
editingLinkId = null;
editingPrefill = new Set();
editingDecode = new Set();
errorMessage = null;
}
type Role = "prefill" | "decode" | "none";
function roleOf(id: string): Role {
if (editingPrefill.has(id)) return "prefill";
if (editingDecode.has(id)) return "decode";
return "none";
}
function setRole(id: string, role: Role) {
const p = new Set(editingPrefill);
const d = new Set(editingDecode);
p.delete(id);
d.delete(id);
if (role === "prefill") p.add(id);
if (role === "decode") d.add(id);
editingPrefill = p;
editingDecode = d;
}
const editingFamilies = $derived.by<string[]>(() => {
const fams = new Set<string>();
for (const id of [...editingPrefill, ...editingDecode]) {
const r = instanceById[id];
if (r && r.baseModel) fams.add(r.baseModel.toLowerCase());
}
return Array.from(fams);
});
const editingMultiNode = $derived.by<string[]>(() => {
const names: string[] = [];
for (const id of [...editingPrefill, ...editingDecode]) {
const r = instanceById[id];
if (r && r.nodeCount > 1) {
names.push(r.baseModel || r.modelId);
}
}
return names;
});
const editingMismatch = $derived(editingFamilies.length > 1);
const canSave = $derived(
editingLinkId !== null &&
editingPrefill.size > 0 &&
editingDecode.size > 0 &&
!saving,
);
async function save() {
if (editingLinkId === null) return;
saving = true;
errorMessage = null;
try {
const prefill = Array.from(editingPrefill);
const decode = Array.from(editingDecode);
if (editingLinkId === "new") {
await createInstanceLink(prefill, decode);
} else {
await updateInstanceLink(editingLinkId, prefill, decode);
}
cancelEdit();
await refreshState();
} catch (err) {
errorMessage = err instanceof Error ? err.message : String(err);
} finally {
saving = false;
}
}
async function remove(linkId: string) {
if (!confirm("Remove this routing?")) return;
try {
await deleteInstanceLink(linkId);
if (editingLinkId === linkId) cancelEdit();
await refreshState();
} catch (err) {
errorMessage = err instanceof Error ? err.message : String(err);
}
}
</script>
<div class="font-mono text-foreground">
<div class="mb-6 space-y-4">
<details open class="group [&_summary::-webkit-details-marker]:hidden">
<summary
class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
>
<span
class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
>▶</span
>
Prefill vs Decode
</summary>
<div class="mt-2 text-white/80 text-sm leading-relaxed">
Prefill is the compute-bound pass that consumes the entire prompt and
builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
tokens sequentially from that cache. The two phases have very different
bottlenecks, so running them on different hardware can be substantially
faster than doing both on one node.
</div>
</details>
<details class="group [&_summary::-webkit-details-marker]:hidden">
<summary
class="cursor-pointer list-none text-exo-yellow text-xs font-mono tracking-widest uppercase flex items-center gap-2 hover:opacity-80 transition-opacity"
>
<span
class="inline-block transition-transform group-open:rotate-90 text-exo-light-gray"
>▶</span
>
Linking Instances
</summary>
<div class="mt-2 text-white/80 text-sm leading-relaxed space-y-2">
<p>
A linked route here tells the cluster: when a request is sent to a
model in that cluster, the decode node (or the least active one if
there are multiple) will handle it. If it decides it must do a lot of
prefill not already cached in the prefix cache, it routes the request
to the prefill node over TCP IP. The prefill node streams the KV cache
back to the decode node which picks up from there.
</p>
<p>
Linked instances must be running the same model family — KV layouts
differ across architectures. More on the <a
class="text-exo-yellow underline underline-offset-2 hover:text-exo-yellow-darker transition-colors"
href="https://blog.exolabs.net/nvidia-dgx-spark/"
target="_blank"
rel="noreferrer noopener">blog</a
>.
</p>
</div>
</details>
</div>
{#if errorMessage}
<div
class="mb-4 px-4 py-3 bg-red-500/10 border border-red-500/40 text-red-300 text-sm"
>
{errorMessage}
</div>
{/if}
<section class="mt-12">
<h2
class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
>
Existing routes
</h2>
{#if linkRows.length === 0}
{#if editingLinkId === null}
<div class="flex items-center justify-between">
<p class="text-exo-light-gray italic text-sm m-0">
No routes yet. Create one to enable remote prefill.
</p>
<button
class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
onclick={startCreate}
>
+ New route
</button>
</div>
{/if}
{:else}
{#if editingLinkId === null}
<div class="flex justify-end mb-3">
<button
class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 transition-colors"
onclick={startCreate}
>
+ New route
</button>
</div>
{/if}
<div
class="bg-exo-dark-gray/60 border border-exo-medium-gray/40 flex flex-col"
>
{#each linkRows as row (row.linkId)}
{#if editingLinkId !== row.linkId}
<article
class="p-4 border-b border-exo-light-gray/25 last:border-b-0"
>
{#if row.multiNode}
<div
class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
>
⚠ Multi-node instance detected. Remote prefill currently only
works on single-node (rank-0) instances. This route will not
function until that's supported.
</div>
{/if}
{#if row.families.length > 1}
<div
class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
>
⚠ Mixed model families: {row.families.join(", ")}
</div>
{/if}
<div
class="grid grid-cols-[1fr_auto_1fr_auto] items-center gap-x-3 gap-y-2"
>
<span
class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-yellow/15 border border-exo-yellow/40 text-exo-yellow"
>Prefill</span
>
<span></span>
<span
class="inline-block justify-self-start text-[10px] font-mono tracking-widest uppercase px-2 py-0.5 bg-exo-medium-gray/40 border border-exo-medium-gray/60 text-foreground"
>Decode</span
>
<span></span>
<div class="min-w-0">
<ul class="list-none p-0 m-0 flex flex-col gap-2">
{#each row.prefill as id (id)}
{@const r = instanceById[id]}
{#if r}
<li
class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
>
<FamilyLogos family={r.family} />
<div class="min-w-0 flex-1">
<div
class="text-exo-yellow text-xs font-mono truncate"
>
{r.baseModel || r.modelId}
</div>
<div
class="text-exo-light-gray text-[11px] truncate"
>
{r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
? ` (${r.nodeCount} nodes)`
: ""}
</div>
<div
class="text-exo-light-gray/40 text-[10px] font-mono truncate"
title={r.id}
>
{r.id.slice(0, 8)}
</div>
</div>
</li>
{/if}
{/each}
</ul>
</div>
<div class="text-exo-yellow/60 text-xl px-2" aria-hidden="true">
</div>
<div class="min-w-0">
<ul class="list-none p-0 m-0 flex flex-col gap-2">
{#each row.decode as id (id)}
{@const r = instanceById[id]}
{#if r}
<li
class="flex items-center gap-2 px-2.5 py-2 bg-exo-medium-gray/20 border border-exo-medium-gray/40"
>
<FamilyLogos family={r.family} />
<div class="min-w-0 flex-1">
<div
class="text-exo-yellow text-xs font-mono truncate"
>
{r.baseModel || r.modelId}
</div>
<div
class="text-exo-light-gray text-[11px] truncate"
>
{r.nodeNames.join(", ") || "?"}{r.nodeCount > 1
? ` (${r.nodeCount} nodes)`
: ""}
</div>
<div
class="text-exo-light-gray/40 text-[10px] font-mono truncate"
title={r.id}
>
{r.id.slice(0, 8)}
</div>
</div>
</li>
{/if}
{/each}
</ul>
</div>
<div class="flex gap-2 pl-3">
<button
class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 hover:text-exo-yellow disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
onclick={() => startEdit(row)}
disabled={editingLinkId !== null}
>
Edit
</button>
<button
class="px-2 py-0.5 text-[11px] font-mono tracking-wider uppercase bg-red-500/15 border border-red-500/40 rounded text-red-300 hover:bg-red-500/25 transition-colors"
onclick={() => remove(row.linkId)}
>
Remove
</button>
</div>
</div>
</article>
{/if}
{/each}
</div>
{/if}
</section>
{#if editingLinkId !== null && instanceRows.length === 0}
<section
class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 px-4 py-2.5 flex items-center justify-between gap-3"
>
<span class="text-exo-light-gray italic text-sm font-mono"
>No instances available.</span
>
<button
class="px-3 py-1 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 rounded text-foreground hover:border-exo-yellow/60 transition-colors"
onclick={cancelEdit}
>
Cancel
</button>
</section>
{:else if editingLinkId !== null}
<section class="mt-6 bg-exo-dark-gray/60 border border-exo-yellow/30 p-5">
<h2
class="text-exo-yellow text-xs font-mono tracking-widest uppercase m-0 mb-3"
>
{editingLinkId === "new" ? "New route" : "Edit route"}
</h2>
{#if editingMismatch}
<div
class="mb-3 px-3 py-2 bg-amber-500/10 border border-amber-500/40 text-amber-300 text-xs tracking-wide"
>
⚠ Selected instances span multiple model families: <strong
>{editingFamilies.join(", ")}</strong
>. Linking across families produces a corrupt KV cache.
</div>
{/if}
{#if editingMultiNode.length > 0}
<div
class="mb-3 px-3 py-2 bg-red-500/10 border border-red-500/40 text-red-300 text-xs tracking-wide"
>
⚠ Multi-node instance(s) selected: <strong
>{editingMultiNode.join(", ")}</strong
>. Remote prefill currently only works on single-node instances. This
route will not function until multi-node support lands.
</div>
{/if}
<p class="text-exo-light-gray text-xs mb-4">
Pick a role for each instance:
<span class="text-exo-yellow">Prefill</span>
serves KV cache,
<span class="text-foreground">Decode</span> consumes it.
</p>
<div
class="grid gap-2.5"
style="grid-template-columns: repeat(auto-fill, minmax(360px, 1fr));"
>
{#each instanceRows as row (row.id)}
{@const role = roleOf(row.id)}
<div
class="border p-3 flex flex-col gap-2.5 transition-colors {role ===
'prefill'
? 'border-exo-yellow/60 bg-exo-dark-gray/60'
: role === 'decode'
? 'border-exo-light-gray/60 bg-exo-dark-gray/60'
: 'border-exo-medium-gray/40 bg-exo-dark-gray/40'}"
>
<div class="flex items-center gap-2">
<FamilyLogos family={row.family} />
<div class="min-w-0 flex-1">
<div class="text-exo-yellow text-xs font-mono truncate">
{row.baseModel || row.modelId}
</div>
<div class="text-exo-light-gray text-[11px] truncate">
{row.nodeNames.join(", ") || "?"}{row.nodeCount > 1
? ` (${row.nodeCount} nodes)`
: ""}
</div>
<div
class="text-exo-light-gray/40 text-[10px] font-mono truncate"
title={row.id}
>
{row.id.slice(0, 8)}
</div>
</div>
{#if row.nodeCount > 1}
<span
class="text-[9px] font-mono tracking-widest uppercase px-1.5 py-0.5 bg-red-500/15 border border-red-500/40 text-red-300"
title="Multi-node instances are not supported by remote prefill yet."
>Unsupported</span
>
{/if}
</div>
<div
class="flex rounded-md overflow-hidden border border-exo-light-gray/40 divide-x divide-exo-light-gray/40"
>
<button
class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
'prefill'
? 'bg-exo-yellow/20 text-exo-yellow'
: 'bg-transparent text-white/80 hover:text-exo-yellow'}"
onclick={() =>
setRole(row.id, role === "prefill" ? "none" : "prefill")}
>Prefill</button
>
<button
class="flex-1 px-2 py-1 text-[11px] font-mono tracking-wider uppercase transition-colors {role ===
'decode'
? 'bg-exo-medium-gray/50 text-foreground'
: 'bg-transparent text-white/80 hover:text-foreground'}"
onclick={() =>
setRole(row.id, role === "decode" ? "none" : "decode")}
>Decode</button
>
</div>
</div>
{/each}
</div>
<div class="flex gap-2 mt-5 justify-end">
<button
class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-yellow/15 border border-exo-yellow/50 text-exo-yellow hover:bg-exo-yellow/25 hover:border-exo-yellow/80 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
onclick={save}
disabled={!canSave}
>
{saving ? "Saving..." : "Save route"}
</button>
<button
class="px-3 py-1.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 border border-exo-medium-gray/60 text-foreground hover:border-exo-yellow/60 disabled:opacity-40 disabled:cursor-not-allowed transition-colors"
onclick={cancelEdit}
disabled={saving}
>
Cancel
</button>
</div>
</section>
{/if}
</div>
+22 -129
View File
@@ -74,12 +74,6 @@ export interface Instance {
};
}
export interface RawInstanceLink {
linkId: string;
prefillInstances: string[];
decodeInstances: string[];
}
// Granular node state types from the new state structure
interface RawNodeIdentity {
modelId?: string;
@@ -174,7 +168,7 @@ export interface ModelDownloadStatus {
export interface PlacementPreview {
model_id: string;
sharding: "Pipeline" | "Tensor";
instance_meta: "MlxRing" | "MlxJaccl";
instance_meta: "MlxRing" | "MlxJaccl" | "Vllm";
instance: unknown | null;
memory_delta_by_node: Record<string, number> | null;
error: string | null;
@@ -229,7 +223,6 @@ interface RawStateResponse {
}
>;
runners?: Record<string, unknown>;
instanceLinks?: Record<string, RawInstanceLink>;
downloads?: Record<string, unknown[]>;
// New granular node state fields
nodeIdentities?: Record<string, RawNodeIdentity>;
@@ -548,8 +541,6 @@ class AppStore {
topologyData = $state<TopologyData | null>(null);
instances = $state<Record<string, unknown>>({});
runners = $state<Record<string, unknown>>({});
instanceLinks = $state<Record<string, RawInstanceLink>>({});
featureFlags = $state<Record<string, boolean>>({});
downloads = $state<Record<string, unknown[]>>({});
nodeDisk = $state<
Record<
@@ -557,6 +548,7 @@ class AppStore {
{ total: { inBytes: number }; available: { inBytes: number } }
>
>({});
vllmAvailable = $state(false);
placementPreviews = $state<PlacementPreview[]>([]);
selectedPreviewModelId = $state<string | null>(null);
isLoadingPreviews = $state(false);
@@ -1283,7 +1275,6 @@ class AppStore {
startPolling() {
this.fetchState();
this.fetchFeatureFlags();
this.fetchInterval = setInterval(() => this.fetchState(), 1000);
}
@@ -1295,16 +1286,6 @@ class AppStore {
this.stopPreviewsPolling();
}
async fetchFeatureFlags() {
try {
const response = await fetch("/v1/feature-flags");
if (!response.ok) return;
this.featureFlags = await response.json();
} catch {
// Silently ignore — defaults to all-disabled.
}
}
async fetchState() {
try {
const response = await fetch("/state");
@@ -1330,11 +1311,6 @@ class AppStore {
if (data.runners) {
this.runners = data.runners;
}
if (data.instanceLinks) {
this.instanceLinks = data.instanceLinks;
} else {
this.instanceLinks = {};
}
if (data.downloads) {
this.downloads = data.downloads;
}
@@ -1357,6 +1333,15 @@ class AppStore {
this.isConnected = true;
}
this.consecutiveFailures = 0;
fetch("/capabilities")
.then((r) => (r.ok ? r.json() : null))
.then((data: { vllm_available?: boolean } | null) => {
this.vllmAvailable = data?.vllm_available ?? false;
})
.catch(() => {
this.vllmAvailable = false;
});
} catch (error) {
this.consecutiveFailures++;
if (
@@ -1695,15 +1680,7 @@ class AppStore {
}
}
}
const out: {
role: string;
content: string;
reasoning_content?: string;
} = { role: m.role, content: msgContent };
if (m.role === "assistant" && m.thinking) {
out.reasoning_content = m.thinking;
}
return out;
return { role: m.role, content: msgContent };
}),
];
@@ -1910,15 +1887,7 @@ class AppStore {
const apiMessages = [
systemPrompt,
...targetConversation.messages.slice(0, -1).map((m) => {
const out: {
role: string;
content: string;
reasoning_content?: string;
} = { role: m.role, content: m.content };
if (m.role === "assistant" && m.thinking) {
out.reasoning_content = m.thinking;
}
return out;
return { role: m.role, content: m.content };
}),
];
@@ -2449,15 +2418,10 @@ class AppStore {
contentParts.push({ type: "text", text: textContent });
}
const out: {
role: string;
content: typeof contentParts;
reasoning_content?: string;
} = { role: m.role, content: contentParts };
if (m.role === "assistant" && m.thinking) {
out.reasoning_content = m.thinking;
}
return out;
return {
role: m.role,
content: contentParts,
};
}
// Text-only message (original path)
@@ -2475,15 +2439,10 @@ class AppStore {
}
}
const out: {
role: string;
content: string;
reasoning_content?: string;
} = { role: m.role, content: msgContent };
if (m.role === "assistant" && m.thinking) {
out.reasoning_content = m.thinking;
}
return out;
return {
role: m.role,
content: msgContent,
};
}),
];
@@ -3332,60 +3291,6 @@ class AppStore {
}
}
async createInstanceLink(
prefillInstances: string[],
decodeInstances: string[],
): Promise<void> {
const response = await fetch("/v1/instance-links", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
prefill_instances: prefillInstances,
decode_instances: decodeInstances,
}),
});
if (!response.ok) {
throw new Error(
`Failed to create instance link: ${response.status} ${await response.text()}`,
);
}
}
async updateInstanceLink(
linkId: string,
prefillInstances: string[],
decodeInstances: string[],
): Promise<void> {
const response = await fetch(
`/v1/instance-links/${encodeURIComponent(linkId)}`,
{
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
prefill_instances: prefillInstances,
decode_instances: decodeInstances,
}),
},
);
if (!response.ok) {
throw new Error(
`Failed to update instance link: ${response.status} ${await response.text()}`,
);
}
}
async deleteInstanceLink(linkId: string): Promise<void> {
const response = await fetch(
`/v1/instance-links/${encodeURIComponent(linkId)}`,
{ method: "DELETE" },
);
if (!response.ok) {
throw new Error(
`Failed to delete instance link: ${response.status} ${await response.text()}`,
);
}
}
/**
* Delete a downloaded model from a specific node
*/
@@ -3484,19 +3389,6 @@ export const prefillProgress = () => appStore.prefillProgress;
export const topologyData = () => appStore.topologyData;
export const instances = () => appStore.instances;
export const runners = () => appStore.runners;
export const instanceLinks = () => appStore.instanceLinks;
export const featureFlags = () => appStore.featureFlags;
export const createInstanceLink = (
prefillInstances: string[],
decodeInstances: string[],
) => appStore.createInstanceLink(prefillInstances, decodeInstances);
export const updateInstanceLink = (
linkId: string,
prefillInstances: string[],
decodeInstances: string[],
) => appStore.updateInstanceLink(linkId, prefillInstances, decodeInstances);
export const deleteInstanceLink = (linkId: string) =>
appStore.deleteInstanceLink(linkId);
export const downloads = () => appStore.downloads;
export const nodeDisk = () => appStore.nodeDisk;
export const placementPreviews = () => appStore.placementPreviews;
@@ -3608,6 +3500,7 @@ export const nodeThunderbolt = () => appStore.nodeThunderbolt;
export const nodeRdmaCtl = () => appStore.nodeRdmaCtl;
export const thunderboltBridgeCycles = () => appStore.thunderboltBridgeCycles;
export const nodeThunderboltBridge = () => appStore.nodeThunderboltBridge;
export const vllmAvailable = () => appStore.vllmAvailable;
// Image generation params
export const imageGenerationParams = () => appStore.getImageGenerationParams();
-55
View File
@@ -1,55 +0,0 @@
export async function copyText(text: string): Promise<boolean> {
if (
typeof window !== "undefined" &&
window.isSecureContext &&
navigator.clipboard?.writeText
) {
try {
await navigator.clipboard.writeText(text);
return true;
} catch {
// fall through to execCommand fallback
}
}
if (typeof document === "undefined") {
return false;
}
const textarea = document.createElement("textarea");
textarea.value = text;
textarea.setAttribute("readonly", "");
textarea.style.position = "fixed";
textarea.style.top = "0";
textarea.style.left = "0";
textarea.style.width = "1px";
textarea.style.height = "1px";
textarea.style.padding = "0";
textarea.style.border = "none";
textarea.style.outline = "none";
textarea.style.boxShadow = "none";
textarea.style.background = "transparent";
textarea.style.opacity = "0";
document.body.appendChild(textarea);
const previousSelection = document.getSelection();
const previousRange =
previousSelection && previousSelection.rangeCount > 0
? previousSelection.getRangeAt(0)
: null;
try {
textarea.focus();
textarea.select();
textarea.setSelectionRange(0, text.length);
return document.execCommand("copy");
} catch {
return false;
} finally {
document.body.removeChild(textarea);
if (previousRange && previousSelection) {
previousSelection.removeAllRanges();
previousSelection.addRange(previousRange);
}
}
}
-44
View File
@@ -1,44 +0,0 @@
// Mirrors src/exo/shared/models/model_cards.py:derive_base_model
const QUANT_SUFFIXES = new RegExp(
"[-_ ](?:MLX|MXFP[0-9]+|NVFP[0-9]+|GPTQ|AWQ|GGUF|fp16|bf16|fp8|int[0-9]+|[0-9]+(?:\\.[0-9]+)?bit|Q[0-9]+(?:_[A-Z0-9]+)?|gs[0-9]+)" +
"(?:[-_ ](?:MLX|Q[0-9]+|Int[0-9]+|[A-Z0-9]+|gs[0-9]+))*$",
"i",
);
function normalize(s: string): string {
return s
.replaceAll("-", " ")
.replaceAll("_", " ")
.replaceAll(" ", " ")
.trim();
}
export function deriveBaseModel(modelId: string): string {
const short = modelId.includes("/")
? (modelId.split("/").pop() ?? modelId)
: modelId;
const stripped = short.replace(QUANT_SUFFIXES, "");
return normalize(stripped);
}
export function baseModelsCompatible(a: string, b: string): boolean {
return deriveBaseModel(a).toLowerCase() === deriveBaseModel(b).toLowerCase();
}
// Mirrors src/exo/shared/models/model_cards.py:derive_family
export function deriveFamily(modelId: string): string {
const short = modelId.includes("/")
? (modelId.split("/").pop() ?? modelId)
: modelId;
const stripped = short
.replace(QUANT_SUFFIXES, "")
.toLowerCase()
.replaceAll("_", "-");
const parts = stripped.split(/[-.]/);
const familyParts: string[] = [];
for (const p of parts) {
if (/^\d+$/.test(p) || /^\d+[bm]?$/i.test(p)) break;
familyParts.push(p);
}
return familyParts.length > 0 ? familyParts.join("-") : stripped;
}
+17 -52
View File
@@ -65,7 +65,7 @@
nodeThunderboltBridge,
nodeIdentities,
isConnected,
featureFlags,
vllmAvailable,
type DownloadProgress,
type PlacementPreview,
} from "$lib/stores/app.svelte";
@@ -297,7 +297,7 @@
const seen = new Set<string>();
const deduped: typeof candidates = [];
for (const m of candidates) {
const key = m.base_model || m.family || m.id;
const key = (m.base_model || m.family || m.id).toLowerCase().replace(/[-_]/g, " ");
if (seen.has(key)) continue;
seen.add(key);
deduped.push(m);
@@ -704,7 +704,7 @@
: 1;
const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
const instanceType =
nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
nodeCount <= 1 && selectedInstanceType !== "Vllm"
? "MlxRing"
: selectedInstanceType;
try {
@@ -787,7 +787,6 @@
quantization?: string;
base_model?: string;
capabilities?: string[];
requires_vllm?: boolean;
}>
>([]);
type ModelMemoryFitStatus =
@@ -937,12 +936,11 @@
// Apply sharding and instance type unconditionally
selectedSharding = defaults.sharding;
selectedInstanceType =
defaults.instanceType === "MlxRing"
? "MlxRing"
: defaults.instanceType === "Vllm"
? "Vllm"
defaults.instanceType === "Vllm"
? "Vllm"
: defaults.instanceType === "MlxRing"
? "MlxRing"
: "MlxJaccl";
userPickedInstanceType = true;
// Apply minNodes if valid (between 1 and maxNodes)
if (
@@ -964,23 +962,6 @@
}
let selectedInstanceType = $state<InstanceMeta>("MlxRing");
let userPickedInstanceType = $state(false);
$effect(() => {
if (!userPickedInstanceType && featureFlags()["vllm_available"]) {
selectedInstanceType = "Vllm";
}
});
const selectedModelRequiresVllm = $derived.by((): boolean => {
const id = selectedPreviewModelId();
if (!id) return false;
const model = models.find((m) => m.id === id);
return model?.requires_vllm === true;
});
$effect(() => {
if (selectedModelRequiresVllm) {
selectedInstanceType = "Vllm";
}
});
let selectedMinNodes = $state<number>(1);
let minNodesInitialized = $state(false);
let launchingModelId = $state<string | null>(null);
@@ -2088,7 +2069,7 @@
let instanceType = "Unknown";
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
else if (instanceTag === "VllmInstance") instanceType = "vLLM";
else if (instanceTag === "VllmInstance") instanceType = "vLLM (CUDA)";
const inst = instance as {
shardAssignments?: {
@@ -5795,18 +5776,14 @@
</div>
<div class="flex gap-2">
<button
disabled={selectedModelRequiresVllm}
onclick={() => {
if (selectedModelRequiresVllm) return;
selectedInstanceType = "MlxRing";
userPickedInstanceType = true;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
: selectedInstanceType === 'MlxRing'
? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
: 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
'MlxRing'
? 'bg-transparent text-exo-yellow border-exo-yellow'
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
>
<span
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5822,18 +5799,14 @@
TCP/IP
</button>
<button
disabled={selectedModelRequiresVllm}
onclick={() => {
if (selectedModelRequiresVllm) return;
selectedInstanceType = "MlxJaccl";
userPickedInstanceType = true;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
: selectedInstanceType === 'MlxJaccl'
? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
: 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
'MlxJaccl'
? 'bg-transparent text-exo-yellow border-exo-yellow'
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
>
<span
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5848,11 +5821,10 @@
</span>
RDMA (Fast)
</button>
{#if featureFlags()["vllm_available"] || selectedModelRequiresVllm}
{#if vllmAvailable()}
<button
onclick={() => {
selectedInstanceType = "Vllm";
userPickedInstanceType = true;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
@@ -5876,13 +5848,6 @@
</button>
{/if}
</div>
{#if selectedModelRequiresVllm}
<div
class="mt-2 text-[11px] font-mono text-orange-300/80"
>
This model requires vLLM.
</div>
{/if}
</div>
<!-- Minimum Devices -->
@@ -1,81 +0,0 @@
<script lang="ts">
import { browser } from "$app/environment";
import HeaderNav from "$lib/components/HeaderNav.svelte";
import PrefillDecodeDisaggregation from "$lib/components/PrefillDecodeDisaggregation.svelte";
import { featureFlags, refreshState } from "$lib/stores/app.svelte";
import { onMount } from "svelte";
type TabId = "prefill-decode";
const tabs: { id: TabId; label: string }[] = [
{ id: "prefill-decode", label: "Prefill / Decode" },
];
let activeTab = $state<TabId>(tabs[0].id);
let flagsLoaded = $state(false);
onMount(() => {
refreshState().finally(() => {
flagsLoaded = true;
});
});
const flags = $derived(featureFlags());
const enabled = $derived(flags["disaggregation"] === true);
$effect(() => {
if (browser && flagsLoaded && !enabled) {
// No advanced features enabled — bounce home.
window.location.hash = "/";
}
});
</script>
<div class="min-h-screen bg-exo-dark-gray flex flex-col">
<HeaderNav />
<main class="flex-1 max-w-[1100px] mx-auto w-full px-4 md:px-6 py-8">
{#if !flagsLoaded}
<div class="text-exo-light-gray/60 text-sm">Loading…</div>
{:else if !enabled}
<div class="text-exo-light-gray/60 text-sm">
No advanced features enabled. Set <code
class="text-exo-yellow font-mono">ENABLE_DISAGGREGATION=true</code
> on the cluster to access prefill/decode disaggregation.
</div>
{:else}
<div class="mb-4">
<h1
class="text-white text-xl md:text-2xl font-semibold tracking-wide mb-2"
>
Advanced
</h1>
<p class="text-exo-light-gray/60 text-sm">
Cluster-level configuration. Most users don't need anything here.
</p>
</div>
<div
class="flex flex-wrap gap-2 mb-6 border-b border-exo-light-gray/10 pb-3"
>
{#each tabs as tab (tab.id)}
<button
onclick={() => (activeTab = tab.id)}
class="px-3 py-1.5 text-xs rounded-md transition-all cursor-pointer
{activeTab === tab.id
? 'bg-exo-yellow/15 text-exo-yellow border border-exo-yellow/30'
: 'text-exo-light-gray/60 hover:text-white/80 border border-transparent hover:border-exo-light-gray/20'}"
>
{tab.label}
</button>
{/each}
</div>
<div class="space-y-4">
{#if activeTab === "prefill-decode"}
<PrefillDecodeDisaggregation />
{/if}
</div>
{/if}
</main>
</div>
+1 -112
View File
@@ -14,7 +14,6 @@
let modelCapabilities = $state<Record<string, string[]>>({});
let modelContextLengths = $state<Record<string, number>>({});
let modelReasoningDialects = $state<Record<string, string>>({});
const runningModels = $derived.by(() => {
const models: string[] = [];
@@ -89,12 +88,10 @@
let codexModel = $state("");
let codexMcpPath = $state("/Users/username");
let openClawModel = $state("");
let piModel = $state("");
$effect(() => {
const def = modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id";
codexModel = def;
openClawModel = def;
piModel = def;
});
const claudeShellCommand = $derived(
@@ -133,7 +130,6 @@
for (const modelId of runningModels) {
const caps = modelCapabilities[modelId] || [];
const ctxLen = modelContextLengths[modelId] || 0;
const dialect = modelReasoningDialects[modelId];
const entry: Record<string, unknown> = { name: modelId };
if (ctxLen > 0) {
entry.limit = { context: ctxLen, output: Math.min(ctxLen, 16384) };
@@ -141,27 +137,6 @@
if (caps.includes("vision")) {
entry.modalities = { input: ["text", "image"], output: ["text"] };
}
// Reasoning round-trip: opencode's `interleaved` field tells the
// openai-compatible adapter to send the assistant's prior
// reasoning_content back in subsequent turns. Emit it for dialects
// whose chat templates use prior reasoning:
// - `tool_conditional` (DeepSeek V3.2 / V4): wrapper preserves all
// reasoning when tools are present.
// - `post_last_user` (Qwen3-Thinking, GLM 4.5+, MiniMax M2.x):
// Jinja template reads reasoning_content for assistant turns since
// the last user message — exactly the tool-chain window.
// - `channel` (gpt-oss / Harmony): the model's Jinja template reads
// `message.thinking` rather than `message.reasoning_content`, but
// the server bridges `reasoning_content` → `thinking` before
// rendering, so the round-trip works through the standard field.
// `suffix` (Kimi): reasoning lives in content; no separate field path.
if (
dialect === "tool_conditional" ||
dialect === "post_last_user" ||
dialect === "channel"
) {
entry.interleaved = { field: "reasoning_content" };
}
models[modelId] = entry;
}
if (Object.keys(models).length === 0) {
@@ -243,55 +218,6 @@
),
);
const piModelsJson = $derived.by(() => {
const models: Record<string, unknown>[] = [];
for (const modelId of runningModels) {
const caps = modelCapabilities[modelId] || [];
const ctxLen = modelContextLengths[modelId] || 0;
const entry: Record<string, unknown> = { id: modelId };
if (caps.includes("vision")) {
entry.input = ["text", "image"];
}
// Mark thinking-capable models so pi surfaces its thinking-level selector
// for them. exo capability strings: "thinking" (model emits reasoning
// content) and "thinking_toggle" (user can turn it on/off).
if (caps.includes("thinking") || caps.includes("thinking_toggle")) {
entry.reasoning = true;
}
if (ctxLen > 0) {
entry.contextWindow = ctxLen;
}
models.push(entry);
}
if (models.length === 0) {
models.push({ id: "your-model-id" });
}
return JSON.stringify(
{
providers: {
exo: {
baseUrl: `${apiUrl}/v1`,
api: "openai-completions",
apiKey: "exo",
compat: {
supportsDeveloperRole: false,
// exo's OpenAI surface takes a boolean `enable_thinking` toggle,
// not graded effort levels, so disable pi's `reasoning_effort`
// parameter and use the matching top-level-boolean format.
supportsReasoningEffort: false,
thinkingFormat: "qwen",
},
models,
},
},
},
null,
2,
);
});
const piShellCommand = $derived(`pi --provider exo --model ${piModel}`);
const ollamaCommand = $derived(
`OLLAMA_HOST=${apiUrl}/ollama ollama run ${modelsBySize.length > 0 ? modelsBySize[0] : "your-model-id"}`,
);
@@ -351,7 +277,6 @@
"OpenCode",
"Codex",
"OpenClaw",
"Pi",
"Open WebUI",
"n8n",
"Firefox",
@@ -373,25 +298,16 @@
try {
const resp = await fetch("/v1/models");
const data = (await resp.json()) as {
data: {
id: string;
capabilities: string[];
context_length: number;
reasoning_dialect?: string;
}[];
data: { id: string; capabilities: string[]; context_length: number }[];
};
const caps: Record<string, string[]> = {};
const ctxs: Record<string, number> = {};
const dialects: Record<string, string> = {};
for (const model of data.data) {
caps[model.id] = model.capabilities || [];
if (model.context_length > 0) ctxs[model.id] = model.context_length;
if (model.reasoning_dialect)
dialects[model.id] = model.reasoning_dialect;
}
modelCapabilities = caps;
modelContextLengths = ctxs;
modelReasoningDialects = dialects;
} catch {
/* ignore */
}
@@ -599,33 +515,6 @@
config={`openclaw doctor --fix${(modelCapabilities[openClawModel] || []).includes("vision") ? `\nopenclaw models set-image exo/${openClawModel}` : ""}\nopenclaw gateway &\nopenclaw dashboard`}
language="bash"
/>
{:else if activeTab === "Pi"}
{#if runningModels.length > 1}
<div class="text-xs">
<span
class="text-exo-light-gray/50 text-[10px] uppercase tracking-wider block mb-1"
>Model</span
>
<select bind:value={piModel} class={selectClass}>
{#each runningModels as model}
<option value={model}>{model.split("/").pop()}</option>
{/each}
</select>
</div>
{/if}
<IntegrationCard
title="Models Config"
subtitle="~/.pi/agent/models.json"
description="Register exo as a custom provider in pi. Create or edit this file, then run pi and pick an exo model via /model. Install pi with: npm install -g @mariozechner/pi-coding-agent"
config={piModelsJson}
/>
<IntegrationCard
title="Shell Command"
subtitle="Run in terminal"
description="Launch pi directly with the exo provider and model selected."
config={piShellCommand}
language="bash"
/>
{:else if activeTab === "Open WebUI"}
<IntegrationCard
title="1. Start Open WebUI"
+1 -1
View File
@@ -81,4 +81,4 @@ Whenever a device produces side effects, it captures those side effects in an `E
## Purity
A significant goal of the current design is to make data flow explicit. Classes should either represent simple data (`FrozenModel`s typically, and `TaggedModel`s for unions) or active `System`s (Erlang `Actor`s), with all transformations of that data being "referentially transparent" - destructure and construct new data, don't mutate in place. We have had varying degrees of success with this, and are still exploring where purity makes sense.
A significant goal of the current design is to make data flow explicit. Classes should either represent simple data (`CamelCaseModel`s typically, and `TaggedModel`s for unions) or active `System`s (Erlang `Actor`s), with all transformations of that data being "referentially transparent" - destructure and construct new data, don't mutate in place. We have had varying degrees of success with this, and are still exploring where purity makes sense.
Generated
-21
View File
@@ -96,26 +96,6 @@
"type": "github"
}
},
"nixglhost": {
"inputs": {
"nixpkgs": [
"nixpkgs"
]
},
"locked": {
"lastModified": 1732211616,
"narHash": "sha256-QZCKJoypcwgS3tDNSWMjlxEBZtOYPW3eXV24rMzKsac=",
"owner": "numtide",
"repo": "nix-gl-host",
"rev": "5269b233f83880a0b433eafe026f0bc0d8f1a4a9",
"type": "github"
},
"original": {
"owner": "numtide",
"repo": "nix-gl-host",
"type": "github"
}
},
"nixpkgs": {
"locked": {
"lastModified": 1775595990,
@@ -207,7 +187,6 @@
"dream2nix": "dream2nix",
"fenix": "fenix",
"flake-parts": "flake-parts",
"nixglhost": "nixglhost",
"nixpkgs": "nixpkgs",
"pyproject-build-systems": "pyproject-build-systems",
"pyproject-nix": "pyproject-nix",
+54 -26
View File
@@ -45,11 +45,6 @@
inputs.uv2nix.follows = "uv2nix";
inputs.nixpkgs.follows = "nixpkgs";
};
nixglhost = {
url = "github:numtide/nix-gl-host";
inputs.nixpkgs.follows = "nixpkgs";
};
};
nixConfig = {
@@ -76,11 +71,14 @@
perSystem = { config, self', pkgs, lib, system, ... }:
let
pkgsArgs = {
pkgsCuda = import ./nix/cuda-pkgs.nix { inherit (inputs) nixpkgs; inherit system; };
in
{
# Allow unfree for metal-toolchain (needed for Darwin Metal packages)
_module.args.pkgs = import inputs.nixpkgs {
inherit system;
config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
overlays = [
inputs.nixglhost.overlays.default
(import ./nix/apple-sdk-overlay.nix)
(final: _: {
macmon = final.rustPlatform.buildRustPackage {
@@ -97,13 +95,6 @@
})
];
};
in
{
# Allow unfree for metal-toolchain (needed for Darwin Metal packages)
_module.args = {
pkgs = import inputs.nixpkgs pkgsArgs;
unfreePkgs = import inputs.nixpkgs (pkgsArgs // { config.allowUnfree = true; });
};
treefmt = {
projectRootFile = "flake.nix";
programs = {
@@ -130,24 +121,59 @@
};
};
packages = {
default = self'.packages.exo;
} //
lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
(
let
uvLock = builtins.fromTOML (builtins.readFile ./uv.lock);
mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package);
uvLockMlxVersion = mlxPackage.version;
uvLockMlxRev = builtins.elemAt (builtins.split "#" mlxPackage.source.git) 2;
in
{
metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { };
mlx = pkgs.callPackage ./nix/mlx.nix {
inherit (self'.packages) metal-toolchain;
inherit uvLockMlxVersion uvLockMlxRev;
};
default = self'.packages.exo;
}
) // lib.optionalAttrs (pkgsCuda != null) {
torch-cuda = pkgsCuda.python313Packages.torch;
vllm-cuda = pkgsCuda.python313Packages.vllm;
# exo with CUDA torch + vLLM — wraps the uv2nix-built package with host driver libs
exo-cuda = pkgs.writeShellApplication {
name = "exo-cuda";
runtimeInputs = [ self'.packages.exo-cuda-unwrapped ];
text = ''
for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do
if [ -e "$dir/libcuda.so.1" ]; then
NVIDIA_LIBS="$dir/libcuda.so.1"
for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do
[ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib"
done
export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}"
break
fi
done
export LD_LIBRARY_PATH="${pkgsCuda.stdenv.cc.cc.lib}/lib:${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
exec exo-cuda "$@"
'';
};
};
devShells.default = with pkgs; pkgs.mkShell {
devShells.default = pkgs.mkShell {
inputsFrom = [ self'.checks.cargo-build ];
packages =
[
with pkgs; [
# FORMATTING
config.treefmt.build.wrapper
# PYTHON
self'.packages.exo.passthru.evenv
self'.packages.python
uv
ruff
basedpyright
# RUST
config.rust.toolchain
@@ -164,16 +190,18 @@
just
jq
]
++ lib.optionals stdenv.isLinux [
unixtools.ifconfig
]
++ lib.optionals stdenv.isDarwin [
macmon
];
OPENSSL_NO_VENDOR = "1";
shellHook = ''
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
${lib.optionalString stdenv.isLinux ''
export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.python}/lib"
${lib.optionalString pkgs.stdenv.isLinux ''
export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH"
''}
'';
};
+14 -15
View File
@@ -18,11 +18,23 @@ check:
sync:
uv sync --all-packages
sync-cuda:
#!/usr/bin/env bash
set -euo pipefail
if command -v nvidia-smi &>/dev/null; then
arch=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | head -1 | tr -d ' ')
export TORCH_CUDA_ARCH_LIST="$arch"
export VLLM_TARGET_DEVICE=cuda
fi
uv pip install "cmake>=3.26.1" ninja "packaging>=24.2" "setuptools>=77.0.3,<81.0.0" "setuptools-scm>=8.0" wheel jinja2
find ~/.cache/uv/git-v0 -name CMakeCache.txt -delete 2>/dev/null || true
uv sync --extra cuda
sync-clean:
uv sync --all-packages --force-reinstall --no-cache
rust-rebuild:
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo run --bin stub_gen
cargo run --bin stub_gen
uv sync --reinstall-package exo_pyo3_bindings
build-dashboard:
@@ -36,23 +48,10 @@ package: build-dashboard
uv run pyinstaller packaging/pyinstaller/exo.spec
rm -rf build
build-app: rust-rebuild sync-clean package
build-app: package
xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
@echo "\nBuild complete. Run with:\n open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
sync-cuda:
#!/usr/bin/env bash
set -euo pipefail
uv sync --extra vllm-cuda13 --extra mlx-cpu --no-install-package vllm
dest=".venv/lib/python3.13/site-packages"
[[ -d $dest/vllm ]] || {
nix build .#exo-cuda-13.passthru.evenv
# will also grab vllm-0.19.1-distinfo
cp -aL result/lib/python3.13/site-packages/vllm* .venv/lib/python3.13/site-packages
chmod -R u+rwX .venv/lib/python3.13/site-packages/vllm*
rm result
}
clean:
rm -rf **/__pycache__
rm -rf target/
+102
View File
@@ -0,0 +1,102 @@
{ nixpkgs, system }:
let
pkgs = import nixpkgs { inherit system; };
in
if system == "aarch64-linux" then
import nixpkgs
{
inherit system;
config = {
allowUnfree = true;
allowBroken = true;
allowUnsupportedSystem = true;
cudaSupport = true;
cudaCapabilities = [ "12.1" ];
};
overlays = [
(final: prev:
let
cudaCompatStub = cfinal: cprev: {
cuda_compat = prev.runCommand "cuda13.0-cuda_compat-stub" { } "mkdir -p $out";
};
in
{
cudaPackages = prev.cudaPackages_13.overrideScope cudaCompatStub // {
override = args:
(prev.cudaPackages_13.override args).overrideScope cudaCompatStub;
};
pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
(pyFinal: pyPrev: {
fastsafetensors = pyFinal.buildPythonPackage {
pname = "fastsafetensors";
version = "0.2.2";
src = prev.fetchFromGitHub {
owner = "foundation-model-stack";
repo = "fastsafetensors";
rev = "v0.2.2";
hash = "";
};
pyproject = true;
build-system = [
pyFinal.setuptools
pyFinal.pybind11
];
buildInputs = [
final.cudaPackages.cuda_cudart
final.cudaPackages.cuda_nvml_dev
];
nativeBuildInputs = [
final.cudaPackages.cuda_nvcc
];
dependencies = [
pyFinal.typer
];
env.CUDA_HOME = "${final.cudaPackages.cuda_nvcc}";
pythonImportsCheck = [ "fastsafetensors" ];
};
cupy = pyPrev.cupy.override {
cudaPackages = final.cudaPackages;
};
bitsandbytes = pyPrev.bitsandbytes.overrideAttrs (old: {
preConfigure = (old.preConfigure or "") + ''
export CXXFLAGS="''${CXXFLAGS:-} -I${final.cudaPackages.cuda_crt}/include"
export CUDAFLAGS="''${CUDAFLAGS:-} -I${final.cudaPackages.cuda_crt}/include"
export CMAKE_CUDA_FLAGS="''${CMAKE_CUDA_FLAGS:-} -I${final.cudaPackages.cuda_crt}/include"
'';
buildInputs = (old.buildInputs or [ ]) ++ [ final.cudaPackages.cuda_crt ];
});
vllm = pyPrev.vllm.overrideAttrs (old: {
buildInputs = (old.buildInputs or [ ]) ++ [
final.cuda_cccl_with_prefix
final.cudaPackages.cuda_crt
];
preConfigure = (old.preConfigure or "") + ''
export CXXFLAGS="''${CXXFLAGS:-} -I${final.cuda_cccl_with_prefix}/include -I${final.cudaPackages.cuda_crt}/include"
export CUDAFLAGS="''${CUDAFLAGS:-} -I${final.cuda_cccl_with_prefix}/include -I${final.cudaPackages.cuda_crt}/include"
'';
});
})
];
magma-cuda-static = prev.magma-cuda-static.overrideAttrs (old: {
postPatch = (old.postPatch or "") + ''
sed -i '/err = cudaGetDeviceProperties( &prop, dev );/a\ int clock_khz = 0; cudaDeviceGetAttribute(\&clock_khz, cudaDevAttrClockRate, dev);' interface_cuda/interface.cpp
sed -i 's/prop\.clockRate/clock_khz/g' interface_cuda/interface.cpp
'';
});
cuda_cccl_with_prefix = prev.runCommand "cuda13.0-cuda_cccl-with-cccl-prefix" { } ''
mkdir -p $out/include
ln -s ${final.cudaPackages.cuda_cccl}/include $out/include/cccl
'';
opencv = prev.opencv.override { enableCuda = false; };
opencv4 = prev.opencv4.override { enableCuda = false; };
})
];
}
else
null
+158
View File
@@ -0,0 +1,158 @@
{ stdenv
, lib
, fetchFromGitHub
, replaceVars
, fetchzip
, cmake
, nlohmann_json
, apple-sdk_26
, metal-toolchain
, runCommand
, fmt
, python313Packages
, uvLockMlxVersion
, uvLockMlxRev
}:
assert stdenv.isDarwin;
let
python = python313Packages.python;
# Static dependencies included directly during compilation
gguf-tools = fetchFromGitHub {
owner = "antirez";
repo = "gguf-tools";
rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
};
metal_cpp = fetchzip {
url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
};
nanobind = fetchFromGitHub {
owner = "wjakob";
repo = "nanobind";
rev = "v2.10.2";
hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
fetchSubmodules = true;
};
mlx = stdenv.mkDerivation rec {
pname = "mlx";
version = uvLockMlxVersion;
pyproject = true;
src = fetchFromGitHub {
owner = "rltakashige";
repo = "mlx-jaccl-fix-small-recv";
rev = uvLockMlxRev;
hash = "sha256-M9x9QBYxwHv2z47qGZNJ4FgJyqLSIZ/3G1fEFQ421Lo=";
};
patches = [
(replaceVars ./darwin-build-fixes.patch {
sdkVersion = apple-sdk_26.version;
metalVersion = metal-toolchain.metalVersion;
})
];
postPatch = ''
substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
--replace-fail "g++" "$CXX"
'';
dontUseCmakeConfigure = true;
enableParallelBuilding = true;
# Allows multiple cores to be used in Python builds.
postUnpack = ''
export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
'';
# Updates the wrong fetcher rev attribute
passthru.skipBulkUpdate = true;
env = {
DEV_RELEASE = 1;
CMAKE_ARGS = toString [
(lib.cmakeBool "USE_SYSTEM_FMT" true)
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
(lib.cmakeBool "MLX_BUILD_CPU" true)
(lib.cmakeBool "MLX_BUILD_METAL" true)
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
(lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}")
(lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}")
];
SDKROOT = apple-sdk_26.passthru.sdkroot;
MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version;
};
build-system = [
python313Packages.setuptools
];
nativeBuildInputs = [
cmake
metal-toolchain
python313Packages.pypaBuildHook
python313Packages.pypaInstallHook
python313Packages.setuptools
python313Packages.typing-extensions
python313Packages.wheel
python313Packages.cmake
python313Packages.ninja
];
buildInputs = [
fmt
gguf-tools
python313Packages.nanobind
python313Packages.pybind11
apple-sdk_26
];
# Tests require Metal GPU access which isn't available in the Nix sandbox.
# To run tests, build with: nix build --option sandbox false .#mlx.passthru.tests.mlxTest
doCheck = false;
pythonImportsCheck = [ "mlx" ];
passthru.tests = {
# Runs example scripts to verify MLX works. Requires --option sandbox false
# since Metal GPU access is needed.
mlxTest =
runCommand "run-mlx-examples"
{
buildInputs = [ mlx ];
nativeBuildInputs = [ python ];
}
''
cp ${src}/examples/python/logistic_regression.py .
${python.interpreter} logistic_regression.py
rm logistic_regression.py
cp ${src}/examples/python/linear_regression.py .
${python.interpreter} linear_regression.py
rm linear_regression.py
touch $out
'';
};
meta = {
homepage = "https://github.com/ml-explore/mlx";
description = "Array framework for Apple silicon";
changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
license = lib.licenses.mit;
platforms = [ "aarch64-darwin" ];
};
};
in
mlx
-26
View File
@@ -1,26 +0,0 @@
diff --git a/setup.py b/setup.py
index 6dc2ed028..bdcc6354a 100644
--- a/setup.py
+++ b/setup.py
@@ -18,6 +18,13 @@ from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext
+if "NIX_ATTRS_JSON_FILE" in os.environ:
+ with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
+ NIX_ATTRS = json.load(f)
+else:
+ NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
+
+
def load_module_from_path(module_name, path):
spec = importlib.util.spec_from_file_location(module_name, path)
module = importlib.util.module_from_spec(spec)
@@ -184,6 +191,7 @@ class cmake_build_ext(build_ext):
cmake_args = [
"-DCMAKE_BUILD_TYPE={}".format(cfg),
"-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
+ *NIX_ATTRS["cmakeFlags"],
]
verbose = envs.VERBOSE
-6
View File
@@ -1,6 +0,0 @@
{
"name": "exo",
"lockfileVersion": 3,
"requires": true,
"packages": {}
}
+9 -12
View File
@@ -1,6 +1,5 @@
# -*- mode: python ; coding: utf-8 -*-
import sys
import importlib.util
import shutil
from pathlib import Path
@@ -57,7 +56,6 @@ HIDDEN_IMPORTS = sorted(
set(
collect_submodules("mlx")
+ _safe_collect("mlx_lm")
+ _safe_collect("mlx_vlm")
+ _safe_collect("transformers")
)
)
@@ -69,19 +67,18 @@ DATAS: list[tuple[str, str]] = [
(str(EXO_SHARED_MODELS_DIR), "exo/shared/models"),
]
if sys.platform == "darwin":
MACMON_PATH = shutil.which("macmon")
if MACMON_PATH is None:
raise SystemExit(
"macmon binary not found in PATH. "
"Install the pinned fork used by exo via: "
"cargo install --git https://github.com/vladkens/macmon "
"--rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d macmon --force"
)
MACMON_PATH = shutil.which("macmon")
if MACMON_PATH is None:
raise SystemExit(
"macmon binary not found in PATH. "
"Install the pinned fork used by exo via: "
"cargo install --git https://github.com/vladkens/macmon "
"--rev a1cd06b6cc0d5e61db24fd8832e74cd992097a7d macmon --force"
)
BINARIES: list[tuple[str, str]] = [
(MACMON_PATH, "."),
] if sys.platform == "darwin" else []
]
a = Analysis(
[str(ENTRYPOINT)],
+32 -106
View File
@@ -1,9 +1,9 @@
[project]
name = "exo"
version = "0.3.70"
version = "0.3.69"
description = "Exo"
readme = "README.md"
requires-python = "==3.13.*"
requires-python = ">=3.13"
dependencies = [
"aiofiles>=24.1.0",
"aiohttp>=3.12.14",
@@ -15,18 +15,22 @@ dependencies = [
"huggingface-hub>=1.8.0",
"psutil>=7.0.0",
"loguru>=0.7.3",
"exo-pyo3-bindings", # rust bindings
"exo_pyo3_bindings", # rust bindings
"anyio==4.11.0",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"mlx; sys_platform == 'darwin'",
"mlx==0.30.6; sys_platform == 'linux'",
"mlx-lm",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
"openai-harmony>=0.0.8",
"httpx>=0.28.1",
"tomlkit>=0.14.0",
"mflux==0.17.2; sys_platform == 'darwin'",
"python-multipart>=0.0.21",
"msgspec>=0.19.0",
"zstandard>=0.23.0",
"transformers>=5.6.2",
"nvidia-ml-py>=13.595.45",
"mlx-vlm>=0.3.11",
"transformers>=5.0.0,<5.4.0",
]
[project.scripts]
@@ -44,31 +48,11 @@ dev = [
]
[project.optional-dependencies]
build = ["nanobind"]
mlx-none = ["anyio"]
mlx = [
"mlx==0.31.2",
"mlx-lm",
"mlx-vlm>=0.3.11",
"mflux==0.17.5",
# pinning vllms versions for consistency.
"torch==2.10.0; sys_platform == 'darwin'",
"torch==2.10.0; sys_platform == 'linux'",
"torchaudio==2.10.0; sys_platform == 'darwin'",
"torchaudio==2.10.0; sys_platform == 'linux'",
"torchvision==0.25.0; sys_platform == 'darwin'",
"torchvision==0.25.0; sys_platform == 'linux'",
]
mlx-cpu = ["exo[mlx]", "mlx-cpu==0.31.2; sys_platform == 'linux'"]
mlx-cuda12 = ["exo[mlx]", "mlx-cuda-12==0.31.1; sys_platform == 'linux'"]
mlx-cuda13 = ["exo[mlx]", "mlx-cuda-13==0.31.1; sys_platform == 'linux'"]
vllm-none = ["anyio"]
vllm-cuda13 = [
"vllm[cuda13, fastsafetensors]; sys_platform == 'linux'",
"torch==2.10.0; sys_platform == 'linux'",
"torchaudio==2.10.0; sys_platform == 'linux'",
"torchvision==0.25.0; sys_platform == 'linux'",
cuda = [
"torch>=2.10.0; sys_platform == 'linux'",
"vllm>=0.13.0; sys_platform == 'linux'",
"mlx-cuda-13==0.30.6; sys_platform == 'linux'",
"fastsafetensors>=0.1.10; sys_platform == 'linux'",
]
###
@@ -79,37 +63,20 @@ vllm-cuda13 = [
members = ["rust/exo_pyo3_bindings", "bench"]
[tool.uv.sources]
exo-pyo3-bindings = { workspace = true }
exo_pyo3_bindings = { workspace = true }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
mflux = { git = "http://github.com/evanev7/mflux", branch = "exo" }
vllm = { git = "http://github.com/evanev7/vllm", branch = "exo2" }
torch = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
]
torchvision = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
]
torchaudio = [
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
{ index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
{ index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine!='x86_64'" },
{ index = "pytorch-cpu", marker = "sys_platform == 'linux' and platform_machine=='x86_64'" }
]
vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
[[tool.uv.index]]
name = "pytorch-cu130"
url = "https://download.pytorch.org/whl/cu130"
explicit = true
[[tool.uv.index]]
name = "pytorch-cu128"
url = "https://download.pytorch.org/whl/cu128"
explicit = true
[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
@@ -124,7 +91,7 @@ build-backend = "uv_build"
###
[tool.basedpyright]
include = ["src", "bench"]
include = [".venv/lib/mlx", ".venv/lib/mlx_lm", "src", "bench"]
typeCheckingMode = "strict"
failOnWarnings = true
@@ -148,17 +115,20 @@ exclude = [
"**/.direnv",
"**/rust",
"**/.github",
"**/vllm_patches",
"**/engines/vllm",
]
stubPath = ".mlx_typings"
extraPaths = [".cuda_typings"]
[[tool.basedpyright.executionEnvironments]]
root = "src/exo/worker/engines/image"
root = "src/exo/worker/runner/vllm_inference"
extraPaths = ["src", ".cuda_typings"]
reportMissingModuleSource = false
[[tool.basedpyright.executionEnvironments]]
root = "src"
###
# uv configuration
###
@@ -167,60 +137,14 @@ root = "src"
[tool.uv]
required-version = ">=0.8.6"
prerelease = "allow"
environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
override-dependencies = ["opencv-python; python_version < '0'"]
conflicts = [
[
{ extra = "mlx-cuda13" },
{ extra = "mlx-cuda12" },
{ extra = "mlx-cpu" },
{ extra = "mlx-none" },
],
[
{ extra = "vllm-cuda13" },
{ extra = "mlx-cuda12" },
{ extra = "vllm-none" },
],
environments = [
"sys_platform == 'linux' and platform_machine == 'x86_64'",
"sys_platform == 'linux' and platform_machine == 'aarch64'",
"sys_platform == 'darwin'",
]
[tool.uv.extra-build-dependencies]
miniaudio = ["setuptools", "cffi", "pycparser"]
mlx = [
"setuptools",
"typing-extensions",
"nanobind",
"pybind11",
"wheel",
"cmake",
"ninja",
]
mlx-lm = ["setuptools"]
mflux = ["uv_build"]
xgrammar = [
"nanobind",
"setuptools",
"scikit-build-core",
"packaging",
"pathspec",
]
rouge-score = ["setuptools"]
sacrebleu = ["setuptools"]
sqlitedict = ["setuptools"]
word2number = ["setuptools"]
vllm = [
"setuptools",
"setuptools-scm",
"scikit-build-core",
"jinja2",
"wheel",
"markupsafe",
"typing-extensions",
"torch",
]
fastsafetensors = ["setuptools", "pybind11"]
torch = ["typing-extensions"]
torchvision = ["torch"]
torchaudio = ["torch"]
###
# ruff configuration
@@ -228,7 +152,9 @@ torchaudio = ["torch"]
[tool.ruff]
extend-exclude = [
"shared/protobufs/**",
"*mlx_typings/**",
"*cuda_typings/**",
"rust/exo_pyo3_bindings/**",
"bench/vendor/**",
]
+188 -363
View File
@@ -1,38 +1,19 @@
{ inputs, ... }:
let
# Load workspace from uv.lock
workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
workspaceRoot = ../.;
};
mkPythonSet = { pkgs, lib, self', members }:
{
perSystem =
{ config, self', pkgs, lib, system, ... }:
let
inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
inherit (pkgs.config) cudaSupport;
inherit (pkgs) cudaPackages;
libmlx_source =
if (builtins.elem "mlx-cuda13" members.exo or [ ]) then "mlx-cuda-13"
else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
else "mlx-cpu";
python = pkgs.python313;
cudaLibs = with cudaPackages; [
cuda_cudart
cuda_cccl
cuda_cupti
cuda_nvrtc
cuda_nvtx
cudnn
libcufile
libcublas
libcufft
libcurand
libcusolver
libcusparse
libcusparse_lt
libnvjitlink
libnvshmem
nccl
];
pkgsCuda = import ../nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; };
# Load workspace from uv.lock
workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
workspaceRoot = inputs.self;
};
# Create overlay from workspace
# Use wheels from PyPI for most packages; we override mlx with our pure Nix Metal build
overlay = workspace.mkPyprojectOverlay { sourcePreference = "wheel"; };
# Override overlay to inject Nix-built components
exoOverlay = final: prev: {
# Replace workspace exo_pyo3_bindings with Nix-built wheel.
# Preserve passthru so mkVirtualEnv can resolve dependency groups.
@@ -52,378 +33,222 @@ let
'';
};
};
buildSystemsOverlay = final: prev:
lib.optionalAttrs isDarwin
{
mlx = prev.mlx.overrideAttrs (old:
let
# Static dependencies included directly during compilation
gguf-tools = pkgs.fetchFromGitHub {
owner = "antirez";
repo = "gguf-tools";
rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
};
metal_cpp = pkgs.fetchzip {
url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
};
python = pkgs.python313;
nanobind = pkgs.fetchFromGitHub {
owner = "wjakob";
repo = "nanobind";
rev = "v2.10.2";
hash = "sha256-io44YhN+VpfHFWyvvLWSanRgbzA0whK8WlDNRi3hahU=";
fetchSubmodules = true;
};
in
{
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ pkgs.cmake self'.packages.metal-toolchain ];
# TODO: non-sdk_26 support
buildInputs = (old.buildInputs or [ ])
++ [ gguf-tools pkgs.fmt pkgs.nlohmann_json pkgs.apple-sdk_26 ];
patches = [
(pkgs.replaceVars ../nix/darwin-build-fixes.patch {
sdkVersion = pkgs.apple-sdk_26.version;
inherit (self'.packages.metal-toolchain) metalVersion;
})
];
postPatch = ''
substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
--replace-fail "g++" "${lib.getExe' pkgs.stdenv.cc "c++"}"
'';
# Overlay to provide build systems and custom packages
buildSystemsOverlay = final: prev: {
# mlx-lm is a git dependency that needs setuptools
mlx-lm = prev.mlx-lm.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
# rouge-score and sacrebleu don't declare setuptools as a build dependency
rouge-score = prev.rouge-score.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
sacrebleu = prev.sacrebleu.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
sqlitedict = prev.sqlitedict.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
word2number = prev.word2number.overrideAttrs (old: {
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
final.setuptools
];
});
} // lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
# Use our pure Nix-built MLX with Metal support (macOS only)
mlx = self'.packages.mlx;
};
DEV_RELEASE = 1;
CMAKE_ARGS = toString ([
(lib.cmakeBool "USE_SYSTEM_FMT" true)
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${pkgs.nlohmann_json.src}")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_NANOBIND" "${nanobind}")
(lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
(lib.cmakeBool "MLX_BUILD_CPU" true)
(lib.cmakeBool "MLX_BUILD_METAL" true)
(lib.cmakeOptionType "string" "CMAKE_INSTALL_LIBDIR" "lib")
(lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
(lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${pkgs.apple-sdk_26.version}")
(lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${pkgs.apple-sdk_26.passthru.sdkroot}")
] ++ lib.optionals (isDarwin && isx86_64) [
(lib.cmakeBool "MLX_ENABLE_X64_MAC" true)
]);
SDKROOT = pkgs.apple-sdk_26.passthru.sdkroot;
MACOSX_DEPLOYMENT_TARGET = pkgs.apple-sdk_26.version;
});
} // lib.optionalAttrs isLinux {
mlx = prev.mlx.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ lib.optionals cudaSupport [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
postInstall = ''
cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
'';
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
} // lib.optionalAttrs cudaSupport {
"${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
});
nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
});
nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
});
nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
});
torch = prev.torch.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
torchaudio = prev.torchaudio.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
buildInputs = old.buildInputs ++ [ cudaPackages.cuda_cudart ];
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
});
torchvision = prev.torchvision.overrideAttrs (old: {
nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
});
torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old: {
buildInputs = old.buildInputs ++ cudaLibs;
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
});
# Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
vllm = prev.vllm.overrideAttrs (old:
let
cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
mkdir -p $out/include
ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
'';
cudaRoot = pkgs.symlinkJoin {
name = "cuda-merged-exo";
paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
};
cutlass = pkgs.fetchFromGitHub {
name = "cutlass-source";
owner = "NVIDIA";
repo = "cutlass";
tag = "v4.2.1";
hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
};
triton-kernels = pkgs.fetchFromGitHub {
owner = "triton-lang";
repo = "triton";
tag = "v3.6.0";
hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
};
cutlass-flashmla = pkgs.fetchFromGitHub {
owner = "NVIDIA";
repo = "cutlass";
rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
};
flashmla = pkgs.stdenv.mkDerivation {
pname = "flashmla";
version = "1.0.0";
src = pkgs.fetchFromGitHub {
name = "FlashMLA-source";
owner = "vllm-project";
repo = "FlashMLA";
rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
};
dontConfigure = true;
buildPhase = ''
rm -rf csrc/cutlass
ln -sf ${cutlass-flashmla} csrc/cutlass
'';
installPhase = ''
cp -rva . $out
'';
};
qutlass = pkgs.fetchFromGitHub {
name = "qutlass-source";
owner = "IST-DASLab";
repo = "qutlass";
rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
};
vllm-flash-attn = pkgs.stdenv.mkDerivation {
pname = "vllm-flash-attn";
version = "2.7.2.post1";
src = pkgs.fetchFromGitHub {
name = "flash-attention-source";
owner = "vllm-project";
repo = "flash-attention";
rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
};
patches = [
(pkgs.fetchpatch {
url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
})
(pkgs.fetchpatch {
url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
})
];
dontConfigure = true;
buildPhase = ''
rm -rf csrc/cutlass
ln -sf ${cutlass} csrc/cutlass
'';
installPhase = ''
cp -rva . $out
'';
};
in
{
patches = (old.patches or [ ]) ++ [ ../nix/vllm-setuppy-cmake.patch ];
nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
pkgs.cmake
pkgs.ninja
pkgs.autoAddDriverRunpath
] ++ lib.optionals cudaSupport [
cudaPackages.cuda_nvcc
# Additional overlay for Linux-specific fixes (type checking env).
# Native wheels have shared lib dependencies we don't need at type-check time.
linuxOverlay = final: prev:
let
ignoreMissing = drv: drv.overrideAttrs { autoPatchelfIgnoreMissingDeps = [ "*" ]; };
nvidiaPackages = lib.filterAttrs (name: _: lib.hasPrefix "nvidia-" name) prev;
in
lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux (
(lib.mapAttrs (_: ignoreMissing) nvidiaPackages) // {
mlx = ignoreMissing prev.mlx;
mlx-cuda-13 = prev.mlx-cuda-13.overrideAttrs (old: {
buildInputs = (old.buildInputs or [ ]) ++ [
final.nvidia-cublas
final.nvidia-cuda-nvrtc
final.nvidia-cudnn-cu13
final.nvidia-nccl-cu13
];
# TODO: vllm rocm/cpu
VLLM_TARGET_DEVICE = "empty";
preConfigure = ''
export MAX_JOBS="$NIX_BUILD_CORES"
preFixup = ''
addAutoPatchelfSearchPath ${final.nvidia-cublas}
addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc}
addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13}
addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13}
'';
# TODO: vllm non cuda13 support, more arch's, etc.
} // lib.optionalAttrs cudaSupport {
buildInputs = cudaLibs ++ [ cudaRoot ];
VLLM_CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
CUDA_HOME = "${cudaRoot}";
CUDAToolkit_ROOT = "${cudaRoot}";
CUDACXX = "${cudaRoot}/bin/nvcc";
VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
VLLM_TARGET_DEVICE = "cuda";
TORCH_CUDA_ARCH_LIST = "12.0;12.1";
TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
CAFFE2_USE_CUDNN = "ON";
CAFFE2_USE_CUFILE = "ON";
CUTLASS_ENABLE_CUBLAS = "ON";
CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
cmakeFlags = [
(lib.cmakeBool "CMAKE_SKIP_INSTALL_RPATH" true)
(lib.cmakeBool "CMAKE_BUILD_WITH_INSTALL_RPATH" true)
(lib.cmakeFeature "CUDA_HOME" "${cudaRoot}")
(lib.cmakeFeature "CUDAToolkit_ROOT" "${cudaRoot}")
(lib.cmakeFeature "CMAKE_CUDA_COMPILER" "${cudaRoot}/bin/nvcc")
(lib.cmakeFeature "CMAKE_PREFIX_PATH" "${cudaRoot}")
(lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}")
(lib.cmakeFeature "FLASH_MLA_SRC_DIR" "${lib.getDev flashmla}")
(lib.cmakeFeature "VLLM_FLASH_ATTN_SRC_DIR" "${lib.getDev vllm-flash-attn}")
(lib.cmakeFeature "QUTLASS_SRC_DIR" "${lib.getDev qutlass}")
(lib.cmakeFeature "TORCH_CUDA_ARCH_LIST" "12.0;12.1")
(lib.cmakeFeature "CUTLASS_NVCC_ARCHS_ENABLED" "${cudaPackages.flags.cmakeCudaArchitecturesString}")
(lib.cmakeFeature "CUDA_TOOLKIT_ROOT_DIR" "${cudaRoot}")
(lib.cmakeFeature "CAFFE2_USE_CUDNN" "ON")
(lib.cmakeFeature "CAFFE2_USE_CUFILE" "ON")
(lib.cmakeFeature "CUTLASS_ENABLE_CUBLAS" "ON")
];
autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
});
torch = ignoreMissing prev.torch;
triton = ignoreMissing prev.triton;
}
);
baseOverlays = [
inputs.pyproject-build-systems.overlays.default
overlay
exoOverlay
buildSystemsOverlay
linuxOverlay
];
} // lib.optionalAttrs (cudaSupport && isx86_64) {
numba = prev.numba.overrideAttrs (old: {
buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
});
};
pyprojectOverlay = workspace.mkPyprojectOverlay {
sourcePreference = "wheel";
dependencies = members;
};
editableOverlay = workspace.mkEditablePyprojectOverlay {
# Use environment variable pointing to editable root directory
root = "$REPO_ROOT";
members = [ "exo" "exo-bench" ];
};
pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
inherit python;
}).overrideScope (
lib.composeManyExtensions [
inputs.pyproject-build-systems.overlays.default
pyprojectOverlay
exoOverlay
buildSystemsOverlay
]
lib.composeManyExtensions baseOverlays
);
# mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
venv = name: (pythonSet.mkVirtualEnv "${name}-venv" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
mkApp = text: name: pkgs.writeShellApplication {
# mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
# mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
venvCollisionPaths = lib.optionals pkgs.stdenv.hostPlatform.isLinux [
"lib/python3.13/site-packages/mlx*"
"lib/python3.13/site-packages/nvidia*"
];
# Exclude bench deps from main env (bench has its own benchVenv)
exoDeps = removeAttrs workspace.deps.default [ "exo-bench" ];
exoVenv = (pythonSet.mkVirtualEnv "exo-env" exoDeps).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
# Virtual environment with dev dependencies for testing
testVenv = (pythonSet.mkVirtualEnv "exo-test-env" (
exoDeps // {
exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env
}
)).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
mkPythonScript = name: path: pkgs.writeShellApplication {
inherit name;
text = "exec " + lib.optionalString cudaSupport "nixglhost " + text;
runtimeInputs = [ exoVenv ];
runtimeEnv = {
EXO_DASHBOARD_DIR = self'.packages.dashboard;
EXO_RESOURCES_DIR = inputs.self + /resources;
};
runtimeInputs = [
(venv name)
pkgs.nix-gl-host
]
++ lib.optionals isDarwin [ pkgs.macmon ];
passthru = {
venv = venv name;
evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
};
text = ''exec python ${path} "$@"'';
};
in
{
inherit venv;
mkPythonScript = path: mkApp ''python ${path} "$@"'';
mkExo = mkApp ''exo "$@"'';
};
in
{
perSystem =
{ self', pkgs, unfreePkgs, lib, ... }:
let
inherit (pkgs.stdenv.hostPlatform) isLinux;
inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
# Virtual environment with dev dependencies for testing
testVenv = (mkPythonSet {
inherit self' pkgs lib; members = {
exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
benchVenv = pythonSet.mkVirtualEnv "exo-bench-env" {
exo-bench = [ ];
};
}).venv "exo-test";
mkBenchScript = (mkPythonSet {
inherit self' pkgs lib; members = {
exo = [ "mlx-cpu" "vllm-none" ];
exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
mkBenchScript = name: path: pkgs.writeShellApplication {
inherit name;
runtimeInputs = [ benchVenv ];
text = ''exec python ${path} "$@"'';
};
}).mkPythonScript;
mkSimplePythonScript = name: path: pkgs.writeShellApplication {
inherit name;
runtimeInputs = [ pkgs.python313 ];
text = ''exec python ${path} "$@"'';
};
cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" "vllm-none" ]; }; };
cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cpu" "vllm-cuda13" ]; }; };
exoPackage = pkgs.runCommand "exo"
{
nativeBuildInputs = [ pkgs.makeWrapper ];
}
''
mkdir -p $out/bin
# Create wrapper script
makeWrapper ${exoVenv}/bin/exo $out/bin/exo \
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
'';
vllmEnv = pkgsCuda.python313.withPackages (ps: [ ps.vllm ps.fastsafetensors ]);
vllmSite = pkgs.runCommand "vllm-site-filtered" { } ''
mkdir -p $out
for pkg in ${vllmEnv}/${python.sitePackages}/*; do
name=$(basename "$pkg")
case "$name" in
anyio*|pydantic*) ;;
*) ln -s "$pkg" "$out/$name" ;;
esac
done
'';
exoCudaDeps = exoDeps // {
mlx-cuda-13 = [ ];
};
exoCudaVenv = (pythonSet.mkVirtualEnv "exo-cuda-env" exoCudaDeps).overrideAttrs {
venvIgnoreCollisions = venvCollisionPaths;
};
exoCudaPackage = pkgs.runCommand "exo-cuda"
{
nativeBuildInputs = [ pkgs.makeWrapper ];
}
''
mkdir -p $out/bin
makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo-cuda \
--set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \
--set EXO_RESOURCES_DIR ${inputs.self + /resources} \
--prefix PYTHONPATH : "${vllmSite}"
'';
in
{
packages = {
exo = mkExo "exo";
# for running tests in ci
exo-test-env = testVenv;
# Python package only available on macOS (requires MLX/Metal)
packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
{
exo = exoPackage;
# Test environment for running pytest outside of Nix sandbox (needs GPU access)
exo-test-env = testVenv;
} // lib.optionalAttrs (pkgsCuda != null) {
exo-cuda-unwrapped = exoCudaPackage;
} // {
inherit python;
exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
# used by ./tests/run_exo_on.sh
exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
} // lib.optionalAttrs isLinux {
exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
};
checks = {
# Ruff linting (works on all platforms)
lint = pkgs.runCommand "ruff-lint" { } ''
export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
${pkgs.ruff}/bin/ruff check ${inputs.self}
touch $out
'';
typecheck = pkgs.runCommand "typecheck" { nativeBuildInputs = [ testVenv ]; } ''
cd ${inputs.self}
basedpyright
touch $out
'';
# Hermetic basedpyright type checking
typecheck = pkgs.runCommand "typecheck"
{
nativeBuildInputs = [
testVenv
pkgs.basedpyright
];
}
''
cd ${inputs.self}
export HOME=$TMPDIR
basedpyright --pythonpath ${testVenv}/bin/python
touch $out
'';
};
};
}
@@ -1,21 +0,0 @@
model_id = "2imi9/gpt-oss-20B-NVFP4A16-BF16"
n_layers = 24
hidden_size = 2880
num_key_value_heads = 8
supports_tensor = false
tasks = ["TextGeneration"]
family = "gpt-oss"
quantization = "nvfp4"
base_model = "GPT-OSS 20B"
capabilities = ["text", "thinking"]
reasoning_dialect = "channel"
context_length = 131072
requires_vllm = true
[storage_size]
in_bytes = 41829514752
[sampling_defaults]
temperature = 1.0
top_p = 1.0
top_k = 0
@@ -8,15 +8,8 @@ family = "deepseek"
quantization = "4bit"
base_model = "DeepSeek V3.1"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 131072
[storage_size]
in_bytes = 405874409472
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/blob/main/generation_config.json
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/discussions/19
[sampling_defaults]
temperature = 0.6
top_p = 0.95
@@ -8,15 +8,8 @@ family = "deepseek"
quantization = "8bit"
base_model = "DeepSeek V3.1"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 131072
[storage_size]
in_bytes = 765577920512
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/blob/main/generation_config.json
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.1/discussions/19
[sampling_defaults]
temperature = 0.6
top_p = 0.95
@@ -8,15 +8,8 @@ family = "deepseek"
quantization = "4bit"
base_model = "DeepSeek V3.2"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "tool_conditional"
context_length = 131072
[storage_size]
in_bytes = 378086226621
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/generation_config.json
# Source: https://docs.vllm.ai/projects/recipes/en/latest/DeepSeek/DeepSeek-V3_2.html
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "deepseek"
quantization = "8bit"
base_model = "DeepSeek V3.2"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "tool_conditional"
context_length = 131072
[storage_size]
in_bytes = 755957120916
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/generation_config.json
# Source: https://docs.vllm.ai/projects/recipes/en/latest/DeepSeek/DeepSeek-V3_2.html
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -1,21 +0,0 @@
model_id = "mlx-community/DeepSeek-V4-Flash"
n_layers = 43
hidden_size = 4096
num_key_value_heads = 1
supports_tensor = true
tasks = ["TextGeneration"]
family = "deepseek"
quantization = "8bit"
base_model = "DeepSeek V4 Flash"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "tool_conditional"
context_length = 1048576
[storage_size]
in_bytes = 155095760030
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash
[sampling_defaults]
temperature = 1.0
top_p = 1.0
@@ -1,21 +0,0 @@
model_id = "mlx-community/DeepSeek-V4-Pro"
n_layers = 61
hidden_size = 7168
num_key_value_heads = 1
supports_tensor = true
tasks = ["TextGeneration"]
family = "deepseek"
quantization = "8bit"
base_model = "DeepSeek V4 Pro"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "tool_conditional"
context_length = 1048576
[storage_size]
in_bytes = 849681803879
# Source: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro
[sampling_defaults]
temperature = 1.0
top_p = 1.0
@@ -8,13 +8,8 @@ family = "glm"
quantization = "8bit"
base_model = "GLM 4.5 Air"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 131072
[storage_size]
in_bytes = 122406567936
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 0.6
top_p = 0.95
@@ -8,13 +8,8 @@ family = "glm"
quantization = "bf16"
base_model = "GLM 4.5 Air"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 131072
[storage_size]
in_bytes = 229780750336
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 0.6
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "4bit"
base_model = "GLM 4.7"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 198556925568
# Source: https://huggingface.co/zai-org/GLM-4.7
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "6bit"
base_model = "GLM 4.7"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 286737579648
# Source: https://huggingface.co/zai-org/GLM-4.7
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "8bit"
base_model = "GLM 4.7"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 396963397248
# Source: https://huggingface.co/zai-org/GLM-4.7
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "4bit"
base_model = "GLM 4.7 Flash"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 19327352832
# Source: https://huggingface.co/zai-org/GLM-4.7-Flash
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "5bit"
base_model = "GLM 4.7 Flash"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 22548578304
# Source: https://huggingface.co/zai-org/GLM-4.7-Flash
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "6bit"
base_model = "GLM 4.7 Flash"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 26843545600
# Source: https://huggingface.co/zai-org/GLM-4.7-Flash
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
@@ -8,15 +8,8 @@ family = "glm"
quantization = "8bit"
base_model = "GLM 4.7 Flash"
capabilities = ["text", "thinking", "thinking_toggle"]
reasoning_dialect = "post_last_user"
context_length = 202752
[storage_size]
in_bytes = 34359738368
# Source: https://huggingface.co/zai-org/GLM-4.7-Flash
# Source: https://unsloth.ai/docs/models/glm-4.7-flash
# Source: https://docs.z.ai/api-reference/llm/chat-completion
[sampling_defaults]
temperature = 1.0
top_p = 0.95
Loaded 100 of 347 files, more files were not shown because too many files have changed in this diff. Show more