diff --git a/.mlx_typings/mlx/nn/layers/base.pyi b/.mlx_typings/mlx/nn/layers/base.pyi index a4abf36b..70557bac 100644 --- a/.mlx_typings/mlx/nn/layers/base.pyi +++ b/.mlx_typings/mlx/nn/layers/base.pyi @@ -200,7 +200,7 @@ class Module(dict): ) -> mx.MX_ARRAY_TREE: # -> dict[Any, Any | dict[Any, Any | dict[Any, Any] | list[Any]] | dict[Any, Any] | list[Any]]: """Return the submodules that do not contain other modules.""" - def update(self, parameters: dict, strict: bool = ...) -> Module: + def update(self, parameters: dict[str, Any], strict: bool = ...) -> Module: """Replace the parameters of this Module with the provided ones in the dict of dicts and lists. diff --git a/.mlx_typings/mlx/utils.pyi b/.mlx_typings/mlx/utils.pyi index 43738ca7..ca07eb76 100644 --- a/.mlx_typings/mlx/utils.pyi +++ b/.mlx_typings/mlx/utils.pyi @@ -7,7 +7,10 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union from mlx.core import MX_ARRAY_TREE def tree_map( - fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ... + fn: Callable[..., Any], + tree: Any, + *rest: Any, + is_leaf: Callable[..., bool] | None = ..., ) -> Any: """Applies ``fn`` to the leaves of the Python tree ``tree`` and returns a new collection with the results. @@ -44,11 +47,11 @@ def tree_map( """ def tree_map_with_path( - fn: Callable, + fn: Callable[..., Any], tree: Any, *rest: Any, - is_leaf: Optional[Callable] = ..., - path: Optional[Any] = ..., + is_leaf: Callable[..., bool] | None = ..., + path: str | None = ..., ) -> Any: """Applies ``fn`` to the path and leaves of the Python tree ``tree`` and returns a new collection with the results. @@ -80,9 +83,9 @@ def tree_map_with_path( def tree_flatten( tree: Any, prefix: str = ..., - is_leaf: Optional[Callable] = ..., - destination: Optional[Union[List[Tuple[str, Any]], Dict[str, Any]]] = ..., -) -> Union[List[Tuple[str, Any]], Dict[str, Any]]: + is_leaf: Callable[..., bool] | None = ..., + destination: list[tuple[str, Any]] | dict[str, Any] | None = ..., +) -> list[tuple[str, Any]] | dict[str, Any]: """Flattens a Python tree to a list of key, value tuples. The keys are using the dot notation to define trees of arbitrary depth and @@ -118,7 +121,7 @@ def tree_flatten( the Python tree. """ -def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any: +def tree_unflatten(tree: list[tuple[str, Any]] | dict[str, Any]) -> Any: """Recreate a Python tree from its flat representation. .. code-block:: python diff --git a/.mlx_typings/mlx_lm/models/glm_moe_dsa.pyi b/.mlx_typings/mlx_lm/models/glm_moe_dsa.pyi new file mode 100644 index 00000000..135f581f --- /dev/null +++ b/.mlx_typings/mlx_lm/models/glm_moe_dsa.pyi @@ -0,0 +1,46 @@ +"""Type stubs for mlx_lm.models.glm_moe_dsa""" + +from dataclasses import dataclass +from typing import Any, Dict, Optional + +from .base import BaseModelArgs +from .deepseek_v32 import Model as DSV32Model + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + vocab_size: int + hidden_size: int + index_head_dim: int + index_n_heads: int + index_topk: int + intermediate_size: int + moe_intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + n_shared_experts: Optional[int] + n_routed_experts: Optional[int] + routed_scaling_factor: float + kv_lora_rank: int + q_lora_rank: int + qk_rope_head_dim: int + v_head_dim: int + qk_nope_head_dim: int + topk_method: str + scoring_func: str + norm_topk_prob: bool + n_group: int + topk_group: int + num_experts_per_tok: int + moe_layer_freq: int + first_k_dense_replace: int + max_position_embeddings: int + rms_norm_eps: float + rope_parameters: Dict[str, Any] + attention_bias: bool + rope_scaling: Dict[str, Any] | None + rope_theta: float | None + +class Model(DSV32Model): + def __init__(self, config: ModelArgs) -> None: ... diff --git a/Cargo.lock b/Cargo.lock index a45bfe9d..8914f5e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -141,12 +141,6 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" -[[package]] -name = "arrayvec" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" - [[package]] name = "asn1-rs" version = "0.7.1" @@ -304,19 +298,6 @@ version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba" -[[package]] -name = "bigdecimal" -version = "0.4.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "560f42649de9fa436b73517378a147ec21f6c997a546581df4b4b31677828934" -dependencies = [ - "autocfg", - "libm", - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "bimap" version = "0.6.3" @@ -516,15 +497,6 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2f421161cb492475f1661ddc9815a745a1c894592070661180fdec3d4872e9c3" -[[package]] -name = "convert_case" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "633458d4ef8c78b72454de2d54fd6ab2e60f9e02be22f3c6104cdc8a4e0fceb9" -dependencies = [ - "unicode-segmentation", -] - [[package]] name = "core-foundation" version = "0.9.4" @@ -746,29 +718,6 @@ dependencies = [ "powerfmt", ] -[[package]] -name = "derive_more" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10b768e943bed7bf2cab53df09f4bc34bfd217cdb57d971e769874c9a6710618" -dependencies = [ - "derive_more-impl", -] - -[[package]] -name = "derive_more-impl" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d286bfdaf75e988b4a78e013ecd79c581e06399ab53fbacd2d916c2f904f30b" -dependencies = [ - "convert_case", - "proc-macro2", - "quote", - "rustc_version", - "syn 2.0.111", - "unicode-xid", -] - [[package]] name = "digest" version = "0.10.7" @@ -939,22 +888,17 @@ name = "exo_pyo3_bindings" version = "0.0.1" dependencies = [ "delegate", - "derive_more", "env_logger", "extend", - "futures", - "impl-trait-for-tuples", + "futures-lite", "libp2p", "log", "networking", - "once_cell", "pin-project", "pyo3", "pyo3-async-runtimes", "pyo3-log", "pyo3-stub-gen", - "thiserror 2.0.17", - "thread_local", "tokio", "util", ] @@ -970,6 +914,12 @@ dependencies = [ "syn 2.0.111", ] +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + [[package]] name = "ff" version = "0.13.1" @@ -1078,7 +1028,10 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad" dependencies = [ + "fastrand", "futures-core", + "futures-io", + "parking", "pin-project-lite", ] @@ -1640,17 +1593,6 @@ dependencies = [ "xmltree", ] -[[package]] -name = "impl-trait-for-tuples" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0eb5a3343abf848c0984fe4604b2b105da9539376e24fc0a3b0007411ae4fd9" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.111", -] - [[package]] name = "indexmap" version = "2.12.1" @@ -1829,12 +1771,6 @@ version = "0.2.178" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091" -[[package]] -name = "libm" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de" - [[package]] name = "libp2p" version = "0.56.0" @@ -2824,16 +2760,13 @@ name = "networking" version = "0.0.1" dependencies = [ "delegate", - "derive_more", "either", "extend", - "futures", + "futures-lite", "futures-timer", - "impl-trait-for-tuples", "keccak-const", "libp2p", "log", - "thiserror 2.0.17", "tokio", "tracing-subscriber", "util", @@ -2918,17 +2851,6 @@ dependencies = [ "num-traits", ] -[[package]] -name = "num-rational" -version = "0.4.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" -dependencies = [ - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "num-traits" version = "0.2.19" @@ -3279,28 +3201,14 @@ version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab53c047fcd1a1d2a8820fe84f05d6be69e9526be40cb03b73f86b6b03e6d87d" dependencies = [ - "bigdecimal", - "either", - "hashbrown 0.16.1", - "indexmap", "indoc", - "inventory", "libc", - "lock_api", "memoffset", - "num-bigint", - "num-complex", - "num-rational", - "num-traits", "once_cell", - "ordered-float", - "parking_lot", "portable-atomic", "pyo3-build-config", "pyo3-ffi", "pyo3-macros", - "rust_decimal", - "smallvec", "unindent", ] @@ -3741,16 +3649,6 @@ dependencies = [ "tokio", ] -[[package]] -name = "rust_decimal" -version = "1.39.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35affe401787a9bd846712274d97654355d21b2a2c092a3139aabe31e9022282" -dependencies = [ - "arrayvec", - "num-traits", -] - [[package]] name = "rustc-hash" version = "1.1.0" @@ -4615,24 +4513,12 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" -[[package]] -name = "unicode-segmentation" -version = "1.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" - [[package]] name = "unicode-width" version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" -[[package]] -name = "unicode-xid" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" - [[package]] name = "unicode_names2" version = "1.3.0" diff --git a/Cargo.toml b/Cargo.toml index 34339ede..ffa9022c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,49 +26,20 @@ opt-level = 3 networking = { path = "rust/networking" } util = { path = "rust/util" } -# Proc-macro authoring tools -syn = "2.0" -quote = "1.0" -proc-macro2 = "1.0" -darling = "0.20" - # Macro dependecies extend = "1.2" delegate = "0.13" -impl-trait-for-tuples = "0.2" -clap = "4.5" -derive_more = { version = "2.0.1", features = ["display"] } -pin-project = "1" # Utility dependencies -itertools = "0.14" -thiserror = "2" -internment = "0.8" -recursion = "0.5" -regex = "1.11" -once_cell = "1.21" -thread_local = "1.1" -bon = "3.4" -generativity = "1.1" -anyhow = "1.0" keccak-const = "0.2" -# Functional generics/lenses frameworks -frunk_core = "0.4" -frunk = "0.4" -frunk_utils = "0.2" -frunk-enum-core = "0.3" - # Async dependencies tokio = "1.46" -futures = "0.3" -futures-util = "0.3" +futures-lite = "2.6.1" futures-timer = "3.0" # Data structures either = "1.15" -ordered-float = "5.0" -ahash = "0.8" # Tracing/logging log = "0.4" diff --git a/README.md b/README.md index 58d41c37..785f6fba 100644 --- a/README.md +++ b/README.md @@ -72,16 +72,30 @@ There are two ways to run exo: ### Run from Source (macOS) +If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly: + +```bash +nix run .#exo +``` + +**Note:** To accept the Cachix binary cache (and avoid the Xcode Metal ToolChain), add to `/etc/nix/nix.conf`: +``` +trusted-users = root (or your username) +experimental-features = nix-command flakes +``` +Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-daemon` + **Prerequisites:** +- [Xcode](https://developer.apple.com/xcode/) (provides the Metal ToolChain required for MLX compilation) - [brew](https://github.com/Homebrew/brew) (for simple package management on macOS) - + ```bash /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)" ``` - [uv](https://github.com/astral-sh/uv) (for Python dependency management) - [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon) - [node](https://github.com/nodejs/node) (for building the dashboard) - + ```bash brew install uv macmon node ``` diff --git a/app/EXO/EXO/ExoProcessController.swift b/app/EXO/EXO/ExoProcessController.swift index bb84f06e..c67a77b5 100644 --- a/app/EXO/EXO/ExoProcessController.swift +++ b/app/EXO/EXO/ExoProcessController.swift @@ -136,11 +136,37 @@ final class ExoProcessController: ObservableObject { return } process.terminationHandler = nil - if process.isRunning { - process.terminate() - } - self.process = nil status = .stopped + + guard process.isRunning else { + self.process = nil + return + } + + let proc = process + self.process = nil + + Task.detached { + proc.interrupt() + + for _ in 0..<50 { + if !proc.isRunning { return } + try? await Task.sleep(nanoseconds: 100_000_000) + } + + if proc.isRunning { + proc.terminate() + } + + for _ in 0..<30 { + if !proc.isRunning { return } + try? await Task.sleep(nanoseconds: 100_000_000) + } + + if proc.isRunning { + kill(proc.processIdentifier, SIGKILL) + } + } } func restart() { diff --git a/bench/bench.toml b/bench/bench.toml new file mode 100644 index 00000000..3b3c8c07 --- /dev/null +++ b/bench/bench.toml @@ -0,0 +1,7 @@ +# Canary benchmark manifest +# +# Lists the suite files to include. Each file defines benchmarks +# with shared constraints, topology, and default args. +include = [ + "single-m3-ultra.toml", +] diff --git a/bench/eval_tool_calls.py b/bench/eval_tool_calls.py new file mode 100644 index 00000000..13cb2537 --- /dev/null +++ b/bench/eval_tool_calls.py @@ -0,0 +1,1104 @@ +# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false +from __future__ import annotations + +import argparse +import contextlib +import json +import os +import sys +import time +import tomllib +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal + +import httpx +from harness import ( + ExoClient, + ExoHttpError, + add_common_instance_args, + instance_id_from_instance, + nodes_used_in_instance, + resolve_model_short_id, + run_planning_phase, + settle_and_fetch_placements, + wait_for_instance_gone, + wait_for_instance_ready, +) + +SCENARIOS_PATH = Path(__file__).parent / "scenarios.toml" + + +@dataclass +class Scenario: + name: str + description: str + messages: list[dict[str, Any]] + tools: list[dict[str, Any]] + expect_tool_call: bool + expected_function: str | None = None + required_arg_keys: list[str] | None = None + tool_result: str | None = None + nested_array_key: str | None = None + required_item_keys: list[str] | None = None + + +def load_scenarios(path: Path) -> list[Scenario]: + with open(path, "rb") as f: + data = tomllib.load(f) + + tools_data = data.get("tools", {}) + all_tools: list[dict[str, Any]] = [] + tool_by_name: dict[str, dict[str, Any]] = {} + for name, defn in tools_data.items(): + tool: dict[str, Any] = { + "type": "function", + "function": { + "name": name, + "description": defn.get("description", ""), + "parameters": { + "type": "object", + "properties": defn.get("properties", {}), + "required": defn.get("required", []), + }, + }, + } + all_tools.append(tool) + tool_by_name[name] = tool + + scenarios: list[Scenario] = [] + for s in data.get("scenarios", []): + if "tools" in s: + scenario_tools = [tool_by_name[t] for t in s["tools"]] + else: + scenario_tools = list(all_tools) + + messages: list[dict[str, Any]] = [] + for msg in s.get("messages", []): + m: dict[str, Any] = {"role": msg["role"]} + if "content" in msg: + m["content"] = msg["content"] + if "tool_calls" in msg: + m["tool_calls"] = [ + { + "id": tc["id"], + "type": "function", + "function": { + "name": tc["name"], + "arguments": json.dumps(tc["arguments"]), + }, + } + for tc in msg["tool_calls"] + ] + if "tool_call_id" in msg: + m["tool_call_id"] = msg["tool_call_id"] + messages.append(m) + + tool_result: str | None = None + if "tool_result" in s: + tool_result = json.dumps(s["tool_result"]) + + scenarios.append( + Scenario( + name=s["name"], + description=s["description"], + messages=messages, + tools=scenario_tools, + expect_tool_call=s["expect_tool_call"], + expected_function=s.get("expected_function"), + required_arg_keys=s.get("required_arg_keys"), + tool_result=tool_result, + nested_array_key=s.get("nested_array_key"), + required_item_keys=s.get("required_item_keys"), + ) + ) + + return scenarios + + +ApiName = Literal["openai", "claude", "responses"] + + +@dataclass +class ParsedResponse: + finish_reason: str # "tool_calls" | "stop" | ... + has_tool_call: bool + tool_call: dict[str, str] | None # {"id": ..., "name": ..., "arguments": ...} + content: str | None + + +@dataclass +class ScenarioResult: + name: str + api: str + phase: str # "tool_call" or "follow_up" + passed: bool + checks: dict[str, bool] = field(default_factory=dict) + error: str | None = None + latency_ms: float = 0.0 + + +def validate_args(args_str: str, required_keys: list[str]) -> tuple[bool, str | None]: + """Parse JSON arguments and check required keys exist.""" + try: + args = json.loads(args_str) + except (json.JSONDecodeError, TypeError) as exc: + return False, f"Invalid JSON: {exc}" + if not isinstance(args, dict): + return False, f"Expected dict, got {type(args).__name__}" + missing = [k for k in required_keys if k not in args] + if missing: + return False, f"Missing keys: {missing}" + return True, None + + +def validate_nested_args( + args_str: str, + array_key: str, + required_item_keys: list[str], +) -> tuple[bool, str | None]: + """Check that args[array_key] is a list of objects with required keys.""" + try: + args = json.loads(args_str) + except (json.JSONDecodeError, TypeError) as exc: + return False, f"Invalid JSON: {exc}" + if not isinstance(args, dict): + return False, f"Expected dict, got {type(args).__name__}" + arr = args.get(array_key) + if not isinstance(arr, list): + return False, f"'{array_key}' is not an array (got {type(arr).__name__})" + if len(arr) == 0: + return False, f"'{array_key}' is empty" + for i, item in enumerate(arr): + if not isinstance(item, dict): + return ( + False, + f"'{array_key}[{i}]' is not an object (got {type(item).__name__})", + ) + missing = [k for k in required_item_keys if k not in item] + if missing: + return False, f"'{array_key}[{i}]' missing keys: {missing}" + return True, None + + +def call_api( + client: httpx.Client, + host: str, + port: int, + path: str, + body: dict[str, Any], + timeout: float, +) -> tuple[dict[str, Any], float]: + """POST to http://{host}:{port}{path}, return (response_json, latency_ms).""" + url = f"http://{host}:{port}{path}" + t0 = time.monotonic() + resp = client.post(url, json=body, timeout=timeout) + latency = (time.monotonic() - t0) * 1000 + resp.raise_for_status() + return resp.json(), latency + + +def _openai_build_request( + model: str, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], +) -> tuple[str, dict[str, Any]]: + """Build request for /v1/chat/completions.""" + body: dict[str, Any] = { + "model": model, + "messages": messages, + "tools": tools, + "max_tokens": 16384, + "temperature": 0.0, + } + return "/v1/chat/completions", body + + +def _openai_parse_response(data: dict[str, Any]) -> ParsedResponse: + """Parse OpenAI Chat Completions response into common format.""" + choice = data["choices"][0] + finish_reason = choice.get("finish_reason", "") + message = choice.get("message", {}) + tool_calls = message.get("tool_calls") + content = message.get("content") + + has_tool_call = isinstance(tool_calls, list) and len(tool_calls) > 0 + tool_call_info: dict[str, str] | None = None + if has_tool_call: + tc = tool_calls[0] + fn = tc.get("function", {}) + tool_call_info = { + "id": tc.get("id", "call_0"), + "name": fn.get("name", ""), + "arguments": fn.get("arguments", "{}"), + } + + return ParsedResponse( + finish_reason=finish_reason, + has_tool_call=has_tool_call, + tool_call=tool_call_info, + content=content, + ) + + +def _openai_build_followup( + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], + model: str, + parsed: ParsedResponse, + tool_result: str, +) -> tuple[str, dict[str, Any]]: + """Build multi-turn follow-up for OpenAI Chat Completions.""" + assert parsed.tool_call is not None + tc = parsed.tool_call + followup_messages: list[dict[str, Any]] = list(messages) + [ + { + "role": "assistant", + "tool_calls": [ + { + "id": tc["id"], + "type": "function", + "function": { + "name": tc["name"], + "arguments": tc["arguments"], + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": tc["id"], + "content": tool_result, + }, + ] + body: dict[str, Any] = { + "model": model, + "messages": followup_messages, + "tools": tools, + "max_tokens": 16384, + "temperature": 0.0, + } + return "/v1/chat/completions", body + + +def _claude_translate_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Translate OpenAI-format tools to Claude format.""" + claude_tools: list[dict[str, Any]] = [] + for tool in tools: + fn = tool["function"] + claude_tools.append( + { + "name": fn["name"], + "description": fn.get("description", ""), + "input_schema": fn.get("parameters", {}), + } + ) + return claude_tools + + +def _claude_translate_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Translate OpenAI-format messages to Claude Messages format.""" + claude_messages: list[dict[str, Any]] = [] + + for msg in messages: + role = msg["role"] + + if role == "user": + claude_messages.append( + { + "role": "user", + "content": msg["content"], + } + ) + elif role == "assistant": + content_blocks: list[dict[str, Any]] = [] + text_content = msg.get("content") + if text_content and isinstance(text_content, str) and text_content.strip(): + content_blocks.append({"type": "text", "text": text_content}) + tool_calls = msg.get("tool_calls") + if tool_calls: + for tc in tool_calls: + fn = tc.get("function", {}) + args_str = fn.get("arguments", "{}") + try: + args_dict = json.loads(args_str) + except (json.JSONDecodeError, TypeError): + args_dict = {} + content_blocks.append( + { + "type": "tool_use", + "id": tc.get("id", "call_0"), + "name": fn.get("name", ""), + "input": args_dict, + } + ) + if not content_blocks: + content_blocks.append({"type": "text", "text": ""}) + claude_messages.append( + { + "role": "assistant", + "content": content_blocks, + } + ) + elif role == "tool": + claude_messages.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", "call_0"), + "content": msg.get("content", ""), + } + ], + } + ) + elif role == "system": + pass + + return claude_messages + + +def _claude_build_request( + model: str, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], +) -> tuple[str, dict[str, Any]]: + """Build request for /v1/messages.""" + claude_messages = _claude_translate_messages(messages) + claude_tools = _claude_translate_tools(tools) + + system_content: str | None = None + for msg in messages: + if msg["role"] == "system": + system_content = msg["content"] + break + + body: dict[str, Any] = { + "model": model, + "messages": claude_messages, + "tools": claude_tools, + "max_tokens": 16384, + "temperature": 0.0, + } + if system_content is not None: + body["system"] = system_content + + return "/v1/messages", body + + +def _claude_parse_response(data: dict[str, Any]) -> ParsedResponse: + """Parse Claude Messages response into common format.""" + stop_reason = data.get("stop_reason", "") + content_blocks = data.get("content", []) + + if stop_reason == "tool_use": + finish_reason = "tool_calls" + elif stop_reason == "end_turn": + finish_reason = "stop" + else: + finish_reason = stop_reason + + tool_call_info: dict[str, str] | None = None + text_parts: list[str] = [] + has_tool_call = False + + for block in content_blocks: + block_type = block.get("type") + if block_type == "tool_use": + has_tool_call = True + if tool_call_info is None: + input_data = block.get("input", {}) + tool_call_info = { + "id": block.get("id", "call_0"), + "name": block.get("name", ""), + "arguments": json.dumps(input_data) + if isinstance(input_data, dict) + else str(input_data), + } + elif block_type == "text": + text = block.get("text", "") + if text.strip(): + text_parts.append(text) + + content = "\n".join(text_parts) if text_parts else None + + return ParsedResponse( + finish_reason=finish_reason, + has_tool_call=has_tool_call, + tool_call=tool_call_info, + content=content, + ) + + +def _claude_build_followup( + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], + model: str, + parsed: ParsedResponse, + tool_result: str, +) -> tuple[str, dict[str, Any]]: + """Build multi-turn follow-up for Claude Messages.""" + assert parsed.tool_call is not None + tc = parsed.tool_call + + try: + args_dict = json.loads(tc["arguments"]) + except (json.JSONDecodeError, TypeError): + args_dict = {} + + claude_messages = _claude_translate_messages(messages) + + claude_messages.append( + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": tc["id"], + "name": tc["name"], + "input": args_dict, + } + ], + } + ) + + claude_messages.append( + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": tc["id"], + "content": tool_result, + } + ], + } + ) + + claude_tools = _claude_translate_tools(tools) + + system_content: str | None = None + for msg in messages: + if msg["role"] == "system": + system_content = msg["content"] + break + + body: dict[str, Any] = { + "model": model, + "messages": claude_messages, + "tools": claude_tools, + "max_tokens": 16384, + "temperature": 0.0, + } + if system_content is not None: + body["system"] = system_content + + return "/v1/messages", body + + +def _responses_translate_input(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Translate OpenAI chat messages to Responses API input items.""" + items: list[dict[str, Any]] = [] + + for msg in messages: + role = msg["role"] + + if role in ("user", "system"): + items.append( + { + "type": "message", + "role": role, + "content": msg["content"], + } + ) + elif role == "assistant": + text_content = msg.get("content") + if text_content and isinstance(text_content, str) and text_content.strip(): + items.append( + { + "type": "message", + "role": "assistant", + "content": text_content, + } + ) + tool_calls = msg.get("tool_calls") + if tool_calls: + for tc in tool_calls: + fn = tc.get("function", {}) + items.append( + { + "type": "function_call", + "call_id": tc.get("id", "call_0"), + "name": fn.get("name", ""), + "arguments": fn.get("arguments", "{}"), + } + ) + elif role == "tool": + items.append( + { + "type": "function_call_output", + "call_id": msg.get("tool_call_id", "call_0"), + "output": msg.get("content", ""), + } + ) + + return items + + +def _responses_build_request( + model: str, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], +) -> tuple[str, dict[str, Any]]: + """Build request for /v1/responses.""" + input_items = _responses_translate_input(messages) + + body: dict[str, Any] = { + "model": model, + "input": input_items, + "tools": tools, + "temperature": 0.0, + "max_output_tokens": 4096, + } + return "/v1/responses", body + + +def _responses_parse_response(data: dict[str, Any]) -> ParsedResponse: + """Parse OpenAI Responses API response into common format.""" + output = data.get("output", []) + + tool_call_info: dict[str, str] | None = None + text_parts: list[str] = [] + has_tool_call = False + + for item in output: + item_type = item.get("type") + if item_type == "function_call": + has_tool_call = True + if tool_call_info is None: + tool_call_info = { + "id": item.get("call_id", "call_0"), + "name": item.get("name", ""), + "arguments": item.get("arguments", "{}"), + } + elif item_type == "message": + msg_content = item.get("content", []) + if isinstance(msg_content, list): + for block in msg_content: + if isinstance(block, dict): + text = block.get("text", "") + if text and text.strip(): + text_parts.append(text) + elif isinstance(msg_content, str) and msg_content.strip(): + text_parts.append(msg_content) + + content = "\n".join(text_parts) if text_parts else None + + if has_tool_call: + finish_reason = "tool_calls" + else: + status = data.get("status", "completed") + finish_reason = "stop" if status == "completed" else status + + return ParsedResponse( + finish_reason=finish_reason, + has_tool_call=has_tool_call, + tool_call=tool_call_info, + content=content, + ) + + +def _responses_build_followup( + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], + model: str, + parsed: ParsedResponse, + tool_result: str, +) -> tuple[str, dict[str, Any]]: + """Build multi-turn follow-up for Responses API.""" + assert parsed.tool_call is not None + tc = parsed.tool_call + + input_items = _responses_translate_input(messages) + + input_items.append( + { + "type": "function_call", + "call_id": tc["id"], + "name": tc["name"], + "arguments": tc["arguments"], + } + ) + + input_items.append( + { + "type": "function_call_output", + "call_id": tc["id"], + "output": tool_result, + } + ) + + body: dict[str, Any] = { + "model": model, + "input": input_items, + "tools": tools, + "temperature": 0.0, + "max_output_tokens": 4096, + } + return "/v1/responses", body + + +ADAPTERS: dict[ApiName, dict[str, Any]] = { + "openai": { + "build_request": _openai_build_request, + "parse_response": _openai_parse_response, + "build_followup": _openai_build_followup, + }, + "claude": { + "build_request": _claude_build_request, + "parse_response": _claude_parse_response, + "build_followup": _claude_build_followup, + }, + "responses": { + "build_request": _responses_build_request, + "parse_response": _responses_parse_response, + "build_followup": _responses_build_followup, + }, +} + + +def run_scenario( + client: httpx.Client, + host: str, + port: int, + model: str, + scenario: Scenario, + api_name: ApiName, + timeout: float, + verbose: bool, +) -> list[ScenarioResult]: + """Run a single scenario against one API adapter. Returns 1-2 results.""" + adapter = ADAPTERS[api_name] + build_request = adapter["build_request"] + parse_response = adapter["parse_response"] + build_followup = adapter["build_followup"] + results: list[ScenarioResult] = [] + + # --- Phase 1: initial request --- + path, body = build_request(model, scenario.messages, scenario.tools) + + if verbose: + print( + f" [{api_name}] request: {path} {json.dumps(body, indent=2)}", + file=sys.stderr, + ) + + try: + data, latency = call_api(client, host, port, path, body, timeout) + except Exception as exc: + results.append( + ScenarioResult( + name=scenario.name, + api=api_name, + phase="tool_call", + passed=False, + error=f"API error: {exc}", + ) + ) + return results + + if verbose: + print( + f" [{api_name}] response: {json.dumps(data, indent=2)}", file=sys.stderr + ) + + parsed = parse_response(data) + checks: dict[str, bool] = {} + + if scenario.expect_tool_call: + checks["finish_reason_tool_calls"] = parsed.finish_reason == "tool_calls" + checks["has_tool_call"] = parsed.has_tool_call + + args_err: str | None = None + if parsed.has_tool_call and parsed.tool_call is not None: + checks["correct_function"] = ( + scenario.expected_function is None + or parsed.tool_call["name"] == scenario.expected_function + ) + if scenario.required_arg_keys: + ok, args_err = validate_args( + parsed.tool_call["arguments"], scenario.required_arg_keys + ) + checks["valid_arguments"] = ok + else: + checks["valid_arguments"] = True + if scenario.nested_array_key and scenario.required_item_keys: + ok, nested_err = validate_nested_args( + parsed.tool_call["arguments"], + scenario.nested_array_key, + scenario.required_item_keys, + ) + checks["valid_nested_structure"] = ok + if not ok: + args_err = nested_err + else: + checks["correct_function"] = False + checks["valid_arguments"] = False + args_err = "No tool call returned" + + passed = all(checks.values()) + error = args_err if not passed else None + else: + checks["finish_reason_stop"] = parsed.finish_reason == "stop" + checks["no_tool_call"] = not parsed.has_tool_call + checks["has_content"] = ( + parsed.content is not None and len(parsed.content.strip()) > 0 + ) + passed = all(checks.values()) + error = ( + None + if passed + else ( + f"finish_reason={parsed.finish_reason}, " + f"tool_call={'yes' if parsed.has_tool_call else 'no'}, " + f"content={'yes' if parsed.content else 'no'}" + ) + ) + + results.append( + ScenarioResult( + name=scenario.name, + api=api_name, + phase="tool_call", + passed=passed, + checks=checks, + error=error, + latency_ms=latency, + ) + ) + + # --- Phase 2: multi-turn follow-up --- + if ( + scenario.tool_result is not None + and parsed.has_tool_call + and parsed.tool_call is not None + ): + followup_path, followup_body = build_followup( + scenario.messages, + scenario.tools, + model, + parsed, + scenario.tool_result, + ) + + if verbose: + print( + f" [{api_name}] follow_up request: {followup_path} {json.dumps(followup_body, indent=2)}", + file=sys.stderr, + ) + + try: + data2, latency2 = call_api( + client, host, port, followup_path, followup_body, timeout + ) + except Exception as exc: + results.append( + ScenarioResult( + name=scenario.name, + api=api_name, + phase="follow_up", + passed=False, + error=f"API error: {exc}", + ) + ) + return results + + if verbose: + print( + f" [{api_name}] follow_up response: {json.dumps(data2, indent=2)}", + file=sys.stderr, + ) + + parsed2 = parse_response(data2) + checks2: dict[str, bool] = {} + checks2["finish_reason_stop"] = parsed2.finish_reason == "stop" + checks2["no_tool_call"] = not parsed2.has_tool_call + checks2["has_content"] = ( + parsed2.content is not None and len(parsed2.content.strip()) > 0 + ) + + passed2 = all(checks2.values()) + error2: str | None = None + if not passed2: + error2 = ( + f"finish_reason={parsed2.finish_reason}, " + f"tool_call={'yes' if parsed2.has_tool_call else 'no'}, " + f"content={'yes' if parsed2.content else 'no'}" + ) + results.append( + ScenarioResult( + name=scenario.name, + api=api_name, + phase="follow_up", + passed=passed2, + checks=checks2, + error=error2, + latency_ms=latency2, + ) + ) + + return results + + +def result_to_dict(result: ScenarioResult) -> dict[str, Any]: + """Convert a ScenarioResult to a JSON-serializable dict.""" + return { + "name": result.name, + "api": result.api, + "phase": result.phase, + "passed": result.passed, + "checks": result.checks, + "error": result.error, + "latency_ms": round(result.latency_ms, 1), + } + + +_MULTI_NODE_PRIORITY: dict[tuple[str, str], int] = { + ("tensor", "jaccl"): 0, + ("pipeline", "jaccl"): 2, + ("pipeline", "ring"): 3, + ("tensor", "ring"): 4, +} +_SINGLE_NODE_PRIORITY = 1 + + +def _placement_sort_key(p: dict[str, Any]) -> tuple[int, int]: + sharding = p.get("sharding", "").lower() + meta = p.get("instance_meta", "").lower() + kind = ( + "tensor" if "tensor" in sharding else "pipeline", + "jaccl" if "jaccl" in meta else "ring", + ) + n_nodes = nodes_used_in_instance(p["instance"]) + if n_nodes == 1: + return (_SINGLE_NODE_PRIORITY, -n_nodes) + priority = _MULTI_NODE_PRIORITY.get(kind, 99) + return (priority, -n_nodes) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Multi-API tool-calling eval for exo", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog="""\ +Examples: + %(prog)s --model mlx-community/Qwen3-30B-A3B-4bit + %(prog)s --model my-model --api openai --repeat 3 + %(prog)s --model my-model --api all --scenarios weather_simple calculator_multi_turn + %(prog)s --model my-model --stdout +""", + ) + add_common_instance_args(parser) + parser.add_argument( + "--api", + choices=["openai", "claude", "responses", "all"], + default="all", + help="Which API adapter(s) to test (default: all)", + ) + parser.add_argument( + "--repeat", + type=int, + default=1, + help="Repeat each scenario N times (default: 1)", + ) + parser.add_argument( + "--scenarios", + nargs="*", + help="Run only these scenarios (by name)", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Print full API responses to stderr", + ) + parser.add_argument( + "--json-out", + default="bench/eval_results.json", + help="Write JSON results to file (default: bench/eval_results.json)", + ) + parser.add_argument( + "--stdout", + action="store_true", + help="Write JSON results to stdout instead of file", + ) + args = parser.parse_args() + + all_scenarios = load_scenarios(SCENARIOS_PATH) + if args.scenarios: + scenarios = [s for s in all_scenarios if s.name in args.scenarios] + if not scenarios: + print( + f"No matching scenarios. Available: {[s.name for s in all_scenarios]}", + file=sys.stderr, + ) + sys.exit(1) + else: + scenarios = all_scenarios + + api_names: list[ApiName] = ( + ["openai", "claude", "responses"] if args.api == "all" else [args.api] + ) + + log = sys.stderr if args.stdout else sys.stdout + exo = ExoClient(args.host, args.port, timeout_s=args.timeout) + _short_id, full_model_id = resolve_model_short_id(exo, args.model) + + selected = settle_and_fetch_placements( + exo, full_model_id, args, settle_timeout=args.settle_timeout + ) + if not selected: + print("No valid placements matched your filters.", file=sys.stderr) + sys.exit(1) + + selected.sort(key=_placement_sort_key) + preview = selected[0] + + settle_deadline = ( + time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None + ) + + print("Planning phase: checking downloads...", file=log) + run_planning_phase( + exo, + full_model_id, + preview, + args.danger_delete_downloads, + args.timeout, + settle_deadline, + ) + + 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) + + print(f"Model: {full_model_id}", file=log) + print(f"Placement: {sharding} / {instance_meta} / {n_nodes} nodes", file=log) + print(f"Endpoint: http://{args.host}:{args.port}", file=log) + print(f"APIs: {', '.join(api_names)}", file=log) + + total_runs = len(scenarios) * args.repeat * len(api_names) + print( + f"Scenarios: {len(scenarios)} x {args.repeat} repeats x {len(api_names)} APIs = {total_runs} runs", + file=log, + ) + print("=" * 72, file=log) + + exo.request_json("POST", "/instance", body={"instance": instance}) + try: + wait_for_instance_ready(exo, instance_id) + except (RuntimeError, TimeoutError) as e: + print(f"Failed to initialize placement: {e}", file=sys.stderr) + with contextlib.suppress(ExoHttpError): + exo.request_json("DELETE", f"/instance/{instance_id}") + sys.exit(1) + + time.sleep(1) + all_results: list[ScenarioResult] = [] + + try: + with httpx.Client() as http_client: + 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) + + 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}") + except ExoHttpError as e: + if e.status != 404: + raise + wait_for_instance_gone(exo, instance_id) + + # --- Summary --- + print(f"\n{'=' * 72}", file=log) + + total = len(all_results) + passed = sum(1 for r in all_results if r.passed) + + tool_call_results = [r for r in all_results if r.phase == "tool_call"] + follow_up_results = [r for r in all_results if r.phase == "follow_up"] + tc_passed = sum(1 for r in tool_call_results if r.passed) + fu_passed = sum(1 for r in follow_up_results if r.passed) + avg_latency = sum(r.latency_ms for r in all_results) / total if total else 0 + + print( + f"Total: {passed}/{total} passed ({100 * passed / total:.0f}%)", file=log + ) + print(f"Tool call: {tc_passed}/{len(tool_call_results)} passed", file=log) + if follow_up_results: + print(f"Follow-up: {fu_passed}/{len(follow_up_results)} passed", file=log) + print(f"Avg latency: {avg_latency:.0f}ms", file=log) + + for api_name in api_names: + api_results = [r for r in all_results if r.api == api_name] + api_passed = sum(1 for r in api_results if r.passed) + print(f" {api_name:>9}: {api_passed}/{len(api_results)} passed", file=log) + + if passed < total: + print("\nFailed:", file=log) + for r in all_results: + if not r.passed: + print(f" - {r.name} [{r.api}/{r.phase}]: {r.error}", file=log) + + json_results = [result_to_dict(r) for r in all_results] + + if args.stdout: + print(json.dumps(json_results, indent=2)) + else: + json_path = args.json_out + parent = os.path.dirname(json_path) + if parent: + os.makedirs(parent, exist_ok=True) + with open(json_path, "w") as f: + json.dump(json_results, f, indent=2) + f.write("\n") + print(f"\nJSON results written to {json_path}", file=log) + + sys.exit(0 if passed == total else 1) + + +if __name__ == "__main__": + main() diff --git a/bench/exo_bench.py b/bench/exo_bench.py index 56bb0421..f6fcd342 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -1,29 +1,48 @@ +# type: ignore #!/usr/bin/env python3 -# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false +"""Tool-calling eval for exo's OpenAI-compatible API. + +Tests whether models correctly: +- Trigger tool calls when appropriate +- Return valid JSON arguments matching function schemas +- Handle multi-turn tool use (call -> result -> final answer) +- Avoid calling tools when unnecessary + +Start exo with a model first, then run: + uv run python tool_call_eval.py --model + uv run python tool_call_eval.py --model --host 10.0.0.5 --port 52415 + uv run python tool_call_eval.py --model --repeat 3 + uv run python tool_call_eval.py --model --scenarios weather_simple calculator_multi_turn +""" + from __future__ import annotations import argparse import contextlib -import http.client import itertools import json -import os import sys import time from collections.abc import Callable from pathlib import Path from statistics import mean from typing import Any -from urllib.parse import urlencode +from harness import ( + ExoClient, + ExoHttpError, + add_common_instance_args, + instance_id_from_instance, + nodes_used_in_instance, + resolve_model_short_id, + run_planning_phase, + settle_and_fetch_placements, + wait_for_instance_gone, + wait_for_instance_ready, +) from loguru import logger from transformers import AutoTokenizer -# Backoff constants for cluster settling retry -_SETTLE_INITIAL_BACKOFF_S = 1.0 -_SETTLE_MAX_BACKOFF_S = 60.0 -_SETTLE_BACKOFF_MULTIPLIER = 2.0 - # Monkey-patch for transformers 5.x compatibility # Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location # which was moved in transformers 5.0.0rc2 @@ -103,154 +122,6 @@ def load_tokenizer_for_bench(model_id: str) -> Any: return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) -class ExoHttpError(RuntimeError): - def __init__(self, status: int, reason: str, body_preview: str): - super().__init__(f"HTTP {status} {reason}: {body_preview}") - self.status = status - - -class ExoClient: - def __init__(self, host: str, port: int, timeout_s: float = 7200.0): - self.host = host - self.port = port - self.timeout_s = timeout_s - - def request_json( - self, - method: str, - path: str, - params: dict[str, Any] | None = None, - body: dict[str, Any] | None = None, - headers: dict[str, str] | None = None, - ) -> Any: - if not path.startswith("/"): - path = "/" + path - if params: - path = path + "?" + urlencode(params) - - conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s) - try: - payload: bytes | None = None - hdrs: dict[str, str] = {"Accept": "application/json"} - - if body is not None: - payload = json.dumps(body).encode("utf-8") - hdrs["Content-Type"] = "application/json" - if headers: - hdrs.update(headers) - - conn.request(method.upper(), path, body=payload, headers=hdrs) - resp = conn.getresponse() - raw = resp.read() - text = raw.decode("utf-8", errors="replace") if raw else "" - - if resp.status >= 400: - raise ExoHttpError(resp.status, resp.reason, text[:300]) - - if not text: - return None - return json.loads(text) - finally: - conn.close() - - def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]: - return self.request_json("POST", "/bench/chat/completions", body=payload) - - -def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]: - if len(instance) != 1: - raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}") - - tag = next(iter(instance)) - inner = instance[tag] - if not isinstance(inner, dict): - raise TypeError(f"payload for {tag} must be dict, got {type(inner)}") - return inner - - -def instance_id_from_instance(instance: dict[str, Any]) -> str: - inner = unwrap_instance(instance) - return str(inner["instanceId"]) - - -def nodes_used_in_instance(instance: dict[str, Any]) -> int: - inner = unwrap_instance(instance) - return len(inner["shardAssignments"]["nodeToRunner"]) - - -def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]: - inner = unwrap_instance(instance) - runner_to_shard = inner["shardAssignments"]["runnerToShard"] - return list(runner_to_shard.keys()) - - -def runner_ready(runner: dict[str, Any]) -> bool: - return "RunnerReady" in runner - - -def runner_failed(runner: dict[str, Any]) -> bool: - return "RunnerFailed" in runner - - -def get_runner_failed_message(runner: dict[str, Any]) -> str | None: - if "RunnerFailed" in runner: - return runner["RunnerFailed"].get("errorMessage") - return None - - -def wait_for_instance_ready( - client: ExoClient, instance_id: str, timeout: float = 24000.0 -) -> None: - start_time = time.time() - instance_existed = False - while time.time() - start_time < timeout: - state = client.request_json("GET", "/state") - instances = state.get("instances", {}) - - if instance_id not in instances: - if instance_existed: - # Instance was deleted after being created - likely due to runner failure - raise RuntimeError( - f"Instance {instance_id} was deleted (runner may have failed)" - ) - time.sleep(0.1) - continue - - instance_existed = True - instance = instances[instance_id] - runner_ids = runner_ids_from_instance(instance) - runners = state.get("runners", {}) - - # Check for failed runners first - for rid in runner_ids: - runner = runners.get(rid, {}) - if runner_failed(runner): - error_msg = get_runner_failed_message(runner) or "Unknown error" - raise RuntimeError(f"Runner {rid} failed: {error_msg}") - - if all(runner_ready(runners.get(rid, {})) for rid in runner_ids): - return - - time.sleep(0.1) - - raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}") - - -def wait_for_instance_gone( - client: ExoClient, instance_id: str, timeout: float = 3.0 -) -> None: - start_time = time.time() - while time.time() - start_time < timeout: - try: - client.request_json("GET", f"/instance/{instance_id}") - time.sleep(0.4) - except ExoHttpError as e: - if e.status == 404: - return - - raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}") - - def format_peak_memory(b: float) -> str: for unit in ["B", "KB", "MB", "GB", "TB"]: if b < 1024.0: @@ -269,184 +140,6 @@ def parse_int_list(values: list[str]) -> list[int]: return items -def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]: - models = client.request_json("GET", "/models") or {} - data = models.get("data") or [] - - for m in data: - if m.get("name").lower() == model_arg.lower(): - short_id = str(m["name"]) - full_id = str(m.get("hugging_face_id") or m["name"]) - return short_id, full_id - - for m in data: - if m.get("hugging_face_id") == model_arg: - short_id = str(m["name"]) - full_id = str(m["hugging_face_id"]) - return short_id, full_id - - raise ValueError(f"Model not found in /models: {model_arg}") - - -def run_planning_phase( - client: ExoClient, - full_model_id: str, - preview: dict[str, Any], - danger_delete: bool, - timeout: float, - settle_deadline: float | None, -) -> None: - """Check disk space and ensure model is downloaded before benchmarking.""" - # Get model size from /models - models = client.request_json("GET", "/models") or {} - model_bytes = 0 - for m in models.get("data", []): - if m.get("hugging_face_id") == full_model_id: - model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024 - break - - if not model_bytes: - logger.warning( - f"Could not determine size for {full_model_id}, skipping disk check" - ) - return - - # Get nodes from preview - inner = unwrap_instance(preview["instance"]) - node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys()) - runner_to_shard = inner["shardAssignments"]["runnerToShard"] - - state = client.request_json("GET", "/state") - downloads = state.get("downloads", {}) - node_disk = state.get("nodeDisk", {}) - - for node_id in node_ids: - node_downloads = downloads.get(node_id, []) - - # Check if model already downloaded on this node - already_downloaded = any( - "DownloadCompleted" in p - and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ - "modelId" - ] - == full_model_id - for p in node_downloads - ) - if already_downloaded: - continue - - # Wait for disk info if settle_deadline is set - disk_info = node_disk.get(node_id, {}) - backoff = _SETTLE_INITIAL_BACKOFF_S - while not disk_info and settle_deadline and time.monotonic() < settle_deadline: - remaining = settle_deadline - time.monotonic() - logger.info( - f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..." - ) - time.sleep(min(backoff, remaining)) - backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S) - state = client.request_json("GET", "/state") - node_disk = state.get("nodeDisk", {}) - disk_info = node_disk.get(node_id, {}) - - if not disk_info: - logger.warning(f"No disk info for {node_id}, skipping space check") - continue - - avail = disk_info.get("available", {}).get("inBytes", 0) - if avail >= model_bytes: - continue - - if not danger_delete: - raise RuntimeError( - f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, " - f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space." - ) - - # Delete from smallest to largest - completed = [ - ( - unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ - "modelId" - ], - p["DownloadCompleted"]["totalBytes"]["inBytes"], - ) - for p in node_downloads - if "DownloadCompleted" in p - ] - for del_model, size in sorted(completed, key=lambda x: x[1]): - logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)") - client.request_json("DELETE", f"/download/{node_id}/{del_model}") - avail += size - if avail >= model_bytes: - break - - if avail < model_bytes: - raise RuntimeError(f"Could not free enough space on {node_id}") - - # Start downloads (idempotent) - for node_id in node_ids: - runner_id = inner["shardAssignments"]["nodeToRunner"][node_id] - shard = runner_to_shard[runner_id] - client.request_json( - "POST", - "/download/start", - body={ - "targetNodeId": node_id, - "shardMetadata": shard, - }, - ) - logger.info(f"Started download on {node_id}") - - # Wait for downloads - start = time.time() - while time.time() - start < timeout: - state = client.request_json("GET", "/state") - downloads = state.get("downloads", {}) - all_done = True - for node_id in node_ids: - done = any( - "DownloadCompleted" in p - and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[ - "modelCard" - ]["modelId"] - == full_model_id - for p in downloads.get(node_id, []) - ) - failed = [ - p["DownloadFailed"]["errorMessage"] - for p in downloads.get(node_id, []) - if "DownloadFailed" in p - and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][ - "modelId" - ] - == full_model_id - ] - if failed: - raise RuntimeError(f"Download failed on {node_id}: {failed[0]}") - if not done: - all_done = False - if all_done: - return - time.sleep(1) - - raise TimeoutError("Downloads did not complete in time") - - -def placement_filter(instance_meta: str, wanted: str) -> bool: - s = (instance_meta or "").lower() - if wanted == "both": - return ("ring" in s) or ("jaccl" in s) - return wanted in s - - -def sharding_filter(sharding: str, wanted: str) -> bool: - s = (sharding or "").lower() - if wanted == "both": - return ("pipeline" in s) or ("tensor" in s) - return wanted in s - - def run_one_completion( client: ExoClient, model_id: str, pp_hint: int, tg: int, prompt_sizer: PromptSizer ) -> tuple[dict[str, Any], int]: @@ -538,76 +231,12 @@ class PromptSizer: return content, tok -def fetch_and_filter_placements( - client: ExoClient, full_model_id: str, args: argparse.Namespace -) -> list[dict[str, Any]]: - previews_resp = client.request_json( - "GET", "/instance/previews", params={"model_id": full_model_id} - ) - previews = previews_resp.get("previews") or [] - - selected: list[dict[str, Any]] = [] - for p in previews: - if p.get("error") is not None: - continue - if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta): - continue - if not sharding_filter(str(p.get("sharding", "")), args.sharding): - continue - - instance = p.get("instance") - if not isinstance(instance, dict): - continue - - n = nodes_used_in_instance(instance) - # Skip tensor ring single node as it is pointless when pipeline ring - if n == 1 and ( - (args.sharding == "both" and "tensor" in p.get("sharding", "").lower()) - or ( - args.instance_meta == "both" - and "jaccl" in p.get("instance_meta", "").lower() - ) - ): - continue - - if ( - args.skip_pipeline_jaccl - and ( - args.instance_meta == "both" - and "jaccl" in p.get("instance_meta", "").lower() - ) - and ( - args.sharding == "both" and "pipeline" in p.get("sharding", "").lower() - ) - ): - continue - - if ( - args.skip_tensor_ring - and ( - args.instance_meta == "both" - and "ring" in p.get("instance_meta", "").lower() - ) - and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower()) - ): - continue - - if args.min_nodes <= n <= args.max_nodes: - selected.append(p) - - return selected - - def main() -> int: ap = argparse.ArgumentParser( prog="exo-bench", description="Benchmark exo model throughput across placement previews.", ) - ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost")) - ap.add_argument( - "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415")) - ) - ap.add_argument("--model", required=True, help="Model short id or huggingface id") + add_common_instance_args(ap) ap.add_argument( "--pp", nargs="+", @@ -620,34 +249,6 @@ def main() -> int: required=True, help="Generation lengths (ints). Accepts commas.", ) - ap.add_argument( - "--max-nodes", - type=int, - default=4, - help="Only consider placements using <= this many nodes.", - ) - ap.add_argument( - "--min-nodes", - type=int, - default=1, - help="Only consider placements using >= this many nodes.", - ) - ap.add_argument( - "--instance-meta", choices=["ring", "jaccl", "both"], default="both" - ) - ap.add_argument( - "--sharding", choices=["pipeline", "tensor", "both"], default="both" - ) - ap.add_argument( - "--skip-pipeline-jaccl", - action="store_true", - help="Skip pipeline+jaccl placements, as it's often pointless.", - ) - ap.add_argument( - "--skip-tensor-ring", - action="store_true", - help="Skip tensor+ring placements, as it's so slow.", - ) ap.add_argument( "--repeat", type=int, default=1, help="Repetitions per (pp,tg) pair." ) @@ -657,9 +258,6 @@ def main() -> int: default=0, help="Warmup runs per placement (uses first pp/tg).", ) - ap.add_argument( - "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)." - ) ap.add_argument( "--json-out", default="bench/results.json", @@ -674,17 +272,6 @@ def main() -> int: action="store_true", help="Force all pp×tg combinations (cartesian product) even when lists have equal length.", ) - ap.add_argument( - "--settle-timeout", - type=float, - default=0, - help="Max seconds to wait for the cluster to produce valid placements (0 = try once).", - ) - ap.add_argument( - "--danger-delete-downloads", - action="store_true", - help="Delete existing models from smallest to largest to make room for benchmark model.", - ) args = ap.parse_args() pp_list = parse_int_list(args.pp) @@ -719,24 +306,10 @@ def main() -> int: logger.error("[exo-bench] tokenizer usable but prompt sizing failed") raise - settle_deadline = ( - time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None + selected = settle_and_fetch_placements( + client, full_model_id, args, settle_timeout=args.settle_timeout ) - selected = fetch_and_filter_placements(client, full_model_id, args) - - if not selected and settle_deadline: - backoff = _SETTLE_INITIAL_BACKOFF_S - while not selected and time.monotonic() < settle_deadline: - remaining = settle_deadline - time.monotonic() - logger.warning( - f"No valid placements yet (cluster may still be settling). " - f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..." - ) - time.sleep(min(backoff, remaining)) - backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S) - selected = fetch_and_filter_placements(client, full_model_id, args) - if not selected: logger.error("No valid placements matched your filters.") return 1 @@ -760,6 +333,10 @@ def main() -> int: if args.dry_run: return 0 + settle_deadline = ( + time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None + ) + logger.info("Planning phase: checking downloads...") run_planning_phase( client, diff --git a/bench/harness.py b/bench/harness.py new file mode 100644 index 00000000..58aa8435 --- /dev/null +++ b/bench/harness.py @@ -0,0 +1,477 @@ +# type: ignore +from __future__ import annotations + +import argparse +import http.client +import json +import os +import time +from typing import Any +from urllib.parse import urlencode + +from loguru import logger + +_SETTLE_INITIAL_BACKOFF_S = 1.0 +_SETTLE_MAX_BACKOFF_S = 60.0 +_SETTLE_BACKOFF_MULTIPLIER = 2.0 + + +class ExoHttpError(RuntimeError): + def __init__(self, status: int, reason: str, body_preview: str): + super().__init__(f"HTTP {status} {reason}: {body_preview}") + self.status = status + + +class ExoClient: + def __init__(self, host: str, port: int, timeout_s: float = 7200.0): + self.host = host + self.port = port + self.timeout_s = timeout_s + + def request_json( + self, + method: str, + path: str, + params: dict[str, Any] | None = None, + body: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, + ) -> Any: + if not path.startswith("/"): + path = "/" + path + if params: + path = path + "?" + urlencode(params) + + conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s) + try: + payload: bytes | None = None + hdrs: dict[str, str] = {"Accept": "application/json"} + + if body is not None: + payload = json.dumps(body).encode("utf-8") + hdrs["Content-Type"] = "application/json" + if headers: + hdrs.update(headers) + + conn.request(method.upper(), path, body=payload, headers=hdrs) + resp = conn.getresponse() + raw = resp.read() + text = raw.decode("utf-8", errors="replace") if raw else "" + + if resp.status >= 400: + raise ExoHttpError(resp.status, resp.reason, text[:300]) + + if not text: + return None + return json.loads(text) + finally: + conn.close() + + def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]: + return self.request_json("POST", "/bench/chat/completions", body=payload) + + +def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]: + if len(instance) != 1: + raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}") + + tag = next(iter(instance)) + inner = instance[tag] + if not isinstance(inner, dict): + raise TypeError(f"payload for {tag} must be dict, got {type(inner)}") + return inner + + +def instance_id_from_instance(instance: dict[str, Any]) -> str: + inner = unwrap_instance(instance) + return str(inner["instanceId"]) + + +def nodes_used_in_instance(instance: dict[str, Any]) -> int: + inner = unwrap_instance(instance) + return len(inner["shardAssignments"]["nodeToRunner"]) + + +def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]: + inner = unwrap_instance(instance) + runner_to_shard = inner["shardAssignments"]["runnerToShard"] + return list(runner_to_shard.keys()) + + +def runner_ready(runner: dict[str, Any]) -> bool: + return "RunnerReady" in runner + + +def runner_failed(runner: dict[str, Any]) -> bool: + return "RunnerFailed" in runner + + +def get_runner_failed_message(runner: dict[str, Any]) -> str | None: + if "RunnerFailed" in runner: + return runner["RunnerFailed"].get("errorMessage") + return None + + +def wait_for_instance_ready( + client: ExoClient, instance_id: str, timeout: float = 24000.0 +) -> None: + start_time = time.time() + instance_existed = False + while time.time() - start_time < timeout: + state = client.request_json("GET", "/state") + instances = state.get("instances", {}) + + if instance_id not in instances: + if instance_existed: + # Instance was deleted after being created - likely due to runner failure + raise RuntimeError( + f"Instance {instance_id} was deleted (runner may have failed)" + ) + time.sleep(0.1) + continue + + instance_existed = True + instance = instances[instance_id] + runner_ids = runner_ids_from_instance(instance) + runners = state.get("runners", {}) + + # Check for failed runners first + for rid in runner_ids: + runner = runners.get(rid, {}) + if runner_failed(runner): + error_msg = get_runner_failed_message(runner) or "Unknown error" + raise RuntimeError(f"Runner {rid} failed: {error_msg}") + + if all(runner_ready(runners.get(rid, {})) for rid in runner_ids): + return + + time.sleep(0.1) + + raise TimeoutError(f"Instance {instance_id} did not become ready within {timeout=}") + + +def wait_for_instance_gone( + client: ExoClient, instance_id: str, timeout: float = 3.0 +) -> None: + start_time = time.time() + while time.time() - start_time < timeout: + try: + client.request_json("GET", f"/instance/{instance_id}") + time.sleep(0.4) + except ExoHttpError as e: + if e.status == 404: + return + raise + + raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}") + + +def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]: + models = client.request_json("GET", "/models") or {} + data = models.get("data") or [] + + for m in data: + if (m.get("name") or "").lower() == model_arg.lower(): + short_id = str(m["name"]) + full_id = str(m.get("hugging_face_id") or m["name"]) + return short_id, full_id + + for m in data: + if m.get("hugging_face_id") == model_arg: + short_id = str(m["name"]) + full_id = str(m["hugging_face_id"]) + return short_id, full_id + + raise ValueError(f"Model not found in /models: {model_arg}") + + +def placement_filter(instance_meta: str, wanted: str) -> bool: + s = (instance_meta or "").lower() + if wanted == "both": + return ("ring" in s) or ("jaccl" in s) + return wanted in s + + +def sharding_filter(sharding: str, wanted: str) -> bool: + s = (sharding or "").lower() + if wanted == "both": + return ("pipeline" in s) or ("tensor" in s) + return wanted in s + + +def fetch_and_filter_placements( + client: ExoClient, full_model_id: str, args: argparse.Namespace +) -> list[dict[str, Any]]: + previews_resp = client.request_json( + "GET", "/instance/previews", params={"model_id": full_model_id} + ) + previews = previews_resp.get("previews") or [] + + selected: list[dict[str, Any]] = [] + for p in previews: + if p.get("error") is not None: + continue + if not placement_filter(str(p.get("instance_meta", "")), args.instance_meta): + continue + if not sharding_filter(str(p.get("sharding", "")), args.sharding): + continue + + instance = p.get("instance") + if not isinstance(instance, dict): + continue + + n = nodes_used_in_instance(instance) + # Skip tensor ring single node as it is pointless when pipeline ring + if n == 1 and ( + (args.sharding == "both" and "tensor" in p.get("sharding", "").lower()) + or ( + args.instance_meta == "both" + and "jaccl" in p.get("instance_meta", "").lower() + ) + ): + continue + + if ( + args.skip_pipeline_jaccl + and ( + args.instance_meta == "both" + and "jaccl" in p.get("instance_meta", "").lower() + ) + and ( + args.sharding == "both" and "pipeline" in p.get("sharding", "").lower() + ) + ): + continue + + if ( + args.skip_tensor_ring + and ( + args.instance_meta == "both" + and "ring" in p.get("instance_meta", "").lower() + ) + and (args.sharding == "both" and "tensor" in p.get("sharding", "").lower()) + ): + continue + + if args.min_nodes <= n <= args.max_nodes: + selected.append(p) + + return selected + + +def settle_and_fetch_placements( + client: ExoClient, + full_model_id: str, + args: argparse.Namespace, + settle_timeout: float = 0, +) -> list[dict[str, Any]]: + selected = fetch_and_filter_placements(client, full_model_id, args) + + if not selected and settle_timeout > 0: + backoff = _SETTLE_INITIAL_BACKOFF_S + deadline = time.monotonic() + settle_timeout + while not selected and time.monotonic() < deadline: + remaining = deadline - time.monotonic() + logger.warning( + f"No valid placements yet (cluster may still be settling). " + f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..." + ) + time.sleep(min(backoff, remaining)) + backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S) + selected = fetch_and_filter_placements(client, full_model_id, args) + + return selected + + +def run_planning_phase( + client: ExoClient, + full_model_id: str, + preview: dict[str, Any], + danger_delete: bool, + timeout: float, + settle_deadline: float | None, +) -> None: + """Check disk space and ensure model is downloaded before benchmarking.""" + # Get model size from /models + models = client.request_json("GET", "/models") or {} + model_bytes = 0 + for m in models.get("data", []): + if m.get("hugging_face_id") == full_model_id: + model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024 + break + + if not model_bytes: + logger.warning( + f"Could not determine size for {full_model_id}, skipping disk check" + ) + return + + # Get nodes from preview + inner = unwrap_instance(preview["instance"]) + node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys()) + runner_to_shard = inner["shardAssignments"]["runnerToShard"] + + state = client.request_json("GET", "/state") + downloads = state.get("downloads", {}) + node_disk = state.get("nodeDisk", {}) + + for node_id in node_ids: + node_downloads = downloads.get(node_id, []) + + # Check if model already downloaded on this node + already_downloaded = any( + "DownloadCompleted" in p + and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ + "modelId" + ] + == full_model_id + for p in node_downloads + ) + if already_downloaded: + continue + + # Wait for disk info if settle_deadline is set + disk_info = node_disk.get(node_id, {}) + backoff = _SETTLE_INITIAL_BACKOFF_S + while not disk_info and settle_deadline and time.monotonic() < settle_deadline: + remaining = settle_deadline - time.monotonic() + logger.info( + f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..." + ) + time.sleep(min(backoff, remaining)) + backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S) + state = client.request_json("GET", "/state") + node_disk = state.get("nodeDisk", {}) + disk_info = node_disk.get(node_id, {}) + + if not disk_info: + logger.warning(f"No disk info for {node_id}, skipping space check") + continue + + avail = disk_info.get("available", {}).get("inBytes", 0) + if avail >= model_bytes: + continue + + if not danger_delete: + raise RuntimeError( + f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, " + f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space." + ) + + # Delete from smallest to largest + completed = [ + ( + unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][ + "modelId" + ], + p["DownloadCompleted"]["totalBytes"]["inBytes"], + ) + for p in node_downloads + if "DownloadCompleted" in p + ] + for del_model, size in sorted(completed, key=lambda x: x[1]): + logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)") + client.request_json("DELETE", f"/download/{node_id}/{del_model}") + avail += size + if avail >= model_bytes: + break + + if avail < model_bytes: + raise RuntimeError(f"Could not free enough space on {node_id}") + + # Start downloads (idempotent) + for node_id in node_ids: + runner_id = inner["shardAssignments"]["nodeToRunner"][node_id] + shard = runner_to_shard[runner_id] + client.request_json( + "POST", + "/download/start", + body={ + "targetNodeId": node_id, + "shardMetadata": shard, + }, + ) + logger.info(f"Started download on {node_id}") + + # Wait for downloads + start = time.time() + while time.time() - start < timeout: + state = client.request_json("GET", "/state") + downloads = state.get("downloads", {}) + all_done = True + for node_id in node_ids: + done = any( + "DownloadCompleted" in p + and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[ + "modelCard" + ]["modelId"] + == full_model_id + for p in downloads.get(node_id, []) + ) + failed = [ + p["DownloadFailed"]["errorMessage"] + for p in downloads.get(node_id, []) + if "DownloadFailed" in p + and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][ + "modelId" + ] + == full_model_id + ] + if failed: + raise RuntimeError(f"Download failed on {node_id}: {failed[0]}") + if not done: + all_done = False + if all_done: + return + time.sleep(1) + + raise TimeoutError("Downloads did not complete in time") + + +def add_common_instance_args(ap: argparse.ArgumentParser) -> None: + ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost")) + ap.add_argument( + "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415")) + ) + ap.add_argument("--model", required=True, help="Model short id or huggingface id") + ap.add_argument( + "--max-nodes", + type=int, + default=4, + help="Only consider placements using <= this many nodes.", + ) + ap.add_argument( + "--min-nodes", + type=int, + default=1, + help="Only consider placements using >= this many nodes.", + ) + ap.add_argument( + "--instance-meta", choices=["ring", "jaccl", "both"], default="both" + ) + ap.add_argument( + "--sharding", choices=["pipeline", "tensor", "both"], default="both" + ) + ap.add_argument( + "--skip-pipeline-jaccl", + action="store_true", + help="Skip pipeline+jaccl placements, as it's often pointless.", + ) + ap.add_argument( + "--skip-tensor-ring", + action="store_true", + help="Skip tensor+ring placements, as it's so slow.", + ) + ap.add_argument( + "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)." + ) + ap.add_argument( + "--settle-timeout", + type=float, + default=0, + help="Max seconds to wait for the cluster to produce valid placements (0 = try once).", + ) + ap.add_argument( + "--danger-delete-downloads", + action="store_true", + help="Delete existing models from smallest to largest to make room for benchmark model.", + ) diff --git a/bench/pyproject.toml b/bench/pyproject.toml index cc277231..a2a450f2 100644 --- a/bench/pyproject.toml +++ b/bench/pyproject.toml @@ -4,6 +4,7 @@ version = "0.1.0" description = "Benchmarking tool for exo distributed inference" requires-python = ">=3.13" dependencies = [ + "httpx>=0.27.0", "loguru>=0.7.3", "transformers>=5.0.0", "huggingface-hub>=0.33.4", diff --git a/bench/scenarios.toml b/bench/scenarios.toml new file mode 100644 index 00000000..892a044a --- /dev/null +++ b/bench/scenarios.toml @@ -0,0 +1,306 @@ +# Tool definitions — each becomes an OpenAI function tool. +# All scenarios get all tools unless they specify a `tools` list. + +[tools.get_current_weather] +description = "Get the current weather in a given location" +required = ["location"] + +[tools.get_current_weather.properties.location] +type = "string" +description = "City and state, e.g. San Francisco, CA" + +[tools.get_current_weather.properties.unit] +type = "string" +enum = ["celsius", "fahrenheit"] +description = "Temperature unit" + +[tools.calculate] +description = "Evaluate a mathematical expression and return the numeric result" +required = ["expression"] + +[tools.calculate.properties.expression] +type = "string" +description = "The math expression to evaluate, e.g. '2 + 3 * 4'" + +[tools.search_products] +description = "Search for products in a catalog by query, category, and price" +required = ["query"] + +[tools.search_products.properties.query] +type = "string" +description = "Search query string" + +[tools.search_products.properties.category] +type = "string" +enum = ["electronics", "clothing", "food", "books"] +description = "Product category to filter by" + +[tools.search_products.properties.max_price] +type = "number" +description = "Maximum price in USD" + +[tools.create_todos] +description = "Create a structured todo list" +required = ["todos"] + +[tools.create_todos.properties.todos] +type = "array" +description = "List of todo items" + +[tools.create_todos.properties.todos.items] +type = "object" +required = ["content", "status", "priority"] + +[tools.create_todos.properties.todos.items.properties.content] +type = "string" +description = "The todo item text" + +[tools.create_todos.properties.todos.items.properties.status] +type = "string" +description = "Status: pending, in_progress, or completed" + +[tools.create_todos.properties.todos.items.properties.priority] +type = "string" +description = "Priority: low, normal, or high" + +# -- Should call a tool -- + +[[scenarios]] +name = "weather_simple" +description = "Basic weather query -> get_current_weather" +expect_tool_call = true +expected_function = "get_current_weather" +required_arg_keys = ["location"] + +[[scenarios.messages]] +role = "user" +content = "What's the weather like in Tokyo right now?" + +[[scenarios]] +name = "calculator_simple" +description = "Math question -> calculate" +expect_tool_call = true +expected_function = "calculate" +required_arg_keys = ["expression"] + +[[scenarios.messages]] +role = "user" +content = "Use the calculator to compute 3847 * 926 + 17293" + +[[scenarios]] +name = "search_with_filters" +description = "Product search with category and price filter" +expect_tool_call = true +expected_function = "search_products" +required_arg_keys = ["query"] + +[[scenarios.messages]] +role = "user" +content = "Find me electronics under $50" + +# -- Multi-turn: tool call then follow-up -- + +[[scenarios]] +name = "weather_multi_turn" +description = "Weather query -> tool result -> natural language summary" +expect_tool_call = true +expected_function = "get_current_weather" +required_arg_keys = ["location"] + +[scenarios.tool_result] +temperature = "18C" +condition = "partly cloudy" +humidity = "65%" +wind = "12 km/h NW" + +[[scenarios.messages]] +role = "user" +content = "What's the weather in Paris?" + +[[scenarios]] +name = "calculator_multi_turn" +description = "Math query -> tool result -> model reports the answer" +expect_tool_call = true +expected_function = "calculate" +required_arg_keys = ["expression"] + +[scenarios.tool_result] +result = 491682 + +[[scenarios.messages]] +role = "user" +content = "Use the calculator to compute 1847 * 263 + 5921" + +[[scenarios]] +name = "search_multi_turn" +description = "Search query -> tool result -> model summarizes products" +expect_tool_call = true +expected_function = "search_products" +required_arg_keys = ["query"] + +[[scenarios.tool_result.results]] +name = "Hands-On Machine Learning" +price = 45.99 +rating = 4.8 + +[[scenarios.tool_result.results]] +name = "Deep Learning with Python" +price = 39.99 +rating = 4.6 + +[[scenarios.messages]] +role = "user" +content = "Search for books about machine learning" + +# -- Sequential tool calls -- + +[[scenarios]] +name = "chained_tool_calls_same" +description = "Thinking + weather(Tokyo) -> result -> model must call weather(London)" +expect_tool_call = true +expected_function = "get_current_weather" +required_arg_keys = ["location"] + +[[scenarios.messages]] +role = "user" +content = "Compare the weather in Tokyo and London." + +[[scenarios.messages]] +role = "assistant" +content = "I'll check both cities. Let me start with Tokyo." + +[[scenarios.messages.tool_calls]] +id = "call_1" +name = "get_current_weather" +arguments = { location = "Tokyo" } + +[[scenarios.messages]] +role = "tool" +tool_call_id = "call_1" +content = '{"temperature": "25C", "condition": "sunny"}' + +[[scenarios]] +name = "chained_tool_calls_different" +description = "Thinking + weather(Berlin) -> result -> model must call calculator" +expect_tool_call = true +expected_function = "calculate" +required_arg_keys = ["expression"] + +[[scenarios.messages]] +role = "user" +content = "What's the weather in Berlin, and also use the calculator to compute 4819 * 37 + 291." + +[[scenarios.messages]] +role = "assistant" +content = "I'll handle both. Let me check Berlin's weather first." + +[[scenarios.messages.tool_calls]] +id = "call_2" +name = "get_current_weather" +arguments = { location = "Berlin" } + +[[scenarios.messages]] +role = "tool" +tool_call_id = "call_2" +content = '{"temperature": "12C", "condition": "rainy"}' + +[[scenarios]] +name = "chained_tool_calls_three" +description = "Two prior thinking+tool calls -> results -> model must make a third" +expect_tool_call = true +expected_function = "get_current_weather" +required_arg_keys = ["location"] + +[[scenarios.messages]] +role = "user" +content = "Compare weather in Tokyo, Paris, and London." + +[[scenarios.messages]] +role = "assistant" +content = "I'll check all three cities. Starting with Tokyo." + +[[scenarios.messages.tool_calls]] +id = "call_3" +name = "get_current_weather" +arguments = { location = "Tokyo" } + +[[scenarios.messages]] +role = "tool" +tool_call_id = "call_3" +content = '{"temperature": "25C", "condition": "sunny"}' + +[[scenarios.messages]] +role = "assistant" +content = "Got Tokyo. Now checking Paris." + +[[scenarios.messages.tool_calls]] +id = "call_4" +name = "get_current_weather" +arguments = { location = "Paris" } + +[[scenarios.messages]] +role = "tool" +tool_call_id = "call_4" +content = '{"temperature": "18C", "condition": "cloudy"}' + +# -- Nested object schema (regression for lossy chat template rendering) -- + +[[scenarios]] +name = "nested_schema_tool_call" +description = "Tool call with nested object array schema -> create_todos" +expect_tool_call = true +expected_function = "create_todos" +required_arg_keys = ["todos"] +nested_array_key = "todos" +required_item_keys = ["content", "status", "priority"] +tools = ["create_todos"] + +[[scenarios.messages]] +role = "user" +content = "Create a todo list with 3 items to learn Python" + +# -- Tool name integrity (regression for harmony token leaking into name) -- + +[tools.glob] +description = "Search for files matching a glob pattern in the codebase" +required = ["pattern"] + +[tools.glob.properties.pattern] +type = "string" +description = "The glob pattern to match files against, e.g. '**/*.py'" + +[tools.glob.properties.path] +type = "string" +description = "The directory to search in" + +[[scenarios]] +name = "tool_name_integrity" +description = "Tool name must not contain harmony tokens like <|channel|>" +expect_tool_call = true +expected_function = "glob" +required_arg_keys = ["pattern"] +tools = ["glob"] + +[[scenarios.messages]] +role = "user" +content = "Find all Python files in the src directory" + +# -- Should NOT call a tool -- + +[[scenarios]] +name = "no_tool_joke" +description = "Joke request should NOT trigger any tool" +expect_tool_call = false + +[[scenarios.messages]] +role = "user" +content = "Tell me a funny joke about cats." + +[[scenarios]] +name = "no_tool_factual" +description = "Factual question answerable from training data" +expect_tool_call = false + +[[scenarios.messages]] +role = "user" +content = "What is the capital of Japan?" diff --git a/bench/single-m3-ultra.toml b/bench/single-m3-ultra.toml new file mode 100644 index 00000000..1d167ac9 --- /dev/null +++ b/bench/single-m3-ultra.toml @@ -0,0 +1,189 @@ +# Single-node M3 Ultra benchmarks +# +# Shared constraints applied to ALL benchmarks in this file. +constraints = [ + "All(MacOsBuild(=25D125))", + "Hosts(=1)", + "All(Chip(m3_ultra))", + "All(GpuCores(=80))", +] + +[topology] +type = "none" + +# Default args merged into each benchmark's args (benchmark-level args win). +[defaults] +pp = [512, 2048, 8192, 16384] +tg = 128 + +[[benchmark]] +model = "mlx-community/Meta-Llama-3.1-70B-Instruct-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/gpt-oss-120b-MXFP4-Q8" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-Flash-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Coder-Next-6bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-30B-A3B-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-0.6B-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-0.6B-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Llama-3.2-1B-Instruct-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Llama-3.2-3B-Instruct-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Llama-3.2-3B-Instruct-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Meta-Llama-3.1-8B-Instruct-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Meta-Llama-3.1-8B-Instruct-bf16" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/gpt-oss-20b-MXFP4-Q8" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-30B-A3B-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-Flash-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-Flash-5bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-Flash-6bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Llama-3.3-70B-Instruct-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Coder-Next-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Coder-Next-5bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Coder-Next-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit" +extra_constraints = ["All(Memory(>=96GiB))"] + +[[benchmark]] +model = "mlx-community/Llama-3.3-70B-Instruct-8bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/llama-3.3-70b-instruct-fp16" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.5-Air-8bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.5-Air-bf16" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-4bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/MiniMax-M2.1-3bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/MiniMax-M2.1-8bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Coder-Next-bf16" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/Step-3.5-Flash-4bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/Step-3.5-Flash-6bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/Step-3.5-Flash-8Bit" +extra_constraints = ["All(Memory(>=256GiB))"] + +[[benchmark]] +model = "mlx-community/DeepSeek-V3.1-4bit" +extra_constraints = ["All(Memory(>=512GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-6bit" +extra_constraints = ["All(Memory(>=512GiB))"] + +[[benchmark]] +model = "mlx-community/GLM-4.7-8bit-gs32" +extra_constraints = ["All(Memory(>=512GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit" +extra_constraints = ["All(Memory(>=512GiB))"] + +[[benchmark]] +model = "mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit" +extra_constraints = ["All(Memory(>=512GiB))"] diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte index 1bc554e8..a7bf7673 100644 --- a/dashboard/src/lib/components/ChatForm.svelte +++ b/dashboard/src/lib/components/ChatForm.svelte @@ -14,6 +14,7 @@ totalTokens, thinkingEnabled as thinkingEnabledStore, setConversationThinking, + stopGeneration, } from "$lib/stores/app.svelte"; import ChatAttachments from "./ChatAttachments.svelte"; import ImageParamsPanel from "./ImageParamsPanel.svelte"; @@ -105,7 +106,7 @@ const modelSupportsThinking = $derived(() => { if (!currentModel) return false; const caps = modelCapabilities[currentModel] || []; - return caps.includes("thinking") && caps.includes("text"); + return caps.includes("thinking_toggle") && caps.includes("text"); }); const isEditOnlyWithoutImage = $derived( @@ -657,86 +658,92 @@ style="min-height: 28px; max-height: 150px;" > - + + {:else} + + {/if} diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte index 7c076459..c716ec3b 100644 --- a/dashboard/src/lib/components/ChatMessages.svelte +++ b/dashboard/src/lib/components/ChatMessages.svelte @@ -3,16 +3,17 @@ messages, currentResponse, isLoading, + prefillProgress, deleteMessage, editAndRegenerate, regenerateLastResponse, regenerateFromToken, setEditingImage, } from "$lib/stores/app.svelte"; - import type { Message } from "$lib/stores/app.svelte"; import type { MessageAttachment } from "$lib/stores/app.svelte"; import MarkdownContent from "./MarkdownContent.svelte"; import TokenHeatmap from "./TokenHeatmap.svelte"; + import PrefillProgressBar from "./PrefillProgressBar.svelte"; import ImageLightbox from "./ImageLightbox.svelte"; interface Props { @@ -25,6 +26,7 @@ const messageList = $derived(messages()); const response = $derived(currentResponse()); const loading = $derived(isLoading()); + const prefill = $derived(prefillProgress()); // Scroll management - user controls scroll, show button when not at bottom const SCROLL_THRESHOLD = 100; @@ -428,6 +430,9 @@ {:else}
+ {#if loading && isLastAssistantMessage(message.id) && prefill && !message.content} + + {/if} {#if message.thinking && message.thinking.trim().length > 0}
= 1000000) { return `${(num / 1000000).toFixed(1)}M`; } else if (num >= 1000) { diff --git a/dashboard/src/lib/components/ImageParamsPanel.svelte b/dashboard/src/lib/components/ImageParamsPanel.svelte index 8da07559..e4536d09 100644 --- a/dashboard/src/lib/components/ImageParamsPanel.svelte +++ b/dashboard/src/lib/components/ImageParamsPanel.svelte @@ -59,13 +59,14 @@ } const sizeOptions: ImageGenerationParams["size"][] = [ + "auto", "512x512", "768x768", "1024x1024", "1024x768", "768x1024", - "1024x1365", - "1365x1024", + "1024x1536", + "1536x1024", ]; const qualityOptions: ImageGenerationParams["quality"][] = [ @@ -176,92 +177,90 @@
- - {#if !isEditMode} -
- SIZE: +
+ SIZE: +
+ +
+ - {params.size} - -
- - - + + +
+
+ + {#if isSizeDropdownOpen} + + + + +
+
+ {#each sizeOptions as size} + + {/each}
- - {#if isSizeDropdownOpen} - - - - -
-
- {#each sizeOptions as size} - - {/each} -
-
- {/if} -
- {/if} + {/if} +
@@ -311,7 +310,7 @@
diff --git a/dashboard/src/lib/components/PrefillProgressBar.svelte b/dashboard/src/lib/components/PrefillProgressBar.svelte new file mode 100644 index 00000000..c07be12c --- /dev/null +++ b/dashboard/src/lib/components/PrefillProgressBar.svelte @@ -0,0 +1,70 @@ + + +
+
+ Processing prompt + + {formatTokenCount(progress.processed)} / {formatTokenCount( + progress.total, + )} tokens + +
+
+
+
+
+ {etaText ?? ""} + {percentage}% +
+
+ + diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index fd24e6c8..dfeec9eb 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -250,6 +250,11 @@ interface RawStateResponse { >; // Thunderbolt bridge cycles (nodes with bridge enabled forming loops) thunderboltBridgeCycles?: string[][]; + // Disk usage per node + nodeDisk?: Record< + string, + { total: { inBytes: number }; available: { inBytes: number } } + >; } export interface MessageAttachment { @@ -273,6 +278,13 @@ export interface TokenData { topLogprobs: TopLogprob[]; } +export interface PrefillProgress { + processed: number; + total: number; + /** Timestamp (performance.now()) when prefill started. */ + startedAt: number; +} + export interface Message { id: string; role: "user" | "assistant" | "system"; @@ -306,13 +318,14 @@ const IMAGE_PARAMS_STORAGE_KEY = "exo-image-generation-params"; export interface ImageGenerationParams { // Basic params size: + | "auto" | "512x512" | "768x768" | "1024x1024" | "1024x768" | "768x1024" - | "1024x1365" - | "1365x1024"; + | "1024x1536" + | "1536x1024"; quality: "low" | "medium" | "high"; outputFormat: "png" | "jpeg"; numImages: number; @@ -336,7 +349,7 @@ export interface EditingImage { } const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = { - size: "1024x1024", + size: "auto", quality: "medium", outputFormat: "png", numImages: 1, @@ -519,6 +532,10 @@ class AppStore { ttftMs = $state(null); // Time to first token in ms tps = $state(null); // Tokens per second totalTokens = $state(0); // Total tokens in current response + prefillProgress = $state(null); + + // Abort controller for stopping generation + private currentAbortController: AbortController | null = null; // Topology state topologyData = $state(null); @@ -1660,11 +1677,12 @@ class AppStore { if (!reader) throw new Error("No response body"); let fullContent = prefixText; + let streamedThinking = ""; const collectedTokens: TokenData[] = [...tokensToKeep]; interface ChatCompletionChunk { choices?: Array<{ - delta?: { content?: string }; + delta?: { content?: string; reasoning_content?: string }; logprobs?: { content?: Array<{ token: string; @@ -1685,6 +1703,7 @@ class AppStore { (parsed) => { const choice = parsed.choices?.[0]; const delta = choice?.delta?.content; + const thinkingDelta = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; @@ -1703,7 +1722,11 @@ class AppStore { } } - if (delta) { + if (thinkingDelta) { + streamedThinking += thinkingDelta; + } + + if (delta || thinkingDelta) { if (firstTokenTime === null) { firstTokenTime = performance.now(); this.ttftMs = firstTokenTime - requestStartTime; @@ -1717,9 +1740,14 @@ class AppStore { this.tps = ((tokenCount - tokensToKeep.length) / elapsed) * 1000; } - fullContent += delta; - const { displayContent, thinkingContent } = + if (delta) { + fullContent += delta; + } + const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(fullContent); + const combinedThinking = [streamedThinking, tagThinking] + .filter(Boolean) + .join("\n\n"); if (this.activeConversationId === targetConversationId) { this.currentResponse = displayContent; @@ -1731,7 +1759,7 @@ class AppStore { messageId, (m) => { m.content = displayContent; - m.thinking = thinkingContent || undefined; + m.thinking = combinedThinking || undefined; m.tokens = [...collectedTokens]; }, ); @@ -1743,11 +1771,14 @@ class AppStore { // Final update if (this.conversationExists(targetConversationId)) { - const { displayContent, thinkingContent } = + const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(fullContent); + const finalThinking = [streamedThinking, tagThinking] + .filter(Boolean) + .join("\n\n"); this.updateConversationMessage(targetConversationId, messageId, (m) => { m.content = displayContent; - m.thinking = thinkingContent || undefined; + m.thinking = finalThinking || undefined; m.tokens = [...collectedTokens]; if (this.ttftMs !== null) m.ttftMs = this.ttftMs; if (this.tps !== null) m.tps = this.tps; @@ -1855,11 +1886,12 @@ class AppStore { } let streamedContent = ""; + let streamedThinking = ""; const collectedTokens: TokenData[] = []; interface ChatCompletionChunk { choices?: Array<{ - delta?: { content?: string }; + delta?: { content?: string; reasoning_content?: string }; logprobs?: { content?: Array<{ token: string; @@ -1880,6 +1912,7 @@ class AppStore { (parsed) => { const choice = parsed.choices?.[0]; const delta = choice?.delta?.content; + const thinkingDelta = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; @@ -1898,10 +1931,19 @@ class AppStore { } } - if (delta) { - streamedContent += delta; - const { displayContent, thinkingContent } = + if (thinkingDelta) { + streamedThinking += thinkingDelta; + } + + if (delta || thinkingDelta) { + if (delta) { + streamedContent += delta; + } + const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); + const combinedThinking = [streamedThinking, tagThinking] + .filter(Boolean) + .join("\n\n"); // Only update currentResponse if target conversation is active if (this.activeConversationId === targetConversationId) { @@ -1914,7 +1956,7 @@ class AppStore { assistantMessage.id, (msg) => { msg.content = displayContent; - msg.thinking = thinkingContent || undefined; + msg.thinking = combinedThinking || undefined; msg.tokens = [...collectedTokens]; }, ); @@ -1926,14 +1968,17 @@ class AppStore { // Final cleanup of the message (if conversation still exists) if (this.conversationExists(targetConversationId)) { - const { displayContent, thinkingContent } = + const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); + const finalThinking = [streamedThinking, tagThinking] + .filter(Boolean) + .join("\n\n"); this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = displayContent; - msg.thinking = thinkingContent || undefined; + msg.thinking = finalThinking || undefined; msg.tokens = [...collectedTokens]; }, ); @@ -2022,6 +2067,7 @@ class AppStore { reader: ReadableStreamDefaultReader, targetConversationId: string, onChunk: (parsed: T) => void, + onEvent?: Record void>, ): Promise { const decoder = new TextDecoder(); let buffer = ""; @@ -2042,6 +2088,24 @@ class AppStore { const trimmed = line.trim(); if (!trimmed) continue; + // Handle SSE comments (": key json") for prefill progress etc. + if (trimmed.startsWith(": ") && onEvent) { + const comment = trimmed.slice(2); + const spaceIdx = comment.indexOf(" "); + if (spaceIdx > 0) { + const key = comment.slice(0, spaceIdx); + if (onEvent[key]) { + try { + const parsed = JSON.parse(comment.slice(spaceIdx + 1)); + onEvent[key](parsed); + } catch { + // Skip malformed JSON in comment + } + } + } + continue; + } + if (trimmed.startsWith("data: ")) { const data = trimmed.slice(6); if (data === "[DONE]") continue; @@ -2273,6 +2337,9 @@ class AppStore { let firstTokenTime: number | null = null; let tokenCount = 0; + const abortController = new AbortController(); + this.currentAbortController = abortController; + const response = await fetch("/v1/chat/completions", { method: "POST", headers: { @@ -2289,6 +2356,7 @@ class AppStore { enable_thinking: enableThinking, }), }), + signal: abortController.signal, }); if (!response.ok) { @@ -2302,10 +2370,11 @@ class AppStore { } let streamedContent = ""; + let streamedThinking = ""; interface ChatCompletionChunk { choices?: Array<{ - delta?: { content?: string }; + delta?: { content?: string; reasoning_content?: string }; logprobs?: { content?: Array<{ token: string; @@ -2326,8 +2395,14 @@ class AppStore { reader, targetConversationId, (parsed) => { + // Clear prefill progress when first token data arrives + if (this.prefillProgress) { + this.prefillProgress = null; + } + const choice = parsed.choices?.[0]; const tokenContent = choice?.delta?.content; + const thinkingContent = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; @@ -2346,7 +2421,11 @@ class AppStore { } } - if (tokenContent) { + if (thinkingContent) { + streamedThinking += thinkingContent; + } + + if (tokenContent || thinkingContent) { // Track first token for TTFT if (firstTokenTime === null) { firstTokenTime = performance.now(); @@ -2363,11 +2442,16 @@ class AppStore { this.tps = (tokenCount / elapsed) * 1000; } - streamedContent += tokenContent; + if (tokenContent) { + streamedContent += tokenContent; + } - // Strip thinking tags for display and extract thinking content - const { displayContent, thinkingContent } = + // Use stripThinkingTags as fallback for any tags still in content + const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); + const combinedThinking = [streamedThinking, tagThinking] + .filter(Boolean) + .join("\n\n"); // Only update currentResponse if target conversation is active if (this.activeConversationId === targetConversationId) { @@ -2380,7 +2464,7 @@ class AppStore { assistantMessage.id, (msg) => { msg.content = displayContent; - msg.thinking = thinkingContent || undefined; + msg.thinking = combinedThinking || undefined; msg.tokens = [...collectedTokens]; }, ); @@ -2388,8 +2472,27 @@ class AppStore { this.persistConversation(targetConversationId); } }, + { + prefill_progress: (data) => { + // TaggedModel wraps as {"PrefillProgressChunk": {...}} + // model_dump_json() uses snake_case (by_alias defaults to False) + const raw = data as Record; + const inner = (raw["PrefillProgressChunk"] ?? raw) as { + processed_tokens: number; + total_tokens: number; + }; + this.prefillProgress = { + processed: inner.processed_tokens, + total: inner.total_tokens, + startedAt: this.prefillProgress?.startedAt ?? performance.now(), + }; + }, + }, ); + // Clear prefill progress after stream ends + this.prefillProgress = null; + // Calculate final TPS if (firstTokenTime !== null && tokenCount > 1) { const totalGenerationTime = performance.now() - firstTokenTime; @@ -2398,14 +2501,17 @@ class AppStore { // Final cleanup of the message (if conversation still exists) if (this.conversationExists(targetConversationId)) { - const { displayContent, thinkingContent } = + const { displayContent, thinkingContent: tagThinking } = this.stripThinkingTags(streamedContent); + const finalThinking = [streamedThinking, tagThinking] + .filter(Boolean) + .join("\n\n"); this.updateConversationMessage( targetConversationId, assistantMessage.id, (msg) => { msg.content = displayContent; - msg.thinking = thinkingContent || undefined; + msg.thinking = finalThinking || undefined; msg.tokens = [...collectedTokens]; // Store performance metrics on the message if (this.ttftMs !== null) { @@ -2420,20 +2526,31 @@ class AppStore { this.persistConversation(targetConversationId); } } catch (error) { - console.error("Error sending message:", error); - this.handleStreamingError( - error, - targetConversationId, - assistantMessage.id, - "Failed to get response", - ); + if (error instanceof DOMException && error.name === "AbortError") { + // User stopped generation — not an error + } else { + console.error("Error sending message:", error); + this.handleStreamingError( + error, + targetConversationId, + assistantMessage.id, + "Failed to get response", + ); + } } finally { + this.currentAbortController = null; + this.prefillProgress = null; this.isLoading = false; this.currentResponse = ""; this.saveConversationsToStorage(); } } + stopGeneration(): void { + this.currentAbortController?.abort(); + this.currentAbortController = null; + } + /** * Generate an image using the image generation API */ @@ -3060,6 +3177,7 @@ export const isLoading = () => appStore.isLoading; export const ttftMs = () => appStore.ttftMs; export const tps = () => appStore.tps; export const totalTokens = () => appStore.totalTokens; +export const prefillProgress = () => appStore.prefillProgress; export const topologyData = () => appStore.topologyData; export const instances = () => appStore.instances; export const runners = () => appStore.runners; @@ -3077,6 +3195,7 @@ export const topologyOnlyMode = () => appStore.getTopologyOnlyMode(); export const chatSidebarVisible = () => appStore.getChatSidebarVisible(); // Actions +export const stopGeneration = () => appStore.stopGeneration(); export const startChat = () => appStore.startChat(); export const sendMessage = ( content: string, diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 663239f6..05714db5 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -153,6 +153,74 @@ }); let tb5InfoDismissed = $state(false); + // Detect Mac Studio nodes using RDMA on en2 (the port next to ethernet — RDMA doesn't work there) + const macStudioEn2RdmaWarning = $derived.by(() => { + const edges = data?.edges; + const ids = tbIdentifiers; + const rdmaCtl = rdmaCtlData; + if (!edges || !ids || !rdmaCtl) return null; + + const affectedConnections: Array<{ + nodeId: string; + nodeName: string; + peerNodeId: string; + peerNodeName: string; + rdmaIface: string; + }> = []; + + const isMacStudio = (node: (typeof data.nodes)[string] | undefined) => + node?.system_info?.model_id === "Mac Studio"; + + for (const edge of edges) { + if (!edge.sourceRdmaIface && !edge.sinkRdmaIface) continue; + + const sourceNode = data?.nodes?.[edge.source]; + if ( + isMacStudio(sourceNode) && + edge.sourceRdmaIface === "rdma_en2" && + rdmaCtl[edge.source]?.enabled + ) { + affectedConnections.push({ + nodeId: edge.source, + nodeName: + sourceNode?.friendly_name || edge.source.slice(0, 8) + "...", + peerNodeId: edge.target, + peerNodeName: + data?.nodes?.[edge.target]?.friendly_name || + edge.target.slice(0, 8) + "...", + rdmaIface: "en2", + }); + } + + const sinkNode = data?.nodes?.[edge.target]; + if ( + isMacStudio(sinkNode) && + edge.sinkRdmaIface === "rdma_en2" && + rdmaCtl[edge.target]?.enabled + ) { + affectedConnections.push({ + nodeId: edge.target, + nodeName: sinkNode?.friendly_name || edge.target.slice(0, 8) + "...", + peerNodeId: edge.source, + peerNodeName: + sourceNode?.friendly_name || edge.source.slice(0, 8) + "...", + rdmaIface: "en2", + }); + } + } + + // Deduplicate by nodeId + const seen = new Set(); + const unique = affectedConnections.filter((c) => { + if (seen.has(c.nodeId)) return false; + seen.add(c.nodeId); + return true; + }); + + return unique.length > 0 ? unique : null; + }); + let macStudioEn2Dismissed = $state(false); + // Helper to get friendly node name from node ID function getNodeName(nodeId: string): string { const node = data?.nodes?.[nodeId]; @@ -1010,10 +1078,8 @@ if (!progress || typeof progress !== "object") return null; const prog = progress as Record; - const totalBytes = getBytes(prog.total_bytes ?? prog.totalBytes); - const downloadedBytes = getBytes( - prog.downloaded_bytes ?? prog.downloadedBytes, - ); + const totalBytes = getBytes(prog.total); + const downloadedBytes = getBytes(prog.downloaded); const speed = (prog.speed as number) ?? 0; const completedFiles = (prog.completed_files as number) ?? (prog.completedFiles as number) ?? 0; @@ -1026,8 +1092,8 @@ for (const [fileName, fileData] of Object.entries(filesObj)) { if (!fileData || typeof fileData !== "object") continue; const fd = fileData as Record; - const fTotal = getBytes(fd.total_bytes ?? fd.totalBytes); - const fDownloaded = getBytes(fd.downloaded_bytes ?? fd.downloadedBytes); + const fTotal = getBytes(fd.total); + const fDownloaded = getBytes(fd.downloaded); files.push({ name: fileName, totalBytes: fTotal, @@ -1152,13 +1218,6 @@ }; } - // Debug: Log downloads data when it changes - $effect(() => { - if (downloadsData && Object.keys(downloadsData).length > 0) { - console.log("[Download Debug] Current downloads:", downloadsData); - } - }); - // Helper to get download status for an instance function getInstanceDownloadStatus( instanceId: string, @@ -1423,7 +1482,6 @@ if (typeof value === "number") return value; if (value && typeof value === "object") { const v = value as Record; - if (typeof v.in_bytes === "number") return v.in_bytes; if (typeof v.inBytes === "number") return v.inBytes; } return 0; @@ -1986,7 +2044,7 @@ {#snippet clusterWarnings()} - {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)} + {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (macStudioEn2RdmaWarning && !macStudioEn2Dismissed)}
{#if tbBridgeCycles.length > 0} {@const cycle = tbBridgeCycles[0]} @@ -2151,12 +2209,260 @@
{/if} + + {#if macStudioEn2RdmaWarning && !macStudioEn2Dismissed} + + {/if}
{/if} {/snippet} {#snippet clusterWarningsCompact()} - {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)} + {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (macStudioEn2RdmaWarning && !macStudioEn2Dismissed)}
{#if tbBridgeCycles.length > 0}
{/if} + {#if macStudioEn2RdmaWarning && !macStudioEn2Dismissed} +
+ + + + BAD RDMA PORT +
+ {/if}
{/if} {/snippet} diff --git a/dashboard/src/routes/downloads/+page.svelte b/dashboard/src/routes/downloads/+page.svelte index 94e10612..a57a18d3 100644 --- a/dashboard/src/routes/downloads/+page.svelte +++ b/dashboard/src/routes/downloads/+page.svelte @@ -1,43 +1,59 @@ @@ -415,253 +339,384 @@
{:else} -
- {#each downloadOverview as node} -
-
-
-
- {node.nodeName} -
-
- {node.nodeId} -
-
- {formatBytes( - node.models - .filter((m) => m.status === "completed") - .reduce((sum, m) => sum + m.totalBytes, 0), - )} models{#if node.diskAvailable != null} - - {formatBytes(node.diskAvailable)} free{/if} -
-
-
+ + + + + {/each} + + + + {#each modelRows as row} + -
-
-
- {model.prettyName ?? model.modelId} -
-
- {model.modelId} -
- {#if model.status !== "completed"} -
- {formatBytes(model.downloadedBytes)} / {formatBytes( - model.totalBytes, - )} -
- {/if} -
+
-
-
-
- -
- {model.status === "completed" - ? `Completed (${formatBytes(model.totalBytes)})` - : `${formatSpeed(model.speed)} • ETA ${formatEta(model.etaMs)}`} - {#if model.status !== "completed"} - {model.files.length} file{model.files.length === 1 - ? "" - : "s"} - {/if} -
- - {#if isExpanded} -
- {#if model.files.length === 0} -
- No file details reported. + {#each nodeColumns as col} + {@const cell = row.cells[col.nodeId] ?? { + kind: "not_present" as const, + }} +
+ {/each} + {/each} - - {/each} + +
-
- {node.models.filter((m) => m.status === "completed") - .length} - / {node.models.length} models -
- - - - {#each node.models as model} - {@const key = `${node.nodeId}|${model.modelId}`} - {@const pct = clampPercent(model.percentage)} - {@const gradient = getBarGradient(pct)} - {@const isExpanded = expanded.has(key)} -
+ {#each nodeColumns as col} +
+
{col.label}
+ {#if col.diskAvailable != null} +
+ {formatBytes(col.diskAvailable)} free +
+ {/if} +
- - {pct.toFixed(1)}% - - {#if model.status !== "completed" && model.shardMetadata} - - {/if} - {#if model.status === "completed"} - - {/if} + {row.modelId} +
+ {/if} + - + 		+ {#if cell.kind === "completed"} +
+ + + + {formatBytes(cell.totalBytes)} +
- {:else} - {#each model.files as f} - {@const fpct = clampPercent(f.percentage)} - {@const fgradient = getBarGradient(fpct)} + {:else if cell.kind === "downloading"} +
+ {clampPercent(cell.percentage).toFixed(1)}%
- {f.name} - = 100 - ? "text-green-400" - : fpct <= 0 - ? "text-red-400" - : "text-exo-yellow"}>{fpct.toFixed(1)}% -
-
-
-
-
- {formatBytes(f.downloadedBytes)} / {formatBytes( - f.totalBytes, - )} - {formatSpeed(f.speed)} • ETA {formatEta( - f.etaMs, - )} -
+ class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300" + style="width: {clampPercent( + cell.percentage, + ).toFixed(1)}%" + >
- {/each} + {formatSpeed(cell.speed)} + + {:else if cell.kind === "pending"} +
+ ... +
+ {:else if cell.kind === "failed"} +
+ + + + {#if row.shardMetadata} + + {/if} +
+ {:else} +
+ -- + {#if row.shardMetadata} + + {/if} +
{/if} - - {/if} - +
{/if}
+ +{#if infoRow} +
(infoRow = null)} + role="presentation" + >
+ +{/if} + diff --git a/flake.nix b/flake.nix index 9c2ca1ef..3bf7d912 100644 --- a/flake.nix +++ b/flake.nix @@ -74,7 +74,6 @@ perSystem = { config, self', inputs', pkgs, lib, system, ... }: let - fenixToolchain = inputs'.fenix.packages.complete; # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) pkgsSwift = import inputs.nixpkgs-swift { inherit system; }; in @@ -115,7 +114,7 @@ packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin ( let uvLock = builtins.fromTOML (builtins.readFile ./uv.lock); - mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx") uvLock.package); + mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package); uvLockMlxVersion = mlxPackage.version; in { diff --git a/nix/mlx.nix b/nix/mlx.nix index f29217b8..4434d3cc 100644 --- a/nix/mlx.nix +++ b/nix/mlx.nix @@ -41,16 +41,16 @@ let mlx = stdenv.mkDerivation rec { pname = "mlx"; - version = let v = "0.30.6"; in + version = let v = "0.30.7.dev20260218+14841977"; in assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix."; v; pyproject = true; src = fetchFromGitHub { - owner = "ml-explore"; - repo = "mlx"; - tag = "v${version}"; - hash = "sha256-avD5EGhwgmPdXLAyQSqTO6AXk/W3ziH+f6AetjK3Sdo="; + owner = "rltakashige"; + repo = "mlx-jaccl-fix-small-recv"; + rev = "1484197707f35186ad3bd614357c7c47fdf86ebc"; + hash = "sha256-FupCMoK/SF/ldfKuvMSAKECcOP8c+ANgkQlPZttDsLk="; }; patches = [ diff --git a/pyproject.toml b/pyproject.toml index def495c7..c4bfa550 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,9 +17,9 @@ dependencies = [ "loguru>=0.7.3", "exo_pyo3_bindings", # rust bindings "anyio==4.11.0", - "mlx==0.30.6; sys_platform == 'darwin'", + "mlx; sys_platform == 'darwin'", "mlx[cpu]==0.30.6; sys_platform == 'linux'", - "mlx-lm==0.30.6", + "mlx-lm==0.30.7", "tiktoken>=0.12.0", # required for kimi k2 tokenizer "hypercorn>=0.18.0", "openai-harmony>=0.0.8", @@ -64,6 +64,7 @@ members = [ [tool.uv.sources] 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/davidmcc73/mlx-lm", branch = "stable" } # Uncomment to use local mlx/mlx-lm development versions: # mlx = { path = "/Users/Shared/mlx", editable=true } @@ -132,7 +133,7 @@ markers = [ env = [ "EXO_TESTS=1" ] -addopts = "-m 'not slow'" +addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py" filterwarnings = [ "ignore:builtin type Swig:DeprecationWarning", ] diff --git a/python/parts.nix b/python/parts.nix index 46b4abdf..9ba703a0 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -58,6 +58,21 @@ 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 + ]; + preFixup = '' + addAutoPatchelfSearchPath ${final.nvidia-cublas} + addAutoPatchelfSearchPath ${final.nvidia-cuda-nvrtc} + addAutoPatchelfSearchPath ${final.nvidia-cudnn-cu13} + addAutoPatchelfSearchPath ${final.nvidia-nccl-cu13} + ''; + autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ]; + }); torch = ignoreMissing prev.torch; triton = ignoreMissing prev.triton; } @@ -74,14 +89,25 @@ linuxOverlay ] ); - exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default; + # 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*" + ]; + + exoVenv = (pythonSet.mkVirtualEnv "exo-env" workspace.deps.default).overrideAttrs { + venvIgnoreCollisions = venvCollisionPaths; + }; # Virtual environment with dev dependencies for testing - testVenv = pythonSet.mkVirtualEnv "exo-test-env" ( + testVenv = (pythonSet.mkVirtualEnv "exo-test-env" ( workspace.deps.default // { exo = [ "dev" ]; # Include pytest, pytest-asyncio, pytest-env } - ); + )).overrideAttrs { + venvIgnoreCollisions = venvCollisionPaths; + }; mkPythonScript = name: path: pkgs.writeShellApplication { inherit name; @@ -132,6 +158,7 @@ exo-test-env = testVenv; } // { exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py); + exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py); exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py); }; diff --git a/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml b/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml index 41784cf6..2c982882 100644 --- a/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml +++ b/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "deepseek" quantization = "4bit" base_model = "DeepSeek V3.1" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 405874409472 diff --git a/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml b/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml index a5d77bcd..4cf99bec 100644 --- a/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml +++ b/resources/inference_model_cards/mlx-community--DeepSeek-V3.1-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "deepseek" quantization = "8bit" base_model = "DeepSeek V3.1" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 765577920512 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml index a7acea44..0f8708bc 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM 4.5 Air" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 122406567936 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml b/resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml index 4258c225..764372fd 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "bf16" base_model = "GLM 4.5 Air" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 229780750336 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml index 0672d664..559cea65 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "4bit" base_model = "GLM 4.7" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 198556925568 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml index bcf1cae4..bdcd0b4e 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-6bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "6bit" base_model = "GLM 4.7" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 286737579648 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml index 0f56c2f7..1d942440 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM 4.7" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 396963397248 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml index 8637cef0..2aad9b5b 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "4bit" base_model = "GLM 4.7 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 19327352832 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml index b9a9da4d..1efed44f 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-5bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "5bit" base_model = "GLM 4.7 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 22548578304 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml index e3cb1fa8..b5ec0fa5 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-6bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "6bit" base_model = "GLM 4.7 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 26843545600 diff --git a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml index bd6df312..4b400aeb 100644 --- a/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml +++ b/resources/inference_model_cards/mlx-community--GLM-4.7-Flash-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "glm" quantization = "8bit" base_model = "GLM 4.7 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 34359738368 diff --git a/resources/inference_model_cards/mlx-community--GLM-5-8bit.toml b/resources/inference_model_cards/mlx-community--GLM-5-8bit.toml new file mode 100644 index 00000000..24beb79f --- /dev/null +++ b/resources/inference_model_cards/mlx-community--GLM-5-8bit.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/GLM-5-8bit-MXFP8" +n_layers = 78 +hidden_size = 6144 +supports_tensor = true +tasks = ["TextGeneration"] +family = "glm" +quantization = "8bit" +base_model = "GLM-5" +capabilities = ["text", "thinking"] + +[storage_size] +in_bytes = 790517400864 diff --git a/resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml b/resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml new file mode 100644 index 00000000..ceb1f74c --- /dev/null +++ b/resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/GLM-5-MXFP4-Q8" +n_layers = 78 +hidden_size = 6144 +supports_tensor = true +tasks = ["TextGeneration"] +family = "glm" +quantization = "MXFP4-Q8" +base_model = "GLM-5" +capabilities = ["text", "thinking"] + +[storage_size] +in_bytes = 405478939008 diff --git a/resources/inference_model_cards/mlx-community--GLM-5-bf16.toml b/resources/inference_model_cards/mlx-community--GLM-5-bf16.toml new file mode 100644 index 00000000..18a7aec6 --- /dev/null +++ b/resources/inference_model_cards/mlx-community--GLM-5-bf16.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/GLM-5" +n_layers = 78 +hidden_size = 6144 +supports_tensor = true +tasks = ["TextGeneration"] +family = "glm" +quantization = "bf16" +base_model = "GLM-5" +capabilities = ["text", "thinking"] + +[storage_size] +in_bytes = 1487822475264 diff --git a/resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml b/resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml index 0a955b04..3e7dedd2 100644 --- a/resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml +++ b/resources/inference_model_cards/mlx-community--Kimi-K2-Thinking.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "kimi" quantization = "" base_model = "Kimi K2" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 706522120192 diff --git a/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml b/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml index 806c6b30..eb73ea09 100644 --- a/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml +++ b/resources/inference_model_cards/mlx-community--Kimi-K2.5.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "kimi" quantization = "" base_model = "Kimi K2.5" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 662498705408 diff --git a/resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml b/resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml index 92ec6746..f740e0d5 100644 --- a/resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml +++ b/resources/inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "minimax" quantization = "3bit" base_model = "MiniMax M2.1" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 100086644736 diff --git a/resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml b/resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml index c1388d2f..6cf55637 100644 --- a/resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml +++ b/resources/inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "minimax" quantization = "8bit" base_model = "MiniMax M2.1" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 242986745856 diff --git a/resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml b/resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml new file mode 100644 index 00000000..d64f29f9 --- /dev/null +++ b/resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/MiniMax-M2.5-4bit" +n_layers = 62 +hidden_size = 3072 +supports_tensor = true +tasks = ["TextGeneration"] +family = "minimax" +quantization = "4bit" +base_model = "MiniMax M2.5" +capabilities = ["text", "thinking"] + +[storage_size] +in_bytes = 128666664960 diff --git a/resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml b/resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml new file mode 100644 index 00000000..c92dcd0d --- /dev/null +++ b/resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/MiniMax-M2.5-6bit" +n_layers = 62 +hidden_size = 3072 +supports_tensor = true +tasks = ["TextGeneration"] +family = "minimax" +quantization = "6bit" +base_model = "MiniMax M2.5" +capabilities = ["text", "thinking"] + +[storage_size] +in_bytes = 185826705408 diff --git a/resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml b/resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml new file mode 100644 index 00000000..b1abb744 --- /dev/null +++ b/resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml @@ -0,0 +1,12 @@ +model_id = "mlx-community/MiniMax-M2.5-8bit" +n_layers = 62 +hidden_size = 3072 +supports_tensor = true +tasks = ["TextGeneration"] +family = "minimax" +quantization = "8bit" +base_model = "MiniMax M2.5" +capabilities = ["text", "thinking"] + +[storage_size] +in_bytes = 242986745856 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml index 7929aaba..03c6c28b 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 0.6B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 342884352 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml index d9fcc368..44177385 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 0.6B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 698351616 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml index ef835c6a..e8ef3494 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 235B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 141733920768 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml index f6e079ab..1ce1fd62 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 235B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 268435456000 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-4bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-4bit.toml index 48a6666f..bcae53f9 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-4bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 30B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 17612931072 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-8bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-8bit.toml index c283396f..ed08db63 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-8bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-30B-A3B-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 30B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 33279705088 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml index 2a3e3c19..48caa3fc 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "4bit" base_model = "Qwen3 Next 80B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 47080074240 diff --git a/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml b/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml index 65d33253..bd4ae549 100644 --- a/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml +++ b/resources/inference_model_cards/mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "qwen" quantization = "8bit" base_model = "Qwen3 Next 80B" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 88814387200 diff --git a/resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml b/resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml index 9c12a5b9..78385661 100644 --- a/resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml +++ b/resources/inference_model_cards/mlx-community--Step-3.5-Flash-4bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "step" quantization = "4bit" base_model = "Step 3.5 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 114572190076 diff --git a/resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml b/resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml index d564498a..ab336bc2 100644 --- a/resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml +++ b/resources/inference_model_cards/mlx-community--Step-3.5-Flash-6bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "step" quantization = "6bit" base_model = "Step 3.5 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 159039627774 diff --git a/resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml b/resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml index afef6d7f..ef2a5074 100644 --- a/resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml +++ b/resources/inference_model_cards/mlx-community--Step-3.5-Flash-8Bit.toml @@ -6,7 +6,7 @@ tasks = ["TextGeneration"] family = "step" quantization = "8bit" base_model = "Step 3.5 Flash" -capabilities = ["text", "thinking"] +capabilities = ["text", "thinking", "thinking_toggle"] [storage_size] in_bytes = 209082699847 diff --git a/rust/clippy.toml b/rust/clippy.toml deleted file mode 100644 index 6d5a6187..00000000 --- a/rust/clippy.toml +++ /dev/null @@ -1,2 +0,0 @@ -# we can manually exclude false-positive lint errors for dual packages (if in dependencies) -#allowed-duplicate-crates = ["hashbrown"] \ No newline at end of file diff --git a/rust/exo_pyo3_bindings/Cargo.toml b/rust/exo_pyo3_bindings/Cargo.toml index 12803ab4..77777055 100644 --- a/rust/exo_pyo3_bindings/Cargo.toml +++ b/rust/exo_pyo3_bindings/Cargo.toml @@ -25,17 +25,17 @@ workspace = true networking = { workspace = true } # interop -pyo3 = { version = "0.27.1", features = [ - # "abi3-py311", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.11 - "nightly", # enables better-supported GIL integration +pyo3 = { version = "0.27.2", features = [ + # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13 + # "nightly", # enables better-supported GIL integration "experimental-async", # async support in #[pyfunction] & #[pymethods] #"experimental-inspect", # inspection of generated binary => easier to automate type-hint generation #"py-clone", # adding Clone-ing of `Py` without GIL (may cause panics - remove if panics happen) - "multiple-pymethods", # allows multiple #[pymethods] sections per class + # "multiple-pymethods", # allows multiple #[pymethods] sections per class # integrations with other libraries - "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational", - "ordered-float", "rust_decimal", "smallvec", + # "arc_lock", "bigdecimal", "either", "hashbrown", "indexmap", "num-bigint", "num-complex", "num-rational", + # "ordered-float", "rust_decimal", "smallvec", # "anyhow", "chrono", "chrono-local", "chrono-tz", "eyre", "jiff-02", "lock_api", "parking-lot", "time", "serde", ] } pyo3-stub-gen = { version = "0.17.2" } @@ -45,33 +45,18 @@ pyo3-log = "0.13.2" # macro dependencies extend = { workspace = true } delegate = { workspace = true } -impl-trait-for-tuples = { workspace = true } -derive_more = { workspace = true } -pin-project = { workspace = true } # async runtime tokio = { workspace = true, features = ["full", "tracing"] } -futures = { workspace = true } +futures-lite = { workspace = true } # utility dependencies -once_cell = "1.21.3" -thread_local = "1.1.9" util = { workspace = true } -thiserror = { workspace = true } -#internment = { workspace = true } -#recursion = { workspace = true } -#generativity = { workspace = true } -#itertools = { workspace = true } - # Tracing -#tracing = "0.1" -#tracing-subscriber = "0.3" -#console-subscriber = "0.1.5" -#tracing-log = "0.2.0" log = { workspace = true } env_logger = "0.11" - # Networking libp2p = { workspace = true, features = ["full"] } +pin-project = "1.1.10" diff --git a/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi b/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi index fa6700ff..3d55c9e8 100644 --- a/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi +++ b/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi @@ -19,7 +19,7 @@ class ConnectionUpdate: Whether this is a connection or disconnection event """ @property - def peer_id(self) -> PeerId: + def peer_id(self) -> builtins.str: r""" Identity of the peer that we have connected to or disconnected from. """ @@ -40,92 +40,22 @@ class Keypair: Identity keypair of a node. """ @staticmethod - def generate_ed25519() -> Keypair: + def generate() -> Keypair: r""" Generate a new Ed25519 keypair. """ @staticmethod - def generate_ecdsa() -> Keypair: + def from_bytes(bytes: bytes) -> Keypair: r""" - Generate a new ECDSA keypair. - """ - @staticmethod - def generate_secp256k1() -> Keypair: - r""" - Generate a new Secp256k1 keypair. - """ - @staticmethod - def from_protobuf_encoding(bytes: bytes) -> Keypair: - r""" - Decode a private key from a protobuf structure and parse it as a `Keypair`. - """ - @staticmethod - def rsa_from_pkcs8(bytes: bytes) -> Keypair: - r""" - Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo` - format (i.e. unencrypted) as defined in [RFC5208]. - - [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5 - """ - @staticmethod - def secp256k1_from_der(bytes: bytes) -> Keypair: - r""" - Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey` - structure as defined in [RFC5915]. - - [RFC5915]: https://tools.ietf.org/html/rfc5915 - """ - @staticmethod - def ed25519_from_bytes(bytes: bytes) -> Keypair: ... - def to_protobuf_encoding(self) -> bytes: - r""" - Encode a private key as protobuf structure. - """ - def to_peer_id(self) -> PeerId: - r""" - Convert the `Keypair` into the corresponding `PeerId`. - """ - -@typing.final -class Multiaddr: - r""" - Representation of a Multiaddr. - """ - @staticmethod - def empty() -> Multiaddr: - r""" - Create a new, empty multiaddress. - """ - @staticmethod - def with_capacity(n: builtins.int) -> Multiaddr: - r""" - Create a new, empty multiaddress with the given capacity. - """ - @staticmethod - def from_bytes(bytes: bytes) -> Multiaddr: - r""" - Parse a `Multiaddr` value from its byte slice representation. - """ - @staticmethod - def from_string(string: builtins.str) -> Multiaddr: - r""" - Parse a `Multiaddr` value from its string representation. - """ - def len(self) -> builtins.int: - r""" - Return the length in bytes of this multiaddress. - """ - def is_empty(self) -> builtins.bool: - r""" - Returns true if the length of this multiaddress is 0. + Construct an Ed25519 keypair from secret key bytes """ def to_bytes(self) -> bytes: r""" - Return a copy of this [`Multiaddr`]'s byte representation. + Get the secret key bytes underlying the keypair """ - def to_string(self) -> builtins.str: + def to_node_id(self) -> builtins.str: r""" - Convert a Multiaddr to a string. + Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`. """ @typing.final @@ -180,37 +110,6 @@ class NoPeersSubscribedToTopicError(builtins.Exception): def __repr__(self) -> builtins.str: ... def __str__(self) -> builtins.str: ... -@typing.final -class PeerId: - r""" - Identifier of a peer of the network. - - The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer - as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md). - """ - @staticmethod - def random() -> PeerId: - r""" - Generates a random peer ID from a cryptographically secure PRNG. - - This is useful for randomly walking on a DHT, or for testing purposes. - """ - @staticmethod - def from_bytes(bytes: bytes) -> PeerId: - r""" - Parses a `PeerId` from bytes. - """ - def to_bytes(self) -> bytes: - r""" - Returns a raw bytes representation of this `PeerId`. - """ - def to_base58(self) -> builtins.str: - r""" - Returns a base-58 encoded string of this `PeerId`. - """ - def __repr__(self) -> builtins.str: ... - def __str__(self) -> builtins.str: ... - @typing.final class ConnectionUpdateType(enum.Enum): r""" diff --git a/rust/exo_pyo3_bindings/src/allow_threading.rs b/rust/exo_pyo3_bindings/src/allow_threading.rs index 3106e535..142ff0c6 100644 --- a/rust/exo_pyo3_bindings/src/allow_threading.rs +++ b/rust/exo_pyo3_bindings/src/allow_threading.rs @@ -2,11 +2,10 @@ //! use pin_project::pin_project; -use pyo3::marker::Ungil; use pyo3::prelude::*; use std::{ future::Future, - pin::{Pin, pin}, + pin::Pin, task::{Context, Poll}, }; @@ -26,15 +25,13 @@ where impl Future for AllowThreads where - F: Future + Ungil, - F::Output: Ungil, + F: Future + Send, + F::Output: Send, { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let waker = cx.waker(); - Python::with_gil(|py| { - py.allow_threads(|| self.project().0.poll(&mut Context::from_waker(waker))) - }) + Python::attach(|py| py.detach(|| self.project().0.poll(&mut Context::from_waker(waker)))) } } diff --git a/rust/exo_pyo3_bindings/src/examples/mod.rs b/rust/exo_pyo3_bindings/src/examples/mod.rs deleted file mode 100644 index bde14199..00000000 --- a/rust/exo_pyo3_bindings/src/examples/mod.rs +++ /dev/null @@ -1,240 +0,0 @@ -//! This module exists to hold examples of some pyo3 patterns that may be too complex to -//! re-create from scratch, but too inhomogenous to create an abstraction/wrapper around. -//! -//! Pattern examples include: -//! - Async task handles: with GC-integrated cleanup -//! - Sync/async callbacks from python: with propper eventloop handling -//! -//! Mutability pattern: https://pyo3.rs/v0.26.0/async-await.html#send--static-constraint -//! - Store mutable fields in tokio's `Mutex` -//! - For async code: take `&self` and `.lock().await` -//! - For sync code: take `&mut self` and `.get_mut()` - -use crate::ext::{PyResultExt as _, ResultExt as _, TokioRuntimeExt as _}; -use futures::FutureExt as _; -use futures::future::BoxFuture; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::{PyModule, PyModuleMethods as _}; -use pyo3::{ - Bound, Py, PyAny, PyErr, PyResult, PyTraverseError, PyVisit, Python, pyclass, pymethods, -}; -use std::time::Duration; -use tokio::sync::mpsc; -use tokio::sync::mpsc::error::TryRecvError; - -fn needs_tokio_runtime() { - tokio::runtime::Handle::current(); -} - -type SyncCallback = Box; -type AsyncCallback = Box BoxFuture<'static, ()> + Send + Sync>; - -enum AsyncTaskMessage { - SyncCallback(SyncCallback), - AsyncCallback(AsyncCallback), -} - -async fn async_task( - sender: mpsc::UnboundedSender<()>, - mut receiver: mpsc::UnboundedReceiver, -) { - log::info!("RUST: async task started"); - - // task state - let mut interval = tokio::time::interval(Duration::from_secs(1)); - - let mut sync_cbs: Vec = vec![]; - let mut async_cbs: Vec = vec![]; - - loop { - tokio::select! { - // handle incoming messages from task-handle - message = receiver.recv() => { - // handle closed channel by exiting - let Some(message) = message else { - log::info!("RUST: channel closed"); - break; - }; - - // dispatch incoming event - match message { - AsyncTaskMessage::SyncCallback(cb) => { - sync_cbs.push(cb); - } - AsyncTaskMessage::AsyncCallback(cb) => { - async_cbs.push(cb); - } - } - } - - // handle all other events - _ = interval.tick() => { - log::info!("RUST: async task tick"); - - // call back all sync callbacks - for cb in &sync_cbs { - cb(); - } - - // call back all async callbacks - for cb in &async_cbs { - cb().await; - } - - // send event on unbounded channel - sender.send(()).expect("handle receiver cannot be closed/dropped"); - } - } - } - - log::info!("RUST: async task stopped"); -} - -// #[gen_stub_pyclass] -#[pyclass(name = "AsyncTaskHandle")] -#[derive(Debug)] -struct PyAsyncTaskHandle { - sender: Option>, - receiver: mpsc::UnboundedReceiver<()>, -} - -#[allow(clippy::expect_used)] -impl PyAsyncTaskHandle { - const fn sender(&self) -> &mpsc::UnboundedSender { - self.sender - .as_ref() - .expect("The sender should only be None after de-initialization.") - } - - const fn sender_mut(&mut self) -> &mpsc::UnboundedSender { - self.sender - .as_mut() - .expect("The sender should only be None after de-initialization.") - } - - const fn new( - sender: mpsc::UnboundedSender, - receiver: mpsc::UnboundedReceiver<()>, - ) -> Self { - Self { - sender: Some(sender), - receiver, - } - } -} - -// #[gen_stub_pymethods] -#[pymethods] -impl PyAsyncTaskHandle { - #[new] - fn py_new(py: Python<'_>) -> PyResult { - use pyo3_async_runtimes::tokio::get_runtime; - - // create communication channel TOWARDS our task - let (h_sender, t_receiver) = mpsc::unbounded_channel::(); - - // create communication channel FROM our task - let (t_sender, h_receiver) = mpsc::unbounded_channel::<()>(); - - // perform necessary setup within tokio context - or it crashes - let () = get_runtime().block_on(async { needs_tokio_runtime() }); - - // spawn tokio task with this thread's task-locals - without this, async callbacks on the new threads will not work!! - _ = get_runtime().spawn_with_scope(py, async move { - async_task(t_sender, t_receiver).await; - }); - Ok(Self::new(h_sender, h_receiver)) - } - - /// NOTE: exceptions in callbacks are silently ignored until end of execution - fn add_sync_callback( - &self, - // #[gen_stub(override_type( - // type_repr="collections.abc.Callable[[], None]", - // imports=("collections.abc") - // ))] - callback: Py, - ) -> PyResult<()> { - // blocking call to async method -> can do non-blocking if needed - self.sender() - .send(AsyncTaskMessage::SyncCallback(Box::new(move || { - _ = Python::with_gil(|py| callback.call0(py).write_unraisable_with(py)); - }))) - .pyerr()?; - Ok(()) - } - - /// NOTE: exceptions in callbacks are silently ignored until end of execution - fn add_async_callback( - &self, - // #[gen_stub(override_type( - // type_repr="collections.abc.Callable[[], collections.abc.Awaitable[None]]", - // imports=("collections.abc") - // ))] - callback: Py, - ) -> PyResult<()> { - // blocking call to async method -> can do non-blocking if needed - self.sender() - .send(AsyncTaskMessage::AsyncCallback(Box::new(move || { - let c = Python::with_gil(|py| callback.clone_ref(py)); - async move { - if let Some(f) = Python::with_gil(|py| { - let coroutine = c.call0(py).write_unraisable_with(py)?; - pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py)) - .write_unraisable_with(py) - }) { - _ = f.await.write_unraisable(); - } - } - .boxed() - }))) - .pyerr()?; - Ok(()) - } - - async fn receive_unit(&mut self) -> PyResult<()> { - self.receiver - .recv() - .await - .ok_or(PyErr::new::( - "cannot receive unit on closed channel", - )) - } - - fn drain_units(&mut self) -> PyResult { - let mut cnt = 0; - loop { - match self.receiver.try_recv() { - Err(TryRecvError::Disconnected) => { - return Err(PyErr::new::( - "cannot receive unit on closed channel", - )); - } - Err(TryRecvError::Empty) => return Ok(cnt), - Ok(()) => { - cnt += 1; - continue; - } - } - } - } - - // #[gen_stub(skip)] - const fn __traverse__(&self, _visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - Ok(()) // This is needed purely so `__clear__` can work - } - - // #[gen_stub(skip)] - fn __clear__(&mut self) { - // TODO: may or may not need to await a "kill-signal" oneshot channel message, - // to ensure that the networking task is done BEFORE exiting the clear function... - // but this may require GIL?? and it may not be safe to call GIL here?? - self.sender = None; // Using Option as a trick to force `sender` channel to be dropped - } -} - -pub fn examples_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - - Ok(()) -} diff --git a/rust/exo_pyo3_bindings/src/ident.rs b/rust/exo_pyo3_bindings/src/ident.rs new file mode 100644 index 00000000..55f40bc6 --- /dev/null +++ b/rust/exo_pyo3_bindings/src/ident.rs @@ -0,0 +1,47 @@ +use crate::ext::ResultExt as _; +use libp2p::identity::Keypair; +use pyo3::types::{PyBytes, PyBytesMethods as _}; +use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; +use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; + +/// Identity keypair of a node. +#[gen_stub_pyclass] +#[pyclass(name = "Keypair", frozen)] +#[repr(transparent)] +pub struct PyKeypair(pub Keypair); + +#[gen_stub_pymethods] +#[pymethods] +#[allow(clippy::needless_pass_by_value)] +impl PyKeypair { + /// Generate a new Ed25519 keypair. + #[staticmethod] + fn generate() -> Self { + Self(Keypair::generate_ed25519()) + } + + /// Construct an Ed25519 keypair from secret key bytes + #[staticmethod] + fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult { + let mut bytes = Vec::from(bytes.as_bytes()); + Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?)) + } + + /// Get the secret key bytes underlying the keypair + fn to_bytes<'py>(&self, py: Python<'py>) -> PyResult> { + let bytes = self + .0 + .clone() + .try_into_ed25519() + .pyerr()? + .secret() + .as_ref() + .to_vec(); + Ok(PyBytes::new(py, &bytes)) + } + + /// Convert the `Keypair` into the corresponding `PeerId` string, which we use as our `NodeId`. + fn to_node_id(&self) -> String { + self.0.public().to_peer_id().to_base58() + } +} diff --git a/rust/exo_pyo3_bindings/src/lib.rs b/rust/exo_pyo3_bindings/src/lib.rs index 4f591b8c..8c1fb6b5 100644 --- a/rust/exo_pyo3_bindings/src/lib.rs +++ b/rust/exo_pyo3_bindings/src/lib.rs @@ -4,28 +4,14 @@ //! //! -// enable Rust-unstable features for convenience -#![feature(trait_alias)] -#![feature(tuple_trait)] -#![feature(unboxed_closures)] -// #![feature(stmt_expr_attributes)] -// #![feature(assert_matches)] -// #![feature(async_fn_in_dyn_trait)] -// #![feature(async_for_loop)] -// #![feature(auto_traits)] -// #![feature(negative_impls)] - -extern crate core; mod allow_threading; -mod examples; -pub(crate) mod networking; -pub(crate) mod pylibp2p; +mod ident; +mod networking; +use crate::ident::PyKeypair; use crate::networking::networking_submodule; -use crate::pylibp2p::ident::ident_submodule; -use crate::pylibp2p::multiaddr::multiaddr_submodule; use pyo3::prelude::PyModule; -use pyo3::prelude::*; +use pyo3::types::PyModuleMethods; use pyo3::{Bound, PyResult, pyclass, pymodule}; use pyo3_stub_gen::define_stub_info_gatherer; @@ -34,24 +20,11 @@ pub(crate) mod r#const { pub const MPSC_CHANNEL_SIZE: usize = 1024; } -/// Namespace for all the type/trait aliases used by this crate. -pub(crate) mod alias { - use std::error::Error; - use std::marker::Tuple; - - pub trait SendFn = - Fn + Send + 'static; - - pub type AnyError = Box; - pub type AnyResult = Result; -} - /// Namespace for crate-wide extension traits/methods pub(crate) mod ext { use crate::allow_threading::AllowThreads; use extend::ext; use pyo3::exceptions::{PyConnectionError, PyRuntimeError}; - use pyo3::marker::Ungil; use pyo3::types::PyBytes; use pyo3::{Py, PyErr, PyResult, Python}; use tokio::runtime::Runtime; @@ -62,7 +35,7 @@ pub(crate) mod ext { #[ext(pub, name = ByteArrayExt)] impl [u8] { fn pybytes(&self) -> Py { - Python::with_gil(|py| PyBytes::new(py, self).unbind()) + Python::attach(|py| PyBytes::new(py, self).unbind()) } } @@ -98,7 +71,7 @@ pub(crate) mod ext { #[ext(pub, name = PyResultExt)] impl PyResult { fn write_unraisable(self) -> Option { - Python::with_gil(|py| self.write_unraisable_with(py)) + Python::attach(|py| self.write_unraisable_with(py)) } fn write_unraisable_with(self, py: Python<'_>) -> Option { @@ -175,24 +148,6 @@ pub(crate) mod ext { } } -pub(crate) mod private { - use std::marker::Sized; - - /// Sealed traits support - pub trait Sealed {} - impl Sealed for T {} -} - -/// A wrapper around [`Py`] that implements [`Clone`] using [`Python::with_gil`]. -#[repr(transparent)] -pub(crate) struct ClonePy(pub Py); - -impl Clone for ClonePy { - fn clone(&self) -> Self { - Python::with_gil(|py| Self(self.0.clone_ref(py))) - } -} - /// A Python module implemented in Rust. The name of this function must match /// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to /// import the module. @@ -204,8 +159,7 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> { // TODO: for now this is all NOT a submodule, but figure out how to make the submodule system // work with maturin, where the types generate correctly, in the right folder, without // too many importing issues... - ident_submodule(m)?; - multiaddr_submodule(m)?; + m.add_class::()?; networking_submodule(m)?; // top-level constructs diff --git a/rust/exo_pyo3_bindings/src/networking.rs b/rust/exo_pyo3_bindings/src/networking.rs index 01024cf9..2fb1d78e 100644 --- a/rust/exo_pyo3_bindings/src/networking.rs +++ b/rust/exo_pyo3_bindings/src/networking.rs @@ -8,12 +8,12 @@ use crate::r#const::MPSC_CHANNEL_SIZE; use crate::ext::{ByteArrayExt as _, FutureExt, PyErrExt as _}; use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt as _}; +use crate::ident::PyKeypair; use crate::pyclass; -use crate::pylibp2p::ident::{PyKeypair, PyPeerId}; use libp2p::futures::StreamExt as _; +use libp2p::gossipsub; use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError}; use libp2p::swarm::SwarmEvent; -use libp2p::{gossipsub, mdns}; use networking::discovery; use networking::swarm::create_swarm; use pyo3::prelude::{PyModule, PyModuleMethods as _}; @@ -25,7 +25,7 @@ use tokio::sync::{Mutex, mpsc, oneshot}; mod exception { use pyo3::types::PyTuple; - use pyo3::{PyErrArguments, exceptions::PyException, prelude::*}; + use pyo3::{exceptions::PyException, prelude::*}; use pyo3_stub_gen::derive::*; #[gen_stub_pyclass] @@ -119,7 +119,7 @@ struct PyConnectionUpdate { /// Identity of the peer that we have connected to or disconnected from. #[pyo3(get)] - peer_id: PyPeerId, + peer_id: String, /// Remote connection's IPv4 address. #[pyo3(get)] @@ -155,7 +155,6 @@ async fn networking_task( ) { use SwarmEvent::*; use ToTask::*; - use mdns::Event::*; use networking::swarm::BehaviourEvent::*; log::info!("RUST: networking task started"); @@ -252,7 +251,7 @@ async fn networking_task( // send connection event to channel (or exit if connection closed) if let Err(e) = connection_update_tx.send(PyConnectionUpdate { update_type: PyConnectionUpdateType::Connected, - peer_id: PyPeerId(peer_id), + peer_id: peer_id.to_base58(), remote_ipv4, remote_tcp_port, }).await { @@ -273,7 +272,7 @@ async fn networking_task( // send disconnection event to channel (or exit if connection closed) if let Err(e) = connection_update_tx.send(PyConnectionUpdate { update_type: PyConnectionUpdateType::Disconnected, - peer_id: PyPeerId(peer_id), + peer_id: peer_id.to_base58(), remote_ipv4, remote_tcp_port, }).await { @@ -485,7 +484,7 @@ impl PyNetworkingHandle { let (tx, rx) = oneshot::channel(); // send off request to subscribe - let data = Python::with_gil(|py| Vec::from(data.as_bytes(py))); + let data = Python::attach(|py| Vec::from(data.as_bytes(py))); self.to_task_tx() .send_py(ToTask::GossipsubPublish { topic, diff --git a/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs b/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs deleted file mode 100644 index 3c27526a..00000000 --- a/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs +++ /dev/null @@ -1,159 +0,0 @@ -use crate::ext::ResultExt as _; -use libp2p::PeerId; -use libp2p::identity::Keypair; -use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _}; -use pyo3::types::PyBytes; -use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; -use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; - -/// Identity keypair of a node. -#[gen_stub_pyclass] -#[pyclass(name = "Keypair", frozen)] -#[repr(transparent)] -pub struct PyKeypair(pub Keypair); - -#[gen_stub_pymethods] -#[pymethods] -#[allow(clippy::needless_pass_by_value)] -impl PyKeypair { - /// Generate a new Ed25519 keypair. - #[staticmethod] - fn generate_ed25519() -> Self { - Self(Keypair::generate_ed25519()) - } - - /// Generate a new ECDSA keypair. - #[staticmethod] - fn generate_ecdsa() -> Self { - Self(Keypair::generate_ecdsa()) - } - - /// Generate a new Secp256k1 keypair. - #[staticmethod] - fn generate_secp256k1() -> Self { - Self(Keypair::generate_secp256k1()) - } - - /// Decode a private key from a protobuf structure and parse it as a `Keypair`. - #[staticmethod] - fn from_protobuf_encoding(bytes: Bound<'_, PyBytes>) -> PyResult { - let bytes = Vec::from(bytes.as_bytes()); - Ok(Self(Keypair::from_protobuf_encoding(&bytes).pyerr()?)) - } - - /// Decode an keypair from a DER-encoded secret key in PKCS#8 `PrivateKeyInfo` - /// format (i.e. unencrypted) as defined in [RFC5208]. - /// - /// [RFC5208]: https://tools.ietf.org/html/rfc5208#section-5 - #[staticmethod] - fn rsa_from_pkcs8(bytes: Bound<'_, PyBytes>) -> PyResult { - let mut bytes = Vec::from(bytes.as_bytes()); - Ok(Self(Keypair::rsa_from_pkcs8(&mut bytes).pyerr()?)) - } - - /// Decode a keypair from a DER-encoded Secp256k1 secret key in an `ECPrivateKey` - /// structure as defined in [RFC5915]. - /// - /// [RFC5915]: https://tools.ietf.org/html/rfc5915 - #[staticmethod] - fn secp256k1_from_der(bytes: Bound<'_, PyBytes>) -> PyResult { - let mut bytes = Vec::from(bytes.as_bytes()); - Ok(Self(Keypair::secp256k1_from_der(&mut bytes).pyerr()?)) - } - - #[staticmethod] - fn ed25519_from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult { - let mut bytes = Vec::from(bytes.as_bytes()); - Ok(Self(Keypair::ed25519_from_bytes(&mut bytes).pyerr()?)) - } - - /// Encode a private key as protobuf structure. - fn to_protobuf_encoding<'py>(&self, py: Python<'py>) -> PyResult> { - let bytes = self.0.to_protobuf_encoding().pyerr()?; - Ok(PyBytes::new(py, &bytes)) - } - - /// Convert the `Keypair` into the corresponding `PeerId`. - fn to_peer_id(&self) -> PyPeerId { - PyPeerId(self.0.public().to_peer_id()) - } - - // /// Hidden constructor for pickling support. TODO: figure out how to do pickling... - // #[gen_stub(skip)] - // #[new] - // fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult { - // Self::from_protobuf_encoding(bytes) - // } - // - // #[gen_stub(skip)] - // fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> { - // *self = Self::from_protobuf_encoding(state)?; - // Ok(()) - // } - // - // #[gen_stub(skip)] - // fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult> { - // self.to_protobuf_encoding(py) - // } - // - // #[gen_stub(skip)] - // pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> { - // Ok((self.to_protobuf_encoding(py)?,)) - // } -} - -/// Identifier of a peer of the network. -/// -/// The data is a `CIDv0` compatible multihash of the protobuf encoded public key of the peer -/// as specified in [specs/peer-ids](https://github.com/libp2p/specs/blob/master/peer-ids/peer-ids.md). -#[gen_stub_pyclass] -#[pyclass(name = "PeerId", frozen)] -#[derive(Debug, Clone)] -#[repr(transparent)] -pub struct PyPeerId(pub PeerId); - -#[gen_stub_pymethods] -#[pymethods] -#[allow(clippy::needless_pass_by_value)] -impl PyPeerId { - /// Generates a random peer ID from a cryptographically secure PRNG. - /// - /// This is useful for randomly walking on a DHT, or for testing purposes. - #[staticmethod] - fn random() -> Self { - Self(PeerId::random()) - } - - /// Parses a `PeerId` from bytes. - #[staticmethod] - fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult { - let bytes = Vec::from(bytes.as_bytes()); - Ok(Self(PeerId::from_bytes(&bytes).pyerr()?)) - } - - /// Returns a raw bytes representation of this `PeerId`. - fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> { - let bytes = self.0.to_bytes(); - PyBytes::new(py, &bytes) - } - - /// Returns a base-58 encoded string of this `PeerId`. - fn to_base58(&self) -> String { - self.0.to_base58() - } - - fn __repr__(&self) -> String { - format!("PeerId({})", self.to_base58()) - } - - fn __str__(&self) -> String { - self.to_base58() - } -} - -pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - - Ok(()) -} diff --git a/rust/exo_pyo3_bindings/src/pylibp2p/mod.rs b/rust/exo_pyo3_bindings/src/pylibp2p/mod.rs deleted file mode 100644 index 8eb1bdc0..00000000 --- a/rust/exo_pyo3_bindings/src/pylibp2p/mod.rs +++ /dev/null @@ -1,8 +0,0 @@ -//! A module for exposing Rust's libp2p datatypes over Pyo3 -//! -//! TODO: right now we are coupled to libp2p's identity, but eventually we want to create our own -//! independent identity type of some kind or another. This may require handshaking. -//! - -pub mod ident; -pub mod multiaddr; diff --git a/rust/exo_pyo3_bindings/src/pylibp2p/multiaddr.rs b/rust/exo_pyo3_bindings/src/pylibp2p/multiaddr.rs deleted file mode 100644 index 4d398b53..00000000 --- a/rust/exo_pyo3_bindings/src/pylibp2p/multiaddr.rs +++ /dev/null @@ -1,81 +0,0 @@ -use crate::ext::ResultExt as _; -use libp2p::Multiaddr; -use pyo3::prelude::{PyBytesMethods as _, PyModule, PyModuleMethods as _}; -use pyo3::types::PyBytes; -use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; -use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; -use std::str::FromStr as _; - -/// Representation of a Multiaddr. -#[gen_stub_pyclass] -#[pyclass(name = "Multiaddr", frozen)] -#[derive(Debug, Clone)] -#[repr(transparent)] -pub struct PyMultiaddr(pub Multiaddr); - -#[gen_stub_pymethods] -#[pymethods] -#[allow(clippy::needless_pass_by_value)] -impl PyMultiaddr { - /// Create a new, empty multiaddress. - #[staticmethod] - fn empty() -> Self { - Self(Multiaddr::empty()) - } - - /// Create a new, empty multiaddress with the given capacity. - #[staticmethod] - fn with_capacity(n: usize) -> Self { - Self(Multiaddr::with_capacity(n)) - } - - /// Parse a `Multiaddr` value from its byte slice representation. - #[staticmethod] - fn from_bytes(bytes: Bound<'_, PyBytes>) -> PyResult { - let bytes = Vec::from(bytes.as_bytes()); - Ok(Self(Multiaddr::try_from(bytes).pyerr()?)) - } - - /// Parse a `Multiaddr` value from its string representation. - #[staticmethod] - fn from_string(string: String) -> PyResult { - Ok(Self(Multiaddr::from_str(&string).pyerr()?)) - } - - /// Return the length in bytes of this multiaddress. - fn len(&self) -> usize { - self.0.len() - } - - /// Returns true if the length of this multiaddress is 0. - fn is_empty(&self) -> bool { - self.0.is_empty() - } - - /// Return a copy of this [`Multiaddr`]'s byte representation. - fn to_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> { - let bytes = self.0.to_vec(); - PyBytes::new(py, &bytes) - } - - /// Convert a Multiaddr to a string. - fn to_string(&self) -> String { - self.0.to_string() - } - - #[gen_stub(skip)] - fn __repr__(&self) -> String { - format!("Multiaddr({})", self.0) - } - - #[gen_stub(skip)] - fn __str__(&self) -> String { - self.to_string() - } -} - -pub fn multiaddr_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - - Ok(()) -} diff --git a/rust/networking/Cargo.toml b/rust/networking/Cargo.toml index 47d61f41..58ca4e0a 100644 --- a/rust/networking/Cargo.toml +++ b/rust/networking/Cargo.toml @@ -19,21 +19,14 @@ either = { workspace = true } # macro dependencies extend = { workspace = true } delegate = { workspace = true } -impl-trait-for-tuples = { workspace = true } -derive_more = { workspace = true } # async tokio = { workspace = true, features = ["full"] } -futures = { workspace = true } +futures-lite = { workspace = true } futures-timer = { workspace = true } # utility dependencies util = { workspace = true } -thiserror = { workspace = true } -#internment = { workspace = true } -#recursion = { workspace = true } -#generativity = { workspace = true } -#itertools = { workspace = true } tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] } keccak-const = { workspace = true } @@ -41,4 +34,4 @@ keccak-const = { workspace = true } log = { workspace = true } # networking -libp2p = { workspace = true, features = ["full"] } \ No newline at end of file +libp2p = { workspace = true, features = ["full"] } diff --git a/rust/networking/examples/chatroom.rs b/rust/networking/examples/chatroom.rs index 3371b46d..99067e9b 100644 --- a/rust/networking/examples/chatroom.rs +++ b/rust/networking/examples/chatroom.rs @@ -1,4 +1,4 @@ -use futures::stream::StreamExt as _; +use futures_lite::StreamExt; use libp2p::{gossipsub, identity, swarm::SwarmEvent}; use networking::{discovery, swarm}; use tokio::{io, io::AsyncBufReadExt as _, select}; @@ -38,19 +38,19 @@ async fn main() { println!("Publish error: {e:?}"); } } - event = swarm.select_next_some() => match event { + event = swarm.next() => match event { // on gossipsub incoming - SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message { + Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Gossipsub(gossipsub::Event::Message { propagation_source: peer_id, message_id: id, message, - })) => println!( + }))) => println!( "\n\nGot message: '{}' with id: {id} from peer: {peer_id}\n\n", String::from_utf8_lossy(&message.data), ), // on discovery - SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) => match e { + Some(SwarmEvent::Behaviour(swarm::BehaviourEvent::Discovery(e)) )=> match e { discovery::Event::ConnectionEstablished { peer_id, connection_id, remote_ip, remote_tcp_port } => { @@ -64,7 +64,7 @@ async fn main() { } // ignore outgoing errors: those are normal - e@SwarmEvent::OutgoingConnectionError { .. } => { log::debug!("Outgoing connection error: {e:?}"); } + e@Some(SwarmEvent::OutgoingConnectionError { .. }) => { log::debug!("Outgoing connection error: {e:?}"); } // otherwise log any other event e => { log::info!("Other event {e:?}"); } diff --git a/rust/networking/examples/chatroom_manual.rs b/rust/networking/examples/chatroom_manual.rs deleted file mode 100644 index 5d92ac86..00000000 --- a/rust/networking/examples/chatroom_manual.rs +++ /dev/null @@ -1,127 +0,0 @@ -// Copyright 2018 Parity Technologies (UK) Ltd. -// -// Permission is hereby granted, free of charge, to any person obtaining a -// copy of this software and associated documentation files (the "Software"), -// to deal in the Software without restriction, including without limitation -// the rights to use, copy, modify, merge, publish, distribute, sublicense, -// and/or sell copies of the Software, and to permit persons to whom the -// Software is furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS -// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING -// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER -// DEALINGS IN THE SOFTWARE. - -use futures::stream::StreamExt; -use libp2p::{ - gossipsub, mdns, noise, - swarm::{NetworkBehaviour, SwarmEvent}, - tcp, yamux, -}; -use std::time::Duration; -use std::{error::Error, hash::Hash}; -use tokio::{io, io::AsyncBufReadExt, select}; -use tracing_subscriber::EnvFilter; - -// We create a custom network behaviour that combines Gossipsub and Mdns. -#[derive(NetworkBehaviour)] -struct MyBehaviour { - gossipsub: gossipsub::Behaviour, - mdns: mdns::tokio::Behaviour, -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - let _ = tracing_subscriber::fmt() - .with_env_filter(EnvFilter::from_default_env()) - .try_init(); - - let mut swarm = libp2p::SwarmBuilder::with_new_identity() - .with_tokio() - .with_tcp( - tcp::Config::default(), - noise::Config::new, - yamux::Config::default, - )? - .with_behaviour(|key| { - // Set a custom gossipsub configuration - let gossipsub_config = gossipsub::ConfigBuilder::default() - .heartbeat_interval(Duration::from_secs(10)) - .validation_mode(gossipsub::ValidationMode::Strict) // This sets the kind of message validation. The default is Strict (enforce message signing) - .build() - .map_err(io::Error::other)?; // Temporary hack because `build` does not return a proper `std::error::Error`. - - // build a gossipsub network behaviour - let gossipsub = gossipsub::Behaviour::new( - gossipsub::MessageAuthenticity::Signed(key.clone()), - gossipsub_config, - )?; - - let mdns = - mdns::tokio::Behaviour::new(mdns::Config::default(), key.public().to_peer_id())?; - Ok(MyBehaviour { gossipsub, mdns }) - })? - .build(); - - println!("Running swarm with identity {}", swarm.local_peer_id()); - - // Create a Gossipsub topic - let topic = gossipsub::IdentTopic::new("test-net"); - // subscribes to our topic - swarm.behaviour_mut().gossipsub.subscribe(&topic)?; - - // Read full lines from stdin - let mut stdin = io::BufReader::new(io::stdin()).lines(); - - // Listen on all interfaces and whatever port the OS assigns - swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?; - - println!("Enter messages via STDIN and they will be sent to connected peers using Gossipsub"); - - // Kick it off - loop { - select! { - Ok(Some(line)) = stdin.next_line() => { - if let Err(e) = swarm - .behaviour_mut().gossipsub - .publish(topic.clone(), line.as_bytes()) { - println!("Publish error: {e:?}"); - } - } - event = swarm.select_next_some() => match event { - SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Discovered(list))) => { - for (peer_id, multiaddr) in list { - println!("mDNS discovered a new peer: {peer_id} on {multiaddr}"); - swarm.behaviour_mut().gossipsub.add_explicit_peer(&peer_id); - } - }, - SwarmEvent::Behaviour(MyBehaviourEvent::Mdns(mdns::Event::Expired(list))) => { - for (peer_id, multiaddr) in list { - println!("mDNS discover peer has expired: {peer_id} on {multiaddr}"); - swarm.behaviour_mut().gossipsub.remove_explicit_peer(&peer_id); - } - }, - SwarmEvent::Behaviour(MyBehaviourEvent::Gossipsub(gossipsub::Event::Message { - propagation_source: peer_id, - message_id: id, - message, - })) => println!( - "Got message: '{}' with id: {id} from peer: {peer_id}", - String::from_utf8_lossy(&message.data), - ), - SwarmEvent::NewListenAddr { address, .. } => { - println!("Local node is listening on {address}"); - } - e => { - println!("Other swarm event: {:?}", e); - } - } - } - } -} diff --git a/rust/networking/src/discovery.rs b/rust/networking/src/discovery.rs index b9a4052c..a6cbd0e0 100644 --- a/rust/networking/src/discovery.rs +++ b/rust/networking/src/discovery.rs @@ -1,8 +1,7 @@ use crate::ext::MultiaddrExt; -use crate::keep_alive; use delegate::delegate; use either::Either; -use futures::FutureExt; +use futures_lite::FutureExt; use futures_timer::Delay; use libp2p::core::transport::PortUse; use libp2p::core::{ConnectedPoint, Endpoint}; @@ -363,7 +362,7 @@ impl NetworkBehaviour for Behaviour { } // retry connecting to all mDNS peers periodically (fails safely if already connected) - if self.retry_delay.poll_unpin(cx).is_ready() { + if self.retry_delay.poll(cx).is_ready() { for (p, mas) in self.mdns_discovered.clone() { for ma in mas { self.dial(p, ma) diff --git a/rust/networking/src/keep_alive.rs b/rust/networking/src/keep_alive.rs deleted file mode 100644 index 881b11d7..00000000 --- a/rust/networking/src/keep_alive.rs +++ /dev/null @@ -1,44 +0,0 @@ -use delegate::delegate; -use libp2p::swarm::handler::ConnectionEvent; -use libp2p::swarm::{ConnectionHandlerEvent, SubstreamProtocol, dummy, handler}; -use std::task::{Context, Poll}; - -/// An implementation of [`ConnectionHandler`] that doesn't handle any protocols, but it keeps -/// the connection alive. -#[derive(Clone)] -#[repr(transparent)] -pub struct ConnectionHandler(dummy::ConnectionHandler); - -impl ConnectionHandler { - pub fn new() -> Self { - ConnectionHandler(dummy::ConnectionHandler) - } -} - -impl handler::ConnectionHandler for ConnectionHandler { - // delegate types and implementation mostly to dummy handler - type FromBehaviour = ::FromBehaviour; - type ToBehaviour = ::ToBehaviour; - type InboundProtocol = - ::InboundProtocol; - type OutboundProtocol = - ::OutboundProtocol; - type InboundOpenInfo = - ::InboundOpenInfo; - type OutboundOpenInfo = - ::OutboundOpenInfo; - - delegate! { - to self.0 { - fn listen_protocol(&self) -> SubstreamProtocol; - fn poll(&mut self, cx: &mut Context<'_>) -> Poll>; - fn on_behaviour_event(&mut self, event: Self::FromBehaviour); - fn on_connection_event(&mut self, event: ConnectionEvent); - } - } - - // specifically override this to force connection to stay alive - fn connection_keep_alive(&self) -> bool { - true - } -} diff --git a/rust/networking/src/lib.rs b/rust/networking/src/lib.rs index 59b83817..6ff8c901 100644 --- a/rust/networking/src/lib.rs +++ b/rust/networking/src/lib.rs @@ -3,19 +3,7 @@ //! this is here as a placeholder documentation //! //! - -// enable Rust-unstable features for convenience -#![feature(trait_alias)] -// #![feature(stmt_expr_attributes)] -// #![feature(unboxed_closures)] -// #![feature(assert_matches)] -// #![feature(async_fn_in_dyn_trait)] -// #![feature(async_for_loop)] -// #![feature(auto_traits)] -// #![feature(negative_impls)] - pub mod discovery; -pub mod keep_alive; pub mod swarm; /// Namespace for all the type/trait aliases used by this crate. @@ -54,11 +42,3 @@ pub(crate) mod ext { } } } - -pub(crate) mod private { - #![allow(dead_code)] - - /// Sealed traits support - pub trait Sealed {} - impl Sealed for T {} -} diff --git a/rust/networking/src/swarm.rs b/rust/networking/src/swarm.rs index 33cf8966..fc734dc3 100644 --- a/rust/networking/src/swarm.rs +++ b/rust/networking/src/swarm.rs @@ -31,7 +31,7 @@ pub fn create_swarm(keypair: identity::Keypair) -> alias::AnyResult { mod transport { use crate::alias; use crate::swarm::{NETWORK_VERSION, OVERRIDE_VERSION_ENV_VAR}; - use futures::{AsyncRead, AsyncWrite}; + use futures_lite::{AsyncRead, AsyncWrite}; use keccak_const::Sha3_256; use libp2p::core::muxing; use libp2p::core::transport::Boxed; diff --git a/rust/parts.nix b/rust/parts.nix index f160da8e..f5272b72 100644 --- a/rust/parts.nix +++ b/rust/parts.nix @@ -1,11 +1,10 @@ { inputs, ... }: { perSystem = - { config, self', inputs', pkgs, lib, ... }: + { inputs', pkgs, lib, ... }: let # Fenix nightly toolchain with all components - fenixPkgs = inputs'.fenix.packages; - rustToolchain = fenixPkgs.complete.withComponents [ + rustToolchain = inputs'.fenix.packages.stable.withComponents [ "cargo" "rustc" "clippy" diff --git a/rust/rust-toolchain.toml b/rust/rust-toolchain.toml deleted file mode 100644 index 271800cb..00000000 --- a/rust/rust-toolchain.toml +++ /dev/null @@ -1,2 +0,0 @@ -[toolchain] -channel = "nightly" \ No newline at end of file diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index a05bd6f8..f2b44495 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -14,6 +14,7 @@ from exo.download.download_utils import ( map_repo_download_progress_to_download_progress_data, ) from exo.download.shard_downloader import ShardDownloader +from exo.shared.constants import EXO_MODELS_DIR from exo.shared.models.model_cards import ModelId from exo.shared.types.commands import ( CancelDownload, @@ -46,6 +47,7 @@ class DownloadCoordinator: download_command_receiver: Receiver[ForwarderDownloadCommand] local_event_sender: Sender[ForwarderEvent] event_index_counter: Iterator[int] + offline: bool = False # Local state download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict) @@ -61,8 +63,13 @@ class DownloadCoordinator: def __post_init__(self) -> None: self.event_sender, self.event_receiver = channel[Event]() + if self.offline: + self.shard_downloader.set_internet_connection(False) self.shard_downloader.on_progress(self._download_progress_callback) + def _model_dir(self, model_id: ModelId) -> str: + return str(EXO_MODELS_DIR / model_id.normalize()) + async def _download_progress_callback( self, callback_shard: ShardMetadata, progress: RepoDownloadProgress ) -> None: @@ -73,7 +80,8 @@ class DownloadCoordinator: completed = DownloadCompleted( shard_metadata=callback_shard, node_id=self.node_id, - total_bytes=progress.total_bytes, + total=progress.total, + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = completed await self.event_sender.send( @@ -93,6 +101,7 @@ class DownloadCoordinator: download_progress=map_repo_download_progress_to_download_progress_data( progress ), + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = ongoing await self.event_sender.send( @@ -101,23 +110,30 @@ class DownloadCoordinator: self._last_progress_time[model_id] = current_time() async def run(self) -> None: - logger.info("Starting DownloadCoordinator") - self._test_internet_connection() + logger.info( + f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}" + ) + if not self.offline: + self._test_internet_connection() async with self._tg as tg: tg.start_soon(self._command_processor) tg.start_soon(self._forward_events) tg.start_soon(self._emit_existing_download_progress) - tg.start_soon(self._check_internet_connection) + if not self.offline: + tg.start_soon(self._check_internet_connection) def _test_internet_connection(self) -> None: - try: - socket.create_connection(("1.1.1.1", 443), timeout=3).close() - self.shard_downloader.set_internet_connection(True) - except OSError: - self.shard_downloader.set_internet_connection(False) - logger.debug( - f"Internet connectivity: {self.shard_downloader.internet_connection}" - ) + # Try multiple endpoints since some ISPs/networks block specific IPs + for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"): + try: + socket.create_connection((host, 443), timeout=3).close() + self.shard_downloader.set_internet_connection(True) + logger.debug(f"Internet connectivity: True (via {host})") + return + except OSError: + continue + self.shard_downloader.set_internet_connection(False) + logger.debug("Internet connectivity: False") async def _check_internet_connection(self) -> None: first_connection = True @@ -170,7 +186,11 @@ class DownloadCoordinator: return # Emit pending status - progress = DownloadPending(shard_metadata=shard, node_id=self.node_id) + progress = DownloadPending( + shard_metadata=shard, + node_id=self.node_id, + model_directory=self._model_dir(model_id), + ) self.download_status[model_id] = progress await self.event_sender.send(NodeDownloadProgress(download_progress=progress)) @@ -183,7 +203,8 @@ class DownloadCoordinator: completed = DownloadCompleted( shard_metadata=shard, node_id=self.node_id, - total_bytes=initial_progress.total_bytes, + total=initial_progress.total, + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = completed await self.event_sender.send( @@ -191,6 +212,20 @@ class DownloadCoordinator: ) return + if self.offline: + logger.warning( + f"Offline mode: model {model_id} is not fully available locally, cannot download" + ) + failed = DownloadFailed( + shard_metadata=shard, + node_id=self.node_id, + error_message=f"Model files not found locally in offline mode: {model_id}", + model_directory=self._model_dir(model_id), + ) + self.download_status[model_id] = failed + await self.event_sender.send(NodeDownloadProgress(download_progress=failed)) + return + # Start actual download self._start_download_task(shard, initial_progress) @@ -206,6 +241,7 @@ class DownloadCoordinator: download_progress=map_repo_download_progress_to_download_progress_data( initial_progress ), + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = status self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status)) @@ -219,6 +255,7 @@ class DownloadCoordinator: shard_metadata=shard, node_id=self.node_id, error_message=str(e), + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = failed await self.event_sender.send( @@ -253,6 +290,7 @@ class DownloadCoordinator: pending = DownloadPending( shard_metadata=current_status.shard_metadata, node_id=self.node_id, + model_directory=self._model_dir(model_id), ) await self.event_sender.send( NodeDownloadProgress(download_progress=pending) @@ -294,12 +332,19 @@ class DownloadCoordinator: status: DownloadProgress = DownloadCompleted( node_id=self.node_id, shard_metadata=progress.shard, - total_bytes=progress.total_bytes, + total=progress.total, + model_directory=self._model_dir( + progress.shard.model_card.model_id + ), ) elif progress.status in ["in_progress", "not_started"]: - if progress.downloaded_bytes_this_session.in_bytes == 0: + if progress.downloaded_this_session.in_bytes == 0: status = DownloadPending( - node_id=self.node_id, shard_metadata=progress.shard + node_id=self.node_id, + shard_metadata=progress.shard, + model_directory=self._model_dir( + progress.shard.model_card.model_id + ), ) else: status = DownloadOngoing( @@ -308,6 +353,9 @@ class DownloadCoordinator: download_progress=map_repo_download_progress_to_download_progress_data( progress ), + model_directory=self._model_dir( + progress.shard.model_card.model_id + ), ) else: continue diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 7974d504..a9da4cc5 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -80,9 +80,9 @@ def map_repo_file_download_progress_to_download_progress_data( repo_file_download_progress: RepoFileDownloadProgress, ) -> DownloadProgressData: return DownloadProgressData( - downloaded_bytes=repo_file_download_progress.downloaded, - downloaded_bytes_this_session=repo_file_download_progress.downloaded_this_session, - total_bytes=repo_file_download_progress.total, + downloaded=repo_file_download_progress.downloaded, + downloaded_this_session=repo_file_download_progress.downloaded_this_session, + total=repo_file_download_progress.total, completed_files=1 if repo_file_download_progress.status == "complete" else 0, total_files=1, speed=repo_file_download_progress.speed, @@ -95,9 +95,9 @@ def map_repo_download_progress_to_download_progress_data( repo_download_progress: RepoDownloadProgress, ) -> DownloadProgressData: return DownloadProgressData( - total_bytes=repo_download_progress.total_bytes, - downloaded_bytes=repo_download_progress.downloaded_bytes, - downloaded_bytes_this_session=repo_download_progress.downloaded_bytes_this_session, + total=repo_download_progress.total, + downloaded=repo_download_progress.downloaded, + downloaded_this_session=repo_download_progress.downloaded_this_session, completed_files=repo_download_progress.completed_files, total_files=repo_download_progress.total_files, speed=repo_download_progress.overall_speed, @@ -142,7 +142,7 @@ async def delete_model(model_id: ModelId) -> bool: async def seed_models(seed_dir: str | Path): - """Move model in resources folder of app to .cache/huggingface/hub""" + """Move models from resources folder to EXO_MODELS_DIR.""" source_dir = Path(seed_dir) dest_dir = await ensure_models_dir() for path in source_dir.iterdir(): @@ -448,12 +448,13 @@ async def download_file_with_retry( target_dir: Path, on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None, on_connection_lost: Callable[[], None] = lambda: None, + skip_internet: bool = False, ) -> Path: n_attempts = 3 for attempt in range(n_attempts): try: return await _download_file( - model_id, revision, path, target_dir, on_progress + model_id, revision, path, target_dir, on_progress, skip_internet ) except HuggingFaceAuthenticationError: raise @@ -487,10 +488,14 @@ async def _download_file( path: str, target_dir: Path, on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None, + skip_internet: bool = False, ) -> Path: target_path = target_dir / path if await aios.path.exists(target_path): + if skip_internet: + return target_path + local_size = (await aios.stat(target_path)).st_size # Try to verify against remote, but allow offline operation @@ -510,6 +515,11 @@ async def _download_file( ) return target_path + if skip_internet: + raise FileNotFoundError( + f"File {path} not found locally and cannot download in offline mode" + ) + await aios.makedirs((target_dir / path).parent, exist_ok=True) length, etag = await file_meta(model_id, revision, path) remote_hash = etag[:-5] if etag.endswith("-gzip") else etag @@ -568,19 +578,20 @@ def calculate_repo_progress( file_progress: dict[str, RepoFileDownloadProgress], all_start_time: float, ) -> RepoDownloadProgress: - all_total_bytes = sum((p.total.in_bytes for p in file_progress.values()), 0) - all_downloaded_bytes = sum( - (p.downloaded.in_bytes for p in file_progress.values()), 0 + all_total = sum((p.total for p in file_progress.values()), Memory.from_bytes(0)) + all_downloaded = sum( + (p.downloaded for p in file_progress.values()), Memory.from_bytes(0) ) - all_downloaded_bytes_this_session = sum( - (p.downloaded_this_session.in_bytes for p in file_progress.values()), 0 + all_downloaded_this_session = sum( + (p.downloaded_this_session for p in file_progress.values()), + Memory.from_bytes(0), ) elapsed_time = time.time() - all_start_time all_speed = ( - all_downloaded_bytes_this_session / elapsed_time if elapsed_time > 0 else 0 + all_downloaded_this_session.in_bytes / elapsed_time if elapsed_time > 0 else 0 ) all_eta = ( - timedelta(seconds=(all_total_bytes - all_downloaded_bytes) / all_speed) + timedelta(seconds=(all_total - all_downloaded).in_bytes / all_speed) if all_speed > 0 else timedelta(seconds=0) ) @@ -599,11 +610,9 @@ def calculate_repo_progress( [p for p in file_progress.values() if p.downloaded == p.total] ), total_files=len(file_progress), - downloaded_bytes=Memory.from_bytes(all_downloaded_bytes), - downloaded_bytes_this_session=Memory.from_bytes( - all_downloaded_bytes_this_session - ), - total_bytes=Memory.from_bytes(all_total_bytes), + downloaded=all_downloaded, + downloaded_this_session=all_downloaded_this_session, + total=all_total, overall_speed=all_speed, overall_eta=all_eta, status=status, @@ -814,6 +823,7 @@ async def download_shard( file, curr_bytes, total_bytes, is_renamed ), on_connection_lost=on_connection_lost, + skip_internet=skip_internet, ) if not skip_download: diff --git a/src/exo/download/shard_downloader.py b/src/exo/download/shard_downloader.py index 9dd8c324..22e85643 100644 --- a/src/exo/download/shard_downloader.py +++ b/src/exo/download/shard_downloader.py @@ -107,9 +107,9 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress( ), completed_files=0, total_files=0, - downloaded_bytes=Memory.from_bytes(0), - downloaded_bytes_this_session=Memory.from_bytes(0), - total_bytes=Memory.from_bytes(0), + downloaded=Memory.from_bytes(0), + downloaded_this_session=Memory.from_bytes(0), + total=Memory.from_bytes(0), overall_speed=0, overall_eta=timedelta(seconds=0), status="complete", diff --git a/src/exo/download/tests/test_offline_mode.py b/src/exo/download/tests/test_offline_mode.py new file mode 100644 index 00000000..15210c3f --- /dev/null +++ b/src/exo/download/tests/test_offline_mode.py @@ -0,0 +1,230 @@ +"""Tests for offline/air-gapped mode.""" + +from collections.abc import AsyncIterator +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import aiofiles +import aiofiles.os as aios +import pytest + +from exo.download.download_utils import ( + _download_file, # pyright: ignore[reportPrivateUsage] + download_file_with_retry, + fetch_file_list_with_cache, +) +from exo.shared.types.common import ModelId +from exo.shared.types.worker.downloads import FileListEntry + + +@pytest.fixture +def model_id() -> ModelId: + return ModelId("test-org/test-model") + + +@pytest.fixture +async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]: + models_dir = tmp_path / "models" + await aios.makedirs(models_dir, exist_ok=True) + with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): + yield models_dir + + +class TestDownloadFileOffline: + """Tests for _download_file with skip_internet=True.""" + + async def test_returns_local_file_without_http_verification( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """When skip_internet=True and file exists locally, return it immediately + without making any HTTP calls (no file_meta verification).""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + local_file = target_dir / "model.safetensors" + async with aiofiles.open(local_file, "wb") as f: + await f.write(b"model weights data") + + with patch( + "exo.download.download_utils.file_meta", + new_callable=AsyncMock, + ) as mock_file_meta: + result = await _download_file( + model_id, + "main", + "model.safetensors", + target_dir, + skip_internet=True, + ) + + assert result == local_file + mock_file_meta.assert_not_called() + + async def test_raises_file_not_found_for_missing_file( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """When skip_internet=True and file does NOT exist locally, + raise FileNotFoundError instead of attempting download.""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + with pytest.raises(FileNotFoundError, match="offline mode"): + await _download_file( + model_id, + "main", + "missing_model.safetensors", + target_dir, + skip_internet=True, + ) + + async def test_returns_local_file_in_subdirectory( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """When skip_internet=True and file exists in a subdirectory, + return it without HTTP calls.""" + target_dir = tmp_path / "downloads" + subdir = target_dir / "transformer" + await aios.makedirs(subdir, exist_ok=True) + + local_file = subdir / "diffusion_pytorch_model.safetensors" + async with aiofiles.open(local_file, "wb") as f: + await f.write(b"weights") + + with patch( + "exo.download.download_utils.file_meta", + new_callable=AsyncMock, + ) as mock_file_meta: + result = await _download_file( + model_id, + "main", + "transformer/diffusion_pytorch_model.safetensors", + target_dir, + skip_internet=True, + ) + + assert result == local_file + mock_file_meta.assert_not_called() + + +class TestDownloadFileWithRetryOffline: + """Tests for download_file_with_retry with skip_internet=True.""" + + async def test_propagates_skip_internet_to_download_file( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """Verify skip_internet is passed through to _download_file.""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + local_file = target_dir / "config.json" + async with aiofiles.open(local_file, "wb") as f: + await f.write(b'{"model_type": "qwen2"}') + + with patch( + "exo.download.download_utils.file_meta", + new_callable=AsyncMock, + ) as mock_file_meta: + result = await download_file_with_retry( + model_id, + "main", + "config.json", + target_dir, + skip_internet=True, + ) + + assert result == local_file + mock_file_meta.assert_not_called() + + async def test_file_not_found_does_not_retry( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """FileNotFoundError from offline mode should not trigger retries.""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + with pytest.raises(FileNotFoundError): + await download_file_with_retry( + model_id, + "main", + "nonexistent.safetensors", + target_dir, + skip_internet=True, + ) + + +class TestFetchFileListOffline: + """Tests for fetch_file_list_with_cache with skip_internet=True.""" + + async def test_uses_cached_file_list( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """When skip_internet=True and cache file exists, use it without network.""" + from pydantic import TypeAdapter + + cache_dir = temp_models_dir / "caches" / model_id.normalize() + await aios.makedirs(cache_dir, exist_ok=True) + + cached_list = [ + FileListEntry(type="file", path="model.safetensors", size=1000), + FileListEntry(type="file", path="config.json", size=200), + ] + cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json" + async with aiofiles.open(cache_file, "w") as f: + await f.write( + TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode() + ) + + with patch( + "exo.download.download_utils.fetch_file_list_with_retry", + new_callable=AsyncMock, + ) as mock_fetch: + result = await fetch_file_list_with_cache( + model_id, "main", skip_internet=True + ) + + assert result == cached_list + mock_fetch.assert_not_called() + + async def test_falls_back_to_local_directory_scan( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """When skip_internet=True and no cache but local files exist, + build file list from local directory.""" + import json + + model_dir = temp_models_dir / model_id.normalize() + await aios.makedirs(model_dir, exist_ok=True) + + async with aiofiles.open(model_dir / "config.json", "w") as f: + await f.write('{"model_type": "qwen2"}') + + index_data = { + "metadata": {}, + "weight_map": {"model.layers.0.weight": "model.safetensors"}, + } + async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: + await f.write(json.dumps(index_data)) + + async with aiofiles.open(model_dir / "model.safetensors", "wb") as f: + await f.write(b"x" * 500) + + with patch( + "exo.download.download_utils.fetch_file_list_with_retry", + new_callable=AsyncMock, + ) as mock_fetch: + result = await fetch_file_list_with_cache( + model_id, "main", skip_internet=True + ) + + mock_fetch.assert_not_called() + paths = {entry.path for entry in result} + assert "config.json" in paths + assert "model.safetensors" in paths + + async def test_raises_when_no_cache_and_no_local_files( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """When skip_internet=True and neither cache nor local files exist, + raise FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="No internet"): + await fetch_file_list_with_cache(model_id, "main", skip_internet=True) diff --git a/src/exo/main.py b/src/exo/main.py index e1ef3a70..27c78165 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -39,12 +39,13 @@ class Node: node_id: NodeId event_index_counter: Iterator[int] + offline: bool _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group) @classmethod async def create(cls, args: "Args") -> "Self": keypair = get_node_id_keypair() - node_id = NodeId(keypair.to_peer_id().to_base58()) + node_id = NodeId(keypair.to_node_id()) session_id = SessionId(master_node_id=node_id, election_clock=0) router = Router.create(keypair) await router.register_topic(topics.GLOBAL_EVENTS) @@ -68,6 +69,7 @@ class Node: download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS), local_event_sender=router.sender(topics.LOCAL_EVENTS), event_index_counter=event_index_counter, + offline=args.offline, ) else: download_coordinator = None @@ -132,10 +134,13 @@ class Node: api, node_id, event_index_counter, + args.offline, ) async def run(self): async with self._tg as tg: + signal.signal(signal.SIGINT, lambda _, __: self.shutdown()) + signal.signal(signal.SIGTERM, lambda _, __: self.shutdown()) tg.start_soon(self.router.run) tg.start_soon(self.election.run) if self.download_coordinator: @@ -147,8 +152,6 @@ class Node: if self.api: tg.start_soon(self.api.run) tg.start_soon(self._elect_loop) - signal.signal(signal.SIGINT, lambda _, __: self.shutdown()) - signal.signal(signal.SIGTERM, lambda _, __: self.shutdown()) def shutdown(self): # if this is our second call to shutdown, just sys.exit @@ -222,6 +225,7 @@ class Node: ), local_event_sender=self.router.sender(topics.LOCAL_EVENTS), event_index_counter=self.event_index_counter, + offline=self.offline, ) self._tg.start_soon(self.download_coordinator.run) if self.worker: @@ -260,6 +264,9 @@ def main(): logger.info("Starting EXO") logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}") + if args.offline: + logger.info("Running in OFFLINE mode — no internet checks, local models only") + # Set FAST_SYNCH override env var for runner subprocesses if args.fast_synch is True: os.environ["EXO_FAST_SYNCH"] = "on" @@ -282,6 +289,7 @@ class Args(CamelCaseModel): tb_only: bool = False no_worker: bool = False no_downloads: bool = False + offline: bool = False fast_synch: bool | None = None # None = auto, True = force on, False = force off @classmethod @@ -329,6 +337,11 @@ class Args(CamelCaseModel): action="store_true", help="Disable the download coordinator (node won't download models)", ) + parser.add_argument( + "--offline", + action="store_true", + help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models", + ) fast_synch_group = parser.add_mutually_exclusive_group() fast_synch_group.add_argument( "--fast-synch", diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py index b86c5ec1..9773806f 100644 --- a/src/exo/master/adapters/chat_completions.py +++ b/src/exo/master/adapters/chat_completions.py @@ -19,7 +19,12 @@ from exo.shared.types.api import ( ToolCall, Usage, ) -from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) from exo.shared.types.common import CommandId from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams @@ -54,7 +59,11 @@ def chat_request_to_text_generation( chat_template_messages.append({"role": "system", "content": content}) else: # Skip messages with no meaningful content - if msg.content is None and msg.thinking is None and msg.tool_calls is None: + if ( + msg.content is None + and msg.reasoning_content is None + and msg.tool_calls is None + ): continue if msg.role in ("user", "assistant", "developer"): @@ -106,6 +115,11 @@ def chunk_to_response( ] ) + if chunk.is_thinking: + delta = ChatCompletionMessage(role="assistant", reasoning_content=chunk.text) + else: + delta = ChatCompletionMessage(role="assistant", content=chunk.text) + return ChatCompletionResponse( id=command_id, created=int(time.time()), @@ -113,7 +127,7 @@ def chunk_to_response( choices=[ StreamingChoiceResponse( index=0, - delta=ChatCompletionMessage(role="assistant", content=chunk.text), + delta=delta, logprobs=logprobs, finish_reason=chunk.finish_reason, ) @@ -123,72 +137,87 @@ def chunk_to_response( async def generate_chat_stream( command_id: CommandId, - chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], + chunk_stream: AsyncGenerator[ + PrefillProgressChunk | ErrorChunk | ToolCallChunk | TokenChunk, None + ], ) -> AsyncGenerator[str, None]: """Generate Chat Completions API streaming events from chunks.""" last_usage: Usage | None = None async for chunk in chunk_stream: - if isinstance(chunk, ErrorChunk): - error_response = ErrorResponse( - error=ErrorInfo( - message=chunk.error_message or "Internal server error", - type="InternalServerError", - code=500, - ) - ) - yield f"data: {error_response.model_dump_json()}\n\n" - yield "data: [DONE]\n\n" - return + match chunk: + case PrefillProgressChunk(): + # Use SSE comment so third-party clients ignore it + yield f": prefill_progress {chunk.model_dump_json()}\n\n" - last_usage = chunk.usage or last_usage - - if isinstance(chunk, ToolCallChunk): - tool_call_deltas = [ - ToolCall( - id=tool.id, - index=i, - function=tool, - ) - for i, tool in enumerate(chunk.tool_calls) - ] - tool_response = ChatCompletionResponse( - id=command_id, - created=int(time.time()), - model=chunk.model, - choices=[ - StreamingChoiceResponse( - index=0, - delta=ChatCompletionMessage( - role="assistant", - tool_calls=tool_call_deltas, - ), - finish_reason="tool_calls", + case ErrorChunk(): + error_response = ErrorResponse( + error=ErrorInfo( + message=chunk.error_message or "Internal server error", + type="InternalServerError", + code=500, ) - ], - usage=last_usage, - ) - yield f"data: {tool_response.model_dump_json()}\n\n" - yield "data: [DONE]\n\n" - return + ) + yield f"data: {error_response.model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + return - chunk_response = chunk_to_response(chunk, command_id) - if chunk.finish_reason is not None: - chunk_response = chunk_response.model_copy(update={"usage": last_usage}) - yield f"data: {chunk_response.model_dump_json()}\n\n" + case ToolCallChunk(): + last_usage = chunk.usage or last_usage - if chunk.finish_reason is not None: - yield "data: [DONE]\n\n" + tool_call_deltas = [ + ToolCall( + id=tool.id, + index=i, + function=tool, + ) + for i, tool in enumerate(chunk.tool_calls) + ] + tool_response = ChatCompletionResponse( + id=command_id, + created=int(time.time()), + model=chunk.model, + choices=[ + StreamingChoiceResponse( + index=0, + delta=ChatCompletionMessage( + role="assistant", + tool_calls=tool_call_deltas, + ), + finish_reason="tool_calls", + ) + ], + usage=last_usage, + ) + yield f"data: {tool_response.model_dump_json()}\n\n" + yield "data: [DONE]\n\n" + return + + case TokenChunk(): + last_usage = chunk.usage or last_usage + + chunk_response = chunk_to_response(chunk, command_id) + if chunk.finish_reason is not None: + chunk_response = chunk_response.model_copy( + update={"usage": last_usage} + ) + yield f"data: {chunk_response.model_dump_json()}\n\n" + + if chunk.finish_reason is not None: + yield "data: [DONE]\n\n" async def collect_chat_response( command_id: CommandId, - chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], ) -> AsyncGenerator[str]: # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason """Collect all token chunks and return a single ChatCompletionResponse.""" text_parts: list[str] = [] + thinking_parts: list[str] = [] tool_calls: list[ToolCall] = [] logprobs_content: list[LogprobsContentItem] = [] model: str | None = None @@ -197,43 +226,52 @@ async def collect_chat_response( last_usage: Usage | None = None async for chunk in chunk_stream: - if isinstance(chunk, ErrorChunk): - error_message = chunk.error_message or "Internal server error" - break + match chunk: + case PrefillProgressChunk(): + continue - if model is None: - model = chunk.model + case ErrorChunk(): + error_message = chunk.error_message or "Internal server error" + break - last_usage = chunk.usage or last_usage - - if isinstance(chunk, TokenChunk): - text_parts.append(chunk.text) - if chunk.logprob is not None: - logprobs_content.append( - LogprobsContentItem( - token=chunk.text, - logprob=chunk.logprob, - top_logprobs=chunk.top_logprobs or [], + case TokenChunk(): + if model is None: + model = chunk.model + last_usage = chunk.usage or last_usage + if chunk.is_thinking: + thinking_parts.append(chunk.text) + else: + text_parts.append(chunk.text) + if chunk.logprob is not None: + logprobs_content.append( + LogprobsContentItem( + token=chunk.text, + logprob=chunk.logprob, + top_logprobs=chunk.top_logprobs or [], + ) ) - ) + if chunk.finish_reason is not None: + finish_reason = chunk.finish_reason - if isinstance(chunk, ToolCallChunk): - tool_calls.extend( - ToolCall( - id=tool.id, - index=i, - function=tool, + case ToolCallChunk(): + if model is None: + model = chunk.model + last_usage = chunk.usage or last_usage + tool_calls.extend( + ToolCall( + id=tool.id, + index=i, + function=tool, + ) + for i, tool in enumerate(chunk.tool_calls) ) - for i, tool in enumerate(chunk.tool_calls) - ) - - if chunk.finish_reason is not None: - finish_reason = chunk.finish_reason + finish_reason = chunk.finish_reason if error_message is not None: raise ValueError(error_message) combined_text = "".join(text_parts) + combined_thinking = "".join(thinking_parts) if thinking_parts else None assert model is not None yield ChatCompletionResponse( @@ -246,6 +284,7 @@ async def collect_chat_response( message=ChatCompletionMessage( role="assistant", content=combined_text, + reasoning_content=combined_thinking, tool_calls=tool_calls if tool_calls else None, ), logprobs=Logprobs(content=logprobs_content) diff --git a/src/exo/master/adapters/claude.py b/src/exo/master/adapters/claude.py index ee7c263a..7a31f6fd 100644 --- a/src/exo/master/adapters/claude.py +++ b/src/exo/master/adapters/claude.py @@ -1,11 +1,17 @@ """Claude Messages API adapter for converting requests/responses.""" import json +import re from collections.abc import AsyncGenerator from typing import Any from exo.shared.types.api import FinishReason, Usage -from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) from exo.shared.types.claude_api import ( ClaudeContentBlock, ClaudeContentBlockDeltaEvent, @@ -23,6 +29,8 @@ from exo.shared.types.claude_api import ( ClaudeStopReason, ClaudeTextBlock, ClaudeTextDelta, + ClaudeThinkingBlock, + ClaudeThinkingDelta, ClaudeToolResultBlock, ClaudeToolUseBlock, ClaudeUsage, @@ -56,6 +64,22 @@ def _extract_tool_result_text(block: ClaudeToolResultBlock) -> str: return "".join(sub_block.text for sub_block in block.content) +# Matches "x-anthropic-billing-header: ...;" (with optional trailing newline) +# or similar telemetry headers that change every request and break KV prefix caching. +_VOLATILE_HEADER_RE = re.compile(r"^x-anthropic-[^\n]*;\n?", re.MULTILINE) + + +def _strip_volatile_headers(text: str) -> str: + """Remove Anthropic billing/telemetry headers from system prompt text. + + Claude Code prepends headers like 'x-anthropic-billing-header: cc_version=...; + cc_entrypoint=...; cch=...;' that contain per-request content hashes. These + change every request and break KV prefix caching (the prefix diverges at ~20 + tokens instead of matching thousands of conversation tokens). + """ + return _VOLATILE_HEADER_RE.sub("", text) + + def claude_request_to_text_generation( request: ClaudeMessagesRequest, ) -> TextGenerationTaskParams: @@ -68,6 +92,8 @@ def claude_request_to_text_generation( instructions = request.system else: instructions = "".join(block.text for block in request.system) + + instructions = _strip_volatile_headers(instructions) chat_template_messages.append({"role": "system", "content": instructions}) # Convert messages to input @@ -80,12 +106,15 @@ def claude_request_to_text_generation( # Process structured content blocks text_parts: list[str] = [] + thinking_parts: list[str] = [] tool_calls: list[dict[str, Any]] = [] tool_results: list[ClaudeToolResultBlock] = [] for block in msg.content: if isinstance(block, ClaudeTextBlock): text_parts.append(block.text) + elif isinstance(block, ClaudeThinkingBlock): + thinking_parts.append(block.thinking) elif isinstance(block, ClaudeToolUseBlock): tool_calls.append( { @@ -101,6 +130,7 @@ def claude_request_to_text_generation( tool_results.append(block) content = "".join(text_parts) + reasoning_content = "".join(thinking_parts) if thinking_parts else None # Build InputMessage from text content if msg.role in ("user", "assistant"): @@ -108,9 +138,14 @@ def claude_request_to_text_generation( # Build chat_template_messages preserving tool structure if tool_calls: - chat_template_messages.append( - {"role": "assistant", "content": content, "tool_calls": tool_calls} - ) + chat_msg: dict[str, Any] = { + "role": "assistant", + "content": content, + "tool_calls": tool_calls, + } + if reasoning_content: + chat_msg["reasoning_content"] = reasoning_content + chat_template_messages.append(chat_msg) elif tool_results: for tr in tool_results: chat_template_messages.append( @@ -121,7 +156,10 @@ def claude_request_to_text_generation( } ) else: - chat_template_messages.append({"role": msg.role, "content": content}) + chat_msg = {"role": msg.role, "content": content} + if reasoning_content: + chat_msg["reasoning_content"] = reasoning_content + chat_template_messages.append(chat_msg) # Convert Claude tool definitions to OpenAI-style function tools tools: list[dict[str, Any]] | None = None @@ -138,6 +176,10 @@ def claude_request_to_text_generation( for tool in request.tools ] + enable_thinking: bool | None = None + if request.thinking is not None: + enable_thinking = request.thinking.type in ("enabled", "adaptive") + return TextGenerationTaskParams( model=request.model, input=input_messages @@ -151,6 +193,7 @@ def claude_request_to_text_generation( stop=request.stop_sequences, stream=request.stream, tools=tools, + enable_thinking=enable_thinking, chat_template_messages=chat_template_messages if chat_template_messages else None, @@ -160,18 +203,24 @@ def claude_request_to_text_generation( async def collect_claude_response( command_id: CommandId, model: str, - chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], ) -> AsyncGenerator[str]: # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason """Collect all token chunks and return a single ClaudeMessagesResponse.""" text_parts: list[str] = [] + thinking_parts: list[str] = [] tool_use_blocks: list[ClaudeToolUseBlock] = [] stop_reason: ClaudeStopReason | None = None last_usage: Usage | None = None error_message: str | None = None async for chunk in chunk_stream: + if isinstance(chunk, PrefillProgressChunk): + continue + if isinstance(chunk, ErrorChunk): error_message = chunk.error_message or "Internal server error" break @@ -190,7 +239,10 @@ async def collect_claude_response( stop_reason = "tool_use" continue - text_parts.append(chunk.text) + if chunk.is_thinking: + thinking_parts.append(chunk.text) + else: + text_parts.append(chunk.text) if chunk.finish_reason is not None: stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason) @@ -199,9 +251,12 @@ async def collect_claude_response( raise ValueError(error_message) combined_text = "".join(text_parts) + combined_thinking = "".join(thinking_parts) # Build content blocks content: list[ClaudeContentBlock] = [] + if combined_thinking: + content.append(ClaudeThinkingBlock(thinking=combined_thinking)) if combined_text: content.append(ClaudeTextBlock(text=combined_text)) content.extend(tool_use_blocks) @@ -230,7 +285,9 @@ async def collect_claude_response( async def generate_claude_stream( command_id: CommandId, model: str, - chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], ) -> AsyncGenerator[str, None]: """Generate Claude Messages API streaming events from TokenChunks.""" # Initial message_start event @@ -244,18 +301,21 @@ async def generate_claude_stream( start_event = ClaudeMessageStartEvent(message=initial_message) yield f"event: message_start\ndata: {start_event.model_dump_json()}\n\n" - # content_block_start for text block at index 0 - block_start = ClaudeContentBlockStartEvent( - index=0, content_block=ClaudeTextBlock(text="") - ) - yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n" - output_tokens = 0 stop_reason: ClaudeStopReason | None = None last_usage: Usage | None = None - next_block_index = 1 # text block is 0, tool blocks start at 1 + next_block_index = 0 + + # Track whether we've started thinking/text blocks + thinking_block_started = False + thinking_block_index = -1 + text_block_started = False + text_block_index = -1 async for chunk in chunk_stream: + if isinstance(chunk, PrefillProgressChunk): + continue + if isinstance(chunk, ErrorChunk): # Close text block and bail break @@ -295,12 +355,45 @@ async def generate_claude_stream( output_tokens += 1 # Count each chunk as one token - # content_block_delta - delta_event = ClaudeContentBlockDeltaEvent( - index=0, - delta=ClaudeTextDelta(text=chunk.text), - ) - yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n" + if chunk.is_thinking: + # Start thinking block on first thinking token + if not thinking_block_started: + thinking_block_started = True + thinking_block_index = next_block_index + next_block_index += 1 + block_start = ClaudeContentBlockStartEvent( + index=thinking_block_index, + content_block=ClaudeThinkingBlock(thinking=""), + ) + yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n" + + delta_event = ClaudeContentBlockDeltaEvent( + index=thinking_block_index, + delta=ClaudeThinkingDelta(thinking=chunk.text), + ) + yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n" + else: + # Close thinking block when transitioning to text + if thinking_block_started and text_block_index == -1: + block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index) + yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" + + # Start text block on first text token + if not text_block_started: + text_block_started = True + text_block_index = next_block_index + next_block_index += 1 + block_start = ClaudeContentBlockStartEvent( + index=text_block_index, + content_block=ClaudeTextBlock(text=""), + ) + yield f"event: content_block_start\ndata: {block_start.model_dump_json()}\n\n" + + delta_event = ClaudeContentBlockDeltaEvent( + index=text_block_index, + delta=ClaudeTextDelta(text=chunk.text), + ) + yield f"event: content_block_delta\ndata: {delta_event.model_dump_json()}\n\n" if chunk.finish_reason is not None: stop_reason = finish_reason_to_claude_stop_reason(chunk.finish_reason) @@ -309,9 +402,22 @@ async def generate_claude_stream( if last_usage is not None: output_tokens = last_usage.completion_tokens - # content_block_stop for text block - block_stop = ClaudeContentBlockStopEvent(index=0) - yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" + # Close any open blocks + if thinking_block_started and text_block_index == -1: + block_stop = ClaudeContentBlockStopEvent(index=thinking_block_index) + yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" + + if text_block_started: + block_stop = ClaudeContentBlockStopEvent(index=text_block_index) + yield f"event: content_block_stop\ndata: {block_stop.model_dump_json()}\n\n" + + if not thinking_block_started and not text_block_started: + empty_start = ClaudeContentBlockStartEvent( + index=0, content_block=ClaudeTextBlock(text="") + ) + yield f"event: content_block_start\ndata: {empty_start.model_dump_json()}\n\n" + empty_stop = ClaudeContentBlockStopEvent(index=0) + yield f"event: content_block_stop\ndata: {empty_stop.model_dump_json()}\n\n" # message_delta message_delta = ClaudeMessageDeltaEvent( diff --git a/src/exo/master/adapters/ollama.py b/src/exo/master/adapters/ollama.py new file mode 100644 index 00000000..75b772d0 --- /dev/null +++ b/src/exo/master/adapters/ollama.py @@ -0,0 +1,456 @@ +from __future__ import annotations + +import json +from collections.abc import AsyncGenerator +from typing import Any + +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) +from exo.shared.types.common import CommandId +from exo.shared.types.ollama_api import ( + OllamaChatRequest, + OllamaChatResponse, + OllamaDoneReason, + OllamaGenerateRequest, + OllamaGenerateResponse, + OllamaMessage, + OllamaToolCall, + OllamaToolFunction, +) +from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams + + +def _map_done_reason( + finish_reason: str | None, +) -> OllamaDoneReason | None: + if finish_reason is None: + return None + if finish_reason == "stop": + return "stop" + if finish_reason == "length": + return "length" + if finish_reason in ("tool_calls", "function_call"): + return "tool_call" + if finish_reason == "error": + return "error" + return "stop" + + +def _try_parse_json(value: str) -> dict[str, Any] | str: + try: + return json.loads(value) # type: ignore + except json.JSONDecodeError: + return value + + +def _build_tool_calls(chunk: ToolCallChunk) -> list[OllamaToolCall]: + tool_calls: list[OllamaToolCall] = [] + for index, tool in enumerate(chunk.tool_calls): + # tool.arguments is always str; try to parse as JSON dict for Ollama format + arguments: dict[str, Any] | str = _try_parse_json(tool.arguments) + tool_calls.append( + OllamaToolCall( + id=tool.id, + type="function", + function=OllamaToolFunction( + name=tool.name, arguments=arguments, index=index + ), + ) + ) + return tool_calls + + +def _get_usage( + chunk: TokenChunk | ToolCallChunk, +) -> tuple[int | None, int | None]: + """Extract (prompt_eval_count, eval_count) from a chunk.""" + if chunk.usage is not None: + return (chunk.usage.prompt_tokens, chunk.usage.completion_tokens) + if chunk.stats is not None: + return (chunk.stats.prompt_tokens, chunk.stats.generation_tokens) + return (None, None) + + +def ollama_request_to_text_generation( + request: OllamaChatRequest, +) -> TextGenerationTaskParams: + """Convert Ollama chat request to exo's internal text generation format.""" + instructions: str | None = None + input_messages: list[InputMessage] = [] + chat_template_messages: list[dict[str, Any]] = [] + tool_message_index = 0 + + for msg in request.messages: + content = msg.content or "" + + if msg.role == "system": + if instructions is None: + instructions = content + else: + instructions = f"{instructions}\n{content}" + chat_template_messages.append({"role": "system", "content": content}) + continue + + if msg.role in ("user", "assistant") and ( + msg.content is not None or msg.thinking is not None or msg.tool_calls + ): + input_messages.append(InputMessage(role=msg.role, content=content)) + + dumped: dict[str, Any] = {"role": msg.role, "content": content} + if msg.thinking is not None: + dumped["thinking"] = msg.thinking + if msg.tool_calls is not None: + tool_calls_list: list[dict[str, Any]] = [] + for tc in msg.tool_calls: + function: dict[str, Any] = { + "name": tc.function.name, + "arguments": ( + json.dumps(tc.function.arguments) + if isinstance(tc.function.arguments, dict) + else tc.function.arguments + ), + } + if tc.function.index is not None: + function["index"] = tc.function.index + tool_call: dict[str, Any] = {"function": function} + if tc.id is not None: + tool_call["id"] = tc.id + if tc.type is not None: + tool_call["type"] = tc.type + tool_calls_list.append(tool_call) + dumped["tool_calls"] = tool_calls_list + if msg.name is not None: + dumped["name"] = msg.name + if msg.role == "tool": + tool_message_index += 1 + tool_call_id = msg.tool_name or msg.name or f"tool_{tool_message_index}" + dumped["tool_call_id"] = tool_call_id + if msg.tool_name is not None: + dumped["tool_name"] = msg.tool_name + chat_template_messages.append(dumped) + + options = request.options + return TextGenerationTaskParams( + model=request.model, + input=input_messages + if input_messages + else [InputMessage(role="user", content="")], + instructions=instructions, + max_output_tokens=options.num_predict if options else None, + temperature=options.temperature if options else None, + top_p=options.top_p if options else None, + top_k=options.top_k if options else None, + stop=options.stop if options else None, + seed=options.seed if options else None, + stream=request.stream, + tools=request.tools, + enable_thinking=request.think, + chat_template_messages=chat_template_messages + if chat_template_messages + else None, + ) + + +async def generate_ollama_chat_stream( + _command_id: CommandId, + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], +) -> AsyncGenerator[str, None]: + """Generate streaming responses in Ollama format (newline-delimited JSON).""" + thinking_parts: list[str] = [] + + async for chunk in chunk_stream: + match chunk: + case PrefillProgressChunk(): + continue + + case ErrorChunk(): + error_response = OllamaChatResponse( + model=str(chunk.model), + message=OllamaMessage( + role="assistant", content=chunk.error_message + ), + done=True, + done_reason="error", + ) + yield f"{error_response.model_dump_json(exclude_none=True)}\n" + return + + case ToolCallChunk(): + prompt_eval, eval_count = _get_usage(chunk) + response = OllamaChatResponse( + model=str(chunk.model), + message=OllamaMessage( + role="assistant", + content="", + tool_calls=_build_tool_calls(chunk), + thinking="".join(thinking_parts) if thinking_parts else None, + ), + done=True, + done_reason="tool_call", + prompt_eval_count=prompt_eval, + eval_count=eval_count, + ) + yield f"{response.model_dump_json(exclude_none=True)}\n" + return + + case TokenChunk(): + done = chunk.finish_reason is not None + + if chunk.is_thinking: + thinking_parts.append(chunk.text) + response = OllamaChatResponse( + model=str(chunk.model), + message=OllamaMessage( + role="assistant", content="", thinking=chunk.text + ), + done=False, + ) + yield f"{response.model_dump_json(exclude_none=True)}\n" + elif done: + prompt_eval, eval_count = _get_usage(chunk) + response = OllamaChatResponse( + model=str(chunk.model), + message=OllamaMessage( + role="assistant", + content=chunk.text, + ), + done=True, + done_reason=_map_done_reason(chunk.finish_reason), + prompt_eval_count=prompt_eval, + eval_count=eval_count, + ) + yield f"{response.model_dump_json(exclude_none=True)}\n" + else: + response = OllamaChatResponse( + model=str(chunk.model), + message=OllamaMessage(role="assistant", content=chunk.text), + done=False, + ) + yield f"{response.model_dump_json(exclude_none=True)}\n" + + if done: + return + + +async def collect_ollama_chat_response( + _command_id: CommandId, + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], +) -> AsyncGenerator[str]: + """Collect streaming chunks into a single non-streaming Ollama response. + + Returns an AsyncGenerator[str] (single yield) for consistency with FastAPI + StreamingResponse cancellation handling. + """ + text_parts: list[str] = [] + thinking_parts: list[str] = [] + tool_calls: list[OllamaToolCall] = [] + model: str | None = None + finish_reason: str | None = None + prompt_eval_count: int | None = None + eval_count: int | None = None + + async for chunk in chunk_stream: + match chunk: + case PrefillProgressChunk(): + continue + + case ErrorChunk(): + raise ValueError(chunk.error_message or "Internal server error") + + case TokenChunk(): + if model is None: + model = str(chunk.model) + if chunk.is_thinking: + thinking_parts.append(chunk.text) + else: + text_parts.append(chunk.text) + if chunk.finish_reason is not None: + finish_reason = chunk.finish_reason + prompt_eval_count, eval_count = _get_usage(chunk) + + case ToolCallChunk(): + if model is None: + model = str(chunk.model) + tool_calls.extend(_build_tool_calls(chunk)) + finish_reason = chunk.finish_reason + prompt_eval_count, eval_count = _get_usage(chunk) + + combined_text = "".join(text_parts) + combined_thinking = "".join(thinking_parts) if thinking_parts else None + assert model is not None + + yield OllamaChatResponse( + model=model, + message=OllamaMessage( + role="assistant", + content=combined_text, + thinking=combined_thinking, + tool_calls=tool_calls if tool_calls else None, + ), + done=True, + done_reason=_map_done_reason(finish_reason), + prompt_eval_count=prompt_eval_count, + eval_count=eval_count, + ).model_dump_json(exclude_none=True) + return + + +# ── /api/generate ── + + +def ollama_generate_request_to_text_generation( + request: OllamaGenerateRequest, +) -> TextGenerationTaskParams: + """Convert Ollama generate request to exo's internal text generation format.""" + chat_template_messages: list[dict[str, Any]] = [] + if request.system: + chat_template_messages.append({"role": "system", "content": request.system}) + chat_template_messages.append({"role": "user", "content": request.prompt}) + + options = request.options + return TextGenerationTaskParams( + model=request.model, + input=[InputMessage(role="user", content=request.prompt)], + instructions=request.system, + max_output_tokens=options.num_predict if options else None, + temperature=options.temperature if options else None, + top_p=options.top_p if options else None, + top_k=options.top_k if options else None, + stop=options.stop if options else None, + seed=options.seed if options else None, + stream=request.stream, + enable_thinking=request.think, + chat_template_messages=chat_template_messages + if chat_template_messages + else None, + ) + + +async def generate_ollama_generate_stream( + _command_id: CommandId, + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], +) -> AsyncGenerator[str, None]: + """Generate streaming responses for /api/generate in Ollama NDJSON format.""" + thinking_parts: list[str] = [] + + async for chunk in chunk_stream: + match chunk: + case PrefillProgressChunk(): + continue + + case ErrorChunk(): + resp = OllamaGenerateResponse( + model=str(chunk.model), + response="", + done=True, + done_reason="error", + ) + yield f"{resp.model_dump_json(exclude_none=True)}\n" + return + + case ToolCallChunk(): + # generate endpoint doesn't support tools; emit as done + prompt_eval, eval_count = _get_usage(chunk) + resp = OllamaGenerateResponse( + model=str(chunk.model), + response="", + done=True, + done_reason="stop", + prompt_eval_count=prompt_eval, + eval_count=eval_count, + ) + yield f"{resp.model_dump_json(exclude_none=True)}\n" + return + + case TokenChunk(): + done = chunk.finish_reason is not None + + if chunk.is_thinking: + thinking_parts.append(chunk.text) + resp = OllamaGenerateResponse( + model=str(chunk.model), + response="", + thinking=chunk.text, + done=False, + ) + yield f"{resp.model_dump_json(exclude_none=True)}\n" + elif done: + prompt_eval, eval_count = _get_usage(chunk) + resp = OllamaGenerateResponse( + model=str(chunk.model), + response=chunk.text, + done=True, + done_reason=_map_done_reason(chunk.finish_reason), + prompt_eval_count=prompt_eval, + eval_count=eval_count, + ) + yield f"{resp.model_dump_json(exclude_none=True)}\n" + else: + resp = OllamaGenerateResponse( + model=str(chunk.model), + response=chunk.text, + done=False, + ) + yield f"{resp.model_dump_json(exclude_none=True)}\n" + + if done: + return + + +async def collect_ollama_generate_response( + _command_id: CommandId, + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], +) -> AsyncGenerator[str]: + """Collect chunks into a single non-streaming /api/generate response.""" + text_parts: list[str] = [] + thinking_parts: list[str] = [] + model: str | None = None + finish_reason: str | None = None + prompt_eval_count: int | None = None + eval_count: int | None = None + + async for chunk in chunk_stream: + match chunk: + case PrefillProgressChunk(): + continue + case ErrorChunk(): + raise ValueError(chunk.error_message or "Internal server error") + case TokenChunk(): + if model is None: + model = str(chunk.model) + if chunk.is_thinking: + thinking_parts.append(chunk.text) + else: + text_parts.append(chunk.text) + if chunk.finish_reason is not None: + finish_reason = chunk.finish_reason + prompt_eval_count, eval_count = _get_usage(chunk) + case ToolCallChunk(): + if model is None: + model = str(chunk.model) + finish_reason = chunk.finish_reason + prompt_eval_count, eval_count = _get_usage(chunk) + + assert model is not None + yield OllamaGenerateResponse( + model=model, + response="".join(text_parts), + thinking="".join(thinking_parts) if thinking_parts else None, + done=True, + done_reason=_map_done_reason(finish_reason), + prompt_eval_count=prompt_eval_count, + eval_count=eval_count, + ).model_dump_json(exclude_none=True) + return diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py index b37b7d54..dc726ba7 100644 --- a/src/exo/master/adapters/responses.py +++ b/src/exo/master/adapters/responses.py @@ -5,7 +5,12 @@ from itertools import count from typing import Any from exo.shared.types.api import Usage -from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk +from exo.shared.types.chunks import ( + ErrorChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) from exo.shared.types.common import CommandId from exo.shared.types.openai_responses import ( FunctionCallInputItem, @@ -24,8 +29,15 @@ from exo.shared.types.openai_responses import ( ResponseOutputItemAddedEvent, ResponseOutputItemDoneEvent, ResponseOutputText, + ResponseReasoningItem, + ResponseReasoningSummaryPartAddedEvent, + ResponseReasoningSummaryPartDoneEvent, + ResponseReasoningSummaryText, + ResponseReasoningSummaryTextDeltaEvent, + ResponseReasoningSummaryTextDoneEvent, ResponsesRequest, ResponsesResponse, + ResponsesStreamEvent, ResponseTextDeltaEvent, ResponseTextDoneEvent, ResponseUsage, @@ -33,6 +45,11 @@ from exo.shared.types.openai_responses import ( from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams +def _format_sse(event: ResponsesStreamEvent) -> str: + """Format a streaming event as an SSE message.""" + return f"event: {event.type}\ndata: {event.model_dump_json()}\n\n" + + def _extract_content(content: str | list[ResponseContentPart]) -> str: """Extract plain text from a content field that may be a string or list of parts.""" if isinstance(content, str): @@ -121,19 +138,26 @@ def responses_request_to_text_generation( async def collect_responses_response( command_id: CommandId, model: str, - chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], ) -> AsyncGenerator[str]: # This is an AsyncGenerator[str] rather than returning a ChatCompletionReponse because # FastAPI handles the cancellation better but wouldn't auto-serialize for some reason """Collect all token chunks and return a single ResponsesResponse.""" response_id = f"resp_{command_id}" item_id = f"item_{command_id}" + reasoning_id = f"rs_{command_id}" accumulated_text = "" + thinking_parts: list[str] = [] function_call_items: list[ResponseFunctionCallItem] = [] last_usage: Usage | None = None error_message: str | None = None async for chunk in chunk_stream: + if isinstance(chunk, PrefillProgressChunk): + continue + if isinstance(chunk, ErrorChunk): error_message = chunk.error_message or "Internal server error" break @@ -144,14 +168,18 @@ async def collect_responses_response( for tool in chunk.tool_calls: function_call_items.append( ResponseFunctionCallItem( - id=f"fc_{tool.id}", - call_id=f"call_{tool.id}", + id=tool.id, + call_id=tool.id, name=tool.name, arguments=tool.arguments, ) ) continue + if chunk.is_thinking: + thinking_parts.append(chunk.text) + continue + accumulated_text += chunk.text if error_message is not None: @@ -166,13 +194,21 @@ async def collect_responses_response( total_tokens=last_usage.total_tokens, ) - output: list[ResponseItem] = [ + output: list[ResponseItem] = [] + if thinking_parts: + output.append( + ResponseReasoningItem( + id=reasoning_id, + summary=[ResponseReasoningSummaryText(text="".join(thinking_parts))], + ) + ) + output.append( ResponseMessageItem( id=item_id, content=[ResponseOutputText(text=accumulated_text)], status="completed", ) - ] + ) output.extend(function_call_items) yield ResponsesResponse( @@ -189,11 +225,14 @@ async def collect_responses_response( async def generate_responses_stream( command_id: CommandId, model: str, - chunk_stream: AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None], + chunk_stream: AsyncGenerator[ + ErrorChunk | ToolCallChunk | TokenChunk | PrefillProgressChunk, None + ], ) -> AsyncGenerator[str, None]: """Generate OpenAI Responses API streaming events from TokenChunks.""" response_id = f"resp_{command_id}" item_id = f"item_{command_id}" + reasoning_id = f"rs_{command_id}" seq = count(1) # response.created @@ -207,42 +246,30 @@ async def generate_responses_stream( created_event = ResponseCreatedEvent( sequence_number=next(seq), response=initial_response ) - yield f"event: response.created\ndata: {created_event.model_dump_json()}\n\n" + yield _format_sse(created_event) # response.in_progress in_progress_event = ResponseInProgressEvent( sequence_number=next(seq), response=initial_response ) - yield f"event: response.in_progress\ndata: {in_progress_event.model_dump_json()}\n\n" - - # response.output_item.added - initial_item = ResponseMessageItem( - id=item_id, - content=[ResponseOutputText(text="")], - status="in_progress", - ) - item_added = ResponseOutputItemAddedEvent( - sequence_number=next(seq), output_index=0, item=initial_item - ) - yield f"event: response.output_item.added\ndata: {item_added.model_dump_json()}\n\n" - - # response.content_part.added - initial_part = ResponseOutputText(text="") - part_added = ResponseContentPartAddedEvent( - sequence_number=next(seq), - item_id=item_id, - output_index=0, - content_index=0, - part=initial_part, - ) - yield f"event: response.content_part.added\ndata: {part_added.model_dump_json()}\n\n" + yield _format_sse(in_progress_event) accumulated_text = "" + accumulated_thinking = "" function_call_items: list[ResponseFunctionCallItem] = [] last_usage: Usage | None = None - next_output_index = 1 # message item is at 0 + next_output_index = 0 + + # Track dynamic block creation + reasoning_started = False + reasoning_output_index = -1 + message_started = False + message_output_index = -1 async for chunk in chunk_stream: + if isinstance(chunk, PrefillProgressChunk): + continue + if isinstance(chunk, ErrorChunk): break @@ -266,7 +293,7 @@ async def generate_responses_stream( output_index=next_output_index, item=fc_item, ) - yield f"event: response.output_item.added\ndata: {fc_added.model_dump_json()}\n\n" + yield _format_sse(fc_added) # response.function_call_arguments.delta args_delta = ResponseFunctionCallArgumentsDeltaEvent( @@ -275,7 +302,7 @@ async def generate_responses_stream( output_index=next_output_index, delta=tool.arguments, ) - yield f"event: response.function_call_arguments.delta\ndata: {args_delta.model_dump_json()}\n\n" + yield _format_sse(args_delta) # response.function_call_arguments.done args_done = ResponseFunctionCallArgumentsDoneEvent( @@ -285,7 +312,7 @@ async def generate_responses_stream( name=tool.name, arguments=tool.arguments, ) - yield f"event: response.function_call_arguments.done\ndata: {args_done.model_dump_json()}\n\n" + yield _format_sse(args_done) # response.output_item.done fc_done_item = ResponseFunctionCallItem( @@ -300,44 +327,205 @@ async def generate_responses_stream( output_index=next_output_index, item=fc_done_item, ) - yield f"event: response.output_item.done\ndata: {fc_item_done.model_dump_json()}\n\n" + yield _format_sse(fc_item_done) function_call_items.append(fc_done_item) next_output_index += 1 continue + if chunk.is_thinking: + # Start reasoning block on first thinking token + if not reasoning_started: + reasoning_started = True + reasoning_output_index = next_output_index + next_output_index += 1 + + # response.output_item.added for reasoning + reasoning_item = ResponseReasoningItem( + id=reasoning_id, + summary=[], + status="in_progress", + ) + rs_added = ResponseOutputItemAddedEvent( + sequence_number=next(seq), + output_index=reasoning_output_index, + item=reasoning_item, + ) + yield _format_sse(rs_added) + + # response.reasoning_summary_part.added + part_added = ResponseReasoningSummaryPartAddedEvent( + sequence_number=next(seq), + item_id=reasoning_id, + output_index=reasoning_output_index, + summary_index=0, + part=ResponseReasoningSummaryText(text=""), + ) + yield _format_sse(part_added) + + accumulated_thinking += chunk.text + + # response.reasoning_summary_text.delta + rs_delta = ResponseReasoningSummaryTextDeltaEvent( + sequence_number=next(seq), + item_id=reasoning_id, + output_index=reasoning_output_index, + summary_index=0, + delta=chunk.text, + ) + yield _format_sse(rs_delta) + continue + + # Close reasoning block when transitioning to text + if reasoning_started and not message_started: + # response.reasoning_summary_text.done + rs_text_done = ResponseReasoningSummaryTextDoneEvent( + sequence_number=next(seq), + item_id=reasoning_id, + output_index=reasoning_output_index, + summary_index=0, + text=accumulated_thinking, + ) + yield _format_sse(rs_text_done) + + # response.reasoning_summary_part.done + rs_part_done = ResponseReasoningSummaryPartDoneEvent( + sequence_number=next(seq), + item_id=reasoning_id, + output_index=reasoning_output_index, + summary_index=0, + part=ResponseReasoningSummaryText(text=accumulated_thinking), + ) + yield _format_sse(rs_part_done) + + # response.output_item.done for reasoning + rs_item_done = ResponseOutputItemDoneEvent( + sequence_number=next(seq), + output_index=reasoning_output_index, + item=ResponseReasoningItem( + id=reasoning_id, + summary=[ResponseReasoningSummaryText(text=accumulated_thinking)], + ), + ) + yield _format_sse(rs_item_done) + + # Start message block on first text token + if not message_started: + message_started = True + message_output_index = next_output_index + next_output_index += 1 + + initial_item = ResponseMessageItem( + id=item_id, + content=[ResponseOutputText(text="")], + status="in_progress", + ) + item_added = ResponseOutputItemAddedEvent( + sequence_number=next(seq), + output_index=message_output_index, + item=initial_item, + ) + yield _format_sse(item_added) + + initial_part = ResponseOutputText(text="") + part_added = ResponseContentPartAddedEvent( + sequence_number=next(seq), + item_id=item_id, + output_index=message_output_index, + content_index=0, + part=initial_part, + ) + yield _format_sse(part_added) + accumulated_text += chunk.text # response.output_text.delta delta_event = ResponseTextDeltaEvent( sequence_number=next(seq), item_id=item_id, - output_index=0, + output_index=message_output_index, content_index=0, delta=chunk.text, ) - yield f"event: response.output_text.delta\ndata: {delta_event.model_dump_json()}\n\n" + yield _format_sse(delta_event) + + # Close reasoning block if it was never followed by text + if reasoning_started and not message_started: + rs_text_done = ResponseReasoningSummaryTextDoneEvent( + sequence_number=next(seq), + item_id=reasoning_id, + output_index=reasoning_output_index, + summary_index=0, + text=accumulated_thinking, + ) + yield _format_sse(rs_text_done) + + rs_part_done = ResponseReasoningSummaryPartDoneEvent( + sequence_number=next(seq), + item_id=reasoning_id, + output_index=reasoning_output_index, + summary_index=0, + part=ResponseReasoningSummaryText(text=accumulated_thinking), + ) + yield _format_sse(rs_part_done) + + rs_item_done = ResponseOutputItemDoneEvent( + sequence_number=next(seq), + output_index=reasoning_output_index, + item=ResponseReasoningItem( + id=reasoning_id, + summary=[ResponseReasoningSummaryText(text=accumulated_thinking)], + ), + ) + yield _format_sse(rs_item_done) + + # If no message block was started, create one now (empty text) + if not message_started: + message_output_index = next_output_index + next_output_index += 1 + + initial_item = ResponseMessageItem( + id=item_id, + content=[ResponseOutputText(text="")], + status="in_progress", + ) + item_added = ResponseOutputItemAddedEvent( + sequence_number=next(seq), + output_index=message_output_index, + item=initial_item, + ) + yield _format_sse(item_added) + + initial_part = ResponseOutputText(text="") + part_added_evt = ResponseContentPartAddedEvent( + sequence_number=next(seq), + item_id=item_id, + output_index=message_output_index, + content_index=0, + part=initial_part, + ) + yield _format_sse(part_added_evt) # response.output_text.done text_done = ResponseTextDoneEvent( sequence_number=next(seq), item_id=item_id, - output_index=0, + output_index=message_output_index, content_index=0, text=accumulated_text, ) - yield f"event: response.output_text.done\ndata: {text_done.model_dump_json()}\n\n" + yield _format_sse(text_done) # response.content_part.done final_part = ResponseOutputText(text=accumulated_text) part_done = ResponseContentPartDoneEvent( sequence_number=next(seq), item_id=item_id, - output_index=0, + output_index=message_output_index, content_index=0, part=final_part, ) - yield f"event: response.content_part.done\ndata: {part_done.model_dump_json()}\n\n" + yield _format_sse(part_done) # response.output_item.done final_message_item = ResponseMessageItem( @@ -346,9 +534,11 @@ async def generate_responses_stream( status="completed", ) item_done = ResponseOutputItemDoneEvent( - sequence_number=next(seq), output_index=0, item=final_message_item + sequence_number=next(seq), + output_index=message_output_index, + item=final_message_item, ) - yield f"event: response.output_item.done\ndata: {item_done.model_dump_json()}\n\n" + yield _format_sse(item_done) # Create usage from usage data if available usage = None @@ -360,7 +550,15 @@ async def generate_responses_stream( ) # response.completed - output: list[ResponseItem] = [final_message_item] + output: list[ResponseItem] = [] + if reasoning_started: + output.append( + ResponseReasoningItem( + id=reasoning_id, + summary=[ResponseReasoningSummaryText(text=accumulated_thinking)], + ) + ) + output.append(final_message_item) output.extend(function_call_items) final_response = ResponsesResponse( id=response_id, @@ -373,4 +571,4 @@ async def generate_responses_stream( completed_event = ResponseCompletedEvent( sequence_number=next(seq), response=final_response ) - yield f"event: response.completed\ndata: {completed_event.model_dump_json()}\n\n" + yield _format_sse(completed_event) diff --git a/src/exo/master/api.py b/src/exo/master/api.py index b8476334..e23ed1d0 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -32,6 +32,14 @@ from exo.master.adapters.claude import ( collect_claude_response, generate_claude_stream, ) +from exo.master.adapters.ollama import ( + collect_ollama_chat_response, + collect_ollama_generate_response, + generate_ollama_chat_stream, + generate_ollama_generate_stream, + ollama_generate_request_to_text_generation, + ollama_request_to_text_generation, +) from exo.master.adapters.responses import ( collect_responses_response, generate_responses_stream, @@ -85,6 +93,7 @@ from exo.shared.types.api import ( ImageGenerationTaskParams, ImageListItem, ImageListResponse, + ImageSize, ModelList, ModelListModel, PlaceInstanceParams, @@ -100,11 +109,13 @@ from exo.shared.types.api import ( TraceRankStats, TraceResponse, TraceStatsResponse, + normalize_image_size, ) from exo.shared.types.chunks import ( ErrorChunk, ImageChunk, InputImageChunk, + PrefillProgressChunk, TokenChunk, ToolCallChunk, ) @@ -138,11 +149,25 @@ from exo.shared.types.events import ( TracesMerged, ) from exo.shared.types.memory import Memory +from exo.shared.types.ollama_api import ( + OllamaChatRequest, + OllamaChatResponse, + OllamaGenerateRequest, + OllamaGenerateResponse, + OllamaModelDetails, + OllamaModelTag, + OllamaPsModel, + OllamaPsResponse, + OllamaShowRequest, + OllamaShowResponse, + OllamaTagsResponse, +) from exo.shared.types.openai_responses import ( ResponsesRequest, ResponsesResponse, ) from exo.shared.types.state import State +from exo.shared.types.worker.downloads import DownloadCompleted from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding from exo.utils.banner import print_startup_banner @@ -218,7 +243,8 @@ class API: ) self._text_generation_queues: dict[ - CommandId, Sender[TokenChunk | ErrorChunk | ToolCallChunk] + CommandId, + Sender[TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk], ] = {} self._image_generation_queues: dict[ CommandId, Sender[ImageChunk | ErrorChunk] @@ -295,6 +321,21 @@ class API: self.app.get("/images/{image_id}")(self.get_image) self.app.post("/v1/messages", response_model=None)(self.claude_messages) self.app.post("/v1/responses", response_model=None)(self.openai_responses) + + # Ollama API + self.app.head("/ollama/")(self.ollama_version) + self.app.head("/ollama/api/version")(self.ollama_version) + self.app.post("/ollama/api/chat", response_model=None)(self.ollama_chat) + self.app.post("/ollama/api/api/chat", response_model=None)(self.ollama_chat) + self.app.post("/ollama/api/v1/chat", response_model=None)(self.ollama_chat) + self.app.post("/ollama/api/generate", response_model=None)(self.ollama_generate) + self.app.get("/ollama/api/tags")(self.ollama_tags) + self.app.get("/ollama/api/api/tags")(self.ollama_tags) + self.app.get("/ollama/api/v1/tags")(self.ollama_tags) + self.app.post("/ollama/api/show")(self.ollama_show) + self.app.get("/ollama/api/ps")(self.ollama_ps) + self.app.get("/ollama/api/version")(self.ollama_version) + self.app.get("/state")(lambda: self.state) self.app.get("/events")(self.stream_events) self.app.post("/download/start")(self.start_download) @@ -524,19 +565,23 @@ class API: async def _token_chunk_stream( self, command_id: CommandId - ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]: + ) -> AsyncGenerator[ + TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk, None + ]: """Yield chunks for a given command until completion. This is the internal low-level stream used by all API adapters. """ try: self._text_generation_queues[command_id], recv = channel[ - ErrorChunk | ToolCallChunk | TokenChunk + TokenChunk | ErrorChunk | ToolCallChunk | PrefillProgressChunk ]() with recv as token_chunks: async for chunk in token_chunks: yield chunk + if isinstance(chunk, PrefillProgressChunk): + continue if chunk.finish_reason is not None: break @@ -563,6 +608,9 @@ class API: stats: GenerationStats | None = None async for chunk in self._token_chunk_stream(command_id): + if isinstance(chunk, PrefillProgressChunk): + continue + if chunk.finish_reason == "error": raise HTTPException( status_code=500, @@ -751,9 +799,11 @@ class API: When stream=True and partial_images > 0, returns a StreamingResponse with SSE-formatted events for partial and final images. """ - payload.model = await self._validate_image_model(ModelId(payload.model)) payload = payload.model_copy( - update={"advanced_params": _ensure_seed(payload.advanced_params)} + update={ + "model": await self._validate_image_model(ModelId(payload.model)), + "advanced_params": _ensure_seed(payload.advanced_params), + } ) command = ImageGeneration( @@ -1009,12 +1059,13 @@ class API: async def bench_image_generations( self, request: Request, payload: BenchImageGenerationTaskParams ) -> BenchImageGenerationResponse: - payload.model = await self._validate_image_model(ModelId(payload.model)) - - payload.stream = False - payload.partial_images = 0 payload = payload.model_copy( - update={"advanced_params": _ensure_seed(payload.advanced_params)} + update={ + "model": await self._validate_image_model(ModelId(payload.model)), + "stream": False, + "partial_images": 0, + "advanced_params": _ensure_seed(payload.advanced_params), + } ) command = ImageGeneration( @@ -1035,7 +1086,7 @@ class API: prompt: str, model: ModelId, n: int, - size: str, + size: ImageSize, response_format: Literal["url", "b64_json"], input_fidelity: Literal["low", "high"], stream: bool, @@ -1105,7 +1156,7 @@ class API: prompt: str = Form(...), model: str = Form(...), n: int = Form(1), - size: str = Form("1024x1024"), + size: str | None = Form(None), response_format: Literal["url", "b64_json"] = Form("b64_json"), input_fidelity: Literal["low", "high"] = Form("low"), stream: str = Form("false"), @@ -1131,7 +1182,7 @@ class API: prompt=prompt, model=ModelId(model), n=n, - size=size, + size=normalize_image_size(size), response_format=response_format, input_fidelity=input_fidelity, stream=stream_bool, @@ -1167,7 +1218,7 @@ class API: prompt: str = Form(...), model: str = Form(...), n: int = Form(1), - size: str = Form("1024x1024"), + size: str | None = Form(None), response_format: Literal["url", "b64_json"] = Form("b64_json"), input_fidelity: Literal["low", "high"] = Form("low"), quality: Literal["high", "medium", "low"] = Form("medium"), @@ -1187,7 +1238,7 @@ class API: prompt=prompt, model=ModelId(model), n=n, - size=size, + size=normalize_image_size(size), response_format=response_format, input_fidelity=input_fidelity, stream=False, @@ -1278,6 +1329,163 @@ class API: media_type="application/json", ) + async def _ollama_root(self) -> JSONResponse: + """Respond to HEAD / from Ollama CLI connectivity checks.""" + return JSONResponse(content="Ollama is running") + + async def ollama_chat( + self, request: Request + ) -> OllamaChatResponse | StreamingResponse: + """Ollama Chat API — accepts JSON regardless of Content-Type.""" + body = await request.body() + payload = OllamaChatRequest.model_validate_json(body) + task_params = ollama_request_to_text_generation(payload) + resolved_model = await self._resolve_and_validate_text_model( + ModelId(task_params.model) + ) + task_params = task_params.model_copy(update={"model": resolved_model}) + + command = TextGeneration(task_params=task_params) + await self._send(command) + + if payload.stream: + return StreamingResponse( + generate_ollama_chat_stream( + command.command_id, + self._token_chunk_stream(command.command_id), + ), + media_type="application/x-ndjson", + headers={ + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) + else: + return StreamingResponse( + collect_ollama_chat_response( + command.command_id, + self._token_chunk_stream(command.command_id), + ), + media_type="application/json", + ) + + async def ollama_generate( + self, request: Request + ) -> OllamaGenerateResponse | StreamingResponse: + """Ollama Generate API — accepts JSON regardless of Content-Type.""" + body = await request.body() + payload = OllamaGenerateRequest.model_validate_json(body) + task_params = ollama_generate_request_to_text_generation(payload) + resolved_model = await self._resolve_and_validate_text_model( + ModelId(task_params.model) + ) + task_params = task_params.model_copy(update={"model": resolved_model}) + + command = TextGeneration(task_params=task_params) + await self._send(command) + + if payload.stream: + return StreamingResponse( + generate_ollama_generate_stream( + command.command_id, + self._token_chunk_stream(command.command_id), + ), + media_type="application/x-ndjson", + headers={ + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + ) + else: + return StreamingResponse( + collect_ollama_generate_response( + command.command_id, + self._token_chunk_stream(command.command_id), + ), + media_type="application/json", + ) + + async def ollama_tags(self) -> OllamaTagsResponse: + """Returns list of models in Ollama tags format. We return the downloaded ones only.""" + + def none_if_empty(value: str) -> str | None: + return value or None + + downloaded_model_ids: set[str] = set() + for node_downloads in self.state.downloads.values(): + for dl in node_downloads: + if isinstance(dl, DownloadCompleted): + downloaded_model_ids.add(dl.shard_metadata.model_card.model_id) + + cards = [ + c for c in await get_model_cards() if c.model_id in downloaded_model_ids + ] + + now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + return OllamaTagsResponse( + models=[ + OllamaModelTag( + name=str(card.model_id), + model=str(card.model_id), + modified_at=now, + size=card.storage_size.in_bytes, + digest="sha256:000000000000", + details=OllamaModelDetails( + family=none_if_empty(card.family), + quantization_level=none_if_empty(card.quantization), + ), + ) + for card in cards + ] + ) + + async def ollama_show(self, request: Request) -> OllamaShowResponse: + """Returns model information in Ollama show format.""" + body = await request.body() + payload = OllamaShowRequest.model_validate_json(body) + model_name = payload.name or payload.model + if not model_name: + raise HTTPException(status_code=400, detail="name or model is required") + try: + card = await ModelCard.load(ModelId(model_name)) + except Exception as exc: + raise HTTPException( + status_code=404, detail=f"Model not found: {model_name}" + ) from exc + + return OllamaShowResponse( + modelfile=f"FROM {card.model_id}", + template="{{ .Prompt }}", + details=OllamaModelDetails( + family=card.family or None, + quantization_level=card.quantization or None, + ), + ) + + async def ollama_ps(self) -> OllamaPsResponse: + """Returns list of running models (active instances).""" + models: list[OllamaPsModel] = [] + seen: set[str] = set() + for instance in self.state.instances.values(): + model_id = str(instance.shard_assignments.model_id) + if model_id in seen: + continue + seen.add(model_id) + models.append( + OllamaPsModel( + name=model_id, + model=model_id, + size=0, + ) + ) + return OllamaPsResponse(models=models) + + async def ollama_version(self) -> dict[str, str]: + """Returns version information for Ollama API compatibility.""" + return {"version": "exo v1.0"} + def _calculate_total_available_memory(self) -> Memory: """Calculate total available memory across all nodes in bytes.""" total_available = Memory() @@ -1287,8 +1495,18 @@ class API: return total_available - async def get_models(self) -> ModelList: - """Returns list of available models.""" + async def get_models(self, status: str | None = Query(default=None)) -> ModelList: + """Returns list of available models, optionally filtered by being downloaded.""" + cards = await get_model_cards() + + if status == "downloaded": + downloaded_model_ids: set[str] = set() + for node_downloads in self.state.downloads.values(): + for dl in node_downloads: + if isinstance(dl, DownloadCompleted): + downloaded_model_ids.add(dl.shard_metadata.model_card.model_id) + cards = [c for c in cards if c.model_id in downloaded_model_ids] + return ModelList( data=[ ModelListModel( @@ -1297,7 +1515,7 @@ class API: name=card.model_id.short(), description="", tags=[], - storage_size_megabytes=int(card.storage_size.in_mb), + storage_size_megabytes=card.storage_size.in_mb, supports_tensor=card.supports_tensor, tasks=[task.value for task in card.tasks], is_custom=is_custom_card(card.model_id), @@ -1306,7 +1524,7 @@ class API: base_model=card.base_model, capabilities=card.capabilities, ) - for card in await get_model_cards() + for card in cards ] ) @@ -1429,7 +1647,6 @@ class API: await queue.send(event.chunk) except BrokenResourceError: self._text_generation_queues.pop(event.command_id, None) - if isinstance(event, TracesMerged): self._save_merged_trace(event) diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index cf31ca78..a0604055 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -141,15 +141,29 @@ def place_instance( if len(selected_cycle) == 1: command.instance_meta = InstanceMeta.MlxRing - # TODO: Single node instances match command.instance_meta: case InstanceMeta.MlxJaccl: + # TODO(evan): shard assignments should contain information about ranks, this is ugly + def get_device_rank(node_id: NodeId) -> int: + runner_id = shard_assignments.node_to_runner[node_id] + shard_metadata = shard_assignments.runner_to_shard.get(runner_id) + assert shard_metadata is not None + return shard_metadata.device_rank + + zero_node_ids = [ + node_id + for node_id in selected_cycle.node_ids + if get_device_rank(node_id) == 0 + ] + assert len(zero_node_ids) == 1 + coordinator_node_id = zero_node_ids[0] + mlx_jaccl_devices = get_mlx_jaccl_devices_matrix( [node_id for node_id in selected_cycle], cycle_digraph, ) mlx_jaccl_coordinators = get_mlx_jaccl_coordinators( - coordinator=selected_cycle.node_ids[0], + coordinator=coordinator_node_id, coordinator_port=random_ephemeral_port(), cycle_digraph=cycle_digraph, node_network=node_network, diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index b20a39cc..eb78c8fa 100644 --- a/src/exo/master/placement_utils.py +++ b/src/exo/master/placement_utils.py @@ -102,22 +102,21 @@ def _allocate_and_validate_layers( layer_allocations = allocate_layers_proportionally( total_layers=model_card.n_layers, memory_fractions=[ - node_memory[node_id].ram_available.in_bytes / total_memory.in_bytes - for node_id in node_ids + node_memory[node_id].ram_available / total_memory for node_id in node_ids ], ) - total_storage_bytes = model_card.storage_size.in_bytes + total_storage = model_card.storage_size total_layers = model_card.n_layers for i, node_id in enumerate(node_ids): node_layers = layer_allocations[i] - required_memory = (total_storage_bytes * node_layers) // total_layers - available_memory = node_memory[node_id].ram_available.in_bytes + required_memory = (total_storage * node_layers) // total_layers + available_memory = node_memory[node_id].ram_available if required_memory > available_memory: raise ValueError( f"Node {i} ({node_id}) has insufficient memory: " - f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, " - f"but only has {available_memory / (1024**3):.2f} GB available" + f"requires {required_memory.in_gb:.2f} GB for {node_layers} layers, " + f"but only has {available_memory.in_gb:.2f} GB available" ) return layer_allocations @@ -342,6 +341,7 @@ def _find_ip_prioritised( other_node_id: NodeId, cycle_digraph: Topology, node_network: Mapping[NodeId, NodeNetworkInfo], + ring: bool, ) -> str | None: """Find an IP address between nodes with prioritization. @@ -354,13 +354,27 @@ def _find_ip_prioritised( ip_to_type = { iface.ip_address: iface.interface_type for iface in other_network.interfaces } - priority = { - "ethernet": 0, - "wifi": 1, - "unknown": 2, - "maybe_ethernet": 3, - "thunderbolt": 4, - } + + # Ring should prioritise fastest connection. As a best-effort, we prioritise TB. + # TODO: Profile and get actual connection speeds. + if ring: + priority = { + "thunderbolt": 0, + "maybe_ethernet": 1, + "ethernet": 2, + "wifi": 3, + "unknown": 4, + } + + # RDMA prefers ethernet coordinator + else: + priority = { + "ethernet": 0, + "wifi": 1, + "unknown": 2, + "maybe_ethernet": 3, + "thunderbolt": 4, + } return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2)) @@ -400,7 +414,7 @@ def get_mlx_ring_hosts_by_node( continue connection_ip = _find_ip_prioritised( - node_id, other_node_id, cycle_digraph, node_network + node_id, other_node_id, cycle_digraph, node_network, ring=True ) if connection_ip is None: raise ValueError( @@ -431,7 +445,9 @@ def get_mlx_jaccl_coordinators( if n == coordinator: return "0.0.0.0" - ip = _find_ip_prioritised(n, coordinator, cycle_digraph, node_network) + ip = _find_ip_prioritised( + n, coordinator, cycle_digraph, node_network, ring=False + ) if ip is not None: return ip diff --git a/src/exo/master/tests/test_claude_tool_use.py b/src/exo/master/tests/test_claude_tool_use.py index 77ce3492..3e814dce 100644 --- a/src/exo/master/tests/test_claude_tool_use.py +++ b/src/exo/master/tests/test_claude_tool_use.py @@ -261,7 +261,7 @@ class TestGenerateClaudeStreamToolUse: parsed = _parse_sse_events(events) - # Two tool block starts (at indices 1 and 2) + # Two tool block starts (at indices 0 and 1 — no text block when only tools) tool_starts = [ e for e in parsed @@ -270,12 +270,11 @@ class TestGenerateClaudeStreamToolUse: == "tool_use" ] assert len(tool_starts) == 2 - assert tool_starts[0]["index"] == 1 - assert tool_starts[1]["index"] == 2 + assert tool_starts[0]["index"] == 0 + assert tool_starts[1]["index"] == 1 - # Two tool block stops (at indices 1 and 2), plus text block stop at 0 + # Two tool block stops (at indices 0 and 1) block_stops = [e for e in parsed if e.get("type") == "content_block_stop"] stop_indices = [e["index"] for e in block_stops] assert 0 in stop_indices assert 1 in stop_indices - assert 2 in stop_indices diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py index 9452108f..fcf71ee4 100644 --- a/src/exo/master/tests/test_master.py +++ b/src/exo/master/tests/test_master.py @@ -42,7 +42,7 @@ from exo.utils.channels import channel @pytest.mark.asyncio async def test_master(): keypair = get_node_id_keypair() - node_id = NodeId(keypair.to_peer_id().to_base58()) + node_id = NodeId(keypair.to_node_id()) session_id = SessionId(master_node_id=node_id, election_clock=0) ge_sender, global_event_receiver = channel[ForwarderEvent]() @@ -75,7 +75,7 @@ async def test_master(): async with anyio.create_task_group() as tg: tg.start_soon(master.run) - sender_node_id = NodeId(f"{keypair.to_peer_id().to_base58()}_sender") + sender_node_id = NodeId(f"{keypair.to_node_id()}_sender") # inject a NodeGatheredInfo event logger.info("inject a NodeGatheredInfo event") await local_event_sender.send( diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py index ad5638e7..cad495ea 100644 --- a/src/exo/master/tests/test_placement.py +++ b/src/exo/master/tests/test_placement.py @@ -80,8 +80,8 @@ def test_get_instance_placements_create_instance( ): # arrange model_card.n_layers = total_layers - model_card.storage_size.in_bytes = sum( - available_memory + model_card.storage_size = Memory.from_bytes( + sum(available_memory) ) # make it exactly fit across all nodes topology = Topology() @@ -349,7 +349,7 @@ def test_tensor_rdma_backend_connectivity_matrix( # arrange topology = Topology() model_card.n_layers = 12 - model_card.storage_size.in_bytes = 1500 + model_card.storage_size = Memory.from_bytes(1500) node_a = NodeId() node_b = NodeId() diff --git a/src/exo/routing/connection_message.py b/src/exo/routing/connection_message.py index 665483ac..0eb68f37 100644 --- a/src/exo/routing/connection_message.py +++ b/src/exo/routing/connection_message.py @@ -30,7 +30,7 @@ class ConnectionMessage(CamelCaseModel): @classmethod def from_update(cls, update: ConnectionUpdate) -> "ConnectionMessage": return cls( - node_id=NodeId(update.peer_id.to_base58()), + node_id=NodeId(update.peer_id), connection_type=ConnectionMessageType.from_update_type(update.update_type), remote_ipv4=update.remote_ipv4, remote_tcp_port=update.remote_tcp_port, diff --git a/src/exo/routing/router.py b/src/exo/routing/router.py index 309f7d0b..d71275b7 100644 --- a/src/exo/routing/router.py +++ b/src/exo/routing/router.py @@ -221,7 +221,7 @@ def get_node_id_keypair( Obtain the :class:`PeerId` by from it. """ # TODO(evan): bring back node id persistence once we figure out how to deal with duplicates - return Keypair.generate_ed25519() + return Keypair.generate() def lock_path(path: str | bytes | PathLike[str] | PathLike[bytes]) -> Path: return Path(str(path) + ".lock") @@ -235,12 +235,12 @@ def get_node_id_keypair( protobuf_encoded = f.read() try: # if decoded successfully, save & return - return Keypair.from_protobuf_encoding(protobuf_encoded) + return Keypair.from_bytes(protobuf_encoded) except ValueError as e: # on runtime error, assume corrupt file logger.warning(f"Encountered error when trying to get keypair: {e}") # if no valid credentials, create new ones and persist with open(path, "w+b") as f: keypair = Keypair.generate_ed25519() - f.write(keypair.to_protobuf_encoding()) + f.write(keypair.to_bytes()) return keypair diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index fd0df265..94869dfe 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -218,11 +218,6 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: key: value for key, value in state.downloads.items() if key != event.node_id } # Clean up all granular node mappings - node_identities = { - key: value - for key, value in state.node_identities.items() - if key != event.node_id - } node_memory = { key: value for key, value in state.node_memory.items() if key != event.node_id } @@ -263,7 +258,6 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State: "downloads": downloads, "topology": topology, "last_seen": last_seen, - "node_identities": node_identities, "node_memory": node_memory, "node_disk": node_disk, "node_system": node_system, diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index f9079854..d02271e9 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -44,7 +44,8 @@ async def _refresh_card_cache(): async for toml_file in path.rglob("*.toml"): try: card = await ModelCard.load_from_path(toml_file) - _card_cache[card.model_id] = card + if card.model_id not in _card_cache: + _card_cache[card.model_id] = card except (ValidationError, TOMLKitError): pass @@ -182,6 +183,7 @@ class ConfigData(BaseModel): def supports_tensor(self) -> bool: return self.architectures in [ ["Glm4MoeLiteForCausalLM"], + ["GlmMoeDsaForCausalLM"], ["DeepseekV32ForCausalLM"], ["DeepseekV3ForCausalLM"], ["Qwen3NextForCausalLM"], diff --git a/src/exo/shared/tests/test_apply/test_apply_node_download.py b/src/exo/shared/tests/test_apply/test_apply_node_download.py index 6b1cb8cc..f9df6e07 100644 --- a/src/exo/shared/tests/test_apply/test_apply_node_download.py +++ b/src/exo/shared/tests/test_apply/test_apply_node_download.py @@ -14,7 +14,7 @@ def test_apply_node_download_progress(): event = DownloadCompleted( node_id=NodeId("node-1"), shard_metadata=shard1, - total_bytes=Memory(), + total=Memory(), ) new_state = apply_node_download_progress( @@ -30,12 +30,12 @@ def test_apply_two_node_download_progress(): event1 = DownloadCompleted( node_id=NodeId("node-1"), shard_metadata=shard1, - total_bytes=Memory(), + total=Memory(), ) event2 = DownloadCompleted( node_id=NodeId("node-1"), shard_metadata=shard2, - total_bytes=Memory(), + total=Memory(), ) state = State(downloads={NodeId("node-1"): [event1]}) diff --git a/src/exo/shared/tests/test_node_id_persistence.py b/src/exo/shared/tests/test_node_id_persistence.py index c067bd3c..ce9e56f5 100644 --- a/src/exo/shared/tests/test_node_id_persistence.py +++ b/src/exo/shared/tests/test_node_id_persistence.py @@ -23,7 +23,7 @@ def _get_keypair_concurrent_subprocess_task( sem.release() # wait to be told to begin simultaneous read ev.wait() - queue.put(get_node_id_keypair().to_protobuf_encoding()) + queue.put(get_node_id_keypair().to_bytes()) def _get_keypair_concurrent(num_procs: int) -> bytes: diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 2756f0d4..23ca9b7b 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -1,9 +1,9 @@ import time from collections.abc import Generator -from typing import Annotated, Any, Literal +from typing import Annotated, Any, Literal, get_args from uuid import uuid4 -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator from exo.shared.models.model_cards import ModelCard, ModelId from exo.shared.types.common import CommandId, NodeId @@ -77,7 +77,7 @@ class ChatCompletionMessage(BaseModel): content: ( str | ChatCompletionMessageText | list[ChatCompletionMessageText] | None ) = None - thinking: str | None = None # Added for GPT-OSS harmony format support + reasoning_content: str | None = None name: str | None = None tool_calls: list[ToolCall] | None = None tool_call_id: str | None = None @@ -262,6 +262,27 @@ class DeleteInstanceResponse(BaseModel): instance_id: InstanceId +ImageSize = Literal[ + "auto", + "512x512", + "768x768", + "1024x768", + "768x1024", + "1024x1024", + "1024x1536", + "1536x1024", +] + + +def normalize_image_size(v: object) -> ImageSize: + """Shared validator for ImageSize fields: maps None → "auto" and rejects invalid values.""" + if v is None: + return "auto" + if v not in get_args(ImageSize): + raise ValueError(f"Invalid size: {v!r}. Must be one of {get_args(ImageSize)}") + return v # pyright: ignore[reportReturnType] + + class AdvancedImageParams(BaseModel): seed: Annotated[int, Field(ge=0)] | None = None num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None @@ -281,7 +302,7 @@ class ImageGenerationTaskParams(BaseModel): partial_images: int | None = 0 quality: Literal["high", "medium", "low"] | None = "medium" response_format: Literal["url", "b64_json"] | None = "b64_json" - size: str | None = "1024x1024" + size: ImageSize = "auto" stream: bool | None = False style: str | None = "vivid" user: str | None = None @@ -289,6 +310,11 @@ class ImageGenerationTaskParams(BaseModel): # Internal flag for benchmark mode - set by API, preserved through serialization bench: bool = False + @field_validator("size", mode="before") + @classmethod + def normalize_size(cls, v: object) -> ImageSize: + return normalize_image_size(v) + class BenchImageGenerationTaskParams(ImageGenerationTaskParams): bench: bool = True @@ -305,13 +331,18 @@ class ImageEditsTaskParams(BaseModel): quality: Literal["high", "medium", "low"] | None = "medium" output_format: Literal["png", "jpeg", "webp"] = "png" response_format: Literal["url", "b64_json"] | None = "b64_json" - size: str | None = "1024x1024" + size: ImageSize = "auto" image_strength: float | None = 0.7 stream: bool = False partial_images: int | None = 0 advanced_params: AdvancedImageParams | None = None bench: bool = False + @field_validator("size", mode="before") + @classmethod + def normalize_size(cls, v: object) -> ImageSize: + return normalize_image_size(v) + def __repr_args__(self) -> Generator[tuple[str, Any], None, None]: for name, value in super().__repr_args__(): # pyright: ignore[reportAny] if name == "image_data": diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py index 5fe9eb1c..64481bab 100644 --- a/src/exo/shared/types/chunks.py +++ b/src/exo/shared/types/chunks.py @@ -27,6 +27,7 @@ class TokenChunk(BaseChunk): stats: GenerationStats | None = None logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None + is_thinking: bool = False class ErrorChunk(BaseChunk): @@ -76,4 +77,13 @@ class InputImageChunk(BaseChunk): yield name, value -GenerationChunk = TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk +class PrefillProgressChunk(BaseChunk): + """Data class for prefill progress events during streaming.""" + + processed_tokens: int + total_tokens: int + + +GenerationChunk = ( + TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk | PrefillProgressChunk +) diff --git a/src/exo/shared/types/claude_api.py b/src/exo/shared/types/claude_api.py index dcf8c515..8eb18497 100644 --- a/src/exo/shared/types/claude_api.py +++ b/src/exo/shared/types/claude_api.py @@ -47,6 +47,14 @@ class ClaudeImageBlock(BaseModel, frozen=True): source: ClaudeImageSource +class ClaudeThinkingBlock(BaseModel, frozen=True): + """Thinking content block in Claude Messages API.""" + + type: Literal["thinking"] = "thinking" + thinking: str + signature: str | None = None + + class ClaudeToolUseBlock(BaseModel, frozen=True): """Tool use content block in Claude Messages API.""" @@ -66,11 +74,17 @@ class ClaudeToolResultBlock(BaseModel, frozen=True): cache_control: dict[str, str] | None = None -ClaudeContentBlock = ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock +ClaudeContentBlock = ( + ClaudeTextBlock | ClaudeImageBlock | ClaudeThinkingBlock | ClaudeToolUseBlock +) # Input content blocks can also include tool_result (sent by user after tool_use) ClaudeInputContentBlock = ( - ClaudeTextBlock | ClaudeImageBlock | ClaudeToolUseBlock | ClaudeToolResultBlock + ClaudeTextBlock + | ClaudeImageBlock + | ClaudeThinkingBlock + | ClaudeToolUseBlock + | ClaudeToolResultBlock ) @@ -82,6 +96,11 @@ class ClaudeMessage(BaseModel, frozen=True): content: str | list[ClaudeInputContentBlock] +class ClaudeThinkingConfig(BaseModel, frozen=True): + type: Literal["enabled", "disabled", "adaptive"] + budget_tokens: int | None = None + + class ClaudeMessagesRequest(BaseModel): """Request body for Claude Messages API.""" @@ -96,6 +115,7 @@ class ClaudeMessagesRequest(BaseModel): top_k: int | None = None tools: list[ClaudeToolDefinition] | None = None metadata: dict[str, str] | None = None + thinking: ClaudeThinkingConfig | None = None # Response types @@ -145,7 +165,7 @@ class ClaudeContentBlockStartEvent(BaseModel, frozen=True): type: Literal["content_block_start"] = "content_block_start" index: int - content_block: ClaudeTextBlock | ClaudeToolUseBlock + content_block: ClaudeTextBlock | ClaudeThinkingBlock | ClaudeToolUseBlock class ClaudeTextDelta(BaseModel, frozen=True): @@ -155,6 +175,13 @@ class ClaudeTextDelta(BaseModel, frozen=True): text: str +class ClaudeThinkingDelta(BaseModel, frozen=True): + """Delta for thinking content block.""" + + type: Literal["thinking_delta"] = "thinking_delta" + thinking: str + + class ClaudeInputJsonDelta(BaseModel, frozen=True): """Delta for tool use input JSON content block.""" @@ -167,7 +194,7 @@ class ClaudeContentBlockDeltaEvent(BaseModel, frozen=True): type: Literal["content_block_delta"] = "content_block_delta" index: int - delta: ClaudeTextDelta | ClaudeInputJsonDelta + delta: ClaudeTextDelta | ClaudeThinkingDelta | ClaudeInputJsonDelta class ClaudeContentBlockStopEvent(BaseModel, frozen=True): diff --git a/src/exo/shared/types/memory.py b/src/exo/shared/types/memory.py index b97fb345..2684d9b9 100644 --- a/src/exo/shared/types/memory.py +++ b/src/exo/shared/types/memory.py @@ -1,10 +1,10 @@ from math import ceil -from typing import Self +from typing import Self, overload -from exo.utils.pydantic_ext import CamelCaseModel +from exo.utils.pydantic_ext import FrozenModel -class Memory(CamelCaseModel): +class Memory(FrozenModel): in_bytes: int = 0 @classmethod @@ -33,12 +33,22 @@ class Memory(CamelCaseModel): return cls(in_bytes=round(val * 1024)) @property - def in_mb(self) -> float: - """The approximate megabytes this memory represents. Setting this property rounds to the nearest byte.""" - return self.in_bytes / (1024**2) + def in_mb(self) -> int: + """The approximate megabytes this memory represents, rounded to nearest MB. Setting this property rounds to the nearest byte.""" + return round(self.in_bytes / (1024**2)) @in_mb.setter - def in_mb(self, val: float): + def in_mb(self, val: int): + """Set the megabytes for this memory.""" + self.in_bytes = val * (1024**2) + + @property + def in_float_mb(self) -> float: + """The megabytes this memory represents as a float. Setting this property rounds to the nearest byte.""" + return self.in_bytes / (1024**2) + + @in_float_mb.setter + def in_float_mb(self, val: float): """Set the megabytes for this memory, rounded to the nearest byte.""" self.in_bytes = round(val * (1024**2)) @@ -57,17 +67,85 @@ class Memory(CamelCaseModel): """The approximate gigabytes this memory represents.""" return self.in_bytes / (1024**3) - def __add__(self, other: "Memory") -> "Memory": - return Memory.from_bytes(self.in_bytes + other.in_bytes) + def __add__(self, other: object) -> "Memory": + if isinstance(other, Memory): + return Memory.from_bytes(self.in_bytes + other.in_bytes) + return NotImplemented - def __lt__(self, other: Self) -> bool: - return self.in_bytes < other.in_bytes + def __radd__(self, other: object) -> "Memory": + if other == 0: + return self + return NotImplemented - def __le__(self, other: Self) -> bool: - return self.in_bytes <= other.in_bytes + def __sub__(self, other: object) -> "Memory": + if isinstance(other, Memory): + return Memory.from_bytes(self.in_bytes - other.in_bytes) + return NotImplemented - def __gt__(self, other: Self) -> bool: - return self.in_bytes > other.in_bytes + def __mul__(self, other: int | float): + return Memory.from_bytes(round(self.in_bytes * other)) - def __ge__(self, other: Self) -> bool: - return self.in_bytes >= other.in_bytes + def __rmul__(self, other: int | float): + return self * other + + @overload + def __truediv__(self, other: "Memory") -> float: ... + @overload + def __truediv__(self, other: int) -> "Memory": ... + @overload + def __truediv__(self, other: float) -> "Memory": ... + def __truediv__(self, other: object) -> "Memory | float": + if isinstance(other, Memory): + return self.in_bytes / other.in_bytes + if isinstance(other, (int, float)): + return Memory.from_bytes(round(self.in_bytes / other)) + return NotImplemented + + def __floordiv__(self, other: object) -> "Memory": + if isinstance(other, (int, float)): + return Memory.from_bytes(int(self.in_bytes // other)) + return NotImplemented + + def __lt__(self, other: object) -> bool: + if isinstance(other, Memory): + return self.in_bytes < other.in_bytes + return NotImplemented + + def __le__(self, other: object) -> bool: + if isinstance(other, Memory): + return self.in_bytes <= other.in_bytes + return NotImplemented + + def __gt__(self, other: object) -> bool: + if isinstance(other, Memory): + return self.in_bytes > other.in_bytes + return NotImplemented + + def __ge__(self, other: object) -> bool: + if isinstance(other, Memory): + return self.in_bytes >= other.in_bytes + return NotImplemented + + def __eq__(self, other: object) -> bool: + if isinstance(other, Memory): + return self.in_bytes == other.in_bytes + return NotImplemented + + def __repr__(self) -> str: + return f"Memory.from_bytes({self.in_bytes})" + + def __str__(self) -> str: + if self.in_gb > 2: + val = self.in_gb + unit = "GiB" + elif self.in_mb > 2: + val = self.in_mb + unit = "MiB" + elif self.in_kb > 3: + val = self.in_kb + unit = "KiB" + else: + val = self.in_bytes + unit = "B" + + return f"{val:.2f} {unit}".rstrip("0").rstrip(".") + f" {unit}" diff --git a/src/exo/shared/types/mlx.py b/src/exo/shared/types/mlx.py index 99fee87a..3cb03195 100644 --- a/src/exo/shared/types/mlx.py +++ b/src/exo/shared/types/mlx.py @@ -4,10 +4,13 @@ from collections.abc import Sequence from mlx_lm.models.cache import ( ArraysCache, + CacheList, KVCache, QuantizedKVCache, RotatingKVCache, ) # This list contains one cache entry per transformer layer -KVCacheType = Sequence[KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache] +KVCacheType = Sequence[ + KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList +] diff --git a/src/exo/shared/types/ollama_api.py b/src/exo/shared/types/ollama_api.py new file mode 100644 index 00000000..9f881fbe --- /dev/null +++ b/src/exo/shared/types/ollama_api.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import time +from typing import Any, Literal + +from pydantic import BaseModel, Field + +from exo.shared.models.model_cards import ModelId + +# https://github.com/ollama/ollama/blob/main/docs/api.md + +OllamaRole = Literal["system", "user", "assistant", "tool"] +OllamaDoneReason = Literal["stop", "length", "tool_call", "error"] + + +class OllamaToolFunction(BaseModel, frozen=True): + name: str + arguments: dict[str, Any] | str + index: int | None = None + + +class OllamaToolCall(BaseModel, frozen=True): + id: str | None = None + type: Literal["function"] | None = None + function: OllamaToolFunction + + +class OllamaMessage(BaseModel, frozen=True): + role: OllamaRole + content: str | None = None + thinking: str | None = None + tool_calls: list[OllamaToolCall] | None = None + name: str | None = None + tool_name: str | None = None + images: list[str] | None = None + + +class OllamaOptions(BaseModel, frozen=True): + num_predict: int | None = None + temperature: float | None = None + top_p: float | None = None + top_k: int | None = None + stop: str | list[str] | None = None + seed: int | None = None + + +class OllamaChatRequest(BaseModel, frozen=True): + model: ModelId + messages: list[OllamaMessage] + stream: bool = True + options: OllamaOptions | None = None + tools: list[dict[str, Any]] | None = None + format: Literal["json"] | dict[str, Any] | None = None + keep_alive: str | int | None = None + think: bool | None = None + + +class OllamaGenerateRequest(BaseModel, frozen=True): + model: ModelId + prompt: str = "" + system: str | None = None + stream: bool = True + options: OllamaOptions | None = None + format: Literal["json"] | dict[str, Any] | None = None + keep_alive: str | int | None = None + think: bool | None = None + raw: bool = False + + +class OllamaGenerateResponse(BaseModel, frozen=True, strict=True): + model: str + created_at: str = Field( + default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + ) + response: str + thinking: str | None = None + done: bool + done_reason: OllamaDoneReason | None = None + total_duration: int | None = None + load_duration: int | None = None + prompt_eval_count: int | None = None + prompt_eval_duration: int | None = None + eval_count: int | None = None + eval_duration: int | None = None + + +class OllamaShowRequest(BaseModel, frozen=True): + name: str | None = None + model: str | None = None + verbose: bool | None = None + + +class OllamaChatResponse(BaseModel, frozen=True, strict=True): + model: str + created_at: str = Field( + default_factory=lambda: time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + ) + message: OllamaMessage + done: bool + done_reason: OllamaDoneReason | None = None + total_duration: int | None = None + load_duration: int | None = None + prompt_eval_count: int | None = None + prompt_eval_duration: int | None = None + eval_count: int | None = None + eval_duration: int | None = None + + +class OllamaModelDetails(BaseModel, frozen=True, strict=True): + format: str | None = None + family: str | None = None + parameter_size: str | None = None + quantization_level: str | None = None + + +class OllamaModelTag(BaseModel, frozen=True, strict=True): + name: str + model: str | None = None + modified_at: str | None = None + size: int | None = None + digest: str | None = None + details: OllamaModelDetails | None = None + + +class OllamaTagsResponse(BaseModel, frozen=True, strict=True): + models: list[OllamaModelTag] + + +class OllamaShowResponse(BaseModel, frozen=True, strict=True): + modelfile: str | None = None + parameters: str | None = None + template: str | None = None + details: OllamaModelDetails | None = None + model_info: dict[str, Any] | None = None + + +class OllamaPsModel(BaseModel, frozen=True, strict=True): + name: str + model: str + size: int + digest: str | None = None + details: OllamaModelDetails | None = None + expires_at: str | None = None + size_vram: int | None = None + + +class OllamaPsResponse(BaseModel, frozen=True, strict=True): + models: list[OllamaPsModel] diff --git a/src/exo/shared/types/openai_responses.py b/src/exo/shared/types/openai_responses.py index e29fb253..331f68e4 100644 --- a/src/exo/shared/types/openai_responses.py +++ b/src/exo/shared/types/openai_responses.py @@ -145,7 +145,23 @@ class ResponseFunctionCallItem(BaseModel, frozen=True): status: ResponseStatus = "completed" -ResponseItem = ResponseMessageItem | ResponseFunctionCallItem +class ResponseReasoningSummaryText(BaseModel, frozen=True): + """Summary text part in a reasoning output item.""" + + type: Literal["summary_text"] = "summary_text" + text: str + + +class ResponseReasoningItem(BaseModel, frozen=True): + """Reasoning output item in response output array.""" + + type: Literal["reasoning"] = "reasoning" + id: str + summary: list[ResponseReasoningSummaryText] = Field(default_factory=list) + status: ResponseStatus = "completed" + + +ResponseItem = ResponseMessageItem | ResponseFunctionCallItem | ResponseReasoningItem class ResponseUsage(BaseModel, frozen=True): @@ -273,6 +289,58 @@ class ResponseFunctionCallArgumentsDoneEvent(BaseModel, frozen=True): arguments: str +class ResponseReasoningSummaryPartAddedEvent(BaseModel, frozen=True): + """Event sent when a reasoning summary part is added.""" + + type: Literal["response.reasoning_summary_part.added"] = ( + "response.reasoning_summary_part.added" + ) + sequence_number: int + item_id: str + output_index: int + summary_index: int + part: ResponseReasoningSummaryText + + +class ResponseReasoningSummaryTextDeltaEvent(BaseModel, frozen=True): + """Event sent for reasoning summary text delta during streaming.""" + + type: Literal["response.reasoning_summary_text.delta"] = ( + "response.reasoning_summary_text.delta" + ) + sequence_number: int + item_id: str + output_index: int + summary_index: int + delta: str + + +class ResponseReasoningSummaryTextDoneEvent(BaseModel, frozen=True): + """Event sent when reasoning summary text is done.""" + + type: Literal["response.reasoning_summary_text.done"] = ( + "response.reasoning_summary_text.done" + ) + sequence_number: int + item_id: str + output_index: int + summary_index: int + text: str + + +class ResponseReasoningSummaryPartDoneEvent(BaseModel, frozen=True): + """Event sent when a reasoning summary part is done.""" + + type: Literal["response.reasoning_summary_part.done"] = ( + "response.reasoning_summary_part.done" + ) + sequence_number: int + item_id: str + output_index: int + summary_index: int + part: ResponseReasoningSummaryText + + class ResponseCompletedEvent(BaseModel, frozen=True): """Event sent when response is completed.""" @@ -292,5 +360,9 @@ ResponsesStreamEvent = ( | ResponseOutputItemDoneEvent | ResponseFunctionCallArgumentsDeltaEvent | ResponseFunctionCallArgumentsDoneEvent + | ResponseReasoningSummaryPartAddedEvent + | ResponseReasoningSummaryTextDeltaEvent + | ResponseReasoningSummaryTextDoneEvent + | ResponseReasoningSummaryPartDoneEvent | ResponseCompletedEvent ) diff --git a/src/exo/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py index 77162e7c..45762846 100644 --- a/src/exo/shared/types/worker/downloads.py +++ b/src/exo/shared/types/worker/downloads.py @@ -10,9 +10,9 @@ from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel class DownloadProgressData(CamelCaseModel): - total_bytes: Memory - downloaded_bytes: Memory - downloaded_bytes_this_session: Memory + total: Memory + downloaded: Memory + downloaded_this_session: Memory completed_files: int total_files: int @@ -26,6 +26,7 @@ class DownloadProgressData(CamelCaseModel): class BaseDownloadProgress(TaggedModel): node_id: NodeId shard_metadata: ShardMetadata + model_directory: str = "" class DownloadPending(BaseDownloadProgress): @@ -33,7 +34,7 @@ class DownloadPending(BaseDownloadProgress): class DownloadCompleted(BaseDownloadProgress): - total_bytes: Memory + total: Memory class DownloadFailed(BaseDownloadProgress): @@ -85,9 +86,9 @@ class RepoDownloadProgress(BaseModel): shard: ShardMetadata completed_files: int total_files: int - downloaded_bytes: Memory - downloaded_bytes_this_session: Memory - total_bytes: Memory + downloaded: Memory + downloaded_this_session: Memory + total: Memory overall_speed: float overall_eta: timedelta status: Literal["not_started", "in_progress", "complete"] diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py index 5f18bf5a..19118e15 100644 --- a/src/exo/shared/types/worker/runner_response.py +++ b/src/exo/shared/types/worker/runner_response.py @@ -28,6 +28,7 @@ class GenerationResponse(BaseRunnerResponse): finish_reason: FinishReason | None = None stats: GenerationStats | None = None usage: Usage | None + is_thinking: bool = False class ImageGenerationResponse(BaseRunnerResponse): @@ -67,3 +68,8 @@ class ToolCallResponse(BaseRunnerResponse): class FinishedResponse(BaseRunnerResponse): pass + + +class PrefillProgressResponse(BaseRunnerResponse): + processed_tokens: int + total_tokens: int diff --git a/src/exo/utils/banner.py b/src/exo/utils/banner.py index 8158ff3e..ac2030e4 100644 --- a/src/exo/utils/banner.py +++ b/src/exo/utils/banner.py @@ -1,5 +1,6 @@ import logging import os +import sys import webbrowser from exo.shared.constants import EXO_CONFIG_HOME @@ -19,7 +20,6 @@ def _mark_first_run_done() -> None: def print_startup_banner(port: int) -> None: - """Print a prominent startup banner with API endpoint information.""" dashboard_url = f"http://localhost:{port}" first_run = _is_first_run() banner = f""" @@ -48,7 +48,7 @@ def print_startup_banner(port: int) -> None: """ - print(banner) + print(banner, file=sys.stderr) if first_run: # Skip browser open when running inside the native macOS app — @@ -59,4 +59,4 @@ def print_startup_banner(port: int) -> None: logger.info("First run detected — opening dashboard in browser") except Exception: logger.debug("Could not auto-open browser", exc_info=True) - _mark_first_run_done() + _mark_first_run_done() \ No newline at end of file diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py index 646ac8f6..6ce4849d 100644 --- a/src/exo/utils/channels.py +++ b/src/exo/utils/channels.py @@ -1,3 +1,4 @@ +import contextlib import multiprocessing as mp from dataclasses import dataclass, field from math import inf @@ -132,7 +133,8 @@ class MpSender[T]: def close(self) -> None: if not self._state.closed.is_set(): self._state.closed.set() - self._state.buffer.put(_MpEndOfStream()) + with contextlib.suppress(Exception): + self._state.buffer.put_nowait(_MpEndOfStream()) self._state.buffer.close() # == unique to Mp channels == @@ -190,7 +192,13 @@ class MpReceiver[T]: try: return self.receive_nowait() except WouldBlock: - item = self._state.buffer.get() + try: + item = self._state.buffer.get() + except (TypeError, OSError): + # Queue pipe can get closed while we are blocked on get(). + # The underlying connection._handle becomes None, causing + # TypeError in read(handle, remaining). + raise ClosedResourceError from None if isinstance(item, _MpEndOfStream): self.close() raise EndOfStream from None @@ -204,6 +212,8 @@ class MpReceiver[T]: def close(self) -> None: if not self._state.closed.is_set(): self._state.closed.set() + with contextlib.suppress(Exception): + self._state.buffer.put_nowait(_MpEndOfStream()) self._state.buffer.close() # == unique to Mp channels == diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py index 8f31369b..a43e54f1 100644 --- a/src/exo/utils/info_gatherer/info_gatherer.py +++ b/src/exo/utils/info_gatherer/info_gatherer.py @@ -388,6 +388,12 @@ class InfoGatherer: if IS_DARWIN: if (macmon_path := shutil.which("macmon")) is not None: tg.start_soon(self._monitor_macmon, macmon_path) + else: + # macmon not installed — fall back to psutil for memory + logger.warning( + "macmon not found, falling back to psutil for memory monitoring" + ) + self.memory_poll_rate = 1 tg.start_soon(self._monitor_system_profiler_thunderbolt_data) tg.start_soon(self._monitor_thunderbolt_bridge_status) tg.start_soon(self._monitor_rdma_ctl_status) diff --git a/src/exo/utils/info_gatherer/net_profile.py b/src/exo/utils/info_gatherer/net_profile.py index 85986ead..3255b207 100644 --- a/src/exo/utils/info_gatherer/net_profile.py +++ b/src/exo/utils/info_gatherer/net_profile.py @@ -108,7 +108,7 @@ async def check_reachable( await send.send((target_ip, expected_node_id)) async with ( - httpx.AsyncClient(timeout=timeout, limits=limits) as client, + httpx.AsyncClient(timeout=timeout, limits=limits, verify=False) as client, create_task_group() as tg, ): for node_id in topology.list_nodes(): diff --git a/src/exo/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py index a59e4eed..2f4c5a4e 100644 --- a/src/exo/worker/engines/image/generate.py +++ b/src/exo/worker/engines/image/generate.py @@ -14,6 +14,7 @@ from exo.shared.types.api import ( ImageEditsTaskParams, ImageGenerationStats, ImageGenerationTaskParams, + ImageSize, ) from exo.shared.types.memory import Memory from exo.shared.types.worker.runner_response import ( @@ -23,9 +24,9 @@ from exo.shared.types.worker.runner_response import ( from exo.worker.engines.image.distributed_model import DistributedImageModel -def parse_size(size_str: str | None) -> tuple[int, int]: +def parse_size(size_str: ImageSize) -> tuple[int, int]: """Parse size parameter like '1024x1024' to (width, height) tuple.""" - if not size_str: + if size_str == "auto": return (1024, 1024) try: @@ -109,6 +110,9 @@ def generate_image( # Decode base64 image data and save to temp file image_path = Path(tmpdir) / "input.png" image_path.write_bytes(base64.b64decode(task.image_data)) + if task.size == "auto": + with Image.open(image_path) as img: + width, height = img.size for image_num in range(num_images): # Increment seed for each image to ensure unique results @@ -162,7 +166,7 @@ def generate_image( else 0.0 ) - peak_memory_gb = mx.get_peak_memory() / (1024**3) + peak_memory = Memory.from_bytes(mx.get_peak_memory()) stats = ImageGenerationStats( seconds_per_step=seconds_per_step, @@ -171,7 +175,7 @@ def generate_image( num_images=num_images, image_width=width, image_height=height, - peak_memory_usage=Memory.from_gb(peak_memory_gb), + peak_memory_usage=peak_memory, ) buffer = io.BytesIO() diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index b9064591..693913c4 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -163,11 +163,14 @@ class PipelineLastLayer(CustomMlxLayer): output, (self.r + 1) % self.s, group=self.group ) if cache is not None: - cache.keys = mx.depends(cache.keys, output) # type: ignore[reportUnknownMemberType] + # CacheList (used by MLA models like DeepSeekV32, GLM MoE DSA) + # doesn't have .keys directly; access via first sub-cache. + _cache = cache[0] if hasattr(cache, "caches") else cache # type: ignore + _cache.keys = mx.depends(_cache.keys, output) # type: ignore if self.is_prefill: mx.eval(output) if cache is not None: - mx.eval(cache.keys) # type: ignore + mx.eval(_cache.keys) # type: ignore if not self.is_prefill: output = mx.distributed.all_gather(output, group=self.group)[ @@ -307,7 +310,9 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T: # Add dependency to last cache entry to ensure distributed ops are evaluated if cache is not None: - cache[-1].state = mx.depends(cache[-1].state, logits) # type: ignore + last = cache[-1] # type: ignore + dep_cache = last[0] if hasattr(last, "caches") else last # type: ignore + dep_cache.keys = mx.depends(dep_cache.keys, logits) # type: ignore return logits @@ -333,7 +338,9 @@ def patch_tensor_model[T](model: T) -> T: # Add dependency to last cache entry to ensure distributed ops are evaluated if cache is not None and len(cache) > 0: # pyright: ignore[reportAny] - cache[-1].state = mx.depends(cache[-1].state, logits) # pyright: ignore[reportAny,reportUnknownMemberType] + last = cache[-1] # pyright: ignore[reportAny] + dep_cache = last[0] if hasattr(last, "caches") else last # pyright: ignore[reportAny] + dep_cache.keys = mx.depends(dep_cache.keys, logits) # pyright: ignore[reportAny,reportUnknownMemberType] return logits @@ -547,10 +554,12 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy): on_timeout: TimeoutCallback | None, ) -> nn.Module: model = cast(DeepseekV3Model, model) + for layer in model.layers: eval_with_timeout( layer.parameters(), timeout_seconds / len(model.layers), on_timeout ) + # Shard the self attention if layer.self_attn.q_lora_rank is None: layer.self_attn.q_proj = self.all_to_sharded_linear( @@ -581,12 +590,18 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy): layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj) layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj) - # Shard the MoE. Shard in place since the MoE should be responsible - # for aggregating the results. + # Shard the MoE. else: - self.all_to_sharded_linear_in_place(layer.mlp.shared_experts.gate_proj) - self.sharded_to_all_linear_in_place(layer.mlp.shared_experts.down_proj) - self.all_to_sharded_linear_in_place(layer.mlp.shared_experts.up_proj) + if getattr(layer.mlp, "shared_experts", None) is not None: + self.all_to_sharded_linear_in_place( + layer.mlp.shared_experts.gate_proj + ) + self.sharded_to_all_linear_in_place( + layer.mlp.shared_experts.down_proj + ) + self.all_to_sharded_linear_in_place( + layer.mlp.shared_experts.up_proj + ) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) @@ -779,8 +794,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy): layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group) # pyright: ignore[reportAttributeAccessIssue,reportArgumentType] - # Shard the MoE. Shard in place since the MoE should be responsible - # for aggregating the results. + # Shard the MoE. self.all_to_sharded_linear_in_place( layer.block_sparse_moe.switch_mlp.gate_proj ) @@ -893,8 +907,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy): layer.self_attn.num_attention_heads //= self.N layer.self_attn.num_key_value_heads //= self.N - # Shard the MoE. Shard in place since the MoE should be responsible - # for aggregating the results. + # Shard the MoE. if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)): self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index c747ba4e..ae6f76fa 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -5,6 +5,7 @@ import mlx.core as mx import psutil from mlx_lm.models.cache import ( ArraysCache, + CacheList, KVCache, QuantizedKVCache, RotatingKVCache, @@ -17,10 +18,22 @@ from exo.worker.engines.mlx import Model from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS from exo.worker.runner.bootstrap import logger -# Fraction of device memory above which LRU eviction kicks in -_DEFAULT_MEMORY_THRESHOLD = 0.9 + +# Fraction of device memory above which LRU eviction kicks in. +# Smaller machines need more aggressive eviction. +def _default_memory_threshold() -> float: + total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb + if total_gb >= 128: + return 0.85 + if total_gb >= 64: + return 0.80 + if total_gb >= 32: + return 0.75 + return 0.70 + + _MEMORY_THRESHOLD = float( - os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD) + os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold()) ) @@ -64,7 +77,7 @@ def has_non_kv_caches(cache: KVCacheType) -> bool: class KVPrefixCache: - def __init__(self, group: mx.distributed.Group | None = None): + def __init__(self, group: mx.distributed.Group | None): self.prompts: list[mx.array] = [] # mx array of tokens (ints) self.caches: list[KVCacheType] = [] self._snapshots: list[list[CacheSnapshot] | None] = [] @@ -156,15 +169,15 @@ class KVPrefixCache: best_length = 0 is_exact = False - # Find best cache + # Find best cache match for i, cached_prompt in enumerate(self.prompts): length = get_prefix_length(prompt_tokens, cached_prompt) + if length >= max_length - 1: + best_index, best_length = i, length + is_exact = True + break if length > best_length: best_index, best_length = i, length - if length == max_length: - is_exact = True - best_index, best_length = i, length - break if best_index is None: return make_kv_cache(model), prompt_tokens, None @@ -172,11 +185,12 @@ class KVPrefixCache: # For exact match: trim to max_length-1 so remaining has the last token # For partial match: trim to best_length, remaining has suffix to prefill # This ensures stream_generate always has at least one token to start with - target = (max_length - 1) if is_exact else best_length + has_ssm = has_non_kv_caches(self.caches[best_index]) + target = (max_length - 1) if is_exact and not has_ssm else best_length restore_pos, restore_snap = self._get_snapshot(best_index, target) # No usable snapshot — need fresh cache - if restore_snap is None and has_non_kv_caches(self.caches[best_index]): + if restore_snap is None and has_ssm: return make_kv_cache(model), prompt_tokens, None prompt_cache = deepcopy(self.caches[best_index]) @@ -257,10 +271,21 @@ def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array: return mx.array(prompt_tokens) +def _entry_length( + c: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList, +) -> int: + # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache). + if hasattr(c, "offset"): + return c.offset + # For CacheList + if hasattr(c, "size"): + return int(c.size()) # type: ignore + return 0 + + def cache_length(cache: KVCacheType) -> int: """Get the number of tokens in a KV cache.""" - # Use .offset attribute which KVCache types have (len() not implemented in older QuantizedKVCache). - return max(getattr(c, "offset", 0) for c in cache) + return max(_entry_length(c) for c in cache) def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int: diff --git a/src/exo/worker/engines/mlx/dsml_encoding.py b/src/exo/worker/engines/mlx/dsml_encoding.py new file mode 100644 index 00000000..d5a5e210 --- /dev/null +++ b/src/exo/worker/engines/mlx/dsml_encoding.py @@ -0,0 +1,72 @@ +import json +import re +from typing import Any + +from mlx_lm.chat_templates import deepseek_v32 + +from exo.shared.types.api import ToolCallItem + +BOS_TOKEN: str = deepseek_v32.bos_token +EOS_TOKEN: str = deepseek_v32.eos_token +DSML_TOKEN: str = deepseek_v32.dsml_token +THINKING_START: str = deepseek_v32.thinking_start_token +THINKING_END: str = deepseek_v32.thinking_end_token +USER_TOKEN = "<\uff5cUser\uff5c>" +ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>" +TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>" +TOOL_CALLS_END = f"" +encode_messages = deepseek_v32.encode_messages + +_INVOKE_PATTERN = re.compile( + rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">" + rf"(.*?)" + rf"", + re.DOTALL, +) + +_PARAM_PATTERN = re.compile( + rf"<{re.escape(DSML_TOKEN)}parameter\s+name=\"([^\"]+)\"\s+string=\"(true|false)\">" + rf"(.*?)" + rf"", + re.DOTALL, +) + + +def parse_dsml_output(text: str) -> list[ToolCallItem] | None: + """Parse DSML function_calls block from model output text. + + Args: + text: The text containing the DSML function_calls block + (including the start/end markers). + + Returns: + List of ToolCallItem, or None if parsing fails. + """ + tool_calls: list[ToolCallItem] = [] + + for invoke_match in _INVOKE_PATTERN.finditer(text): + func_name = invoke_match.group(1) + invoke_body = invoke_match.group(2) + + args: dict[str, Any] = {} + for param_match in _PARAM_PATTERN.finditer(invoke_body): + param_name = param_match.group(1) + is_string = param_match.group(2) == "true" + param_value = param_match.group(3) + + if is_string: + args[param_name] = param_value + else: + try: + args[param_name] = json.loads(param_value) + except (json.JSONDecodeError, ValueError): + args[param_name] = param_value + + tool_calls.append( + ToolCallItem( + name=func_name, + arguments=json.dumps(args), + ) + ) + + return tool_calls if tool_calls else None diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index bc38c61c..6e926205 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -48,7 +48,11 @@ from exo.worker.runner.bootstrap import logger generation_stream = mx.new_stream(mx.default_device()) -_MIN_PREFIX_HIT_TO_UPDATE = 1000 +_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 + + +class PrefillCancelled(BaseException): + """Raised when prefill is cancelled via the progress callback.""" def prefill( @@ -57,6 +61,8 @@ def prefill( sampler: Callable[[mx.array], mx.array], prompt_tokens: mx.array, cache: KVCacheType, + group: mx.distributed.Group | None, + on_prefill_progress: Callable[[int, int], None] | None, ) -> tuple[float, int, list[CacheSnapshot]]: """Prefill the KV cache with prompt tokens. @@ -64,7 +70,7 @@ def prefill( then trims off the extra generated token. Returns: - tokens_per_sec + (tokens_per_sec, num_tokens, snapshots) """ num_tokens = len(prompt_tokens) if num_tokens == 0: @@ -75,6 +81,7 @@ def prefill( has_ssm = has_non_kv_caches(cache) snapshots: list[CacheSnapshot] = [] + # TODO(evan): kill the callbacks/runner refactor def progress_callback(processed: int, total: int) -> None: elapsed = time.perf_counter() - start_time tok_per_sec = processed / elapsed if elapsed > 0 else 0 @@ -84,23 +91,33 @@ def prefill( if has_ssm: snapshots.append(snapshot_ssm_states(cache)) + if on_prefill_progress is not None: + on_prefill_progress(processed, total) + set_pipeline_prefill(model, is_prefill=True) + mx_barrier(group) + logger.info("Starting prefill") + # Use max_tokens=1 because max_tokens=0 does not work. # We just throw away the generated token - we only care about filling the cache - for _ in stream_generate( - model=model, - tokenizer=tokenizer, - prompt=prompt_tokens, - max_tokens=1, - sampler=sampler, - prompt_cache=cache, - prefill_step_size=8192, - kv_group_size=KV_GROUP_SIZE, - kv_bits=KV_BITS, - prompt_progress_callback=progress_callback, - ): - break # Stop after first iteration - cache is now filled + try: + for _ in stream_generate( + model=model, + tokenizer=tokenizer, + prompt=prompt_tokens, + max_tokens=1, + sampler=sampler, + prompt_cache=cache, + prefill_step_size=4096, + kv_group_size=KV_GROUP_SIZE, + kv_bits=KV_BITS, + prompt_progress_callback=progress_callback, + ): + break # Stop after first iteration - cache is now filled + except PrefillCancelled: + set_pipeline_prefill(model, is_prefill=False) + raise set_pipeline_prefill(model, is_prefill=False) @@ -129,7 +146,7 @@ def prefill( def warmup_inference( model: Model, tokenizer: TokenizerWrapper, - group: mx.distributed.Group | None = None, + group: mx.distributed.Group | None, ) -> int: content = "Prompt to warm up the inference engine. Repeat this." @@ -251,8 +268,9 @@ def mlx_generate( tokenizer: TokenizerWrapper, task: TextGenerationTaskParams, prompt: str, - kv_prefix_cache: KVPrefixCache | None = None, - group: mx.distributed.Group | None = None, + kv_prefix_cache: KVPrefixCache | None, + group: mx.distributed.Group | None, + on_prefill_progress: Callable[[int, int], None] | None = None, ) -> Generator[GenerationResponse]: # Ensure that generation stats only contains peak memory for this generation mx.reset_peak_memory() @@ -305,9 +323,6 @@ def mlx_generate( ) max_stop_len = max((len(s) for s in stop_sequences), default=0) - mx_barrier(group) - logger.info("Ready to prefill") - # Prefill cache with all tokens except the last one prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( model, @@ -315,6 +330,8 @@ def mlx_generate( sampler, prompt_tokens[:-1], caches, + group, + on_prefill_progress, ) cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None @@ -331,6 +348,7 @@ def mlx_generate( think_start = tokenizer.think_start think_end = tokenizer.think_end + logger.info("Starting decode") mx_barrier(group) for completion_tokens, out in enumerate( @@ -438,9 +456,14 @@ def mlx_generate( full_prompt_tokens = mx.concatenate( [all_prompt_tokens, generated_tokens_array] ) + hit_ratio = ( + prefix_hit_length / len(all_prompt_tokens) + if len(all_prompt_tokens) > 0 + else 0.0 + ) if ( matched_index is not None - and prefix_hit_length >= _MIN_PREFIX_HIT_TO_UPDATE + and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE ): kv_prefix_cache.update_kv_cache( matched_index, diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 6aceb53c..360a018b 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -1,5 +1,6 @@ import json import os +import re import sys import time from pathlib import Path @@ -231,11 +232,11 @@ def shard_and_load( # Estimate timeout based on model size (5x default for large queued workloads) base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300")) - model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3) - timeout_seconds = base_timeout + model_size_gb + model_size = get_weights_size(shard_metadata) + timeout_seconds = base_timeout + model_size.in_gb logger.info( f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s " - f"(model size: {model_size_gb:.1f}GB)" + f"(model size: {model_size.in_gb:.1f}GB)" ) match shard_metadata: @@ -285,11 +286,15 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None: model_id_lower = model_id.lower() if "kimi-k2" in model_id_lower: return [163586] - elif "glm-4.7-flash" in model_id_lower: + elif "glm-5" in model_id_lower or "glm-4.7" in model_id_lower: + # For GLM-5 and GLM-4.7 # 154820: <|endoftext|>, 154827: <|user|>, 154829: <|observation|> return [154820, 154827, 154829] elif "glm" in model_id_lower: + # For GLM-4.5 and older return [151336, 151329, 151338] + elif "gpt-oss" in model_id_lower: + return [200002, 200012] return None @@ -353,7 +358,13 @@ def load_tokenizer_for_model_id( return list(hf_tokenizer.model.encode(text, allowed_special="all")) # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType] hf_tokenizer.encode = _patched_encode - return TokenizerWrapper(hf_tokenizer, eos_token_ids=eos_token_ids) + return TokenizerWrapper( + hf_tokenizer, + eos_token_ids=eos_token_ids, + tool_call_start="<|tool_calls_section_begin|>", + tool_call_end="<|tool_calls_section_end|>", + tool_parser=_parse_kimi_tool_calls, + ) tokenizer = load_tokenizer( model_path, @@ -397,6 +408,69 @@ def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None: func["arguments"] = json.loads(args) +def _collect_nested_property_names(schema: dict[str, Any]) -> set[str]: + names: set[str] = set() + properties: dict[str, Any] = schema.get("properties", {}) # type: ignore[reportAny] + for prop_spec in properties.values(): # pyright: ignore[reportAny] + if not isinstance(prop_spec, dict): + continue + if prop_spec.get("type") == "array": # type: ignore[reportAny] + items: dict[str, Any] | None = prop_spec.get("items") # type: ignore[reportAny] + if isinstance(items, dict) and items.get("type") == "object": # type: ignore[reportAny] + inner_props: dict[str, Any] = items.get("properties", {}) # type: ignore[reportAny] + for k in inner_props: # pyright: ignore[reportUnknownVariableType] + names.add(str(k)) # pyright: ignore[reportUnknownArgumentType] + names.update(_collect_nested_property_names(items)) # pyright: ignore[reportUnknownArgumentType] + return names + + +def _schemas_lost_in_prompt(prompt: str, tools: list[dict[str, Any]]) -> bool: + """Return True if nested property names from any tool schema are absent.""" + for tool in tools: + fn: dict[str, Any] = tool.get("function", {}) # type: ignore + params: dict[str, Any] = fn.get("parameters", {}) # type: ignore + nested = _collect_nested_property_names(params) + if nested and not all(name in prompt for name in nested): + return True + return False + + +_LOSSY_TEMPLATE_PATTERN = re.compile( + r"""inner_type\s*==\s*["']object \| object["']\s*or\s*inner_type\|length\s*>\s*\d+""", +) + + +def _patch_lossy_chat_template(template: str) -> str | None: + """Patch chat templates that collapse nested object schemas to ``any[]``. + + Some templates (e.g., GPT-OSS) have a guard like:: + + inner_type == "object | object" or inner_type|length > 50 + + The length check silently drops complex array-of-object schemas. + We remove the length guard, keeping only the object-union check. + Returns the patched template, or *None* if no patch was needed. + """ + patched, n = _LOSSY_TEMPLATE_PATTERN.subn( + lambda m: m.group(0).split(" or ")[0], # keep only the object-union check + template, + ) + return patched if n > 0 else None + + +def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool: + if "deepseek-v3.2" not in task_params.model.lower(): + return False + # Use DSML encoding when tools are provided or tool results are in the conversation + if task_params.tools: + return True + if task_params.chat_template_messages: + return any( + msg.get("role") == "tool" for msg in task_params.chat_template_messages + ) + return False + + def apply_chat_template( tokenizer: TokenizerWrapper, task_params: TextGenerationTaskParams, @@ -408,7 +482,6 @@ def apply_chat_template( When chat_template_messages is available (from Chat Completions API), uses those directly to preserve tool_calls, thinking, and other fields. - Otherwise builds messages from the task params input/instructions. """ formatted_messages: list[dict[str, Any]] = [] if task_params.chat_template_messages is not None: @@ -436,6 +509,19 @@ def apply_chat_template( partial_assistant_content = cast(str, formatted_messages[-1].get("content", "")) formatted_messages = formatted_messages[:-1] + if _needs_dsml_encoding(task_params): + from exo.worker.engines.mlx.dsml_encoding import encode_messages + + prompt = encode_messages( + messages=formatted_messages, + thinking_mode="thinking" if task_params.enable_thinking else "chat", + tools=task_params.tools, + ) + if partial_assistant_content: + prompt += partial_assistant_content + logger.info(prompt) + return prompt + extra_kwargs: dict[str, Any] = {} if task_params.enable_thinking is not None: # Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking". @@ -443,14 +529,28 @@ def apply_chat_template( extra_kwargs["enable_thinking"] = task_params.enable_thinking extra_kwargs["thinking"] = task_params.enable_thinking + patched_template: str | None = None + if task_params.tools: + original_template: str | None = getattr(tokenizer, "chat_template", None) + if isinstance(original_template, str): + patched_template = _patch_lossy_chat_template(original_template) + if patched_template is not None: + logger.info( + "Patched lossy chat template (removed inner_type length guard)" + ) + prompt: str = tokenizer.apply_chat_template( formatted_messages, tokenize=False, add_generation_prompt=True, tools=task_params.tools, + **({"chat_template": patched_template} if patched_template is not None else {}), **extra_kwargs, ) + if task_params.tools and _schemas_lost_in_prompt(prompt, task_params.tools): + logger.warning("Chat template lost nested tool schemas even after patching") + if partial_assistant_content: prompt += partial_assistant_content @@ -542,18 +642,17 @@ def set_wired_limit_for_model(model_size: Memory): if not mx.metal.is_available(): return - model_bytes = model_size.in_bytes - max_rec_size = int(mx.metal.device_info()["max_recommended_working_set_size"]) - if model_bytes > 0.9 * max_rec_size: - model_mb = model_bytes // 2**20 - max_rec_mb = max_rec_size // 2**20 + max_rec_size = Memory.from_bytes( + int(mx.metal.device_info()["max_recommended_working_set_size"]) + ) + if model_size > 0.9 * max_rec_size: logger.warning( - f"Generating with a model that requires {model_mb} MB " - f"which is close to the maximum recommended size of {max_rec_mb} " + f"Generating with a model that requires {model_size.in_float_mb:.1f} MB " + f"which is close to the maximum recommended size of {max_rec_size.in_float_mb:.1f} " "MB. This can be slow. See the documentation for possible work-arounds: " "https://github.com/ml-explore/mlx-lm/tree/main#large-models" ) - mx.set_wired_limit(max_rec_size) + mx.set_wired_limit(max_rec_size.in_bytes) logger.info(f"Wired limit set to {max_rec_size}.") @@ -585,3 +684,41 @@ def mx_barrier(group: Group | None): mx.array(1.0), group=group, stream=mx.default_stream(mx.Device(mx.cpu)) ) ) + + +def _parse_kimi_tool_calls(text: str): + import regex as re + + # kimi has a fixed function naming scheme, with a json formatted arg + # functions.multiply:0<|tool_call_argument_begin|>{"a": 2, "b": 3} + _func_name_regex = re.compile( + r"^\s*((?:functions\.)?(.+?):\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL + ) + _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL) + _tool_call_split_regex = re.compile( + r"<\|tool_call_begin\|>(.*?)<\|tool_call_end\|>", re.DOTALL + ) + + def _parse_single_tool(text: str) -> dict[str, Any]: + func_name_match = _func_name_regex.search(text) + if func_name_match is None: + raise ValueError("No tool call found.") + tool_call_id = func_name_match.group(1) # e.g. "functions.get_weather:0" + func_name = func_name_match.group(2) # e.g. "get_weather" + + func_args_match = _func_arg_regex.search(text) + if func_args_match is None: + raise ValueError("No tool call arguments found.") + func_args = func_args_match.group(1) + try: + arg_dct = json.loads(func_args) # pyright: ignore[reportAny] + except Exception: + arg_dct = None + + return dict(id=tool_call_id, name=func_name, arguments=arg_dct) + + tool_matches = _tool_call_split_regex.findall(text) + if tool_matches: + return [_parse_single_tool(match) for match in tool_matches] # pyright: ignore[reportAny] + else: + return [_parse_single_tool(text)] diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index af105652..3bafedaa 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -241,6 +241,11 @@ class Worker: cancelled_task_id=cancelled_task_id, runner_id=runner_id ): await self.runners[runner_id].cancel_task(cancelled_task_id) + await self.event_sender.send( + TaskStatusUpdated( + task_id=task.task_id, task_status=TaskStatus.Complete + ) + ) case ImageEdits() if task.task_params.total_input_chunks > 0: # Assemble image from chunks and inject into task cmd_id = task.command_id diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index ad6c78f6..47293d8c 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -1,28 +1,34 @@ import base64 -import json import math import resource import time from collections.abc import Generator from functools import cache -from typing import Any, Callable, Literal +from typing import TYPE_CHECKING, Literal import mlx.core as mx +from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model from mlx_lm.models.gpt_oss import Model as GptOssModel from mlx_lm.tokenizer_utils import TokenizerWrapper from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs] HarmonyEncodingName, + HarmonyError, # pyright: ignore[reportUnknownVariableType] Role, StreamableParser, load_harmony_encoding, ) -from pydantic import ValidationError from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED from exo.shared.models.model_cards import ModelId, ModelTask from exo.shared.tracing import clear_trace_buffer, get_trace_buffer from exo.shared.types.api import ImageGenerationStats -from exo.shared.types.chunks import ErrorChunk, ImageChunk, TokenChunk, ToolCallChunk +from exo.shared.types.chunks import ( + ErrorChunk, + ImageChunk, + PrefillProgressChunk, + TokenChunk, + ToolCallChunk, +) from exo.shared.types.common import CommandId from exo.shared.types.events import ( ChunkGenerated, @@ -82,7 +88,11 @@ from exo.worker.engines.image import ( ) from exo.worker.engines.mlx import Model from exo.worker.engines.mlx.cache import KVPrefixCache -from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference +from exo.worker.engines.mlx.generator.generate import ( + PrefillCancelled, + mlx_generate, + warmup_inference, +) from exo.worker.engines.mlx.utils_mlx import ( apply_chat_template, detect_thinking_prompt_suffix, @@ -93,6 +103,8 @@ from exo.worker.engines.mlx.utils_mlx import ( ) from exo.worker.runner.bootstrap import logger +from .tool_parsers import ToolParser, make_mlx_parser + def _is_primary_output_node(shard_metadata: ShardMetadata) -> bool: """Check if this node is the primary output node for image generation. @@ -138,6 +150,7 @@ def main( inference_model: Model | None = None image_model: DistributedImageModel | None = None tokenizer = None + tool_parser: ToolParser | None = None group = None kv_prefix_cache: KVPrefixCache | None = None check_for_cancel_every: int | None = None @@ -203,8 +216,17 @@ def main( bound_instance, group, on_timeout=on_model_load_timeout ) logger.info( - f"model has_tool_calling={tokenizer.has_tool_calling}" + f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}" ) + if tokenizer.has_tool_calling: + assert tokenizer.tool_call_start + assert tokenizer.tool_call_end + assert tokenizer.tool_parser # pyright: ignore[reportAny] + tool_parser = make_mlx_parser( + tokenizer.tool_call_start, + tokenizer.tool_call_end, + tokenizer.tool_parser, # pyright: ignore[reportAny] + ) kv_prefix_cache = KVPrefixCache(group) elif ( @@ -287,6 +309,34 @@ def main( assert tokenizer assert check_for_cancel_every + # Define callback to send prefill progress events + # and check for cancellation between prefill chunks. + # TODO(evan): kill the callbacks/runner refactor + # Specifically the part that this is literally duplicated code. + def on_prefill_progress( + processed: int, + total: int, + _task_id: TaskId = task.task_id, + _group: mx.distributed.Group | None = group, + ) -> None: + if device_rank == 0: + event_sender.send( + ChunkGenerated( + command_id=command_id, + chunk=PrefillProgressChunk( + model=shard_metadata.model_card.model_id, + processed_tokens=processed, + total_tokens=total, + ), + ) + ) + cancelled_tasks.update(cancel_receiver.collect()) + want_to_cancel = (_task_id in cancelled_tasks) or ( + TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks + ) + if mx_any(want_to_cancel, _group): + raise PrefillCancelled() + try: _check_for_debug_prompts(task_params) @@ -300,42 +350,29 @@ def main( task=task_params, prompt=prompt, kv_prefix_cache=kv_prefix_cache, + on_prefill_progress=on_prefill_progress, group=group, ) - # For other thinking models (GLM, etc.), check if we need to - # prepend the thinking tag that was consumed by the chat template - if detect_thinking_prompt_suffix(prompt, tokenizer): + if tokenizer.has_thinking: mlx_generator = parse_thinking_models( - mlx_generator, tokenizer - ) - - # Kimi-K2 has tool call sections - we don't care about them - if "kimi" in shard_metadata.model_card.model_id.lower(): - mlx_generator = filter_kimi_tokens(mlx_generator) - patch_kimi_tokenizer(tokenizer) - - # GLM models need patched parser (upstream has bug with None regex match) - elif "glm" in shard_metadata.model_card.model_id.lower(): - patch_glm_tokenizer(tokenizer) - - # GPT-OSS specific parsing to match other model formats. - elif isinstance(inference_model, GptOssModel): - mlx_generator = parse_gpt_oss(mlx_generator) - - if tokenizer.has_tool_calling and not isinstance( - inference_model, GptOssModel - ): - assert tokenizer.tool_call_start - assert tokenizer.tool_call_end - assert tokenizer.tool_parser # pyright: ignore[reportAny] - mlx_generator = parse_tool_calls( mlx_generator, - tokenizer.tool_call_start, - tokenizer.tool_call_end, - tokenizer.tool_parser, # pyright: ignore[reportAny] + tokenizer, + # For other thinking models (GLM, etc.), check if we need to + # prepend the thinking tag that was consumed by the chat template + starts_in_thinking=detect_thinking_prompt_suffix( + prompt, tokenizer + ), ) + # Model-specific output parsing for tool calls. + if isinstance(inference_model, GptOssModel): + mlx_generator = parse_gpt_oss(mlx_generator) + elif isinstance(inference_model, DeepseekV32Model): + mlx_generator = parse_deepseek_v32(mlx_generator) + elif tool_parser: + mlx_generator = parse_tool_calls(mlx_generator, tool_parser) + completion_tokens = 0 tokens_since_last_cancel_check = 0 for response in mlx_generator: @@ -384,6 +421,7 @@ def main( stats=response.stats, logprob=response.logprob, top_logprobs=response.top_logprobs, + is_thinking=response.is_thinking, ), ) ) @@ -401,6 +439,8 @@ def main( ) ) + except PrefillCancelled: + logger.info(f"Prefill cancelled for task {task.task_id}") # can we make this more explicit? except Exception as e: if device_rank == 0: @@ -548,6 +588,13 @@ def main( case Shutdown(): current_status = RunnerShuttingDown() logger.info("runner shutting down") + if not TYPE_CHECKING: + del inference_model, image_model, tokenizer, group + mx.clear_cache() + import gc + + gc.collect() + event_sender.send( RunnerStatusUpdated( runner_id=runner_id, runner_status=current_status @@ -572,12 +619,8 @@ def main( event_sender.send( RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status) ) - if isinstance(current_status, RunnerShutdown): - del inference_model, image_model, tokenizer, group - mx.clear_cache() - import gc - gc.collect() + if isinstance(current_status, RunnerShutdown): break @@ -587,21 +630,8 @@ def get_gpt_oss_encoding(): return encoding -def filter_kimi_tokens( - responses: Generator[GenerationResponse | ToolCallResponse], -) -> Generator[GenerationResponse]: - for resp in responses: - assert isinstance(resp, GenerationResponse) - if ( - resp.text == "<|tool_calls_section_begin|>" - or resp.text == "<|tool_calls_section_end|>" - ): - continue - yield resp - - def parse_gpt_oss( - responses: Generator[GenerationResponse | ToolCallResponse], + responses: Generator[GenerationResponse], ) -> Generator[GenerationResponse | ToolCallResponse]: encoding = get_gpt_oss_encoding() stream = StreamableParser(encoding, role=Role.ASSISTANT) @@ -611,17 +641,31 @@ def parse_gpt_oss( for response in responses: assert isinstance(response, GenerationResponse) - stream.process(response.token) + try: + stream.process(response.token) + except HarmonyError: + logger.error("Encountered critical Harmony Error, returning early") + return delta = stream.last_content_delta ch = stream.current_channel recipient = stream.current_recipient + # Debug: log every token with state + logger.debug( + f"parse_gpt_oss token={response.token} text={response.text!r} " + f"recipient={recipient!r} ch={ch!r} delta={delta!r} " + f"state={stream.state} current_tool={current_tool_name!r}" + ) + if recipient != current_tool_name: if current_tool_name is not None: prefix = "functions." if current_tool_name.startswith(prefix): current_tool_name = current_tool_name[len(prefix) :] + logger.info( + f"parse_gpt_oss yielding tool call: name={current_tool_name!r}" + ) yield ToolCallResponse( tool_calls=[ ToolCallItem( @@ -642,44 +686,208 @@ def parse_gpt_oss( if ch == "analysis" and not thinking: thinking = True - yield response.model_copy(update={"text": ""}) if ch != "analysis" and thinking: thinking = False - yield response.model_copy(update={"text": ""}) if delta: - yield response.model_copy(update={"text": delta}) + yield response.model_copy(update={"text": delta, "is_thinking": thinking}) if response.finish_reason is not None: - if thinking: - yield response.model_copy(update={"text": ""}) yield response -def parse_thinking_models( - responses: Generator[GenerationResponse | ToolCallResponse], - tokenizer: TokenizerWrapper, +def parse_deepseek_v32( + responses: Generator[GenerationResponse], ) -> Generator[GenerationResponse | ToolCallResponse]: + """Parse DeepSeek V3.2 DSML tool calls from the generation stream. + + Uses accumulated-text matching (not per-token marker checks) because + DSML markers like <｜DSML｜function_calls> may span multiple tokens. + Also handles ... blocks for thinking mode. """ - For models that inject thinking tags in the prompt (like GLM-4.7), - prepend the thinking tag to the output stream so the frontend - can properly parse thinking content. + from exo.worker.engines.mlx.dsml_encoding import ( + THINKING_END, + THINKING_START, + TOOL_CALLS_END, + TOOL_CALLS_START, + parse_dsml_output, + ) + + accumulated = "" + in_tool_call = False + thinking = False + # Tokens buffered while we detect the start of a DSML block + pending_buffer: list[GenerationResponse] = [] + # Text accumulated during a tool call block + tool_call_text = "" + + for response in responses: + assert isinstance(response, GenerationResponse) + + # ── Handle thinking tags ── + if not thinking and THINKING_START in response.text: + thinking = True + # Yield any text before the tag + before = response.text[: response.text.index(THINKING_START)] + if before: + yield response.model_copy(update={"text": before}) + continue + + if thinking and THINKING_END in response.text: + thinking = False + # Yield any text after the tag + after = response.text[ + response.text.index(THINKING_END) + len(THINKING_END) : + ] + if after: + yield response.model_copy(update={"text": after, "is_thinking": False}) + continue + + if thinking: + yield response.model_copy(update={"is_thinking": True}) + continue + + # ── Handle tool call accumulation ── + if in_tool_call: + tool_call_text += response.text + if TOOL_CALLS_END in tool_call_text: + # Parse the accumulated DSML block + parsed = parse_dsml_output(tool_call_text) + if parsed is not None: + logger.info(f"parsed DSML tool calls: {parsed}") + yield ToolCallResponse( + tool_calls=parsed, + usage=response.usage, + stats=response.stats, + ) + else: + logger.warning( + f"DSML tool call parsing failed for: {tool_call_text}" + ) + yield response.model_copy(update={"text": tool_call_text}) + in_tool_call = False + tool_call_text = "" + continue + + # EOS reached before end marker — yield buffered text as-is + if response.finish_reason is not None: + logger.info("DSML tool call parsing interrupted by EOS") + yield response.model_copy(update={"text": tool_call_text}) + in_tool_call = False + tool_call_text = "" + continue + + # ── Detect start of tool call block ── + accumulated += response.text + + if TOOL_CALLS_START in accumulated: + # The start marker might be split across pending_buffer + current token + start_idx = accumulated.index(TOOL_CALLS_START) + # Yield any pending tokens that are purely before the marker + pre_text = accumulated[:start_idx] + if pre_text: + # Flush pending buffer tokens that contributed text before the marker + for buf_resp in pending_buffer: + if pre_text: + chunk = buf_resp.text + if len(chunk) <= len(pre_text): + yield buf_resp + pre_text = pre_text[len(chunk) :] + else: + yield buf_resp.model_copy(update={"text": pre_text}) + pre_text = "" + pending_buffer = [] + tool_call_text = accumulated[start_idx:] + accumulated = "" + + # Check if the end marker is already present (entire tool call in one token) + if TOOL_CALLS_END in tool_call_text: + parsed = parse_dsml_output(tool_call_text) + if parsed is not None: + logger.info(f"parsed DSML tool calls: {parsed}") + yield ToolCallResponse( + tool_calls=parsed, + usage=response.usage, + stats=response.stats, + ) + else: + logger.warning( + f"DSML tool call parsing failed for: {tool_call_text}" + ) + yield response.model_copy(update={"text": tool_call_text}) + tool_call_text = "" + else: + in_tool_call = True + continue + + # Check if accumulated text might be the start of a DSML marker + # Buffer tokens if we see a partial match at the end + if _could_be_dsml_prefix(accumulated): + pending_buffer.append(response) + continue + + # No partial match — flush all pending tokens and the current one + for buf_resp in pending_buffer: + yield buf_resp + pending_buffer = [] + accumulated = "" + yield response + + # Flush any remaining pending buffer at generator end + for buf_resp in pending_buffer: + yield buf_resp + + +def _could_be_dsml_prefix(text: str) -> bool: + """Check if the end of text could be the start of a DSML function_calls marker. + + We look for suffixes of text that are prefixes of the TOOL_CALLS_START pattern. + This allows us to buffer tokens until we can determine if a tool call is starting. """ - first = True + from exo.worker.engines.mlx.dsml_encoding import TOOL_CALLS_START + + # Only check the last portion of text that could overlap with the marker + max_check = len(TOOL_CALLS_START) + tail = text[-max_check:] if len(text) > max_check else text + + # Check if any suffix of tail is a prefix of TOOL_CALLS_START + for i in range(len(tail)): + suffix = tail[i:] + if TOOL_CALLS_START.startswith(suffix): + return True + return False + + +def parse_thinking_models( + responses: Generator[GenerationResponse], + tokenizer: TokenizerWrapper, + starts_in_thinking: bool = True, +) -> Generator[GenerationResponse]: + """Route thinking tokens via is_thinking flag. + + Swallows think tag tokens, sets is_thinking on all others. + Always yields tokens with finish_reason to avoid hanging the chunk stream. + """ + in_thinking = starts_in_thinking for response in responses: if isinstance(response, ToolCallResponse): yield response continue - if first: - first = False - yield response.model_copy( - update={ - "text": tokenizer.think_start, - "token": tokenizer.think_start_id, - } - ) - yield response + + is_think_tag = ( + tokenizer.think_end is not None and response.text == tokenizer.think_end + ) or ( + tokenizer.think_start is not None and response.text == tokenizer.think_start + ) + + if is_think_tag: + in_thinking = response.text != tokenizer.think_end + # Never swallow finish_reason — the chunk stream needs it to terminate. + if response.finish_reason is not None: + yield response.model_copy(update={"text": "", "is_thinking": False}) + continue + yield response.model_copy(update={"is_thinking": in_thinking}) def _send_image_chunk( @@ -781,221 +989,55 @@ def _process_image_response( def parse_tool_calls( - responses: Generator[GenerationResponse | ToolCallResponse], - tool_call_start: str, - tool_call_end: str, - tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]], + responses: Generator[GenerationResponse], tool_parser: ToolParser ) -> Generator[GenerationResponse | ToolCallResponse]: in_tool_call = False tool_call_text_parts: list[str] = [] for response in responses: - assert isinstance(response, GenerationResponse) - # assumption: the tool call start is one token - if response.text == tool_call_start: + if response.text.startswith(tool_parser.start_parsing): in_tool_call = True - continue - # assumption: the tool call end is one token - if in_tool_call and response.text == tool_call_end: - try: - # tool_parser returns an arbitrarily nested python dictionary - # we actually don't want the python dictionary, we just want to - # parse the top level { function: ..., arguments: ... } structure - # as we're just gonna hand it back to the api anyway - parsed = tool_parser("".join(tool_call_text_parts).strip()) - logger.info(f"parsed {tool_call_text_parts=} into {parsed=}") - if isinstance(parsed, list): - tools = [_validate_single_tool(tool) for tool in parsed] - else: - tools = [_validate_single_tool(parsed)] - yield ToolCallResponse( - tool_calls=tools, usage=response.usage, stats=response.stats - ) - - except ( - json.JSONDecodeError, - ValidationError, - ValueError, - AttributeError, - ) as e: - # ValueError: our parsers raise this for malformed tool calls - # AttributeError: upstream parsers (e.g. glm47) may raise this when regex doesn't match - logger.opt(exception=e).warning("tool call parsing failed") - # assumption: talking about tool calls, not making a tool call - response.text = ( - tool_call_start + "".join(tool_call_text_parts) + tool_call_end - ) - yield response - - in_tool_call = False - tool_call_text_parts = [] - continue if in_tool_call: tool_call_text_parts.append(response.text) + if response.text.endswith(tool_parser.end_parsing): + # parse the actual tool calls from the tool call text + parsed = tool_parser.parse_tool_calls( + "".join(tool_call_text_parts).strip() + ) + logger.info(f"parsed {tool_call_text_parts=} into {parsed=}") + if parsed is not None: + yield ToolCallResponse( + tool_calls=parsed, usage=response.usage, stats=response.stats + ) + else: + logger.warning( + f"tool call parsing failed for text {''.join(tool_call_text_parts)}" + ) + response.text = "".join(tool_call_text_parts) + yield response + + in_tool_call = False + tool_call_text_parts = [] + continue + if response.finish_reason is not None: logger.info( - "toll call parsing interrupted, yield partial tool call as text" + "tool call parsing interrupted, yield partial tool call as text" ) - yield GenerationResponse( - text=tool_call_start + "".join(tool_call_text_parts), - token=0, - finish_reason=response.finish_reason, - usage=response.usage, - stats=response.stats, + response = response.model_copy( + update={ + "text": "".join(tool_call_text_parts), + "token": 0, + } ) + yield response + continue + # fallthrough yield response -def patch_kimi_tokenizer(tokenizer: TokenizerWrapper): - """ - Version of to-be-upstreamed kimi-k2 tool parser - """ - import ast - import json - from typing import Any - - import regex as re - - # kimi has a fixed function naming scheme, with a json formatted arg - # functions.multiply:0 <|tool_call_argument_begin|> {"a": 2, "b": 3} - # Also needs to handle tools like call_0<|tool_call_argument_begin|>{"filePath": "..."} - _func_name_regex = re.compile( - r"^\s*(.+)[:](\d+)\s*<\|tool_call_argument_begin\|>", re.DOTALL - ) - _func_arg_regex = re.compile(r"<\|tool_call_argument_begin\|>\s*(.*)\s*", re.DOTALL) - - # kimi has a tool_calls_section - we're leaving this up to the caller to handle - tool_call_start = "<|tool_call_begin|>" - tool_call_end = "<|tool_call_end|>" - - def _deserialize(value: str) -> Any: # pyright: ignore[reportAny] - try: - return json.loads(value) # pyright: ignore[reportAny] - except Exception: - pass - - try: - return ast.literal_eval(value) # pyright: ignore[reportAny] - except Exception: - pass - return value - - def parse_tool_call(text: str, tools: Any | None = None): - func_name_match = _func_name_regex.search(text) - if func_name_match is None: - raise ValueError(f"Could not parse function name from tool call: {text!r}") - original_func_name = func_name_match.group(1) - tool_id = func_name_match.group(2) - # strip off the `functions.` prefix, if it exists. - func_name = original_func_name[original_func_name.find(".") + 1 :] - - func_args_match = _func_arg_regex.search(text) - if func_args_match is None: - raise ValueError(f"Could not parse function args from tool call: {text!r}") - func_args = func_args_match.group(1) - # the args should be valid json - no need to check against our tools to deserialize - arg_dct = _deserialize(func_args) # pyright: ignore[reportAny] - - return dict( - id=f"{original_func_name}:{tool_id}", - name=func_name, - arguments=arg_dct, # pyright: ignore[reportAny] - ) - - tokenizer._tool_call_start = tool_call_start - tokenizer._tool_call_end = tool_call_end - tokenizer._tool_parser = parse_tool_call - - -def patch_glm_tokenizer(tokenizer: TokenizerWrapper): - """ - Fixed version of mlx_lm's glm47 tool parser that handles regex match failures. - """ - import ast - import json - from typing import Any - - import regex as re - - _func_name_regex = re.compile(r"^(.*?)", re.DOTALL) - _func_arg_regex = re.compile( - r"(.*?)(?:\n|\s)*(.*?)(?:|(?=)|$)", - re.DOTALL, - ) - - tool_call_start = "" - tool_call_end = "" - - def _is_string_type( - tool_name: str, - arg_name: str, - tools: list[Any] | None, - ) -> bool: - if tools is None: - return False - for tool in tools: # pyright: ignore[reportAny] - func = tool["function"] # pyright: ignore[reportAny] - if func["name"] == tool_name: - params = func["parameters"] # pyright: ignore[reportAny] - if params is None: - return False - props = params.get("properties", {}) # pyright: ignore[reportAny] - arg_props = props.get(arg_name, {}) # pyright: ignore[reportAny] - arg_type = arg_props.get("type", None) # pyright: ignore[reportAny] - return arg_type == "string" # pyright: ignore[reportAny] - return False - - def _deserialize(value: str) -> Any: # pyright: ignore[reportAny] - try: - return json.loads(value) # pyright: ignore[reportAny] - except Exception: - pass - try: - return ast.literal_eval(value) # pyright: ignore[reportAny] - except Exception: - pass - return value - - def parse_tool_call(text: str, tools: list[Any] | None = None): - func_name_match = _func_name_regex.search(text) - if func_name_match is None: - raise ValueError(f"Could not parse function name from tool call: {text!r}") - func_name = func_name_match.group(1) - - pairs = _func_arg_regex.findall(text) - arg_dct: dict[str, Any] = {} - for key, value in pairs: # pyright: ignore[reportAny] - arg_key = key.strip() # pyright: ignore[reportAny] - arg_val = value.strip() # pyright: ignore[reportAny] - if not _is_string_type(func_name, arg_key, tools): # pyright: ignore[reportAny] - arg_val = _deserialize(arg_val) # pyright: ignore[reportAny] - arg_dct[arg_key] = arg_val - return dict(name=func_name, arguments=arg_dct) - - tokenizer._tool_call_start = tool_call_start - tokenizer._tool_call_end = tool_call_end - tokenizer._tool_parser = parse_tool_call - - -def _validate_single_tool(obj: dict[str, Any]) -> ToolCallItem: - if ( - ((name := obj.get("name")) is not None) - and ((args := obj.get("arguments")) is not None) - and isinstance(name, str) - ): - raw_id: object = obj.get("id") - extra = {"id": str(raw_id)} if raw_id is not None else {} - return ToolCallItem( - **extra, - name=name, - arguments=json.dumps(args), - ) - else: - raise ValidationError - - EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL" EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM" EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT" diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 5d39a881..e8a06a77 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -100,10 +100,10 @@ class RunnerSupervisor: logger.info("Runner supervisor shutting down") self._ev_recv.close() self._task_sender.close() - self._event_sender.close() - self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK")) + with contextlib.suppress(ClosedResourceError): + self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK")) self._cancel_sender.close() - self.runner_process.join(1) + self.runner_process.join(5) if not self.runner_process.is_alive(): logger.info("Runner process succesfully terminated") return @@ -180,6 +180,7 @@ class RunnerSupervisor: await self._check_runner(e) for tid in self.pending: self.pending[tid].set() + self._event_sender.close() def __del__(self) -> None: if self.runner_process.is_alive(): @@ -191,7 +192,7 @@ class RunnerSupervisor: logger.info("Checking runner's status") if self.runner_process.is_alive(): logger.info("Runner was found to be alive, attempting to join process") - await to_thread.run_sync(self.runner_process.join, 1) + await to_thread.run_sync(self.runner_process.join, 5) rc = self.runner_process.exitcode logger.info(f"RunnerSupervisor exited with exit code {rc}") if rc == 0: @@ -208,10 +209,15 @@ class RunnerSupervisor: logger.opt(exception=e).error(f"Runner terminated ({cause})") - await self._event_sender.send( - RunnerStatusUpdated( - runner_id=self.bound_instance.bound_runner_id, - runner_status=RunnerFailed(error_message=f"Terminated ({cause})"), + try: + await self._event_sender.send( + RunnerStatusUpdated( + runner_id=self.bound_instance.bound_runner_id, + runner_status=RunnerFailed(error_message=f"Terminated ({cause})"), + ) + ) + except (ClosedResourceError, BrokenResourceError): + logger.warning( + "Event sender already closed, unable to report runner failure" ) - ) self.shutdown() diff --git a/src/exo/worker/runner/tool_parsers.py b/src/exo/worker/runner/tool_parsers.py new file mode 100644 index 00000000..88cbf0be --- /dev/null +++ b/src/exo/worker/runner/tool_parsers.py @@ -0,0 +1,72 @@ +import json +from dataclasses import dataclass +from typing import Any, Callable + +from exo.shared.types.api import ToolCallItem + + +@dataclass +class ToolParser: + start_parsing: str + end_parsing: str + parse_tool_calls: Callable[[str], list[ToolCallItem] | None] + + +def make_mlx_parser( + tool_call_start: str, + tool_call_end: str, + tool_parser: Callable[[str], dict[str, Any] | list[dict[str, Any]]], +) -> ToolParser: + def parse_tool_calls(text: str) -> list[ToolCallItem] | None: + try: + text = text.removeprefix(tool_call_start) + text = text.removesuffix(tool_call_end) + parsed = tool_parser(text) + if isinstance(parsed, list): + return [ToolCallItem.model_validate(_flatten(p)) for p in parsed] + else: + return [ToolCallItem.model_validate(_flatten(parsed))] + + except Exception: + return None + + return ToolParser( + start_parsing=tool_call_start, + end_parsing=tool_call_end, + parse_tool_calls=parse_tool_calls, + ) + + +# TODO / example code: +def _parse_json_calls(text: str) -> list[ToolCallItem] | None: + try: + text = text.removeprefix("") + text = text.removesuffix("") + top_level = { + k: json.dumps(v) if isinstance(v, (dict, list)) else v + for k, v in json.loads(text).items() # pyright: ignore[reportAny] + } + return [ToolCallItem.model_validate(top_level)] + except Exception: + return None + + +def _flatten(p: dict[str, Any]) -> dict[str, str]: + return { + k: json.dumps(v) if isinstance(v, (dict, list)) else str(v) # pyright: ignore[reportAny] + for k, v in p.items() # pyright: ignore[reportAny] + } + + +json_tool_parser = ToolParser( + start_parsing="", + end_parsing="", + parse_tool_calls=_parse_json_calls, +) + + +def infer_tool_parser(chat_template: str) -> ToolParser | None: + """Attempt to auto-infer a tool parser from the chat template.""" + if "" in chat_template and "tool_call.name" in chat_template: + return json_tool_parser + return None diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py index 9e897141..2f9e1fd3 100644 --- a/src/exo/worker/tests/unittests/test_mlx/conftest.py +++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py @@ -123,7 +123,12 @@ def run_gpt_oss_pipeline_device( generated_text = "" for response in mlx_generate( - model=model, tokenizer=tokenizer, task=task, prompt=prompt + model=model, + tokenizer=tokenizer, + task=task, + prompt=prompt, + kv_prefix_cache=None, + group=group, ): generated_text += response.text if response.finish_reason is not None: @@ -194,6 +199,8 @@ def run_gpt_oss_tensor_parallel_device( tokenizer=tokenizer, task=task, prompt=prompt, + kv_prefix_cache=None, + group=group, ): generated_text += response.text if response.finish_reason is not None: diff --git a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py index 2f360166..5c1893e1 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py @@ -88,12 +88,12 @@ class TestKVPrefix: return tokenizer def test_starts_empty(self, mock_tokenizer): - cache = KVPrefixCache() + cache = KVPrefixCache(None) assert len(cache.prompts) == 0 assert len(cache.caches) == 0 def test_clear_empties_cache(self, mock_tokenizer): - cache = KVPrefixCache() + cache = KVPrefixCache(None) cache.prompts.append(mx.array([1, 2, 3])) cache.caches.append([KVCache()]) cache.clear() @@ -101,7 +101,7 @@ class TestKVPrefix: assert len(cache.caches) == 0 def test_clear_on_empty_cache(self, mock_tokenizer): - cache = KVPrefixCache() + cache = KVPrefixCache(None) cache.clear() assert len(cache.prompts) == 0 @@ -142,7 +142,9 @@ class TestKVPrefixCacheWithModel: tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) - _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache) + _, _, snapshots = prefill( + model, tokenizer, make_sampler(0.0), tokens, cache, group=None + ) # Cache should now hold the prompt tokens minus one assert cache_length(cache) == len(tokens) - 1 @@ -161,9 +163,11 @@ class TestKVPrefixCacheWithModel: tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) - _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache) + _, _, snapshots = prefill( + model, tokenizer, make_sampler(0.0), tokens, cache, group=None + ) - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(tokens, cache, snapshots) assert len(kv_prefix_cache.prompts) == 1 @@ -176,9 +180,11 @@ class TestKVPrefixCacheWithModel: ) assert matched_index == 0 - # Exact match returns only last token - assert len(remaining_tokens) == 1 - assert mx.array_equal(remaining_tokens, tokens[-1:]) + # Exact match returns last token(s) — for models with SSM/rotating caches, + # snapshot availability constrains how far back we can trim, so remaining + # may be 1 or 2 tokens depending on the model. + assert len(remaining_tokens) >= 1 + assert mx.array_equal(remaining_tokens, tokens[-len(remaining_tokens) :]) def test_add_and_get_prefix_match(self, model_and_tokenizer): """get_kv_cache with a longer prompt sharing prefix should return partial match.""" @@ -194,10 +200,10 @@ class TestKVPrefixCacheWithModel: cache = make_kv_cache(model) _, _, snapshots = prefill( - model, tokenizer, make_sampler(0.0), short_tokens, cache + model, tokenizer, make_sampler(0.0), short_tokens, cache, group=None ) - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(short_tokens, cache, snapshots) # Query with longer prompt that shares the chat template prefix @@ -238,9 +244,11 @@ class TestKVPrefixCacheWithModel: tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) - _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache) + _, _, snapshots = prefill( + model, tokenizer, make_sampler(0.0), tokens, cache, group=None + ) - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(tokens, cache, snapshots) stored_length = cache_length(kv_prefix_cache.caches[0]) @@ -276,9 +284,11 @@ class TestKVPrefixCacheWithModel: tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) - _, _, snapshots = prefill(model, tokenizer, make_sampler(0.0), tokens, cache) + _, _, snapshots = prefill( + model, tokenizer, make_sampler(0.0), tokens, cache, group=None + ) - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) kv_prefix_cache.add_kv_cache(tokens, cache, snapshots) stored_length = cache_length(kv_prefix_cache.caches[0]) @@ -301,7 +311,7 @@ class TestKVPrefixCacheWithModel: """mlx_generate should save the cache after generation completes.""" model, tokenizer = model_and_tokenizer - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Hello")], @@ -318,6 +328,7 @@ class TestKVPrefixCacheWithModel: task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, + group=None, ): generated_tokens += 1 @@ -331,7 +342,7 @@ class TestKVPrefixCacheWithModel: """Second mlx_generate call with same prompt should get a prefix hit from stored cache.""" model, tokenizer = model_and_tokenizer - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Reuse test")], @@ -347,6 +358,7 @@ class TestKVPrefixCacheWithModel: task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, + group=None, ): pass @@ -368,7 +380,7 @@ class TestKVPrefixCacheWithModel: """With a prompt > 1000 tokens, second generation should update the cache entry in-place.""" model, tokenizer = model_and_tokenizer - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE base_text = "The quick brown fox jumps over the lazy dog. " @@ -395,6 +407,7 @@ class TestKVPrefixCacheWithModel: task=task1, prompt=prompt1, kv_prefix_cache=kv_prefix_cache, + group=None, ): pass first_gen_time = time.perf_counter() - t0 @@ -427,6 +440,7 @@ class TestKVPrefixCacheWithModel: task=task2, prompt=prompt2, kv_prefix_cache=kv_prefix_cache, + group=None, ): pass second_gen_time = time.perf_counter() - t0 @@ -447,7 +461,7 @@ class TestKVPrefixCacheWithModel: """After mlx_generate saves a cache, a second generation must not corrupt the stored copy.""" model, tokenizer = model_and_tokenizer - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) task = TextGenerationTaskParams( model=DEFAULT_GPT_OSS_MODEL_ID, input=[InputMessage(role="user", content="Immutable test")], @@ -462,6 +476,7 @@ class TestKVPrefixCacheWithModel: task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, + group=None, ): pass @@ -474,6 +489,7 @@ class TestKVPrefixCacheWithModel: task=task, prompt=prompt, kv_prefix_cache=kv_prefix_cache, + group=None, ): pass @@ -484,7 +500,7 @@ class TestKVPrefixCacheWithModel: """Under memory pressure, adding a new cache entry evicts the least recently used one.""" model, tokenizer = model_and_tokenizer - kv_prefix_cache = KVPrefixCache() + kv_prefix_cache = KVPrefixCache(None) # Add three cache entries with different prompts prompts = ["First entry", "Second entry", "Third entry"] @@ -497,7 +513,7 @@ class TestKVPrefixCacheWithModel: prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) - prefill(model, tokenizer, make_sampler(0.0), tokens, cache) + prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=None) kv_prefix_cache.add_kv_cache(tokens, cache) # Stagger _last_used so LRU order is deterministic kv_prefix_cache._last_used[i] = float(i) @@ -522,7 +538,7 @@ class TestKVPrefixCacheWithModel: prompt = apply_chat_template(tokenizer, task) tokens = encode_prompt(tokenizer, prompt) cache = make_kv_cache(model) - prefill(model, tokenizer, make_sampler(0.0), tokens, cache) + prefill(model, tokenizer, make_sampler(0.0), tokens, cache, group=None) kv_prefix_cache.add_kv_cache(tokens, cache) # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py new file mode 100644 index 00000000..5f97c8b2 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py @@ -0,0 +1,297 @@ +import copy +import gc +import importlib +import json +import shutil +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast + +import mlx.core as mx +import mlx.nn as nn +import pytest +from mlx.utils import tree_flatten, tree_unflatten +from mlx_lm.tokenizer_utils import TokenizerWrapper + +from exo.shared.types.common import ModelId +from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams +from exo.worker.engines.mlx import Model +from exo.worker.engines.mlx.cache import KVPrefixCache +from exo.worker.engines.mlx.generator.generate import mlx_generate +from exo.worker.engines.mlx.utils_mlx import ( + apply_chat_template, + load_tokenizer_for_model_id, +) + +HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub" + +# ── Config reduction ──────────────────────────────────────────────────────── # + +_REDUCE = { + "num_hidden_layers": 4, + "hidden_size": 256, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "intermediate_size": 512, + "moe_intermediate_size": 128, + "num_experts": 4, + "num_experts_per_tok": 2, + "n_routed_experts": 4, + "num_local_experts": 4, + "num_nextn_predict_layers": 0, + "first_k_dense_replace": 0, + "linear_num_key_heads": 2, + "linear_num_value_heads": 2, + "num_attention_groups": 4, +} + + +def _reduce_dict(cfg: dict[str, Any]) -> dict[str, Any]: + result = dict(cfg) + for key, val in _REDUCE.items(): + if key in result: + result[key] = val + return result + + +def _reduce_config(cfg: dict[str, Any]) -> dict[str, Any]: + result = _reduce_dict(cfg) + n_layers = cast(int, result.get("num_hidden_layers", 4)) + + if "text_config" in result and isinstance(result["text_config"], dict): + result["text_config"] = _reduce_dict( + cast(dict[str, Any], result["text_config"]) + ) + tc: dict[str, Any] = result["text_config"] + if "num_nextn_predict_layers" in tc: + tc["num_nextn_predict_layers"] = 0 + + if "layer_types" in result and isinstance(result["layer_types"], list): + result["layer_types"] = result["layer_types"][:n_layers] + + if "attention_other_setting" in result and isinstance( + result["attention_other_setting"], dict + ): + aos: dict[str, Any] = dict( + cast(dict[str, Any], result["attention_other_setting"]) + ) + if "num_attention_heads" in aos: + aos["num_attention_heads"] = result.get("num_attention_heads", 4) + if "num_attention_groups" in aos: + aos["num_attention_groups"] = result.get( + "num_attention_groups", cast(int, aos["num_attention_groups"]) + ) + result["attention_other_setting"] = aos + + if "moe_layers_enum" in result and isinstance(result["moe_layers_enum"], str): + indices = [int(x) for x in result["moe_layers_enum"].split(",") if x.strip()] + valid = [i for i in indices if i < n_layers] + result["moe_layers_enum"] = ",".join(str(i) for i in valid) if valid else "" + + return result + + +# ── Helpers ───────────────────────────────────────────────────────────────── # + + +def _find_snapshot(hub_name: str) -> Path | None: + model_dir = HF_CACHE / f"models--mlx-community--{hub_name}" + snaps = model_dir / "snapshots" + if not snaps.exists(): + return None + children = sorted(snaps.iterdir()) + return children[0] if children else None + + +def _copy_tokenizer(src: Path, dst: Path) -> None: + for f in src.iterdir(): + name = f.name + if ( + "tokeniz" in name.lower() + or "tiktoken" in name.lower() + or name.startswith("vocab") + or name.endswith(".jinja") + or "tool_declaration" in name + ) and f.is_file(): + shutil.copy2(f, dst / name) + + +def _build_model(module_name: str, cfg: dict[str, Any]) -> Model: + mod = importlib.import_module(f"mlx_lm.models.{module_name}") + args = mod.ModelArgs.from_dict(cfg) # pyright: ignore[reportAny] + model: nn.Module = mod.Model(args) # pyright: ignore[reportAny] + flat = cast(list[tuple[str, mx.array]], tree_flatten(model.parameters())) + random_weights = [ + (k, mx.random.normal(shape=v.shape, dtype=mx.float16)) for k, v in flat + ] + model.update(cast(dict[str, Any], tree_unflatten(random_weights))) + mx.eval(model.parameters()) + return cast(Model, model) + + +def _collect_tokens( + model: Model, + tokenizer: TokenizerWrapper, + task: TextGenerationTaskParams, + prompt: str, + kv_prefix_cache: KVPrefixCache | None, +) -> list[int]: + tokens: list[int] = [] + for resp in mlx_generate( + model=model, + tokenizer=tokenizer, + task=task, + prompt=prompt, + kv_prefix_cache=kv_prefix_cache, + group=None, + ): + tokens.append(resp.token) + if resp.finish_reason is not None: + break + return tokens + + +# ── Architecture definitions ──────────────────────────────────────────────── # + + +@dataclass(frozen=True) +class ArchSpec: + name: str + hub_name: str + module: str + tokenizer_hub: str | None = None # fallback for models without bundled tokenizer + + +ARCHITECTURES: list[ArchSpec] = [ + ArchSpec("llama", "Llama-3.2-1B-Instruct-4bit", "llama"), + ArchSpec("glm_moe_dsa", "GLM-5-MXFP4-Q8", "glm_moe_dsa"), + ArchSpec( + "glm4_moe", "GLM-4.5-Air-8bit", "glm4_moe", tokenizer_hub="GLM-4.7-8bit-gs32" + ), + ArchSpec( + "glm4_moe_lite", + "GLM-4.7-Flash-8bit", + "glm4_moe_lite", + tokenizer_hub="GLM-4.7-8bit-gs32", + ), + ArchSpec("glm4_moe_47", "GLM-4.7-8bit-gs32", "glm4_moe"), + ArchSpec("qwen3", "Qwen3-4B-Instruct-2507-4bit", "qwen3"), + ArchSpec("qwen3_moe", "Qwen3-30B-A3B-4bit", "qwen3_moe"), + ArchSpec("qwen3_next", "Qwen3-Next-80B-A3B-Thinking-4bit", "qwen3_next"), + ArchSpec("minimax", "MiniMax-M2.1-3bit", "minimax"), + ArchSpec("gpt_oss", "gpt-oss-20b-MXFP4-Q8", "gpt_oss"), + ArchSpec("step3p5", "Step-3.5-Flash-4bit", "step3p5"), + ArchSpec("kimi_k25", "Kimi-K2.5", "kimi_k25"), +] + + +def _arch_available(spec: ArchSpec) -> bool: + snap = _find_snapshot(spec.hub_name) + if snap is None: + return False + if spec.tokenizer_hub is not None: + return _find_snapshot(spec.tokenizer_hub) is not None + return True + + +def _make_task() -> TextGenerationTaskParams: + return TextGenerationTaskParams( + model=ModelId("test"), + input=[ + InputMessage( + role="user", + content="Use the calculator to compute 1847 * 263 + 5921", + ) + ], + max_output_tokens=20, + temperature=0.0, + tools=[ + { + "type": "function", + "function": { + "name": "calculate", + "description": "Evaluate a mathematical expression", + "parameters": { + "type": "object", + "properties": {"expression": {"type": "string"}}, + "required": ["expression"], + }, + }, + } + ], + ) + + +# ── Test class ────────────────────────────────────────────────────────────── # + + +@pytest.mark.slow +class TestPrefixCacheArchitectures: + """Verify prefix cache produces identical output to fresh generation for every architecture.""" + + @pytest.fixture(autouse=True) + def _cleanup(self): + yield + mx.clear_cache() + gc.collect() + + @pytest.mark.parametrize( + "spec", + ARCHITECTURES, + ids=[a.name for a in ARCHITECTURES], + ) + def test_prefix_cache_exact_hit(self, spec: ArchSpec) -> None: + if not _arch_available(spec): + pytest.skip(f"Model {spec.hub_name} not cached locally") + + snapshot = _find_snapshot(spec.hub_name) + assert snapshot is not None + + tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_test_{spec.name}_")) + try: + # Build reduced config + with open(snapshot / "config.json") as f: + cfg = cast(dict[str, Any], json.load(f)) + reduced = _reduce_config(copy.deepcopy(cfg)) + (tmpdir / "config.json").write_text(json.dumps(reduced)) + + # Copy tokenizer + tok_src = snapshot + if spec.tokenizer_hub is not None: + alt = _find_snapshot(spec.tokenizer_hub) + if alt is not None: + tok_src = alt + _copy_tokenizer(tok_src, tmpdir) + + # Load tokenizer and model + model_id = ModelId(f"mlx-community/{spec.hub_name}") + tokenizer = load_tokenizer_for_model_id(model_id, tmpdir) + mx.random.seed(0) + model = _build_model(spec.module, reduced) + + task = _make_task() + prompt = apply_chat_template(tokenizer=tokenizer, task_params=task) + + # Run 1: fresh + mx.random.seed(42) + fresh = _collect_tokens(model, tokenizer, task, prompt, None) + assert len(fresh) > 0, "Fresh generation produced no tokens" + + # Run 2: populate cache + kv = KVPrefixCache(None) + mx.random.seed(42) + populate = _collect_tokens(model, tokenizer, task, prompt, kv) + + # Run 3: exact cache hit + mx.random.seed(42) + cached = _collect_tokens(model, tokenizer, task, prompt, kv) + + assert fresh == populate, ( + f"Fresh vs populate mismatch: {fresh[:5]} vs {populate[:5]}" + ) + assert fresh == cached, ( + f"Fresh vs cached mismatch: {fresh[:5]} vs {cached[:5]}" + ) + finally: + shutil.rmtree(tmpdir, ignore_errors=True) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py index a6f406a5..ce06ab86 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py +++ b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py @@ -343,8 +343,16 @@ async def test_kimi_tokenizer_specifically(): @pytest.mark.asyncio async def test_glm_tokenizer_specifically(): """Test GLM tokenizer with its specific EOS tokens.""" + + def contains(card: ModelCard, x: str): + return x in card.model_id.lower() + glm_model_cards = [ - card for card in await get_model_cards() if "glm" in card.model_id.lower() + card + for card in await get_model_cards() + if contains(card, "glm") + and not contains(card, "-5") + and not contains(card, "4.7") ] if not glm_model_cards: diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py index 9c318517..abcb4939 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py +++ b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py @@ -90,14 +90,10 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting(): global_download_status = { NODE_A: [ - DownloadCompleted( - shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory() - ) + DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) ], NODE_B: [ - DownloadCompleted( - shard_metadata=shard2, node_id=NODE_B, total_bytes=Memory() - ) + DownloadCompleted(shard_metadata=shard2, node_id=NODE_B, total=Memory()) ], } @@ -138,9 +134,7 @@ def test_plan_does_not_request_download_when_shard_already_downloaded(): # Global state shows shard is downloaded for NODE_A global_download_status: dict[NodeId, list[DownloadProgress]] = { NODE_A: [ - DownloadCompleted( - shard_metadata=shard, node_id=NODE_A, total_bytes=Memory() - ) + DownloadCompleted(shard_metadata=shard, node_id=NODE_A, total=Memory()) ], NODE_B: [], } @@ -187,9 +181,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally(): global_download_status = { NODE_A: [ - DownloadCompleted( - shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory() - ) + DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) ], NODE_B: [], # NODE_B has no downloads completed yet } @@ -207,14 +199,10 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally(): global_download_status = { NODE_A: [ - DownloadCompleted( - shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory() - ) + DownloadCompleted(shard_metadata=shard1, node_id=NODE_A, total=Memory()) ], NODE_B: [ - DownloadCompleted( - shard_metadata=shard2, node_id=NODE_B, total_bytes=Memory() - ) + DownloadCompleted(shard_metadata=shard2, node_id=NODE_B, total=Memory()) ], # NODE_B has no downloads completed yet } diff --git a/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py b/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py new file mode 100644 index 00000000..a59383e5 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py @@ -0,0 +1,967 @@ +import json +from collections.abc import Generator +from typing import Any + +from exo.shared.types.worker.runner_response import ( + GenerationResponse, + ToolCallResponse, +) +from exo.worker.engines.mlx.dsml_encoding import ( + ASSISTANT_TOKEN, + BOS_TOKEN, + DSML_TOKEN, + EOS_TOKEN, + THINKING_END, + THINKING_START, + TOOL_CALLS_END, + TOOL_CALLS_START, + USER_TOKEN, + encode_messages, + parse_dsml_output, +) +from exo.worker.runner.runner import parse_deepseek_v32 + +# ── Shared fixtures ────────────────────────────────────────────── + +_WEATHER_TOOLS: list[dict[str, Any]] = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a given city", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string", "description": "The city name"}, + "units": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + "description": "Temperature units", + }, + }, + "required": ["city"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "description": "Get the current time in a timezone", + "parameters": { + "type": "object", + "properties": { + "timezone": {"type": "string"}, + }, + "required": ["timezone"], + }, + }, + }, +] + + +def _simulate_tokens( + texts: list[str], + finish_on_last: bool = True, +) -> Generator[GenerationResponse]: + """Simulate a model producing tokens from a list of text strings.""" + for i, text in enumerate(texts): + is_last = i == len(texts) - 1 + yield GenerationResponse( + text=text, + token=i, + finish_reason="stop" if (is_last and finish_on_last) else None, + usage=None, + ) + + +# ── Test: Standard text response (no tool calls) ──────────────── + + +class TestE2EStandardResponse: + """Model generates a plain text response — no tool calling involved.""" + + def test_plain_text_passthrough(self): + """Simulate model producing: 'The weather in NYC is 72°F and sunny.'""" + # Step 1: Encode the prompt (with tools available) + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What's the weather in NYC?"}, + ] + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + + # Verify prompt structure + assert BOS_TOKEN in prompt + assert "## Tools" in prompt + assert "get_weather" in prompt + assert f"{USER_TOKEN}What's the weather in NYC?{ASSISTANT_TOKEN}" in prompt + + # Step 2: Simulate model response — plain text tokens (no DSML) + model_tokens = [ + "The weather", + " in NYC", + " is 72", + "°F", + " and sunny", + ".", + ] + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + + # Step 3: Verify all tokens pass through as GenerationResponse + gen_results = [r for r in results if isinstance(r, GenerationResponse)] + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 0 + assert len(gen_results) == 6 + full_text = "".join(r.text for r in gen_results) + assert full_text == "The weather in NYC is 72°F and sunny." + assert gen_results[-1].finish_reason == "stop" + + +# ── Test: Tool call response ───────────────────────────────────── + + +class TestE2EToolCallResponse: + """Model generates a DSML tool call — realistic token boundaries.""" + + def test_realistic_tool_call_tokens(self): + """Simulate model generating a get_weather tool call with realistic token splits. + + Real models split DSML markers across tokens unpredictably. + This simulates how DeepSeek V3.2 actually tokenizes DSML output. + """ + # Step 1: Encode prompt + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What's the weather in San Francisco?"}, + ] + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + assert "get_weather" in prompt + + # Step 2: Simulate realistic token-by-token model output + # The model first produces some text, then a DSML tool call block + model_tokens = [ + "I'll check the weather for you.", + "\n\n", + f"<{DSML_TOKEN}", # marker split across tokens + "function_calls>\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">', + "San Francisco", + f"\n", + f'<{DSML_TOKEN}parameter name="units" string="false">', + '"celsius"', + f"\n", + f"\n", + f"", + ] + + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + + # Step 3: Verify + gen_results = [r for r in results if isinstance(r, GenerationResponse)] + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + # Should have text tokens before tool call + one ToolCallResponse + assert len(tool_results) == 1 + assert len(tool_results[0].tool_calls) == 1 + + tc = tool_results[0].tool_calls[0] + assert tc.name == "get_weather" + args = json.loads(tc.arguments) # pyright: ignore[reportAny] + assert args["city"] == "San Francisco" + assert args["units"] == "celsius" + + # The text before the tool call should still be yielded + text_before = "".join(r.text for r in gen_results if not r.is_thinking) + assert "check the weather" in text_before + + def test_multiple_tool_calls_in_one_block(self): + """Model generates two tool calls in a single function_calls block.""" + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Weather in NYC and time in EST?"}, + ] + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + assert "get_weather" in prompt + assert "get_time" in prompt + + # Simulate model output with two invocations + model_tokens = [ + "Let me check both.\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">NYC\n', + f"\n", + f'<{DSML_TOKEN}invoke name="get_time">\n', + f'<{DSML_TOKEN}parameter name="timezone" string="true">EST\n', + f"\n", + TOOL_CALLS_END, + ] + + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 1 + assert len(tool_results[0].tool_calls) == 2 + assert tool_results[0].tool_calls[0].name == "get_weather" + assert tool_results[0].tool_calls[1].name == "get_time" + + args0 = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] + args1 = json.loads(tool_results[0].tool_calls[1].arguments) # pyright: ignore[reportAny] + assert args0 == {"city": "NYC"} + assert args1 == {"timezone": "EST"} + + +# ── Test: Multi-turn tool use flow ─────────────────────────────── + + +class TestE2EMultiTurnToolUse: + """Full multi-turn: user asks → model calls tool → tool result → model answers.""" + + def test_encode_multi_turn_with_tool_results(self): + """Verify the prompt for turn 2 (after tool results) is correctly encoded.""" + # Turn 1: user asks, model calls tool + # Turn 2: tool result provided, model answers + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are a weather assistant."}, + {"role": "user", "content": "What's the weather in NYC?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "NYC"}', + }, + } + ], + }, + {"role": "tool", "content": '{"temperature": 72, "condition": "sunny"}'}, + ] + + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + + # Verify multi-turn structure + assert BOS_TOKEN in prompt + assert "You are a weather assistant." in prompt + assert "## Tools" in prompt + + # The assistant's tool call should be encoded as DSML + assert TOOL_CALLS_START in prompt + assert f'<{DSML_TOKEN}invoke name="get_weather">' in prompt + assert EOS_TOKEN in prompt + + # The tool result should be wrapped in function_results + assert "" in prompt + assert "" in prompt + assert "72" in prompt + assert "" in prompt + + # Now simulate model answering after seeing the tool result + model_tokens = [ + "The current", + " weather in NYC", + " is 72°F", + " and sunny.", + ] + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + + gen_results = [r for r in results if isinstance(r, GenerationResponse)] + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 0 + full_text = "".join(r.text for r in gen_results) + assert full_text == "The current weather in NYC is 72°F and sunny." + + def test_multi_tool_results_encoding(self): + """Verify encoding when model called two tools and both return results.""" + messages: list[dict[str, Any]] = [ + {"role": "user", "content": "Weather and time?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "LA"}', + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "arguments": '{"timezone": "PST"}', + }, + }, + ], + }, + {"role": "tool", "content": "85F, clear skies"}, + {"role": "tool", "content": "3:42 PM PST"}, + ] + + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + + # Should have one function_results block with two results + assert prompt.count("") == 1 + assert prompt.count("") == 1 + assert "85F, clear skies" in prompt + assert "3:42 PM PST" in prompt + + +# ── Test: Thinking + tool call ─────────────────────────────────── + + +class TestE2EThinkingAndToolCall: + """Model uses thinking mode, reasons, then makes a tool call.""" + + def test_thinking_then_tool_call(self): + """Model thinks first, then produces a DSML tool call block.""" + messages: list[dict[str, Any]] = [ + {"role": "user", "content": "What's the weather?"}, + ] + prompt = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + # Thinking mode: prompt should end with + assert prompt.endswith(THINKING_START) + + # Simulate: model outputs , thinks, closes thinking, then tool call. + # In the full pipeline, parse_thinking_models handles the case where + # is in the prompt. Here we test parse_deepseek_v32 directly, + # which detects / markers in the stream. + model_tokens = [ + THINKING_START, + "The user wants weather", + " information. I should use", + " the get_weather tool.", + THINKING_END, + "\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">', + "San Francisco", + f"\n", + f"\n", + TOOL_CALLS_END, + ] + + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + + gen_results = [r for r in results if isinstance(r, GenerationResponse)] + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + # Should have thinking tokens + tool call + thinking_results = [r for r in gen_results if r.is_thinking] + + assert len(thinking_results) >= 1 + thinking_text = "".join(r.text for r in thinking_results) + assert "get_weather tool" in thinking_text + + assert len(tool_results) == 1 + assert tool_results[0].tool_calls[0].name == "get_weather" + args = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] + assert args["city"] == "San Francisco" + + def test_thinking_prompt_encoding(self): + """Verify thinking mode affects prompt encoding correctly.""" + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "Be thorough."}, + {"role": "user", "content": "What's the weather?"}, + ] + + # With thinking enabled + prompt_think = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + assert prompt_think.endswith(THINKING_START) + + # With thinking disabled + prompt_no_think = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="chat" + ) + assert prompt_no_think.endswith(THINKING_END) + + # Both should have the same tool definitions + assert "get_weather" in prompt_think + assert "get_weather" in prompt_no_think + + +# ── Test: Round-trip encode → parse ────────────────────────────── + + +class TestE2ERoundTrip: + """Verify that DSML we encode can be parsed back correctly.""" + + def test_encoded_tool_call_is_parseable(self): + """Encode an assistant tool call message, then parse the DSML output.""" + messages: list[dict[str, Any]] = [ + {"role": "user", "content": "Weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Tokyo", "units": "celsius"}', + }, + } + ], + }, + ] + + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + + # Extract the DSML function_calls block from the prompt + start = prompt.index(TOOL_CALLS_START) + end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END) + dsml_block = prompt[start:end] + + # Parse it back + parsed = parse_dsml_output(dsml_block) + assert parsed is not None + assert len(parsed) == 1 + assert parsed[0].name == "get_weather" + args = json.loads(parsed[0].arguments) # pyright: ignore[reportAny] + assert args["city"] == "Tokyo" + assert args["units"] == "celsius" + + def test_encoded_multi_tool_call_round_trips(self): + """Encode multiple tool calls, verify they parse back correctly.""" + messages: list[dict[str, Any]] = [ + {"role": "user", "content": "Both please"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "arguments": '{"timezone": "CET"}', + }, + }, + ], + }, + ] + + prompt = encode_messages(messages, thinking_mode="chat", tools=_WEATHER_TOOLS) + + start = prompt.index(TOOL_CALLS_START) + end = prompt.index(TOOL_CALLS_END) + len(TOOL_CALLS_END) + dsml_block = prompt[start:end] + + parsed = parse_dsml_output(dsml_block) + assert parsed is not None + assert len(parsed) == 2 + assert parsed[0].name == "get_weather" + assert parsed[1].name == "get_time" + assert json.loads(parsed[0].arguments) == {"city": "Paris"} + assert json.loads(parsed[1].arguments) == {"timezone": "CET"} + + +# ── Test: Edge cases with realistic token boundaries ───────────── + + +class TestE2EEdgeCases: + """Edge cases that occur in real model inference.""" + + def test_dsml_marker_split_at_fullwidth_pipe(self): + """The fullwidth pipe character ｜ might be its own token.""" + # This is a realistic tokenization: the DSML marker is split at the ｜ chars + model_tokens = [ + "Let me help.\n\n", + "<\uff5c", # start of ｜DSML｜ + "DSML\uff5c", # rest of DSML token + "function_calls>\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">NYC\n', + f"\n", + TOOL_CALLS_END, + ] + + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 1 + assert tool_results[0].tool_calls[0].name == "get_weather" + + def test_tool_call_with_nested_json_object(self): + """Model passes a complex JSON object as a non-string parameter.""" + dsml_block = ( + f"{TOOL_CALLS_START}\n" + f'<{DSML_TOKEN}invoke name="create_event">\n' + f'<{DSML_TOKEN}parameter name="title" string="true">Team Standup\n' + f'<{DSML_TOKEN}parameter name="config" string="false">' + f'{{"recurring": true, "days": ["mon", "wed", "fri"], "time": "09:00"}}' + f"\n" + f"\n" + f"{TOOL_CALLS_END}" + ) + + # Feed as single token (model might produce it all at once after prefill) + results = list(parse_deepseek_v32(_simulate_tokens([dsml_block]))) + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 1 + tc = tool_results[0].tool_calls[0] + assert tc.name == "create_event" + args = json.loads(tc.arguments) # pyright: ignore[reportAny] + assert args["title"] == "Team Standup" + assert args["config"]["recurring"] is True + assert args["config"]["days"] == ["mon", "wed", "fri"] + + def test_text_with_angle_brackets_not_mistaken_for_dsml(self): + """Angle brackets in normal text should not trigger DSML buffering.""" + model_tokens = [ + "The formula is ", + "", + " where x > 0", + " and y < 100.", + ] + + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + gen_results = [r for r in results if isinstance(r, GenerationResponse)] + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 0 + full_text = "".join(r.text for r in gen_results) + assert "formula" in full_text + assert "" in full_text + + def test_empty_model_response(self): + """Model produces only EOS (empty response).""" + model_tokens = [""] + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + gen_results = [r for r in results if isinstance(r, GenerationResponse)] + assert len(gen_results) == 1 + assert gen_results[0].text == "" + assert gen_results[0].finish_reason == "stop" + + +# ── Test: Full EPDP spec round-trip ────────────────────────────── + + +class TestE2EFullRoundTrip: + """Full round-trip matching the vLLM EPDP spec. + + Simulates the complete multi-turn flow: + Turn 1: user asks → think → tool call → tool result → think → answer + Turn 2: user asks again → old reasoning stripped → think → answer + """ + + def test_single_tool_full_flow_with_thinking(self): + """Complete flow: user → think → tool call → tool result → think → answer. + + This is the core EPDP flow from the vLLM spec. + """ + # ── Turn 1.1: User asks, encode prompt ── + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are a weather assistant."}, + {"role": "user", "content": "How's the weather in Hangzhou?"}, + ] + prompt_1 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + assert prompt_1.endswith(THINKING_START) + assert "## Tools" in prompt_1 + assert "get_weather" in prompt_1 + + # ── Turn 1.1: Model thinks, then calls tool ── + model_tokens_1 = [ + THINKING_START, + "The user wants to know the weather in Hangzhou.", + " I need to use the get_weather tool.", + THINKING_END, + "\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou\n', + f"\n", + TOOL_CALLS_END, + ] + results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1))) + + # Verify: thinking tokens + tool call + gen_1 = [r for r in results_1 if isinstance(r, GenerationResponse)] + tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)] + thinking_1 = [r for r in gen_1 if r.is_thinking] + + assert len(thinking_1) >= 1 + assert "get_weather tool" in "".join(r.text for r in thinking_1) + assert len(tool_1) == 1 + assert tool_1[0].tool_calls[0].name == "get_weather" + tc_args = json.loads(tool_1[0].tool_calls[0].arguments) # pyright: ignore[reportAny] + assert tc_args == {"city": "Hangzhou"} + + # ── Turn 1.2: Add assistant response + tool result to messages ── + messages.append( + { + "role": "assistant", + "content": "", + "reasoning_content": "The user wants to know the weather in Hangzhou. I need to use the get_weather tool.", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Hangzhou"}', + }, + } + ], + } + ) + messages.append( + { + "role": "tool", + "content": '{"temperature": "7~13°C", "condition": "Cloudy"}', + } + ) + + # Encode prompt for turn 1.2 + prompt_2 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + + # Verify: prompt has the full conversation structure + assert TOOL_CALLS_START in prompt_2 # assistant's encoded tool call + assert EOS_TOKEN in prompt_2 # assistant turn ends with EOS + assert "" in prompt_2 + assert "" in prompt_2 + assert "Cloudy" in prompt_2 + assert "" in prompt_2 + # After tool results with thinking enabled → appended + assert prompt_2.endswith(THINKING_START) + # The assistant's reasoning_content should appear (it's after last_user_idx) + assert "get_weather tool" in prompt_2 + + # ── Turn 1.2: Model thinks about results, then answers ── + model_tokens_2 = [ + THINKING_START, + "The weather in Hangzhou is Cloudy, 7~13°C.", + " I'll tell the user.", + THINKING_END, + "The weather in Hangzhou is currently cloudy with temperatures between 7°C and 13°C.", + ] + results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2))) + + gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)] + tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)] + thinking_2 = [r for r in gen_2 if r.is_thinking] + non_thinking_2 = [r for r in gen_2 if not r.is_thinking] + + assert len(tool_2) == 0 # No more tool calls + assert len(thinking_2) >= 1 + assert "Cloudy" in "".join(r.text for r in thinking_2) + assert len(non_thinking_2) >= 1 + final_text = "".join(r.text for r in non_thinking_2) + assert "7°C" in final_text + assert "13°C" in final_text + + def test_multi_tool_full_flow(self): + """Flow with two tools: user → think → 2 tool calls → 2 results → think → answer.""" + # ── Initial prompt ── + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You help with weather and time."}, + {"role": "user", "content": "Weather in NYC and time in EST?"}, + ] + prompt_1 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + assert prompt_1.endswith(THINKING_START) + + # ── Model thinks, calls both tools ── + model_tokens_1 = [ + THINKING_START, + "Two requests: weather and time. I'll call both.", + THINKING_END, + "\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">NYC\n', + f"\n", + f'<{DSML_TOKEN}invoke name="get_time">\n', + f'<{DSML_TOKEN}parameter name="timezone" string="true">EST\n', + f"\n", + TOOL_CALLS_END, + ] + results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1))) + tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)] + + assert len(tool_1) == 1 + assert len(tool_1[0].tool_calls) == 2 + assert tool_1[0].tool_calls[0].name == "get_weather" + assert tool_1[0].tool_calls[1].name == "get_time" + + # ── Add assistant + both tool results ── + messages.append( + { + "role": "assistant", + "content": "", + "reasoning_content": "Two requests: weather and time. I'll call both.", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "NYC"}', + }, + }, + { + "type": "function", + "function": { + "name": "get_time", + "arguments": '{"timezone": "EST"}', + }, + }, + ], + } + ) + messages.append({"role": "tool", "content": "72°F, sunny"}) + messages.append({"role": "tool", "content": "2:30 PM EST"}) + + prompt_2 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + + # Verify multi-tool result encoding + # Count is 2: 1 in _TOOLS_SYSTEM_TEMPLATE example + 1 in conversation + assert prompt_2.count("") == 2 + assert prompt_2.count("") == 2 + assert "72°F, sunny" in prompt_2 + assert "2:30 PM EST" in prompt_2 + assert prompt_2.endswith(THINKING_START) + + # ── Model thinks about results, answers ── + model_tokens_2 = [ + THINKING_START, + "Got both results. Weather is 72°F sunny, time is 2:30 PM.", + THINKING_END, + "In NYC it's currently 72°F and sunny. The time in EST is 2:30 PM.", + ] + results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2))) + + tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)] + gen_2 = [r for r in results_2 if isinstance(r, GenerationResponse)] + non_thinking_2 = [r for r in gen_2 if not r.is_thinking] + + assert len(tool_2) == 0 + final_text = "".join(r.text for r in non_thinking_2) + assert "72°F" in final_text + assert "2:30 PM" in final_text + + def test_two_user_turns_reasoning_stripped(self): + """Turn 2: old reasoning_content is stripped from history. + + Per the vLLM spec, clear_reasoning_content is called between user turns + to save bandwidth. Our _drop_old_thinking handles this. + """ + # Full turn 1 conversation (already completed) + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Weather in Hangzhou?"}, + { + "role": "assistant", + "content": "", + "reasoning_content": "I need to call get_weather for Hangzhou.", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Hangzhou"}', + }, + } + ], + }, + {"role": "tool", "content": "Cloudy 7~13°C"}, + { + "role": "assistant", + "content": "The weather in Hangzhou is cloudy, 7-13°C.", + "reasoning_content": "The tool returned cloudy weather. I'll summarize.", + }, + # Turn 2: user asks again + {"role": "user", "content": "What about Beijing?"}, + ] + + prompt = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + + # Old reasoning_content from turn 1 assistants should be STRIPPED + # (they're before the last user message at index 5) + assert "I need to call get_weather" not in prompt + assert "tool returned cloudy" not in prompt + + # But the assistant's content and tool calls should still be there + assert "cloudy, 7-13°C" in prompt + assert TOOL_CALLS_START in prompt + + # Prompt ends with for the new turn + assert prompt.endswith(THINKING_START) + + # ── Turn 2: Model thinks, calls tool for Beijing ── + model_tokens = [ + THINKING_START, + "Now the user wants Beijing weather.", + THINKING_END, + "\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">Beijing\n', + f"\n", + TOOL_CALLS_END, + ] + results = list(parse_deepseek_v32(_simulate_tokens(model_tokens))) + tool_results = [r for r in results if isinstance(r, ToolCallResponse)] + + assert len(tool_results) == 1 + assert tool_results[0].tool_calls[0].name == "get_weather" + args = json.loads(tool_results[0].tool_calls[0].arguments) # pyright: ignore[reportAny] + assert args == {"city": "Beijing"} + + def test_chained_tool_calls_loop(self): + """Model calls tool, gets result, calls another tool, gets result, answers. + + This simulates the inner while loop from the vLLM spec where the model + may need multiple sub-turns of tool calling before it has enough info. + """ + # ── Sub-turn 1: user asks, model calls get_time ── + messages: list[dict[str, Any]] = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "What's the weather in Hangzhou tomorrow?"}, + ] + + prompt_1 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + assert prompt_1.endswith(THINKING_START) + + # Model first calls get_time to figure out the date + model_tokens_1 = [ + THINKING_START, + "I need the current date first to calculate tomorrow.", + THINKING_END, + "\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_time">\n', + f'<{DSML_TOKEN}parameter name="timezone" string="true">Asia/Shanghai\n', + f"\n", + TOOL_CALLS_END, + ] + results_1 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_1))) + tool_1 = [r for r in results_1 if isinstance(r, ToolCallResponse)] + assert len(tool_1) == 1 + assert tool_1[0].tool_calls[0].name == "get_time" + + # ── Sub-turn 2: add tool result, model calls get_weather ── + messages.append( + { + "role": "assistant", + "content": "", + "reasoning_content": "I need the current date first to calculate tomorrow.", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_time", + "arguments": '{"timezone": "Asia/Shanghai"}', + }, + } + ], + } + ) + messages.append({"role": "tool", "content": "2025-12-01 14:30 CST"}) + + prompt_2 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + assert "2025-12-01 14:30 CST" in prompt_2 + assert prompt_2.endswith(THINKING_START) + + # Model now knows the date, calls get_weather + model_tokens_2 = [ + THINKING_START, + "Today is 2025-12-01, so tomorrow is 2025-12-02.", + " Now I can check weather for Hangzhou.", + THINKING_END, + "\n\n", + TOOL_CALLS_START, + "\n", + f'<{DSML_TOKEN}invoke name="get_weather">\n', + f'<{DSML_TOKEN}parameter name="city" string="true">Hangzhou\n', + f"\n", + TOOL_CALLS_END, + ] + results_2 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_2))) + tool_2 = [r for r in results_2 if isinstance(r, ToolCallResponse)] + assert len(tool_2) == 1 + assert tool_2[0].tool_calls[0].name == "get_weather" + + # ── Sub-turn 3: add weather result, model answers ── + messages.append( + { + "role": "assistant", + "content": "", + "reasoning_content": "Today is 2025-12-01, so tomorrow is 2025-12-02. Now I can check weather for Hangzhou.", + "tool_calls": [ + { + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Hangzhou"}', + }, + } + ], + } + ) + messages.append({"role": "tool", "content": "Sunny, 5~12°C"}) + + prompt_3 = encode_messages( + messages, tools=_WEATHER_TOOLS, thinking_mode="thinking" + ) + # Should have both function_results blocks (one per tool round) + # Count is 3: 1 in _TOOLS_SYSTEM_TEMPLATE example + 2 in conversation + assert prompt_3.count("") == 3 + assert prompt_3.count("") == 3 + assert "2025-12-01 14:30 CST" in prompt_3 + assert "Sunny, 5~12°C" in prompt_3 + assert prompt_3.endswith(THINKING_START) + + # Model finally answers + model_tokens_3 = [ + THINKING_START, + "I have the weather for tomorrow in Hangzhou.", + THINKING_END, + "Tomorrow in Hangzhou will be sunny with temperatures between 5°C and 12°C.", + ] + results_3 = list(parse_deepseek_v32(_simulate_tokens(model_tokens_3))) + + tool_3 = [r for r in results_3 if isinstance(r, ToolCallResponse)] + gen_3 = [r for r in results_3 if isinstance(r, GenerationResponse)] + non_thinking_3 = [r for r in gen_3 if not r.is_thinking] + + assert len(tool_3) == 0 # No more tool calls — loop ends + final_text = "".join(r.text for r in non_thinking_3) + assert "sunny" in final_text.lower() + assert "5°C" in final_text + assert "12°C" in final_text diff --git a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py index 38a0a921..864e424b 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py +++ b/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py @@ -148,6 +148,7 @@ class MockTokenizer: tool_call_start = None tool_call_end = None has_tool_calling = False + has_thinking = False class MockGroup: diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py new file mode 100644 index 00000000..080f0389 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py @@ -0,0 +1,173 @@ +from collections.abc import Generator + +from exo.shared.types.worker.runner_response import ( + GenerationResponse, + ToolCallResponse, +) +from exo.worker.runner.runner import parse_gpt_oss + +# Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer. +# These are stable since they come from the model's vocabulary. +_CHANNEL = 200005 # <|channel|> +_START = 200006 # <|start|> +_MESSAGE = 200008 # <|message|> +_CALL = 200012 # <|call|> +_END = 200007 # <|end|> +_ASSISTANT = 173781 # "assistant" + +# fmt: off +# " to=functions.get_current_weather<|channel|>commentary json<|message|>{\"location\": \"Tokyo\"}<|call|>" +FORMAT_A_TOKENS: list[tuple[int, str]] = [ + (316, " to"), + (28, "="), + (44580, "functions"), + (775, ".get"), + (23981, "_current"), + (170154, "_weather"), + (_CHANNEL, "<|channel|>"), + (12606, "comment"), + (815, "ary"), + (5701, " json"), + (_MESSAGE, "<|message|>"), + (10848, '{"'), + (7693, "location"), + (1243, '":'), + (392, ' "'), + (173844, "Tokyo"), + (18583, '"}'), + (_CALL, "<|call|>"), +] + +# "<|channel|>commentary to=functions.get_current_weather json<|message|>{\"location\": \"Tokyo\"}<|call|>" +FORMAT_B_TOKENS: list[tuple[int, str]] = [ + (_CHANNEL, "<|channel|>"), + (12606, "comment"), + (815, "ary"), + (316, " to"), + (28, "="), + (44580, "functions"), + (775, ".get"), + (23981, "_current"), + (170154, "_weather"), + (5701, " json"), + (_MESSAGE, "<|message|>"), + (10848, '{"'), + (7693, "location"), + (1243, '":'), + (392, ' "'), + (173844, "Tokyo"), + (18583, '"}'), + (_CALL, "<|call|>"), +] + +# "<|channel|>analysis<|message|>Let me think...<|end|><|start|>assistant<|channel|>commentary to=functions.X ..." +# Full analysis-then-tool-call as the model actually generates it. +THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [ + (_CHANNEL, "<|channel|>"), + (35644, "analysis"), + (_MESSAGE, "<|message|>"), + (12845, "Let"), + (668, " me"), + (2411, " think"), + (1078, " about"), + (495, " this"), + (13, "."), + (_END, "<|end|>"), + # Model generates a new message header for the tool call: + (_START, "<|start|>"), + (_ASSISTANT, "assistant"), + *FORMAT_B_TOKENS, +] +# fmt: on + + +def _make_gen_responses( + tokens: list[tuple[int, str]], +) -> list[GenerationResponse]: + """Build GenerationResponse list from (token_id, text) pairs.""" + responses: list[GenerationResponse] = [] + for i, (tid, text) in enumerate(tokens): + is_last = i == len(tokens) - 1 + responses.append( + GenerationResponse( + text=text, + token=tid, + finish_reason="stop" if is_last else None, + usage=None, + ) + ) + return responses + + +def _collect( + tokens: list[tuple[int, str]], +) -> list[GenerationResponse | ToolCallResponse]: + """Feed tokens through parse_gpt_oss and collect all yielded responses.""" + + def _gen() -> Generator[GenerationResponse, None, None]: + yield from _make_gen_responses(tokens) + + return list(parse_gpt_oss(_gen())) + + +def _get_tool_call( + results: list[GenerationResponse | ToolCallResponse], +) -> ToolCallResponse: + """Extract the single ToolCallResponse from results.""" + tool_calls = [r for r in results if isinstance(r, ToolCallResponse)] + assert len(tool_calls) == 1, f"Expected 1 ToolCallResponse, got {len(tool_calls)}" + return tool_calls[0] + + +class TestParseGptOssRecipientPlacement: + """Both Harmony recipient placements must produce identical tool calls.""" + + def test_format_a_yields_tool_call(self): + results = _collect(FORMAT_A_TOKENS) + tc = _get_tool_call(results) + assert tc.tool_calls[0].name == "get_current_weather" + assert '"location"' in tc.tool_calls[0].arguments + assert "Tokyo" in tc.tool_calls[0].arguments + + def test_format_b_yields_tool_call(self): + results = _collect(FORMAT_B_TOKENS) + tc = _get_tool_call(results) + assert tc.tool_calls[0].name == "get_current_weather" + assert '"location"' in tc.tool_calls[0].arguments + assert "Tokyo" in tc.tool_calls[0].arguments + + def test_both_formats_produce_identical_tool_calls(self): + tc_a = _get_tool_call(_collect(FORMAT_A_TOKENS)) + tc_b = _get_tool_call(_collect(FORMAT_B_TOKENS)) + assert tc_a.tool_calls[0].name == tc_b.tool_calls[0].name + assert tc_a.tool_calls[0].arguments == tc_b.tool_calls[0].arguments + + +class TestParseGptOssThinkingThenToolCall: + """Analysis (thinking) followed by a tool call must yield both.""" + + def test_thinking_then_tool_call(self): + results = _collect(THINKING_THEN_TOOL_TOKENS) + + # Thinking tokens should have is_thinking=True and no tags + thinking_responses = [ + r for r in results if isinstance(r, GenerationResponse) and r.is_thinking + ] + thinking_text = "".join(r.text for r in thinking_responses) + assert "Let me think about this." in thinking_text + assert "" not in thinking_text + assert "" not in thinking_text + + # Non-thinking tokens should have is_thinking=False + non_thinking = [ + r + for r in results + if isinstance(r, GenerationResponse) and not r.is_thinking + ] + non_thinking_text = "".join(r.text for r in non_thinking) + assert "" not in non_thinking_text + + # And the tool call + tc = _get_tool_call(results) + assert tc.tool_calls[0].name == "get_current_weather" + assert "Tokyo" in tc.tool_calls[0].arguments diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py index 31f4822e..8a23a18c 100644 --- a/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py +++ b/src/exo/worker/tests/unittests/test_runner/test_parse_tool_calls.py @@ -5,12 +5,13 @@ from typing import Any from exo.shared.types.worker.runner_response import GenerationResponse, ToolCallResponse from exo.worker.runner.runner import parse_tool_calls +from exo.worker.runner.tool_parsers import make_mlx_parser def _make_responses( texts: list[str], finish_on_last: bool = True, -) -> Generator[GenerationResponse | ToolCallResponse]: +) -> Generator[GenerationResponse]: """Create a sequence of GenerationResponses from text strings.""" for i, text in enumerate(texts): is_last = i == len(texts) - 1 @@ -22,10 +23,13 @@ def _make_responses( ) -def _dummy_parser(text: str) -> dict[str, Any]: +def _dummier_parser(text: str) -> dict[str, Any]: return {"name": "test_fn", "arguments": {"arg": text}} +_dummy_parser = make_mlx_parser("", "", _dummier_parser) + + class TestParseToolCalls: """Tests for parse_tool_calls generator.""" @@ -35,8 +39,6 @@ class TestParseToolCalls: results = list( parse_tool_calls( _make_responses(texts, finish_on_last=False), - "", - "", _dummy_parser, ) ) @@ -50,8 +52,6 @@ class TestParseToolCalls: results = list( parse_tool_calls( _make_responses(texts), - "", - "", _dummy_parser, ) ) @@ -76,9 +76,7 @@ class TestParseToolCalls: results = list( parse_tool_calls( _make_responses(texts, finish_on_last=False), - "", - "", - _failing_parser, + make_mlx_parser("", "", _failing_parser), ) ) diff --git a/tests/eval_tool_calls.sh b/tests/eval_tool_calls.sh new file mode 100755 index 00000000..1b6bd3fe --- /dev/null +++ b/tests/eval_tool_calls.sh @@ -0,0 +1,55 @@ +#!/usr/bin/env bash + +[ $# -lt 1 ] && { + echo "Usage: $0 host1 [host2 ...]" + exit 1 +} + +[ -z "$(git status --porcelain)" ] || { + echo "Uncommitted changes" + exit 1 +} + +commit=$(git rev-parse HEAD) +git fetch -q origin +git branch -r --contains "$commit" | grep -qE '^\s*origin/' || { + echo "Not pushed to origin" + exit 1 +} +hosts=("$@") +cleanup() { + for host in "${hosts[@]}"; do + ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" & + done + sleep 1 + jobs -pr | xargs -r kill 2>/dev/null || true +} +trap 'cleanup' EXIT INT TERM + +for host; do + ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ + "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" & +done +wait +for host; do + ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \ + "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null & +done + +for host; do + echo "Waiting for $host..." 1>&2 + until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done +done + +echo "Waiting 30s for cluster setup" 1>&2 +sleep 30 +echo "EXO loaded" 1>&2 +eval_runner="${hosts[0]}" +mkdir -p "./bench/$commit" +nix run .#exo-get-all-models-on-cluster -- "$eval_runner" | while IFS= read -r model; do + echo "running eval for $model" 1>&2 + ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$eval_runner@$eval_runner" \ + "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-eval-tool-calls -- --model $model --stdout" \ + >>"./bench/$commit/${model//\//--}-eval.json" + echo +done diff --git a/tmp/config_examples/claude_code.sh b/tmp/config_examples/claude_code.sh new file mode 100755 index 00000000..685d0a27 --- /dev/null +++ b/tmp/config_examples/claude_code.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# Run Claude Code against a local exo cluster! (Here, GPT OSS 120B) +ANTHROPIC_BASE_URL="http://localhost:52415/" \ + ANTHROPIC_AUTH_TOKEN="dummy" \ + ANTHROPIC_MODEL="mlx-community/gpt-oss-120b-MXFP4-Q8" \ + ANTHROPIC_SMALL_FAST_MODEL="mlx-community/gpt-oss-120b-MXFP4-Q8" \ + CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 \ + claude diff --git a/uv.lock b/uv.lock index 627e6951..b14a1f69 100644 --- a/uv.lock +++ b/uv.lock @@ -377,8 +377,8 @@ dependencies = [ { name = "hypercorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", extra = ["cpu"], marker = "sys_platform == 'linux'" }, + { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" }, + { name = "mlx", version = "0.30.7.dev20260218+14841977", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }, marker = "sys_platform == 'darwin'" }, { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -416,9 +416,9 @@ requires-dist = [ { name = "hypercorn", specifier = ">=0.18.0" }, { name = "loguru", specifier = ">=0.7.3" }, { name = "mflux", specifier = "==0.15.5" }, - { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.6" }, + { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" }, { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" }, - { name = "mlx-lm", specifier = "==0.30.6" }, + { name = "mlx-lm", specifier = "==0.30.7" }, { name = "msgspec", specifier = ">=0.19.0" }, { name = "openai-harmony", specifier = ">=0.0.8" }, { name = "pillow", specifier = ">=11.0,<12.0" }, @@ -447,6 +447,7 @@ name = "exo-bench" version = "0.1.0" source = { editable = "bench" } dependencies = [ + { name = "httpx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -456,6 +457,7 @@ dependencies = [ [package.metadata] requires-dist = [ + { name = "httpx", specifier = ">=0.27.0" }, { name = "huggingface-hub", specifier = ">=0.33.4" }, { name = "jinja2", specifier = ">=3.1.0" }, { name = "loguru", specifier = ">=0.7.3" }, @@ -1020,8 +1022,8 @@ dependencies = [ { name = "fonttools", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", extra = ["cuda13"], marker = "sys_platform == 'linux'" }, + { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" }, + { name = "mlx", version = "0.30.7.dev20260218+14841977", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }, marker = "sys_platform == 'darwin'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1048,18 +1050,12 @@ wheels = [ name = "mlx" version = "0.30.6" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, +resolution-markers = [ + "sys_platform == 'linux'", ] wheels = [ - { url = "https://files.pythonhosted.org/packages/ae/5b/e460e144a34d5529e010056cccf50b538d56ed001473bc6b246018fd58cb/mlx-0.30.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:ed86f8bffc174c2f259ca589ea25464c96cf69d1bb457074a2bf2ef53737e54f", size = 573515, upload-time = "2026-02-06T03:45:23.405Z" }, - { url = "https://files.pythonhosted.org/packages/60/25/69833fefb9a3fef30b56792b1bcd022496c4fea83e45411d289b77ef7546/mlx-0.30.6-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:c52294958269e20f300639a17c1900ca8fc737d859ddda737f9811e94bd040e5", size = 573516, upload-time = "2026-02-06T03:45:24.618Z" }, - { url = "https://files.pythonhosted.org/packages/9c/6a/7e7fbeebc5cb51b6a5eba96b263a6298707bcbdc059f4b0b73e088bc3dea/mlx-0.30.6-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:b5b6636f7c49a4d86d8ec82643b972f45a144a7a9f3a967b27b2e6e22cf71e6a", size = 573592, upload-time = "2026-02-06T03:45:25.928Z" }, { url = "https://files.pythonhosted.org/packages/93/06/280f6f2ba80520a7109730425eda0d966658793aa0d02d8be8d351f75253/mlx-0.30.6-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:67e6c9e30a9faeacc209917ef5523177cf9b086914b6b5d83ff886e4294b727d", size = 622011, upload-time = "2026-02-06T03:45:28.165Z" }, { url = "https://files.pythonhosted.org/packages/fe/35/f872afbee9c079cc69924d9e9c46f5663adb7da58cba3511db082dd307c1/mlx-0.30.6-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:47db8b16fcb6f6c5a47c0bdb24ed377b41237017ac93aa6cb6aa206c9bdf82e4", size = 663650, upload-time = "2026-02-06T03:45:30.315Z" }, - { url = "https://files.pythonhosted.org/packages/60/23/361dc7a5797634e4d7e9bdd6564c6b28f9b1246672632def2f91bf066b18/mlx-0.30.6-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:78804a89dcff4a838f7c2da72392fe87a523e95122a3c840e53df019122aad45", size = 575028, upload-time = "2026-02-06T03:45:31.549Z" }, - { url = "https://files.pythonhosted.org/packages/a8/69/1854484d414171586814dfbe8def95f75c4ea2c7341ba13ba8ee675f7c62/mlx-0.30.6-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:ec13584ab069665cc7ad34a05494d9291cd623aef6ae96be48875fc87cfc25d6", size = 575026, upload-time = "2026-02-06T03:45:33.072Z" }, - { url = "https://files.pythonhosted.org/packages/6b/b8/3adbc441924209a7e4c568308b2a0b54bd09aee6a68db5bae85304791e54/mlx-0.30.6-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:b2c5e8a090a753ef99a1380a4d059c983083f36198864f6df9faaf1223d083df", size = 575041, upload-time = "2026-02-06T03:45:34.814Z" }, { url = "https://files.pythonhosted.org/packages/3f/54/9d9e06804fb2088202a2cdf60458e00b221f71420bea285720b60f9e82b5/mlx-0.30.6-cp314-cp314-manylinux_2_35_aarch64.whl", hash = "sha256:9ceddede4af0de31d1f6b3099f70e5469d60cd7c546975dedbdbeab3519cab3f", size = 624002, upload-time = "2026-02-06T03:45:36Z" }, { url = "https://files.pythonhosted.org/packages/42/92/3140a15a50cb1f9267a6552171e1dfa577861de53e093124bc43707f2a0e/mlx-0.30.6-cp314-cp314-manylinux_2_35_x86_64.whl", hash = "sha256:4a6ffd2d16728cf95f63a1b555d7c2eaeea686a0e6b73228bd265411cb5d77a4", size = 663569, upload-time = "2026-02-06T03:45:37.242Z" }, ] @@ -1072,6 +1068,14 @@ cuda13 = [ { name = "mlx-cuda-13", marker = "sys_platform == 'linux'" }, ] +[[package]] +name = "mlx" +version = "0.30.7.dev20260218+14841977" +source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" } +resolution-markers = [ + "sys_platform == 'darwin'", +] + [[package]] name = "mlx-cpu" version = "0.30.6" @@ -1098,30 +1102,20 @@ wheels = [ [[package]] name = "mlx-lm" -version = "0.30.6" +version = "0.30.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.30.7.dev20260218+14841977", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#1484197707f35186ad3bd614357c7c47fdf86ebc" }, marker = "sys_platform == 'darwin'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/76/cb/815deddc8699b1f694d7e1f9cbed52934c03a8b49432c8add72932bb2f0b/mlx_lm-0.30.6.tar.gz", hash = "sha256:807e042d7040268f1b19190b7eaefd8b2efbff5590a65460974ad4225b91dda1", size = 271733, upload-time = "2026-02-04T21:27:45.741Z" } +sdist = { url = "https://files.pythonhosted.org/packages/66/0d/56542e2ae13ec6f542d3977d7cff89a205d4f6c5122e0ce23f33265f61c9/mlx_lm-0.30.7.tar.gz", hash = "sha256:e5f31ac58d9f2381f28e1ba639ff903e64f7cff1bdc245c0bc97f72264be329c", size = 275764, upload-time = "2026-02-12T18:41:11.86Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/5f/01d281f1fa8a1521d5936659beb4f5ab1f32b463d059263cf9d4cef969d9/mlx_lm-0.30.6-py3-none-any.whl", hash = "sha256:a7405bd581eacc4bf8209d7a6b7f23629585a0d7c6740c2a97e51fee35b3b0e1", size = 379451, upload-time = "2026-02-04T21:27:43.222Z" }, -] - -[[package]] -name = "mlx-metal" -version = "0.30.6" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f3/85/44406b521f920248fad621334d4dc15e77660a494edf890e7cbee33bf38d/mlx_metal-0.30.6-py3-none-macosx_14_0_arm64.whl", hash = "sha256:ea6d0c973def9a5b4f652cc77036237db3f88c9d0af63701d76b5fddde99b820", size = 38437818, upload-time = "2026-02-06T03:44:56.19Z" }, - { url = "https://files.pythonhosted.org/packages/d0/cb/10a516995f7d0c154b0d7e633c54b51e96977a86a355105b6474cfcbe0d0/mlx_metal-0.30.6-py3-none-macosx_15_0_arm64.whl", hash = "sha256:0f8cb94634d07e06a372d6ad9a090f38a18bab1ff19a140aede60eacf707bb94", size = 38433701, upload-time = "2026-02-06T03:44:59.678Z" }, - { url = "https://files.pythonhosted.org/packages/4c/7d/70cb272f7373c334709f210ed8420511fc9d64d05a7a646c0b3b94c29c04/mlx_metal-0.30.6-py3-none-macosx_26_0_arm64.whl", hash = "sha256:d761ae26304f2c4b454eeea7f612a56919d9e5e57dbb1dc0788f8e34aa6f41c2", size = 47718448, upload-time = "2026-02-06T03:45:03.133Z" }, + { url = "https://files.pythonhosted.org/packages/1e/17/a41c798a3d9cbdc47f39c6db5bba4c2cd199203ead26bf911cb03b644070/mlx_lm-0.30.7-py3-none-any.whl", hash = "sha256:17442a4bf01c4c2d3bca1e647712fe44f19890c3f1eadc8589d389e57b44b9bf", size = 386591, upload-time = "2026-02-12T18:41:10.236Z" }, ] [[package]]