From 86735ece782877d8976ebcc8b42f632253c903e9 Mon Sep 17 00:00:00 2001 From: Evan Date: Mon, 16 Feb 2026 12:20:58 +0000 Subject: [PATCH 01/45] begins begins --- src/exo/master/adapters/responses.py | 4 +- .../worker/engines/mlx/generator/generate.py | 2 +- src/exo/worker/engines/mlx/utils_mlx.py | 46 ++- src/exo/worker/runner/runner.py | 293 ++++-------------- src/exo/worker/runner/tool_parsers.py | 72 +++++ .../test_runner/test_parse_tool_calls.py | 16 +- 6 files changed, 179 insertions(+), 254 deletions(-) create mode 100644 src/exo/worker/runner/tool_parsers.py diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py index b37b7d54..5e059fee 100644 --- a/src/exo/master/adapters/responses.py +++ b/src/exo/master/adapters/responses.py @@ -144,8 +144,8 @@ 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, ) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index bc38c61c..ffa2f5c0 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -306,7 +306,7 @@ def mlx_generate( max_stop_len = max((len(s) for s in stop_sequences), default=0) mx_barrier(group) - logger.info("Ready to prefill") + logger.info("Starting prefill") # Prefill cache with all tokens except the last one prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 6aceb53c..3ed65ecc 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -353,7 +353,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, @@ -585,3 +591,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/runner/runner.py b/src/exo/worker/runner/runner.py index ad6c78f6..e55456d3 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -1,11 +1,10 @@ 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 Literal import mlx.core as mx from mlx_lm.models.gpt_oss import Model as GptOssModel @@ -16,7 +15,6 @@ from openai_harmony import ( # pyright: ignore[reportMissingTypeStubs] 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 @@ -93,6 +91,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 +138,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 +204,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 ( @@ -310,31 +320,11 @@ def main( 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): + if 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] - ) + elif tool_parser: + mlx_generator = parse_tool_calls(mlx_generator, tool_parser) completion_tokens = 0 tokens_since_last_cancel_check = 0 @@ -587,21 +577,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) @@ -658,9 +635,9 @@ def parse_gpt_oss( def parse_thinking_models( - responses: Generator[GenerationResponse | ToolCallResponse], + responses: Generator[GenerationResponse], tokenizer: TokenizerWrapper, -) -> Generator[GenerationResponse | ToolCallResponse]: +) -> Generator[GenerationResponse]: """ For models that inject thinking tags in the prompt (like GLM-4.7), prepend the thinking tag to the output stream so the frontend @@ -781,221 +758,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/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_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), ) ) From 8392e78afe152d6311f5465f70d7ae366fc9902c Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Tue, 17 Feb 2026 10:52:05 +0000 Subject: [PATCH 02/45] bench: add spec for automatic canary benchmarks (#1483) Adds all the models that can fit onto a single M3 Ultra for single machine benchmarks. Fixes the macOS version, GPU spec, and chip type for maximum reproducibility. Specifies the minimum memory accordingly for each type of model, using the smallest machine available (the smallest M3 Ultra is 96GiB). Test plan: - Running this with some code that makes machines of this spec available and stores the results. It works. This will become part of a larger testing/stability strategy once we've collected more of the data. --- bench/bench.toml | 7 ++ bench/single-m3-ultra.toml | 189 +++++++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+) create mode 100644 bench/bench.toml create mode 100644 bench/single-m3-ultra.toml 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/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))"] From c01b6fff21aeeea843cd2bc0ca715e870452f1e7 Mon Sep 17 00:00:00 2001 From: Evan Date: Tue, 17 Feb 2026 11:37:15 +0000 Subject: [PATCH 03/45] eprint banner our banner was being printed to stdout but should be printed to stderr as its essentially a log message --- src/exo/utils/banner.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/exo/utils/banner.py b/src/exo/utils/banner.py index eb6d7b08..ffdb5458 100644 --- a/src/exo/utils/banner.py +++ b/src/exo/utils/banner.py @@ -1,5 +1,7 @@ +import sys + + def print_startup_banner(port: int) -> None: - """Print a prominent startup banner with API endpoint information.""" dashboard_url = f"http://localhost:{port}" banner = f""" ╔═══════════════════════════════════════════════════════════════════════╗ @@ -27,4 +29,4 @@ def print_startup_banner(port: int) -> None: """ - print(banner) + print(banner, file=sys.stderr) From 6d1ca6689bbf13e323b3d2e303a42281f7d70649 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Tue, 17 Feb 2026 11:48:28 +0000 Subject: [PATCH 04/45] don't time out node identities (#1493) currently nodes leaving and rejoining the cluster can lose their identity. We have no need to delete this data on node timing out, so let's just persist it. --- src/exo/shared/apply.py | 6 ------ 1 file changed, 6 deletions(-) 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, From d6301ed593733de84ed00b87032413b736510cf3 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Tue, 17 Feb 2026 06:31:47 -0800 Subject: [PATCH 05/45] =?UTF-8?q?dashboard:=20redesign=20downloads=20page?= =?UTF-8?q?=20as=20model=C3=97node=20table=20(#1465)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation The current downloads page uses a node-centric card grid layout that is messy and hard to read — the same model across different nodes appears in separate cards, and deep nesting wastes space. This makes it difficult to quickly see which models are on which nodes. ## Changes Rewrote the downloads page (`dashboard/src/routes/downloads/+page.svelte`) from a card grid to a clean table layout: - **Rows** = models (unique across all nodes) - **Columns** = nodes (with disk free shown in header) - **Cells** show status at a glance: - ✅ Green checkmark + size for completed downloads - 🟡 Yellow percentage + mini progress bar + speed for active downloads - `...` for pending downloads - ❌ Red X for failed downloads - `--` for models not present on a node - Delete/download action buttons appear on row hover - Model name column is sticky on horizontal scroll (for many-node clusters) - Models sorted by number of nodes with completed downloads - Imported shared utilities from `$lib/utils/downloads` instead of inline re-implementations ### Backend: model directory in download events - Added `model_directory` field to `BaseDownloadProgress` so all download status events include the on-disk path - Added `_model_dir()` helper to `DownloadCoordinator` to compute the path from `EXO_MODELS_DIR` - Dashboard uses this to show file location and enable "open in Finder" for completed downloads ### Info modal - Clicking a model name opens an info modal showing card details (family, quantization, capabilities, storage size, layer count, tensor parallelism support) ### Other fixes - Fixed model name truncation in the table - Excluded `tests/start_distributed_test.py` from pytest collection (CLI script that calls `sys.exit()` at import time) ## Test Plan - [x] `uv run basedpyright` — 0 errors - [x] `uv run ruff check` — all passed - [x] `nix fmt` — clean - [x] `uv run pytest` — 188 passed, 1 skipped 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 --- dashboard/src/routes/downloads/+page.svelte | 993 +++++++++++--------- pyproject.toml | 2 +- src/exo/download/coordinator.py | 28 +- src/exo/shared/types/worker/downloads.py | 1 + 4 files changed, 557 insertions(+), 467 deletions(-) diff --git a/dashboard/src/routes/downloads/+page.svelte b/dashboard/src/routes/downloads/+page.svelte index 94e10612..91719dc1 100644 --- a/dashboard/src/routes/downloads/+page.svelte +++ b/dashboard/src/routes/downloads/+page.svelte @@ -1,43 +1,59 @@ @@ -415,253 +349,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/pyproject.toml b/pyproject.toml index def495c7..5d8d79a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -132,7 +132,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/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index a05bd6f8..db13ccef 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, @@ -63,6 +64,9 @@ class DownloadCoordinator: self.event_sender, self.event_receiver = channel[Event]() 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: @@ -74,6 +78,7 @@ class DownloadCoordinator: shard_metadata=callback_shard, node_id=self.node_id, total_bytes=progress.total_bytes, + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = completed await self.event_sender.send( @@ -93,6 +98,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( @@ -170,7 +176,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)) @@ -184,6 +194,7 @@ class DownloadCoordinator: shard_metadata=shard, node_id=self.node_id, total_bytes=initial_progress.total_bytes, + model_directory=self._model_dir(model_id), ) self.download_status[model_id] = completed await self.event_sender.send( @@ -206,6 +217,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 +231,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 +266,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) @@ -295,11 +309,18 @@ class DownloadCoordinator: node_id=self.node_id, shard_metadata=progress.shard, total_bytes=progress.total_bytes, + 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: 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 +329,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/shared/types/worker/downloads.py b/src/exo/shared/types/worker/downloads.py index 77162e7c..a29edcbf 100644 --- a/src/exo/shared/types/worker/downloads.py +++ b/src/exo/shared/types/worker/downloads.py @@ -26,6 +26,7 @@ class DownloadProgressData(CamelCaseModel): class BaseDownloadProgress(TaggedModel): node_id: NodeId shard_metadata: ShardMetadata + model_directory: str = "" class DownloadPending(BaseDownloadProgress): From db79c350c1bc4af0c56c42d630d2b723d9a18a83 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Tue, 17 Feb 2026 09:03:54 -0800 Subject: [PATCH 06/45] Fix graceful process shutdown in macOS app (#1372) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation Fixes #1370 When the macOS app stops exo, GPU/system memory isn't released. This happens because: 1. The macOS app calls `process.terminate()` (SIGTERM) but the Python process only registers a graceful shutdown handler for SIGINT, not SIGTERM. SIGTERM's default Python behavior raises `SystemExit` which bypasses the cleanup cascade (runner subprocess MLX cleanup via `mx.clear_cache()`, channel closing, etc.). 2. The app doesn't wait for the process to actually finish cleanup — it immediately nils out the process reference. ## Changes **`src/exo/main.py`**: Register SIGTERM handler alongside SIGINT so the graceful shutdown cascade (`Node.shutdown()` → cancel task group → worker/runner cleanup → `mx.clear_cache()` + `gc.collect()`) runs regardless of which signal is received. **`app/EXO/EXO/ExoProcessController.swift`**: Replace immediate `process.terminate()` with escalating shutdown per @Evanev7's suggestion: 1. Send SIGINT via `process.interrupt()` — triggers the registered Python handler for graceful cleanup 2. Wait up to 5 seconds for the process to exit 3. If still running, escalate to SIGTERM via `process.terminate()` 4. Wait up to 3 seconds 5. If still running, force kill via SIGKILL The escalation runs in a detached `Task` so the UI updates immediately (status → stopped) without blocking. ## Why It Works The root cause is that SIGTERM wasn't triggering the graceful shutdown path. By registering a SIGTERM handler in Python and sending SIGINT first from the macOS app, the process gets a chance to run the full cleanup cascade: cancelling the task group, shutting down runners (which call `del model; mx.clear_cache(); gc.collect()`), closing channels, and flushing logs. The escalation to SIGTERM and SIGKILL ensures the process always terminates even if graceful shutdown hangs. ## Test Plan ### Manual Testing - Start exo via macOS app, load a model, run inference - Stop via the toggle switch, verify memory is released without requiring a system restart - Test rapid stop/start (restart) to ensure no race conditions ### Automated Testing - `uv run basedpyright` — 0 errors - `uv run ruff check` — passes - `nix fmt` — no changes --------- Co-authored-by: Claude Opus 4.5 Co-authored-by: Evan Quiney --- app/EXO/EXO/ExoProcessController.swift | 34 +++++++++++++++++++++++--- src/exo/main.py | 4 +-- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/app/EXO/EXO/ExoProcessController.swift b/app/EXO/EXO/ExoProcessController.swift index 704b9d4f..7566674b 100644 --- a/app/EXO/EXO/ExoProcessController.swift +++ b/app/EXO/EXO/ExoProcessController.swift @@ -126,11 +126,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/src/exo/main.py b/src/exo/main.py index e1ef3a70..1d358975 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -136,6 +136,8 @@ class Node: 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 +149,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 From a962a28afc7deb65ecb99d1f1cf83be1c5c10dee Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Tue, 17 Feb 2026 09:48:19 -0800 Subject: [PATCH 07/45] Add MetaInstance declarative layer (#1447) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation Users currently manage instances directly, which means if a node disconnects or connections break, the instance dies and nothing recreates it. MetaInstance is a declarative primitive: "ensure an instance matching these parameters always exists." The reconciler watches for unhealthy or missing backing instances and re-places them automatically. ## Changes - **MetaInstance type** (`meta_instance.py`): declarative constraint with `model_id`, `min_nodes`, optional `node_ids`, and `sharding` - **Reconciler** (`reconcile.py`): `find_unsatisfied_meta_instances` checks which MetaInstances lack a healthy backing instance, `try_place_for_meta_instance` creates one - **Master loop** (`main.py`): periodically reconciles unsatisfied MetaInstances; immediate placement on `CreateMetaInstance` command - **API** (`api.py`): `create_meta_instance` / `delete_meta_instance` / `GET /meta_instances` endpoints; delete cascades to backing instances with task cancellation - **Binding via `meta_instance_id` on Instance** (`instances.py`): no separate binding event or backing map — the instance carries its parent MetaInstance ID directly, eliminating race conditions in the reconciler - **Dashboard**: sidebar shows MetaInstances with their backing instance status; orphan instances (created directly) still shown separately - **Tests**: constraint matching, connection health, unsatisfied detection, exclusive binding, cascade delete with task cancellation ### Recent improvements - **fix: cancel active tasks on cascade delete** — `DeleteMetaInstance` now emits `TaskStatusUpdated(Cancelled)` for any Pending/Running tasks on backing instances before emitting `InstanceDeleted`. Previously, cascade-deleting backing instances left orphaned task references in state. - **Lifecycle logging** — added `logger.info`/`logger.warning` for: `CreateMetaInstance` (model, min_nodes, sharding), `DeleteMetaInstance` (with cascade count), reconciler placement success/failure, and retry decisions with attempt counts in `InstanceHealthReconciler`. - **GET `/meta_instances` endpoint** — lists all meta-instances without needing to fetch full state. - **2 regression tests** — `test_cascade_delete_cancels_active_tasks` and `test_cascade_delete_skips_completed_tasks` verify the cascade-delete event sequence. ## Why It Works Putting `meta_instance_id` on `BaseInstance` makes binding inherent to instance creation. When the reconciler creates an instance for a MetaInstance, it tags it via `model_copy`. When the instance is deleted, the binding disappears with it. This avoids the two bugs that a separate binding mechanism would introduce: 1. Stale exclusion sets — the reconciler loop can't accidentally bind two MetaInstances to the same instance 2. Delete ordering race — no window between deleting an instance and its binding where the reconciler could re-place ## Test Plan ### Manual Testing - Created MetaInstance via dashboard, verified instance placed - Verified delete cascades (deleting MetaInstance removes backing instance) - Verified orphan instances still work independently ### Automated Testing - 30 tests in `test_meta_instance_edge_cases.py`: lifecycle, retry logic, error handling, concurrent operations, cascade delete with task cancellation - 24 tests in `test_reconcile.py`: constraint matching, connection health (single/multi-node, edge removal, IP changes), unsatisfied detection, exclusive binding, idempotency - All 261 tests pass - basedpyright 0 errors, ruff clean, dashboard builds --------- Co-authored-by: Claude Opus 4.6 --- .../src/lib/components/ChatSidebar.svelte | 6 +- dashboard/src/lib/components/ModelCard.svelte | 6 +- dashboard/src/lib/stores/app.svelte.ts | 27 +- dashboard/src/routes/+page.svelte | 782 +++++++++++++++--- src/exo/download/coordinator.py | 12 +- src/exo/main.py | 2 +- src/exo/master/api.py | 78 +- src/exo/master/main.py | 255 ++++-- src/exo/master/placement.py | 20 +- src/exo/master/placement_utils.py | 6 +- src/exo/master/process_managers/__init__.py | 12 + .../process_managers/instance_health.py | 62 ++ .../master/process_managers/meta_instance.py | 92 +++ .../master/process_managers/node_timeout.py | 27 + src/exo/master/reconcile.py | 244 ++++++ .../tests/test_meta_instance_edge_cases.py | 778 +++++++++++++++++ src/exo/master/tests/test_placement_utils.py | 12 +- src/exo/master/tests/test_reconcile.py | 742 +++++++++++++++++ src/exo/shared/apply.py | 126 ++- src/exo/shared/types/api.py | 22 +- src/exo/shared/types/commands.py | 13 +- src/exo/shared/types/common.py | 4 + src/exo/shared/types/events.py | 80 +- src/exo/shared/types/meta_instance.py | 25 + src/exo/shared/types/state.py | 4 +- src/exo/shared/types/tasks.py | 2 +- src/exo/shared/types/worker/instances.py | 3 +- src/exo/utils/channels.py | 4 +- src/exo/worker/engines/mlx/utils_mlx.py | 5 + src/exo/worker/main.py | 34 +- src/exo/worker/plan.py | 23 +- src/exo/worker/runner/bootstrap.py | 13 + src/exo/worker/runner/runner.py | 4 +- src/exo/worker/runner/runner_supervisor.py | 199 ++++- .../test_runner/test_event_ordering.py | 27 +- 35 files changed, 3457 insertions(+), 294 deletions(-) create mode 100644 src/exo/master/process_managers/__init__.py create mode 100644 src/exo/master/process_managers/instance_health.py create mode 100644 src/exo/master/process_managers/meta_instance.py create mode 100644 src/exo/master/process_managers/node_timeout.py create mode 100644 src/exo/master/reconcile.py create mode 100644 src/exo/master/tests/test_meta_instance_edge_cases.py create mode 100644 src/exo/master/tests/test_reconcile.py create mode 100644 src/exo/shared/types/meta_instance.py diff --git a/dashboard/src/lib/components/ChatSidebar.svelte b/dashboard/src/lib/components/ChatSidebar.svelte index b721b033..6a822ddf 100644 --- a/dashboard/src/lib/components/ChatSidebar.svelte +++ b/dashboard/src/lib/components/ChatSidebar.svelte @@ -185,11 +185,7 @@ let instanceType: string | null = null; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; - else if ( - instanceTag === "MlxIbvInstance" || - instanceTag === "MlxJacclInstance" - ) - instanceType = "MLX RDMA"; + else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; let sharding: string | null = null; const inst = instance as { diff --git a/dashboard/src/lib/components/ModelCard.svelte b/dashboard/src/lib/components/ModelCard.svelte index 9046d2a6..561c325b 100644 --- a/dashboard/src/lib/components/ModelCard.svelte +++ b/dashboard/src/lib/components/ModelCard.svelte @@ -21,7 +21,7 @@ } | null; nodes?: Record; sharding?: "Pipeline" | "Tensor"; - runtime?: "MlxRing" | "MlxIbv" | "MlxJaccl"; + runtime?: "MlxRing" | "MlxJaccl"; onLaunch?: () => void; tags?: string[]; apiPreview?: PlacementPreview | null; @@ -348,7 +348,7 @@ // Debug mode state const isDebugMode = $derived(debugMode()); const topology = $derived(topologyData()); - const isRdma = $derived(runtime === "MlxIbv" || runtime === "MlxJaccl"); + const isRdma = $derived(runtime === "MlxJaccl"); // Get interface name for an IP from node data function getInterfaceForIp(nodeId: string, ip?: string): string | null { @@ -575,7 +575,7 @@ > {runtime === "MlxRing" ? "MLX Ring" - : runtime === "MlxIbv" || runtime === "MlxJaccl" + : runtime === "MlxJaccl" ? "MLX RDMA" : runtime} diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index e5dbf902..ebb2d0df 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -168,7 +168,7 @@ export interface ModelDownloadStatus { export interface PlacementPreview { model_id: string; sharding: "Pipeline" | "Tensor"; - instance_meta: "MlxRing" | "MlxIbv" | "MlxJaccl"; + instance_meta: "MlxRing" | "MlxJaccl"; instance: unknown | null; memory_delta_by_node: Record | null; error: string | null; @@ -219,7 +219,6 @@ interface RawStateResponse { string, { MlxRingInstance?: Instance; - MlxIbvInstance?: Instance; MlxJacclInstance?: Instance; } >; @@ -250,6 +249,20 @@ interface RawStateResponse { >; // Thunderbolt bridge cycles (nodes with bridge enabled forming loops) thunderboltBridgeCycles?: string[][]; + // MetaInstances (declarative instance constraints) + metaInstances?: Record; +} + +export interface MetaInstanceData { + metaInstanceId: string; + modelId: string; + sharding: string; + instanceMeta: string; + minNodes: number; + nodeIds: string[] | null; + placementError: string | null; + consecutiveFailures: number; + lastFailureError: string | null; } export interface MessageAttachment { @@ -537,6 +550,7 @@ class AppStore { previewNodeFilter = $state>(new Set()); lastUpdate = $state(null); nodeIdentities = $state>({}); + metaInstances = $state>({}); thunderboltBridgeCycles = $state([]); nodeThunderbolt = $state< Record< @@ -895,11 +909,7 @@ class AppStore { let instanceType: string | null = null; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; - else if ( - instanceTag === "MlxIbvInstance" || - instanceTag === "MlxJacclInstance" - ) - instanceType = "MLX RDMA"; + else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; let sharding: string | null = null; const inst = instance as { @@ -1273,6 +1283,8 @@ class AppStore { this.nodeThunderbolt = data.nodeThunderbolt ?? {}; // RDMA ctl status per node this.nodeRdmaCtl = data.nodeRdmaCtl ?? {}; + // MetaInstances + this.metaInstances = data.metaInstances ?? {}; // Thunderbolt bridge cycles this.thunderboltBridgeCycles = data.thunderboltBridgeCycles ?? []; // Thunderbolt bridge status per node @@ -3044,6 +3056,7 @@ export const tps = () => appStore.tps; export const totalTokens = () => appStore.totalTokens; export const topologyData = () => appStore.topologyData; export const instances = () => appStore.instances; +export const metaInstances = () => appStore.metaInstances; export const runners = () => appStore.runners; export const downloads = () => appStore.downloads; export const nodeDisk = () => appStore.nodeDisk; diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 2fdeb8ab..5f7dcf04 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -44,11 +44,13 @@ toggleChatSidebarVisible, nodeThunderbolt, nodeRdmaCtl, + metaInstances, thunderboltBridgeCycles, nodeThunderboltBridge, nodeIdentities, type DownloadProgress, type PlacementPreview, + type MetaInstanceData, } from "$lib/stores/app.svelte"; import HeaderNav from "$lib/components/HeaderNav.svelte"; import { fade, fly } from "svelte/transition"; @@ -68,7 +70,70 @@ const debugEnabled = $derived(debugMode()); const topologyOnlyEnabled = $derived(topologyOnlyMode()); const sidebarVisible = $derived(chatSidebarVisible()); + const metaInstancesData = $derived(metaInstances()); const tbBridgeCycles = $derived(thunderboltBridgeCycles()); + + // Get status for a MetaInstance that has no backing instance yet + function getMetaInstancePlacingStatus(metaInstanceId: string) { + const meta = metaInstancesData[metaInstanceId]; + const placementError = meta?.placementError; + const failures = meta?.consecutiveFailures ?? 0; + const lastError = meta?.lastFailureError; + + if (placementError) { + return { + statusText: "PLACEMENT FAILED", + statusClass: "failed", + isDownloading: false as const, + isFailed: true, + progress: null, + perNode: [] as Array<{ + nodeId: string; + nodeName: string; + progress: DownloadProgress; + }>, + perNodeStatus: [] as PerNodeRunnerStatus[], + errorMessage: placementError, + }; + } + + if (failures > 0) { + const retryPosition = ((failures - 1) % 3) + 1; + const isRecreated = failures % 3 === 0; + return { + statusText: isRecreated ? "PLACING" : `RETRYING (${retryPosition}/3)`, + statusClass: "starting", + isDownloading: false as const, + isFailed: false, + progress: null, + perNode: [] as Array<{ + nodeId: string; + nodeName: string; + progress: DownloadProgress; + }>, + perNodeStatus: [] as PerNodeRunnerStatus[], + errorMessage: isRecreated + ? `Instance re-created due to failure: ${lastError}` + : `Previous failure: ${lastError}`, + }; + } + + return { + statusText: "PLACING", + statusClass: "starting", + isDownloading: false as const, + isFailed: false, + progress: null, + perNode: [] as Array<{ + nodeId: string; + nodeName: string; + progress: DownloadProgress; + }>, + perNodeStatus: [] as PerNodeRunnerStatus[], + errorMessage: null, + }; + } + const tbBridgeData = $derived(nodeThunderboltBridge()); const identitiesData = $derived(nodeIdentities()); const tbIdentifiers = $derived(nodeThunderbolt()); @@ -114,6 +179,17 @@ }); let tb5InfoDismissed = $state(false); + // Detect [jaccl] RDMA driver errors from MetaInstance failure errors + const jacclError = $derived.by(() => { + for (const mi of Object.values(metaInstancesData)) { + if (mi.lastFailureError?.includes("[jaccl]")) { + return mi.lastFailureError; + } + } + return null; + }); + let jacclDismissedError = $state(null); + // Helper to get friendly node name from node ID function getNodeName(nodeId: string): string { const node = data?.nodes?.[nodeId]; @@ -224,7 +300,7 @@ return model.tasks.includes("ImageToImage"); } let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline"); - type InstanceMeta = "MlxRing" | "MlxIbv" | "MlxJaccl"; + type InstanceMeta = "MlxRing" | "MlxJaccl"; // Launch defaults persistence const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults"; @@ -481,7 +557,7 @@ const matchesSelectedRuntime = (runtime: InstanceMeta): boolean => selectedInstanceType === "MlxRing" ? runtime === "MlxRing" - : runtime === "MlxIbv" || runtime === "MlxJaccl"; + : runtime === "MlxJaccl" || runtime === "MlxJaccl"; // Helper to check if a model can be launched (has valid placement with >= minNodes) function canModelFit(modelId: string): boolean { @@ -697,39 +773,30 @@ launchingModelId = modelId; try { - // Use the specific preview if provided, otherwise fall back to filtered preview const preview = specificPreview ?? filteredPreview(); - let instanceData: unknown; + // Extract node IDs from the preview the user is seeing + const previewNodeIds = preview?.memory_delta_by_node + ? Object.keys(preview.memory_delta_by_node) + : nodeFilter.size > 0 + ? Array.from(nodeFilter) + : undefined; - if (preview?.instance) { - // Use the instance from the preview - instanceData = preview.instance; - } else { - // Fallback: GET placement from API - const placementResponse = await fetch( - `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=${selectedMinNodes}`, - ); - - if (!placementResponse.ok) { - const errorText = await placementResponse.text(); - console.error("Failed to get placement:", errorText); - return; - } - - instanceData = await placementResponse.json(); - } - - // POST the instance to create it - const response = await fetch("/instance", { + const response = await fetch("/meta_instance", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ instance: instanceData }), + body: JSON.stringify({ + model_id: modelId, + sharding: preview?.sharding ?? selectedSharding, + instance_meta: preview?.instance_meta ?? selectedInstanceType, + min_nodes: selectedMinNodes, + node_ids: previewNodeIds, + }), }); if (!response.ok) { const errorText = await response.text(); - console.error("Failed to launch instance:", errorText); + console.error("Failed to create meta instance:", errorText); } else { // Always auto-select the newly launched model so the user chats to what they just launched setSelectedChatModel(modelId); @@ -752,7 +819,7 @@ setTimeout(scrollToBottom, 1000); } } catch (error) { - console.error("Error launching instance:", error); + console.error("Error creating meta instance:", error); } finally { launchingModelId = null; } @@ -954,15 +1021,18 @@ nodeName: string; progress: DownloadProgress; }>; + perNodeStatus: PerNodeRunnerStatus[]; } { if (!downloadsData || Object.keys(downloadsData).length === 0) { + const statusInfo = deriveInstanceStatus(instanceWrapped); return { isDownloading: false, - isFailed: false, - errorMessage: null, + isFailed: statusInfo.statusText === "FAILED", + errorMessage: statusInfo.errorMessage, progress: null, - statusText: "RUNNING", + statusText: statusInfo.statusText, perNode: [], + perNodeStatus: statusInfo.perNodeStatus, }; } @@ -976,6 +1046,7 @@ progress: null, statusText: "PREPARING", perNode: [], + perNodeStatus: [], }; } @@ -1044,6 +1115,7 @@ progress: null, statusText: "FAILED", perNode: [], + perNodeStatus: [], }; } } @@ -1084,10 +1156,11 @@ return { isDownloading: false, isFailed: statusInfo.statusText === "FAILED", - errorMessage: null, + errorMessage: statusInfo.errorMessage, progress: null, statusText: statusInfo.statusText, perNode: [], + perNodeStatus: statusInfo.perNodeStatus, }; } @@ -1111,92 +1184,223 @@ }, statusText: "DOWNLOADING", perNode, + perNodeStatus: [], }; } // Derive instance status from runners // Get color class for a status function getStatusColor(statusText: string): string { - switch (statusText) { - case "FAILED": - return "text-red-400"; - case "SHUTDOWN": - return "text-gray-400"; - case "DOWNLOADING": - return "text-blue-400"; - case "LOADING": - case "WARMING UP": - case "WAITING": - case "INITIALIZING": - return "text-yellow-400"; - case "RUNNING": - return "text-teal-400"; - case "READY": - case "LOADED": - return "text-green-400"; - default: - return "text-exo-light-gray"; - } + if (statusText === "FAILED" || statusText === "PLACEMENT FAILED") + return "text-red-400"; + if (statusText.startsWith("RETRYING")) return "text-orange-400"; + if (statusText === "SHUTDOWN") return "text-gray-400"; + if (statusText === "DOWNLOADING") return "text-blue-400"; + if ( + statusText.startsWith("LOADING") || + statusText.startsWith("WARMING UP") || + statusText === "WAITING" || + statusText === "INITIALIZING" + ) + return "text-yellow-400"; + if (statusText === "RUNNING") return "text-teal-400"; + if (statusText === "READY" || statusText === "LOADED") + return "text-green-400"; + return "text-exo-light-gray"; + } + + const RUNNER_STATUS_MAP: Record = { + RunnerWaitingForInitialization: "WaitingForInitialization", + RunnerInitializingBackend: "InitializingBackend", + RunnerWaitingForModel: "WaitingForModel", + RunnerLoading: "Loading", + RunnerLoaded: "Loaded", + RunnerWarmingUp: "WarmingUp", + RunnerReady: "Ready", + RunnerRunning: "Running", + RunnerShutdown: "Shutdown", + RunnerFailed: "Failed", + }; + + // Friendly labels for display + const RUNNER_STATUS_DISPLAY: Record = { + WaitingForInitialization: "Initializing", + InitializingBackend: "Initializing", + WaitingForModel: "Waiting", + Loading: "Loading", + Loaded: "Loaded", + WarmingUp: "Warming Up", + Ready: "Ready", + Running: "Running", + Shutdown: "Shutdown", + Failed: "Failed", + }; + + interface PerNodeRunnerStatus { + nodeId: string; + nodeName: string; + status: string; // friendly display status } function deriveInstanceStatus(instanceWrapped: unknown): { statusText: string; statusClass: string; + perNodeStatus: PerNodeRunnerStatus[]; + errorMessage: string | null; } { const [, instance] = getTagged(instanceWrapped); if (!instance || typeof instance !== "object") { - return { statusText: "PREPARING", statusClass: "inactive" }; + return { + statusText: "PREPARING", + statusClass: "inactive", + perNodeStatus: [], + errorMessage: null, + }; } const inst = instance as { - shardAssignments?: { runnerToShard?: Record }; + shardAssignments?: { + runnerToShard?: Record; + nodeToRunner?: Record; + }; }; + const nodeToRunner = inst.shardAssignments?.nodeToRunner || {}; const runnerIds = Object.keys(inst.shardAssignments?.runnerToShard || {}); + const totalNodes = runnerIds.length; - const statuses = runnerIds - .map((rid) => { - const r = runnersData[rid]; - if (!r) return null; - const [kind] = getTagged(r); - const statusMap: Record = { - RunnerWaitingForInitialization: "WaitingForInitialization", - RunnerInitializingBackend: "InitializingBackend", - RunnerWaitingForModel: "WaitingForModel", - RunnerLoading: "Loading", - RunnerLoaded: "Loaded", - RunnerWarmingUp: "WarmingUp", - RunnerReady: "Ready", - RunnerRunning: "Running", - RunnerShutdown: "Shutdown", - RunnerFailed: "Failed", - }; - return kind ? statusMap[kind] || null : null; - }) - .filter((s): s is string => s !== null); + // Build per-node status and extract error messages from RunnerFailed + const perNodeStatus: PerNodeRunnerStatus[] = []; + const statuses: string[] = []; + const failedErrors: string[] = []; + for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) { + const r = runnersData[runnerId]; + let status: string | null = null; + if (r) { + const [kind, runnerData] = getTagged(r); + status = kind ? RUNNER_STATUS_MAP[kind] || null : null; + // Extract error message from RunnerFailed + if ( + kind === "RunnerFailed" && + runnerData && + typeof runnerData === "object" + ) { + const rd = runnerData as { errorMessage?: string }; + if (rd.errorMessage) + failedErrors.push(`${getNodeName(nodeId)}: ${rd.errorMessage}`); + } + } + if (status) { + statuses.push(status); + perNodeStatus.push({ + nodeId, + nodeName: getNodeName(nodeId), + status: RUNNER_STATUS_DISPLAY[status] || status, + }); + } + } const has = (s: string) => statuses.includes(s); + const count = (s: string) => statuses.filter((v) => v === s).length; if (statuses.length === 0) - return { statusText: "PREPARING", statusClass: "inactive" }; - if (has("Failed")) return { statusText: "FAILED", statusClass: "failed" }; + return { + statusText: "PREPARING", + statusClass: "inactive", + perNodeStatus, + errorMessage: null, + }; + if (has("Failed")) + return { + statusText: "FAILED", + statusClass: "failed", + perNodeStatus, + errorMessage: failedErrors.length > 0 ? failedErrors.join("; ") : null, + }; if (has("Shutdown")) - return { statusText: "SHUTDOWN", statusClass: "inactive" }; - if (has("Loading")) - return { statusText: "LOADING", statusClass: "starting" }; - if (has("WarmingUp")) - return { statusText: "WARMING UP", statusClass: "starting" }; - if (has("Running")) - return { statusText: "RUNNING", statusClass: "running" }; - if (has("Ready")) return { statusText: "READY", statusClass: "loaded" }; - if (has("Loaded")) return { statusText: "LOADED", statusClass: "loaded" }; - if (has("WaitingForModel")) - return { statusText: "WAITING", statusClass: "starting" }; - if (has("InitializingBackend")) - return { statusText: "INITIALIZING", statusClass: "starting" }; - if (has("WaitingForInitialization")) - return { statusText: "INITIALIZING", statusClass: "starting" }; + return { + statusText: "SHUTDOWN", + statusClass: "inactive", + perNodeStatus, + errorMessage: null, + }; - return { statusText: "RUNNING", statusClass: "active" }; + // For loading/warming states, show node progress when multi-node + if (has("Loading")) { + const readyCount = count("Ready") + count("Running") + count("Loaded"); + const statusText = + totalNodes > 1 + ? `LOADING (${readyCount}/${totalNodes} nodes ready)` + : "LOADING"; + return { + statusText, + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; + } + if (has("WarmingUp")) { + const readyCount = count("Ready") + count("Running"); + const statusText = + totalNodes > 1 + ? `WARMING UP (${readyCount}/${totalNodes} nodes ready)` + : "WARMING UP"; + return { + statusText, + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; + } + + if (has("Running")) + return { + statusText: "RUNNING", + statusClass: "running", + perNodeStatus, + errorMessage: null, + }; + if (has("Ready")) + return { + statusText: "READY", + statusClass: "loaded", + perNodeStatus, + errorMessage: null, + }; + if (has("Loaded")) + return { + statusText: "LOADED", + statusClass: "loaded", + perNodeStatus, + errorMessage: null, + }; + if (has("WaitingForModel")) + return { + statusText: "WAITING", + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; + if (has("InitializingBackend")) + return { + statusText: "INITIALIZING", + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; + if (has("WaitingForInitialization")) + return { + statusText: "INITIALIZING", + statusClass: "starting", + perNodeStatus, + errorMessage: null, + }; + + return { + statusText: "RUNNING", + statusClass: "active", + perNodeStatus, + errorMessage: null, + }; } function getBytes(value: unknown): number { @@ -1255,6 +1459,75 @@ } } + async function deleteMetaInstance(metaInstanceId: string) { + const meta = metaInstancesData[metaInstanceId]; + const modelId = meta?.modelId ?? "unknown"; + if (!confirm(`Delete model ${modelId}?`)) return; + + const wasSelected = selectedChatModel() === modelId; + + try { + const response = await fetch(`/meta_instance/${metaInstanceId}`, { + method: "DELETE", + headers: { "Content-Type": "application/json" }, + }); + + if (!response.ok) { + console.error("Failed to delete meta instance:", response.status); + } else if (wasSelected) { + // Switch to another available model or clear selection + const remainingInstances = Object.entries(instanceData).filter( + ([id]) => id !== getBackingInstanceId(metaInstanceId), + ); + if (remainingInstances.length > 0) { + const [, lastInstance] = + remainingInstances[remainingInstances.length - 1]; + const newModelId = getInstanceModelId(lastInstance); + if ( + newModelId && + newModelId !== "Unknown" && + newModelId !== "Unknown Model" + ) { + setSelectedChatModel(newModelId); + } else { + setSelectedChatModel(""); + } + } else { + setSelectedChatModel(""); + } + } + } catch (error) { + console.error("Error deleting meta instance:", error); + } + } + + // Find the backing Instance ID for a MetaInstance by scanning instances + function getBackingInstanceId(metaInstanceId: string): string | null { + for (const [id, inst] of Object.entries(instanceData)) { + const [, inner] = getTagged(inst); + if ( + inner && + typeof inner === "object" && + (inner as Record).metaInstanceId === metaInstanceId + ) { + return id; + } + } + return null; + } + + // Get orphan Instance IDs (not backing any MetaInstance) + function getOrphanInstanceIds(): string[] { + return Object.keys(instanceData).filter((id) => { + const [, inner] = getTagged(instanceData[id]); + return ( + !inner || + typeof inner !== "object" || + !(inner as Record).metaInstanceId + ); + }); + } + // Helper to unwrap tagged unions like { MlxRingInstance: {...} } function getTagged(obj: unknown): [string | null, unknown] { if (!obj || typeof obj !== "object") return [null, null]; @@ -1295,11 +1568,7 @@ // Instance type from tag let instanceType = "Unknown"; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; - else if ( - instanceTag === "MlxIbvInstance" || - instanceTag === "MlxJacclInstance" - ) - instanceType = "MLX RDMA"; + else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; const inst = instance as { shardAssignments?: { @@ -1647,7 +1916,51 @@ } const nodeCount = $derived(data ? Object.keys(data.nodes).length : 0); - const instanceCount = $derived(Object.keys(instanceData).length); + const metaInstanceCount = $derived(Object.keys(metaInstancesData).length); + const orphanInstanceIds = $derived(getOrphanInstanceIds()); + const instanceCount = $derived(metaInstanceCount + orphanInstanceIds.length); + + // Unified display items: MetaInstances first, then orphan Instances + interface DisplayItem { + id: string; // MetaInstance ID or Instance ID (used as key and displayed) + modelId: string; + instance: unknown | null; // The backing/orphan instance (tagged union) or null if placing + instanceId: string | null; // The actual Instance ID (for topology hover) + isMetaInstance: boolean; + sharding: string | null; // From MetaInstance constraints (used when instance is null) + instanceMeta: string | null; // From MetaInstance constraints (used when instance is null) + } + + const unifiedDisplayItems = $derived.by((): DisplayItem[] => { + const items: DisplayItem[] = []; + // MetaInstances + for (const [metaId, meta] of Object.entries(metaInstancesData)) { + const backingId = getBackingInstanceId(metaId); + items.push({ + id: metaId, + modelId: meta.modelId, + instance: backingId ? instanceData[backingId] : null, + instanceId: backingId, + isMetaInstance: true, + sharding: meta.sharding, + instanceMeta: meta.instanceMeta, + }); + } + // Orphan Instances + for (const orphanId of getOrphanInstanceIds()) { + const inst = instanceData[orphanId]; + items.push({ + id: orphanId, + modelId: getInstanceModelId(inst), + instance: inst, + instanceId: orphanId, + isMetaInstance: false, + sharding: null, + instanceMeta: null, + }); + } + return items; + }); // Helper to get the number of nodes in a placement preview function getPreviewNodeCount(preview: PlacementPreview): number { @@ -1765,8 +2078,71 @@ {#snippet clusterWarnings()} - {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)} + {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (jacclError && jacclError !== jacclDismissedError)}
+ {#if jacclError && jacclError !== jacclDismissedError} + + {/if} + {#if tbBridgeCycles.length > 0} {@const cycle = tbBridgeCycles[0]} {@const serviceName = getTbBridgeServiceName(cycle)} @@ -1935,8 +2311,29 @@ {/snippet} {#snippet clusterWarningsCompact()} - {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)} + {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (jacclError && jacclError !== jacclDismissedError)}
+ {#if jacclError && jacclError !== jacclDismissedError} +
+ + + + JACCL ERROR +
+ {/if} {#if tbBridgeCycles.length > 0}
- {#each Object.entries(instanceData) as [id, instance]} - {@const downloadInfo = getInstanceDownloadStatus( - id, - instance, - )} + {#each unifiedDisplayItems as item (item.id)} + {@const id = item.id} + {@const instance = item.instance} + {@const downloadInfo = instance + ? getInstanceDownloadStatus(item.instanceId ?? id, instance) + : getMetaInstancePlacingStatus(id)} + {@const metaData = item.isMetaInstance + ? metaInstancesData[id] + : null} + {@const retryError = + metaData?.lastFailureError && !downloadInfo.isFailed + ? metaData.consecutiveFailures > 0 + ? `(${((metaData.consecutiveFailures - 1) % 3) + 1}/3) ${metaData.lastFailureError}` + : metaData.lastFailureError + : null} {@const statusText = downloadInfo.statusText} {@const isDownloading = downloadInfo.isDownloading} - {@const isFailed = statusText === "FAILED"} + {@const isFailed = + statusText === "FAILED" || + statusText === "PLACEMENT FAILED"} {@const isLoading = - statusText === "LOADING" || - statusText === "WARMING UP" || - statusText === "WAITING"} + statusText.startsWith("LOADING") || + statusText.startsWith("WARMING UP") || + statusText === "WAITING" || + statusText === "PLACING" || + statusText.startsWith("RETRYING")} {@const isReady = statusText === "READY" || statusText === "LOADED"} {@const isRunning = statusText === "RUNNING"} - {@const instanceModelId = getInstanceModelId(instance)} - {@const instanceInfo = getInstanceInfo(instance)} - {@const instanceConnections = - getInstanceConnections(instance)} + {@const instanceModelId = item.modelId} + {@const instanceInfo = instance + ? getInstanceInfo(instance) + : { + instanceType: + item.instanceMeta === "MlxRing" + ? "MLX Ring" + : item.instanceMeta === "MlxJaccl" + ? "MLX RDMA" + : "Unknown", + sharding: item.sharding ?? "Unknown", + nodeNames: [] as string[], + nodeIds: [] as string[], + nodeCount: 0, + }} + {@const instanceConnections = instance + ? getInstanceConnections(instance) + : []}
(hoveredInstanceId = id)} + onmouseenter={() => + (hoveredInstanceId = item.instanceId ?? id)} onmouseleave={() => (hoveredInstanceId = null)} onclick={() => { if ( @@ -2438,7 +2864,10 @@ >
@@ -2884,21 +3337,21 @@
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index db13ccef..f661c878 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -314,7 +314,17 @@ class DownloadCoordinator: ), ) elif progress.status in ["in_progress", "not_started"]: - if progress.downloaded_bytes_this_session.in_bytes == 0: + if ( + progress.downloaded_bytes.in_bytes + >= progress.total_bytes.in_bytes + > 0 + ): + status = DownloadCompleted( + node_id=self.node_id, + shard_metadata=progress.shard, + total_bytes=progress.total_bytes, + ) + elif progress.downloaded_bytes_this_session.in_bytes == 0: status = DownloadPending( node_id=self.node_id, shard_metadata=progress.shard, diff --git a/src/exo/main.py b/src/exo/main.py index 1d358975..8f0c5a41 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -254,7 +254,7 @@ def main(): target = min(max(soft, 65535), hard) resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard)) - mp.set_start_method("spawn") + mp.set_start_method("spawn", force=True) # TODO: Refactor the current verbosity system logger_setup(EXO_LOG, args.verbosity) logger.info("Starting EXO") diff --git a/src/exo/master/api.py b/src/exo/master/api.py index b8476334..91c74f41 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -71,8 +71,11 @@ from exo.shared.types.api import ( ChatCompletionResponse, CreateInstanceParams, CreateInstanceResponse, + CreateMetaInstanceParams, + CreateMetaInstanceResponse, DeleteDownloadResponse, DeleteInstanceResponse, + DeleteMetaInstanceResponse, ErrorInfo, ErrorResponse, FinishReason, @@ -115,8 +118,10 @@ from exo.shared.types.claude_api import ( from exo.shared.types.commands import ( Command, CreateInstance, + CreateMetaInstance, DeleteDownload, DeleteInstance, + DeleteMetaInstance, DownloadCommand, ForwarderCommand, ForwarderDownloadCommand, @@ -129,7 +134,7 @@ from exo.shared.types.commands import ( TaskFinished, TextGeneration, ) -from exo.shared.types.common import CommandId, Id, NodeId, SessionId +from exo.shared.types.common import CommandId, Id, MetaInstanceId, NodeId, SessionId from exo.shared.types.events import ( ChunkGenerated, Event, @@ -138,6 +143,7 @@ from exo.shared.types.events import ( TracesMerged, ) from exo.shared.types.memory import Memory +from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.openai_responses import ( ResponsesRequest, ResponsesResponse, @@ -276,6 +282,9 @@ class API: self.app.get("/instance/previews")(self.get_placement_previews) self.app.get("/instance/{instance_id}")(self.get_instance) self.app.delete("/instance/{instance_id}")(self.delete_instance) + self.app.get("/meta_instances")(self.list_meta_instances) + self.app.post("/meta_instance")(self.create_meta_instance) + self.app.delete("/meta_instance/{meta_instance_id}")(self.delete_meta_instance) self.app.get("/models")(self.get_models) self.app.get("/v1/models")(self.get_models) self.app.post("/models/add")(self.add_custom_model) @@ -305,12 +314,27 @@ class API: self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw) async def place_instance(self, payload: PlaceInstanceParams): + model_card = await ModelCard.load(payload.model_id) command = PlaceInstance( - model_card=await ModelCard.load(payload.model_id), + model_card=model_card, sharding=payload.sharding, instance_meta=payload.instance_meta, min_nodes=payload.min_nodes, ) + + # Validate placement before sending — fail fast with a clear error + # instead of silently dropping the command in the master. + try: + get_instance_placements( + command, + topology=self.state.topology, + current_instances=self.state.instances, + node_memory=self.state.node_memory, + node_network=self.state.node_network, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + await self._send(command) return CreateInstanceResponse( @@ -522,6 +546,44 @@ class API: instance_id=instance_id, ) + def list_meta_instances(self) -> dict[MetaInstanceId, MetaInstance]: + return dict(self.state.meta_instances) + + async def create_meta_instance( + self, payload: CreateMetaInstanceParams + ) -> CreateMetaInstanceResponse: + meta_instance = MetaInstance( + model_id=payload.model_id, + sharding=payload.sharding, + instance_meta=payload.instance_meta, + min_nodes=payload.min_nodes, + node_ids=payload.node_ids, + ) + command = CreateMetaInstance(meta_instance=meta_instance) + await self._send(command) + return CreateMetaInstanceResponse( + message="Command received.", + command_id=command.command_id, + meta_instance_id=meta_instance.meta_instance_id, + ) + + async def delete_meta_instance( + self, meta_instance_id: MetaInstanceId + ) -> DeleteMetaInstanceResponse: + meta = self.state.meta_instances.get(meta_instance_id) + if not meta: + raise HTTPException(status_code=404, detail="MetaInstance not found") + + # Command processor handles cascade-deleting backing instances + command = DeleteMetaInstance(meta_instance_id=meta_instance_id) + await self._send(command) + + return DeleteMetaInstanceResponse( + message="Command received.", + command_id=command.command_id, + meta_instance_id=meta_instance_id, + ) + async def _token_chunk_stream( self, command_id: CommandId ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]: @@ -541,10 +603,10 @@ class API: break except anyio.get_cancelled_exc_class(): - command = TaskCancelled(cancelled_command_id=command_id) + cancel_command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( - ForwarderCommand(origin=self.node_id, command=command) + ForwarderCommand(origin=self.node_id, command=cancel_command) ) raise finally: @@ -884,10 +946,10 @@ class API: del image_metadata[key] except anyio.get_cancelled_exc_class(): - command = TaskCancelled(cancelled_command_id=command_id) + cancel_command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( - ForwarderCommand(origin=self.node_id, command=command) + ForwarderCommand(origin=self.node_id, command=cancel_command) ) raise finally: @@ -970,10 +1032,10 @@ class API: return (images, stats if capture_stats else None) except anyio.get_cancelled_exc_class(): - command = TaskCancelled(cancelled_command_id=command_id) + cancel_command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( - ForwarderCommand(origin=self.node_id, command=command) + ForwarderCommand(origin=self.node_id, command=cancel_command) ) raise finally: diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 9c7cf578..405f6495 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -1,4 +1,5 @@ -from datetime import datetime, timedelta, timezone +from collections.abc import Sequence +from datetime import datetime, timezone import anyio from anyio.abc import TaskGroup @@ -12,11 +13,22 @@ from exo.master.placement import ( get_transition_events, place_instance, ) +from exo.master.process_managers import ProcessManager +from exo.master.process_managers.instance_health import InstanceHealthReconciler +from exo.master.process_managers.meta_instance import MetaInstanceReconciler +from exo.master.process_managers.node_timeout import NodeTimeoutReconciler +from exo.master.reconcile import ( + find_unsatisfied_meta_instances, + try_place_for_meta_instance, +) from exo.shared.apply import apply from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED +from exo.shared.models.model_cards import ModelCard from exo.shared.types.commands import ( CreateInstance, + CreateMetaInstance, DeleteInstance, + DeleteMetaInstance, ForwarderCommand, ForwarderDownloadCommand, ImageEdits, @@ -36,8 +48,12 @@ from exo.shared.types.events import ( IndexedEvent, InputChunkReceived, InstanceDeleted, + JacclSideChannelData, + JacclSideChannelGathered, + MetaInstanceCreated, + MetaInstanceDeleted, + MetaInstancePlacementFailed, NodeGatheredInfo, - NodeTimedOut, TaskCreated, TaskDeleted, TaskStatusUpdated, @@ -60,7 +76,8 @@ from exo.shared.types.tasks import ( TextGeneration as TextGenerationTask, ) from exo.shared.types.worker.instances import InstanceId -from exo.utils.channels import Receiver, Sender, channel +from exo.shared.types.worker.runners import RunnerId +from exo.utils.channels import Receiver, Sender from exo.utils.event_buffer import MultiSourceBuffer @@ -84,16 +101,16 @@ class Master: self.local_event_receiver = local_event_receiver self.global_event_sender = global_event_sender self.download_command_sender = download_command_sender - send, recv = channel[Event]() - self.event_sender: Sender[Event] = send - self._loopback_event_receiver: Receiver[Event] = recv - self._loopback_event_sender: Sender[ForwarderEvent] = ( - local_event_receiver.clone_sender() - ) self._multi_buffer = MultiSourceBuffer[NodeId, Event]() self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master") self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} self._expected_ranks: dict[TaskId, set[int]] = {} + self._jaccl_pending: dict[InstanceId, dict[int, dict[RunnerId, bytes]]] = {} + self._process_managers: Sequence[ProcessManager] = [ + InstanceHealthReconciler(), + NodeTimeoutReconciler(), + MetaInstanceReconciler(), + ] async def run(self): logger.info("Starting Master") @@ -102,15 +119,12 @@ class Master: async with self._tg as tg: tg.start_soon(self._event_processor) tg.start_soon(self._command_processor) - tg.start_soon(self._loopback_processor) - tg.start_soon(self._plan) + tg.start_soon(self._reconcile) finally: self._event_log.close() self.global_event_sender.close() self.local_event_receiver.close() self.command_receiver.close() - self._loopback_event_sender.close() - self._loopback_event_receiver.close() async def shutdown(self): logger.info("Stopping Master") @@ -292,6 +306,86 @@ class Master: ) ) generated_events.extend(transition_events) + case CreateMetaInstance(): + logger.info( + f"Creating MetaInstance for {command.meta_instance.model_id}" + f" (min_nodes={command.meta_instance.min_nodes}," + f" sharding={command.meta_instance.sharding})" + ) + # Apply immediately so self.state is fresh across + # the await below and the reconciler won't race. + await self._apply_and_broadcast( + MetaInstanceCreated(meta_instance=command.meta_instance) + ) + # Immediate placement attempt for responsiveness + model_card = await ModelCard.load( + command.meta_instance.model_id + ) + # Re-check: reconciler may have satisfied it during the await + meta_id = command.meta_instance.meta_instance_id + still_unsatisfied = any( + m.meta_instance_id == meta_id + for m in find_unsatisfied_meta_instances( + self.state.meta_instances, + self.state.instances, + self.state.topology, + ) + ) + if still_unsatisfied: + result = try_place_for_meta_instance( + command.meta_instance, + model_card, + self.state.topology, + self.state.instances, + self.state.node_memory, + self.state.node_network, + self.state.tasks, + ) + generated_events.extend(result.events) + if result.error is not None: + generated_events.append( + MetaInstancePlacementFailed( + meta_instance_id=meta_id, + reason=result.error, + ) + ) + case DeleteMetaInstance(): + backing_count = sum( + 1 + for inst in self.state.instances.values() + if inst.meta_instance_id == command.meta_instance_id + ) + logger.info( + f"Deleting MetaInstance {command.meta_instance_id}" + f" (cascade-deleting {backing_count} backing instance(s))" + ) + generated_events.append( + MetaInstanceDeleted( + meta_instance_id=command.meta_instance_id + ) + ) + # Cascade-delete backing instances atomically, + # cancelling any active tasks first. + for iid, inst in self.state.instances.items(): + if inst.meta_instance_id == command.meta_instance_id: + for task in self.state.tasks.values(): + if ( + task.instance_id == iid + and task.task_status + in ( + TaskStatus.Pending, + TaskStatus.Running, + ) + ): + generated_events.append( + TaskStatusUpdated( + task_status=TaskStatus.Cancelled, + task_id=task.task_id, + ) + ) + generated_events.append( + InstanceDeleted(instance_id=iid) + ) case PlaceInstance(): placement = place_instance( command, @@ -323,16 +417,19 @@ class Master: ) case TaskCancelled(): if ( - task_id := self.command_task_mapping.get( - command.cancelled_command_id - ) - ) is not None: + command.cancelled_command_id + in self.command_task_mapping + ): generated_events.append( - TaskStatusUpdated( - task_status=TaskStatus.Cancelled, - task_id=task_id, + TaskDeleted( + task_id=self.command_task_mapping[ + command.cancelled_command_id + ] ) ) + del self.command_task_mapping[ + command.cancelled_command_id + ] case TaskFinished(): generated_events.append( TaskDeleted( @@ -341,9 +438,10 @@ class Master: ] ) ) - self.command_task_mapping.pop( - command.finished_command_id, None - ) + if command.finished_command_id in self.command_task_mapping: + del self.command_task_mapping[ + command.finished_command_id + ] case RequestEventLog(): # We should just be able to send everything, since other buffers will ignore old messages # rate limit to 1000 at a time @@ -354,31 +452,32 @@ class Master: ): await self._send_event(IndexedEvent(idx=i, event=event)) for event in generated_events: - await self.event_sender.send(event) + await self._apply_and_broadcast(event) except ValueError as e: logger.opt(exception=e).warning("Error in command processor") - # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands - async def _plan(self) -> None: + async def _apply_and_broadcast(self, event: Event) -> None: + """Apply event to state, persist to disk, and broadcast to workers. + + State is updated synchronously (before any await), so callers can + rely on ``self.state`` reflecting this event immediately after the + call. Python's cooperative scheduling guarantees no interleaving + between the state read and write. + """ + logger.debug(f"Master indexing event: {str(event)[:100]}") + indexed = IndexedEvent(event=event, idx=len(self._event_log)) + self.state = apply(self.state, indexed) + event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage] + self._event_log.append(event) + await self._send_event(indexed) + + async def _reconcile(self) -> None: while True: - # kill broken instances - connected_node_ids = set(self.state.topology.list_nodes()) - for instance_id, instance in self.state.instances.items(): - for node_id in instance.shard_assignments.node_to_runner: - if node_id not in connected_node_ids: - await self.event_sender.send( - InstanceDeleted(instance_id=instance_id) - ) - break - - # time out dead nodes - for node_id, time in self.state.last_seen.items(): - now = datetime.now(tz=timezone.utc) - if now - time > timedelta(seconds=30): - logger.info(f"Manually removing node {node_id} due to inactivity") - await self.event_sender.send(NodeTimedOut(node_id=node_id)) - - await anyio.sleep(10) + for pm in self._process_managers: + events = await pm.reconcile(self.state) + for event in events: + await self._apply_and_broadcast(event) + await anyio.sleep(1) async def _event_processor(self) -> None: with self.local_event_receiver as local_events: @@ -396,32 +495,15 @@ class Master: await self._handle_traces_collected(event) continue - logger.debug(f"Master indexing event: {str(event)[:100]}") - indexed = IndexedEvent(event=event, idx=len(self._event_log)) - self.state = apply(self.state, indexed) + if isinstance(event, JacclSideChannelData): + await self._apply_and_broadcast(event) + await self._handle_jaccl_side_channel(event) + continue - event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage] if isinstance(event, NodeGatheredInfo): event.when = str(datetime.now(tz=timezone.utc)) - self._event_log.append(event) - await self._send_event(indexed) - - async def _loopback_processor(self) -> None: - # this would ideally not be necessary. - # this is WAY less hacky than how I was working around this before - local_index = 0 - with self._loopback_event_receiver as events: - async for event in events: - await self._loopback_event_sender.send( - ForwarderEvent( - origin=NodeId(f"master_{self.node_id}"), - origin_idx=local_index, - session=self.session_id, - event=event, - ) - ) - local_index += 1 + await self._apply_and_broadcast(event) # This function is re-entrant, take care! async def _send_event(self, event: IndexedEvent): @@ -453,10 +535,49 @@ class Master: for trace_data in self._pending_traces[task_id].values(): all_trace_data.extend(trace_data) - await self.event_sender.send( + await self._apply_and_broadcast( TracesMerged(task_id=task_id, traces=all_trace_data) ) del self._pending_traces[task_id] if task_id in self._expected_ranks: del self._expected_ranks[task_id] + + async def _handle_jaccl_side_channel(self, event: JacclSideChannelData) -> None: + """Accumulate SideChannel contributions; when all runners for an instance + have submitted for the same sequence, emit JacclSideChannelGathered.""" + iid = event.instance_id + seq = event.sequence + + if iid not in self._jaccl_pending: + self._jaccl_pending[iid] = {} + if seq not in self._jaccl_pending[iid]: + self._jaccl_pending[iid][seq] = {} + self._jaccl_pending[iid][seq][event.runner_id] = event.data + + instance = self.state.instances.get(iid) + if instance is None: + logger.warning(f"JacclSideChannelData for unknown instance {iid}") + return + + expected_runners = set(instance.shard_assignments.runner_to_shard.keys()) + submitted = set(self._jaccl_pending[iid][seq].keys()) + + logger.info( + f"JACCL side channel: instance={iid} seq={seq} " + f"submitted={len(submitted)}/{len(expected_runners)}" + ) + + if submitted >= expected_runners: + gathered = dict(self._jaccl_pending[iid][seq]) + del self._jaccl_pending[iid][seq] + if not self._jaccl_pending[iid]: + del self._jaccl_pending[iid] + + await self._apply_and_broadcast( + JacclSideChannelGathered( + instance_id=iid, + sequence=seq, + gathered_data=gathered, + ) + ) diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index cf31ca78..ab886c3e 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -6,11 +6,11 @@ from typing import Sequence from exo.master.placement_utils import ( Cycle, filter_cycles_by_memory, + get_largest_cycles, get_mlx_jaccl_coordinators, get_mlx_jaccl_devices_matrix, get_mlx_ring_hosts_by_node, get_shard_assignments, - get_smallest_cycles, ) from exo.shared.models.model_cards import ModelId from exo.shared.topology import Topology @@ -106,23 +106,27 @@ def place_instance( "Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)" ) - smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory) + largest_cycles = get_largest_cycles(cycles_with_sufficient_memory) - smallest_rdma_cycles = [ - cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle) + largest_rdma_cycles = [ + cycle for cycle in largest_cycles if topology.is_rdma_cycle(cycle) ] - if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []: - smallest_cycles = smallest_rdma_cycles + if command.instance_meta == InstanceMeta.MlxJaccl: + if not largest_rdma_cycles: + raise ValueError( + "Requested RDMA (MlxJaccl) but no RDMA-connected cycles available" + ) + largest_cycles = largest_rdma_cycles cycles_with_leaf_nodes: list[Cycle] = [ cycle - for cycle in smallest_cycles + for cycle in largest_cycles if any(topology.node_is_leaf(node_id) for node_id in cycle) ] selected_cycle = max( - cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles, + cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else largest_cycles, key=lambda cycle: sum( (node_memory[node_id].ram_available for node_id in cycle), start=Memory(), diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index b20a39cc..d47c40c0 100644 --- a/src/exo/master/placement_utils.py +++ b/src/exo/master/placement_utils.py @@ -37,11 +37,11 @@ def filter_cycles_by_memory( return filtered_cycles -def get_smallest_cycles( +def get_largest_cycles( cycles: list[Cycle], ) -> list[Cycle]: - min_nodes = min(len(cycle) for cycle in cycles) - return [cycle for cycle in cycles if len(cycle) == min_nodes] + max_nodes = max(len(cycle) for cycle in cycles) + return [cycle for cycle in cycles if len(cycle) == max_nodes] def allocate_layers_proportionally( diff --git a/src/exo/master/process_managers/__init__.py b/src/exo/master/process_managers/__init__.py new file mode 100644 index 00000000..f6a3ee0a --- /dev/null +++ b/src/exo/master/process_managers/__init__.py @@ -0,0 +1,12 @@ +from collections.abc import Sequence +from typing import Protocol, runtime_checkable + +from exo.shared.types.events import Event +from exo.shared.types.state import State + + +@runtime_checkable +class ProcessManager(Protocol): + """A reconciliation step that examines state and returns corrective events.""" + + async def reconcile(self, state: State) -> Sequence[Event]: ... diff --git a/src/exo/master/process_managers/instance_health.py b/src/exo/master/process_managers/instance_health.py new file mode 100644 index 00000000..f5f8e922 --- /dev/null +++ b/src/exo/master/process_managers/instance_health.py @@ -0,0 +1,62 @@ +from collections.abc import Sequence +from typing import final + +from loguru import logger + +from exo.master.reconcile import instance_connections_healthy, instance_runners_failed +from exo.shared.types.events import Event, InstanceDeleted, InstanceRetrying +from exo.shared.types.state import State + +MAX_INSTANCE_RETRIES = 3 + + +@final +class InstanceHealthReconciler: + """Delete instances whose network connections are broken or whose runners have all failed.""" + + async def reconcile(self, state: State) -> Sequence[Event]: + events: list[Event] = [] + for instance_id, instance in state.instances.items(): + if not instance_connections_healthy(instance, state.topology): + events.append( + InstanceDeleted( + instance_id=instance_id, + failure_error="Network connection lost", + ) + ) + continue + + is_failed, error_message = instance_runners_failed( + instance, state.runners, state.node_identities + ) + if is_failed: + # Retry within the same instance if backed by a MetaInstance + mid = instance.meta_instance_id + mi = state.meta_instances.get(mid) if mid else None + if mid and mi and mi.consecutive_failures < MAX_INSTANCE_RETRIES: + logger.info( + f"Instance {instance_id} failed (attempt" + f" {mi.consecutive_failures + 1}/{MAX_INSTANCE_RETRIES})," + f" retrying: {error_message}" + ) + events.append( + InstanceRetrying( + instance_id=instance_id, + meta_instance_id=mid, + failure_error=error_message or "Runner failed", + ) + ) + else: + if mid and mi: + logger.warning( + f"Instance {instance_id} exceeded retry limit" + f" ({MAX_INSTANCE_RETRIES}), deleting:" + f" {error_message}" + ) + events.append( + InstanceDeleted( + instance_id=instance_id, + failure_error=error_message, + ) + ) + return events diff --git a/src/exo/master/process_managers/meta_instance.py b/src/exo/master/process_managers/meta_instance.py new file mode 100644 index 00000000..93037ea8 --- /dev/null +++ b/src/exo/master/process_managers/meta_instance.py @@ -0,0 +1,92 @@ +from collections.abc import Sequence +from typing import final + +import anyio +from loguru import logger + +from exo.master.reconcile import ( + find_unsatisfied_meta_instances, + try_place_for_meta_instance, +) +from exo.shared.models.model_cards import ModelCard +from exo.shared.types.events import Event, InstanceCreated, MetaInstancePlacementFailed +from exo.shared.types.state import State +from exo.shared.types.worker.instances import Instance, InstanceId + +MODEL_CARD_LOAD_TIMEOUT_SECONDS = 10 + + +@final +class MetaInstanceReconciler: + """Place instances for unsatisfied MetaInstances.""" + + async def reconcile(self, state: State) -> Sequence[Event]: + all_events: list[Event] = [] + # Local copy for intermediate tracking — so placement of B + # sees A's instance and doesn't double-place on same resources. + current_instances: dict[InstanceId, Instance] = dict(state.instances) + + unsatisfied = find_unsatisfied_meta_instances( + state.meta_instances, + current_instances, + state.topology, + ) + for meta_instance in unsatisfied: + try: + with anyio.fail_after(MODEL_CARD_LOAD_TIMEOUT_SECONDS): + model_card = await ModelCard.load(meta_instance.model_id) + except TimeoutError: + logger.warning( + f"ModelCard.load timed out for {meta_instance.model_id}, skipping this cycle" + ) + continue + except Exception as exc: + logger.warning( + f"ModelCard.load failed for {meta_instance.model_id}: {exc}" + ) + error = f"Failed to load model card: {exc}" + if meta_instance.placement_error != error: + all_events.append( + MetaInstancePlacementFailed( + meta_instance_id=meta_instance.meta_instance_id, + reason=error, + ) + ) + continue + + result = try_place_for_meta_instance( + meta_instance, + model_card, + state.topology, + current_instances, + state.node_memory, + state.node_network, + state.tasks, + ) + # Update local instance map so next placement sees this one + for event in result.events: + if isinstance(event, InstanceCreated): + logger.info( + f"MetaInstance reconciler placed instance" + f" {event.instance.instance_id} for" + f" {meta_instance.model_id}" + ) + current_instances[event.instance.instance_id] = event.instance + all_events.extend(result.events) + + # Emit placement failure if error differs from what's already in state + if ( + result.error is not None + and meta_instance.placement_error != result.error + ): + logger.warning( + f"MetaInstance placement failed for" + f" {meta_instance.model_id}: {result.error}" + ) + all_events.append( + MetaInstancePlacementFailed( + meta_instance_id=meta_instance.meta_instance_id, + reason=result.error, + ) + ) + return all_events diff --git a/src/exo/master/process_managers/node_timeout.py b/src/exo/master/process_managers/node_timeout.py new file mode 100644 index 00000000..98045c25 --- /dev/null +++ b/src/exo/master/process_managers/node_timeout.py @@ -0,0 +1,27 @@ +from collections.abc import Sequence +from datetime import datetime, timedelta, timezone +from typing import final + +from loguru import logger + +from exo.shared.types.events import Event, NodeTimedOut +from exo.shared.types.state import State + +_DEFAULT_TIMEOUT = timedelta(seconds=30) + + +@final +class NodeTimeoutReconciler: + """Time out nodes that haven't been seen recently.""" + + def __init__(self, timeout: timedelta = _DEFAULT_TIMEOUT) -> None: + self.timeout = timeout + + async def reconcile(self, state: State) -> Sequence[Event]: + now = datetime.now(tz=timezone.utc) + events: list[Event] = [] + for node_id, last_seen in state.last_seen.items(): + if now - last_seen > self.timeout: + logger.info(f"Removing node {node_id} due to inactivity") + events.append(NodeTimedOut(node_id=node_id)) + return events diff --git a/src/exo/master/reconcile.py b/src/exo/master/reconcile.py new file mode 100644 index 00000000..4ca968b8 --- /dev/null +++ b/src/exo/master/reconcile.py @@ -0,0 +1,244 @@ +from collections.abc import Mapping, Sequence +from typing import NamedTuple + +from loguru import logger + +from exo.master.placement import get_transition_events, place_instance +from exo.shared.models.model_cards import ModelCard +from exo.shared.topology import Topology +from exo.shared.types.commands import PlaceInstance +from exo.shared.types.common import MetaInstanceId, NodeId +from exo.shared.types.events import Event +from exo.shared.types.meta_instance import MetaInstance +from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo +from exo.shared.types.tasks import Task, TaskId +from exo.shared.types.topology import RDMAConnection, SocketConnection +from exo.shared.types.worker.instances import ( + BaseInstance, + Instance, + InstanceId, + MlxJacclInstance, + MlxRingInstance, +) +from exo.shared.types.worker.runners import ( + RunnerFailed, + RunnerId, + RunnerShutdown, + RunnerStatus, +) + + +class PlacementResult(NamedTuple): + """Result of a placement attempt: events to apply and optional error reason.""" + + events: Sequence[Event] + error: str | None + + +def _get_ring_order(instance: BaseInstance) -> list[NodeId]: + """Reconstruct ring order from shard device_rank.""" + node_ranks: list[tuple[NodeId, int]] = [] + for node_id, runner_id in instance.shard_assignments.node_to_runner.items(): + shard = instance.shard_assignments.runner_to_shard[runner_id] + node_ranks.append((node_id, shard.device_rank)) + node_ranks.sort(key=lambda x: x[1]) + return [node_id for node_id, _ in node_ranks] + + +def _ring_connections_healthy(instance: MlxRingInstance, topology: Topology) -> bool: + """Check that the specific IPs used by a ring instance still exist in the topology.""" + ring = _get_ring_order(instance) + n = len(ring) + for node in ring: + hosts = instance.hosts_by_node[node] + for idx in range(n): + host = hosts[idx] + if host.ip in ("0.0.0.0", "198.51.100.1"): + continue # self or placeholder + # Real connection: node → ring[idx]. Check specific IP. + connections = topology.get_all_connections_between(node, ring[idx]) + if not any( + isinstance(c, SocketConnection) + and c.sink_multiaddr.ip_address == host.ip + for c in connections + ): + return False + return True + + +def _jaccl_connections_healthy(instance: MlxJacclInstance, topology: Topology) -> bool: + """Check that the specific RDMA interfaces used by a JACCL instance still exist.""" + ring = _get_ring_order(instance) + n = len(ring) + for i in range(n): + for j in range(n): + iface = instance.jaccl_devices[i][j] + if iface is None: + continue + connections = topology.get_all_connections_between(ring[i], ring[j]) + if not any( + isinstance(c, RDMAConnection) and c.source_rdma_iface == iface + for c in connections + ): + return False + return True + + +def instance_connections_healthy(instance: Instance, topology: Topology) -> bool: + """Check that an instance's nodes and specific connections are still in the topology.""" + instance_nodes = set(instance.shard_assignments.node_to_runner.keys()) + if not all(topology.contains_node(n) for n in instance_nodes): + return False + if len(instance_nodes) <= 1: + return True + match instance: + case MlxRingInstance(): + return _ring_connections_healthy(instance, topology) + case MlxJacclInstance(): + return _jaccl_connections_healthy(instance, topology) + + +def instance_runners_failed( + instance: Instance, + runners: Mapping[RunnerId, RunnerStatus], + node_identities: Mapping[NodeId, NodeIdentity], +) -> tuple[bool, str | None]: + """Check if an instance's runners have all reached terminal failure states. + + Returns ``(True, error_message)`` when ALL runners are terminal + (``RunnerFailed`` or ``RunnerShutdown``) and at least one is ``RunnerFailed``. + + Returns ``(False, None)`` when runners are still active, haven't reported + yet, or all gracefully shut down (no ``RunnerFailed``). + """ + instance_runner_ids = set(instance.shard_assignments.node_to_runner.values()) + + if not instance_runner_ids: + return False, None + + # Build reverse mapping: runner_id -> node_id + runner_to_node: dict[RunnerId, NodeId] = { + runner_id: node_id + for node_id, runner_id in instance.shard_assignments.node_to_runner.items() + } + + has_any_failed = False + error_messages: list[str] = [] + + for runner_id in instance_runner_ids: + status = runners.get(runner_id) + if status is None: + # Runner hasn't reported yet — instance is still starting + return False, None + if isinstance(status, RunnerFailed): + has_any_failed = True + if status.error_message: + node_id = runner_to_node.get(runner_id) + name = ( + node_identities[node_id].friendly_name + if node_id and node_id in node_identities + else node_id or "unknown" + ) + error_messages.append(f"{name}: {status.error_message}") + elif isinstance(status, RunnerShutdown): + pass # Terminal but not a failure indicator on its own + else: + # Runner is still active (connecting, loading, running, etc.) + return False, None + + if has_any_failed: + return True, "; ".join(error_messages) if error_messages else "Runner failed" + + # All runners are Shutdown but none Failed — graceful shutdown, not a failure + return False, None + + +def instance_satisfies_meta_instance( + meta_instance: MetaInstance, + instance: Instance, +) -> bool: + """Check if a single instance satisfies a meta-instance's constraints. + + This is a pure constraint check (model, min_nodes, node_ids). + Use ``instance_connections_healthy`` separately for topology health. + """ + if instance.shard_assignments.model_id != meta_instance.model_id: + return False + + instance_nodes = set(instance.shard_assignments.node_to_runner.keys()) + + if len(instance_nodes) < meta_instance.min_nodes: + return False + + return meta_instance.node_ids is None or set(meta_instance.node_ids).issubset( + instance_nodes + ) + + +def find_unsatisfied_meta_instances( + meta_instances: Mapping[MetaInstanceId, MetaInstance], + instances: Mapping[InstanceId, Instance], + topology: Topology, +) -> Sequence[MetaInstance]: + """Return meta-instances that have no healthy backing instance.""" + unsatisfied: list[MetaInstance] = [] + for meta_id, meta_instance in meta_instances.items(): + has_healthy_backing = any( + instance.meta_instance_id == meta_id + and instance_connections_healthy(instance, topology) + for instance in instances.values() + ) + if not has_healthy_backing: + unsatisfied.append(meta_instance) + return unsatisfied + + +def try_place_for_meta_instance( + meta_instance: MetaInstance, + model_card: ModelCard, + topology: Topology, + current_instances: Mapping[InstanceId, Instance], + node_memory: Mapping[NodeId, MemoryUsage], + node_network: Mapping[NodeId, NodeNetworkInfo], + tasks: Mapping[TaskId, Task], +) -> PlacementResult: + """Try to place an instance satisfying the meta-instance constraints. + + Returns a :class:`PlacementResult` with events on success, or an error + reason on failure. + """ + command = PlaceInstance( + model_card=model_card, + sharding=meta_instance.sharding, + instance_meta=meta_instance.instance_meta, + min_nodes=meta_instance.min_nodes, + ) + try: + target_instances = place_instance( + command, + topology, + current_instances, + node_memory, + node_network, + required_nodes=( + set(meta_instance.node_ids) if meta_instance.node_ids else None + ), + ) + # Tag the new instance with meta_instance_id + new_instance_ids = set(target_instances.keys()) - set(current_instances.keys()) + if new_instance_ids: + new_id = next(iter(new_instance_ids)) + target_instances[new_id] = target_instances[new_id].model_copy( + update={"meta_instance_id": meta_instance.meta_instance_id} + ) + return PlacementResult( + events=list( + get_transition_events(current_instances, target_instances, tasks) + ), + error=None, + ) + except ValueError as e: + logger.debug( + f"MetaInstance placement not possible for {meta_instance.model_id}: {e}" + ) + return PlacementResult(events=[], error=str(e)) diff --git a/src/exo/master/tests/test_meta_instance_edge_cases.py b/src/exo/master/tests/test_meta_instance_edge_cases.py new file mode 100644 index 00000000..26554834 --- /dev/null +++ b/src/exo/master/tests/test_meta_instance_edge_cases.py @@ -0,0 +1,778 @@ +"""Edge-case and regression tests for MetaInstance lifecycle, concurrent operations, and error handling.""" + +import pytest + +from exo.master.process_managers.instance_health import ( + MAX_INSTANCE_RETRIES, + InstanceHealthReconciler, +) +from exo.master.process_managers.meta_instance import MetaInstanceReconciler +from exo.master.reconcile import ( + find_unsatisfied_meta_instances, + instance_connections_healthy, + instance_runners_failed, + instance_satisfies_meta_instance, +) +from exo.shared.apply import apply +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.topology import Topology +from exo.shared.types.common import Host, MetaInstanceId, NodeId +from exo.shared.types.events import ( + IndexedEvent, + InstanceCreated, + InstanceDeleted, + InstanceRetrying, + MetaInstanceCreated, + MetaInstanceDeleted, + MetaInstancePlacementFailed, + TaskStatusUpdated, +) +from exo.shared.types.memory import Memory +from exo.shared.types.meta_instance import MetaInstance +from exo.shared.types.multiaddr import Multiaddr +from exo.shared.types.profiling import NodeIdentity +from exo.shared.types.state import State +from exo.shared.types.tasks import LoadModel, TaskId, TaskStatus +from exo.shared.types.topology import Connection, SocketConnection +from exo.shared.types.worker.instances import ( + InstanceId, + MlxRingInstance, +) +from exo.shared.types.worker.runners import ( + RunnerFailed, + RunnerId, + RunnerReady, + ShardAssignments, +) +from exo.shared.types.worker.shards import PipelineShardMetadata + +# --- Helpers (copied from test_reconcile.py for independence) --- + + +def _model_card(model_id: str = "test-org/test-model") -> ModelCard: + return ModelCard( + model_id=ModelId(model_id), + storage_size=Memory.from_kb(1000), + n_layers=10, + hidden_size=30, + supports_tensor=True, + tasks=[ModelTask.TextGeneration], + ) + + +def _topology(*node_ids: str, connect: bool = True) -> Topology: + t = Topology() + nodes = [NodeId(n) for n in node_ids] + for n in nodes: + t.add_node(n) + if connect and len(nodes) > 1: + for i in range(len(nodes)): + j = (i + 1) % len(nodes) + t.add_connection( + Connection( + source=nodes[i], + sink=nodes[j], + edge=SocketConnection( + sink_multiaddr=Multiaddr( + address=f"/ip4/10.0.0.{j + 1}/tcp/50000" + ) + ), + ) + ) + t.add_connection( + Connection( + source=nodes[j], + sink=nodes[i], + edge=SocketConnection( + sink_multiaddr=Multiaddr( + address=f"/ip4/10.0.0.{i + 1}/tcp/50000" + ) + ), + ) + ) + return t + + +def _meta_instance( + model_id: str = "test-org/test-model", + *, + min_nodes: int = 1, + node_ids: list[NodeId] | None = None, + meta_instance_id: MetaInstanceId | None = None, + consecutive_failures: int = 0, + last_failure_error: str | None = None, + placement_error: str | None = None, +) -> MetaInstance: + return MetaInstance( + meta_instance_id=meta_instance_id or MetaInstanceId(), + model_id=ModelId(model_id), + min_nodes=min_nodes, + node_ids=node_ids, + consecutive_failures=consecutive_failures, + last_failure_error=last_failure_error, + placement_error=placement_error, + ) + + +def _instance( + model_id: str = "test-org/test-model", + node_ids: list[str] | None = None, + instance_id: InstanceId | None = None, + meta_instance_id: MetaInstanceId | None = None, +) -> tuple[InstanceId, MlxRingInstance]: + iid = instance_id or InstanceId() + nodes = node_ids or ["node-a"] + n = len(nodes) + mc = _model_card(model_id) + ephemeral_port = 50000 + node_to_runner = {NodeId(nd): RunnerId() for nd in nodes} + runner_to_shard = { + runner_id: PipelineShardMetadata( + model_card=mc, + device_rank=i, + world_size=n, + start_layer=0, + end_layer=mc.n_layers, + n_layers=mc.n_layers, + ) + for i, runner_id in enumerate(node_to_runner.values()) + } + hosts_by_node: dict[NodeId, list[Host]] = {} + for r, node_str in enumerate(nodes): + hosts: list[Host] = [] + for idx in range(n): + if idx == r: + hosts.append(Host(ip="0.0.0.0", port=ephemeral_port)) + elif n > 1 and idx in ((r - 1) % n, (r + 1) % n): + hosts.append(Host(ip=f"10.0.0.{idx + 1}", port=ephemeral_port)) + else: + hosts.append(Host(ip="198.51.100.1", port=0)) + hosts_by_node[NodeId(node_str)] = hosts + return iid, MlxRingInstance( + instance_id=iid, + shard_assignments=ShardAssignments( + model_id=ModelId(model_id), + runner_to_shard=runner_to_shard, + node_to_runner=node_to_runner, + ), + hosts_by_node=hosts_by_node, + ephemeral_port=ephemeral_port, + meta_instance_id=meta_instance_id, + ) + + +# ============================================================================= +# 1. MetaInstance lifecycle edge cases +# ============================================================================= + + +def test_meta_instance_model_is_frozen(): + """MetaInstance should be immutable (frozen model).""" + meta = _meta_instance() + try: + meta.model_id = ModelId("something-else") + raise AssertionError("Should have raised") + except Exception: + pass # Expected — frozen model + + +def test_meta_instance_created_then_deleted_roundtrip(): + """Create and delete a MetaInstance through apply — state should be clean.""" + state = State() + meta = _meta_instance() + state = apply( + state, IndexedEvent(idx=0, event=MetaInstanceCreated(meta_instance=meta)) + ) + assert meta.meta_instance_id in state.meta_instances + state = apply( + state, + IndexedEvent( + idx=1, event=MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id) + ), + ) + assert meta.meta_instance_id not in state.meta_instances + assert len(state.meta_instances) == 0 + + +def test_delete_nonexistent_meta_instance_is_safe(): + """Deleting a MetaInstance that doesn't exist should not crash.""" + state = State() + event = MetaInstanceDeleted(meta_instance_id=MetaInstanceId("nonexistent")) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert len(new_state.meta_instances) == 0 + + +def test_placement_failed_for_nonexistent_meta_instance_is_safe(): + """MetaInstancePlacementFailed for unknown ID should not crash.""" + state = State() + event = MetaInstancePlacementFailed( + meta_instance_id=MetaInstanceId("nonexistent"), + reason="test", + ) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert len(new_state.meta_instances) == 0 + + +def test_multiple_meta_instances_for_same_model(): + """Multiple MetaInstances for the same model are tracked independently.""" + state = State() + meta_a = _meta_instance("test-org/model-x") + meta_b = _meta_instance("test-org/model-x") + state = apply( + state, IndexedEvent(idx=0, event=MetaInstanceCreated(meta_instance=meta_a)) + ) + state = apply( + state, IndexedEvent(idx=1, event=MetaInstanceCreated(meta_instance=meta_b)) + ) + assert len(state.meta_instances) == 2 + assert meta_a.meta_instance_id in state.meta_instances + assert meta_b.meta_instance_id in state.meta_instances + + +# ============================================================================= +# 2. Retry logic edge cases +# ============================================================================= + + +def test_retry_counter_resets_on_successful_instance_creation(): + """When a new instance is created for a meta-instance, failures should reset.""" + meta = _meta_instance(consecutive_failures=2, last_failure_error="old") + _, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State(meta_instances={meta.meta_instance_id: meta}) + state = apply(state, IndexedEvent(idx=0, event=InstanceCreated(instance=inst))) + mi = state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 0 + # last_failure_error is preserved (for UI display) + assert mi.last_failure_error == "old" + + +async def test_retry_count_increments_through_full_cycle(): + """Walk through MAX_INSTANCE_RETRIES worth of retries, then verify delete.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + topology = _topology("node-a") + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + topology=topology, + ) + + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + for idx, i in enumerate(range(MAX_INSTANCE_RETRIES)): + # Simulate runners failing + state_with_runners = state.model_copy( + update={"runners": {runner_ids[0]: RunnerFailed(error_message=f"fail-{i}")}} + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state_with_runners) + assert len(events) == 1 + assert isinstance(events[0], InstanceRetrying), f"iteration {i}" + state = apply(state, IndexedEvent(idx=idx, event=events[0])) + + # After MAX_INSTANCE_RETRIES retries, failure counter should be at max + mi = state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == MAX_INSTANCE_RETRIES + + # Next failure should result in deletion + state_with_runners = state.model_copy( + update={"runners": {runner_ids[0]: RunnerFailed(error_message="final")}} + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state_with_runners) + assert len(events) == 1 + assert isinstance(events[0], InstanceDeleted) + + +async def test_health_reconciler_respects_exact_limit(): + """At exactly MAX_INSTANCE_RETRIES, reconciler should delete, not retry.""" + meta = _meta_instance(consecutive_failures=MAX_INSTANCE_RETRIES) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, + topology=_topology("node-a"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceDeleted) + + +async def test_health_reconciler_at_limit_minus_one_retries(): + """At MAX_INSTANCE_RETRIES - 1, reconciler should still retry.""" + meta = _meta_instance(consecutive_failures=MAX_INSTANCE_RETRIES - 1) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, + topology=_topology("node-a"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceRetrying) + + +# ============================================================================= +# 3. Error handling edge cases +# ============================================================================= + + +def test_runners_failed_with_empty_error_message(): + """RunnerFailed with empty error_message should still report as failed.""" + _, inst = _instance(node_ids=["node-a"]) + runners = { + rid: RunnerFailed(error_message="") + for rid in inst.shard_assignments.node_to_runner.values() + } + is_failed, error = instance_runners_failed(inst, runners, {}) + assert is_failed is True + # Empty error message means we get the fallback + assert error == "Runner failed" + + +def test_runners_failed_with_none_error_message(): + """RunnerFailed with None error_message should still report as failed.""" + _, inst = _instance(node_ids=["node-a"]) + runners = { + rid: RunnerFailed(error_message=None) + for rid in inst.shard_assignments.node_to_runner.values() + } + is_failed, error = instance_runners_failed(inst, runners, {}) + assert is_failed is True + assert error == "Runner failed" + + +def test_runners_failed_collects_all_error_messages(): + """With multiple failed runners, all error messages should be collected.""" + _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + runners = { + runner_ids[0]: RunnerFailed(error_message="OOM on GPU 0"), + runner_ids[1]: RunnerFailed(error_message="OOM on GPU 1"), + runner_ids[2]: RunnerFailed(error_message="OOM on GPU 2"), + } + is_failed, error = instance_runners_failed(inst, runners, {}) + assert is_failed is True + assert error is not None + assert "OOM on GPU 0" in error + assert "OOM on GPU 1" in error + assert "OOM on GPU 2" in error + + +def test_runners_failed_includes_friendly_name(): + """Error messages should include node friendly names when available.""" + _, inst = _instance(node_ids=["node-a"]) + node_id = NodeId("node-a") + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + runners = {runner_ids[0]: RunnerFailed(error_message="OOM")} + identities = {node_id: NodeIdentity(friendly_name="My Mac Studio")} + is_failed, error = instance_runners_failed(inst, runners, identities) + assert is_failed is True + assert error is not None + assert "My Mac Studio" in error + + +def test_instance_retrying_for_missing_instance_is_safe(): + """InstanceRetrying for an instance not in state should not crash. + + NOTE: When the instance is missing, the handler returns early WITHOUT + incrementing the MetaInstance failure counter. This means stale retry + events for already-deleted instances are silently dropped. This is + acceptable since the InstanceDeleted handler already increments failures. + """ + meta = _meta_instance() + state = State(meta_instances={meta.meta_instance_id: meta}) + event = InstanceRetrying( + instance_id=InstanceId("nonexistent"), + meta_instance_id=meta.meta_instance_id, + failure_error="crash", + ) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + # Does not crash, but failure count is NOT incremented (early return) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 0 + + +# ============================================================================= +# 4. Backward compatibility +# ============================================================================= + + +def test_instance_without_meta_instance_id_works(): + """Instances created without meta_instance_id should still function normally.""" + _, inst = _instance(node_ids=["node-a"]) + assert inst.meta_instance_id is None + topology = _topology("node-a") + assert instance_connections_healthy(inst, topology) is True + + +def test_instance_deleted_without_meta_does_not_affect_meta_instances(): + """Deleting an instance without meta_instance_id should not affect meta_instances.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"]) # no meta_instance_id + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + event = InstanceDeleted(instance_id=iid, failure_error="crash") + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 0 # unchanged + + +def test_satisfies_ignores_meta_instance_id_binding(): + """instance_satisfies_meta_instance checks constraints only, not binding.""" + meta = _meta_instance() + _, inst = _instance(node_ids=["node-a"]) # no meta_instance_id set + # Should match on constraints (model, min_nodes) regardless of binding + assert instance_satisfies_meta_instance(meta, inst) is True + + +def test_find_unsatisfied_uses_binding_not_constraints(): + """find_unsatisfied checks meta_instance_id binding, not just constraint matching.""" + meta = _meta_instance() + # Instance matches constraints but is NOT bound to this meta_instance + iid, inst = _instance(node_ids=["node-a"]) + topology = _topology("node-a") + result = find_unsatisfied_meta_instances( + {meta.meta_instance_id: meta}, {iid: inst}, topology + ) + # Should be unsatisfied because instance.meta_instance_id != meta.meta_instance_id + assert list(result) == [meta] + + +# ============================================================================= +# 5. Concurrent / multi-instance scenarios +# ============================================================================= + + +async def test_health_reconciler_handles_multiple_failing_instances(): + """Multiple instances failing simultaneously should each get their own event.""" + meta_a = _meta_instance() + meta_b = _meta_instance() + iid_a, inst_a = _instance( + node_ids=["node-a"], meta_instance_id=meta_a.meta_instance_id + ) + iid_b, inst_b = _instance( + node_ids=["node-b"], meta_instance_id=meta_b.meta_instance_id + ) + runner_ids_a = list(inst_a.shard_assignments.node_to_runner.values()) + runner_ids_b = list(inst_b.shard_assignments.node_to_runner.values()) + state = State( + meta_instances={ + meta_a.meta_instance_id: meta_a, + meta_b.meta_instance_id: meta_b, + }, + instances={iid_a: inst_a, iid_b: inst_b}, + runners={ + runner_ids_a[0]: RunnerFailed(error_message="OOM"), + runner_ids_b[0]: RunnerFailed(error_message="OOM"), + }, + topology=_topology("node-a", "node-b"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 2 + # Both should be InstanceRetrying since failures < MAX + assert all(isinstance(e, InstanceRetrying) for e in events) + instance_ids = {e.instance_id for e in events} # type: ignore[union-attr] + assert instance_ids == {iid_a, iid_b} + + +async def test_health_reconciler_mixed_healthy_and_failing(): + """Only failing instances should produce events; healthy ones should not.""" + meta_healthy = _meta_instance() + meta_failing = _meta_instance() + iid_h, inst_h = _instance( + node_ids=["node-a"], meta_instance_id=meta_healthy.meta_instance_id + ) + iid_f, inst_f = _instance( + node_ids=["node-b"], meta_instance_id=meta_failing.meta_instance_id + ) + runner_ids_h = list(inst_h.shard_assignments.node_to_runner.values()) + runner_ids_f = list(inst_f.shard_assignments.node_to_runner.values()) + state = State( + meta_instances={ + meta_healthy.meta_instance_id: meta_healthy, + meta_failing.meta_instance_id: meta_failing, + }, + instances={iid_h: inst_h, iid_f: inst_f}, + runners={ + runner_ids_h[0]: RunnerReady(), + runner_ids_f[0]: RunnerFailed(error_message="crash"), + }, + topology=_topology("node-a", "node-b"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceRetrying) + assert events[0].instance_id == iid_f + + +async def test_meta_instance_reconciler_empty_state(): + """MetaInstanceReconciler with no meta_instances should produce no events.""" + state = State() + reconciler = MetaInstanceReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 0 + + +# ============================================================================= +# 6. Placement error tracking +# ============================================================================= + + +def test_placement_failed_sets_error(): + """MetaInstancePlacementFailed should set placement_error on the MetaInstance.""" + meta = _meta_instance() + state = State(meta_instances={meta.meta_instance_id: meta}) + event = MetaInstancePlacementFailed( + meta_instance_id=meta.meta_instance_id, + reason="Not enough memory", + ) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.placement_error == "Not enough memory" + + +def test_instance_created_clears_placement_error(): + """InstanceCreated should clear placement_error on the MetaInstance.""" + meta = _meta_instance(placement_error="Not enough memory") + _, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State(meta_instances={meta.meta_instance_id: meta}) + state = apply(state, IndexedEvent(idx=0, event=InstanceCreated(instance=inst))) + mi = state.meta_instances[meta.meta_instance_id] + assert mi.placement_error is None + + +def test_placement_error_does_not_increment_failures(): + """Placement failures should only set placement_error, not increment consecutive_failures.""" + meta = _meta_instance() + state = State(meta_instances={meta.meta_instance_id: meta}) + event = MetaInstancePlacementFailed( + meta_instance_id=meta.meta_instance_id, + reason="No resources", + ) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 0 + assert mi.placement_error == "No resources" + + +# ============================================================================= +# 7. State serialization roundtrip +# ============================================================================= + + +def test_state_with_meta_instances_serializes(): + """State with meta_instances should serialize and deserialize correctly.""" + meta = _meta_instance(consecutive_failures=2, last_failure_error="test") + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + json_str = state.model_dump_json() + restored = State.model_validate_json(json_str) + assert meta.meta_instance_id in restored.meta_instances + mi = restored.meta_instances[meta.meta_instance_id] + assert mi.model_id == meta.model_id + assert mi.consecutive_failures == 2 + assert mi.last_failure_error == "test" + assert iid in restored.instances + assert restored.instances[iid].meta_instance_id == meta.meta_instance_id + + +# ============================================================================= +# 8. MetaInstanceReconciler error handling +# ============================================================================= + + +async def test_meta_instance_reconciler_model_load_error_emits_placement_failed( + monkeypatch: "pytest.MonkeyPatch", +): + """When ModelCard.load raises, reconciler emits MetaInstancePlacementFailed.""" + import exo.master.process_managers.meta_instance as mi_mod + + meta = _meta_instance() + topo = _topology("node-a") + state = State( + meta_instances={meta.meta_instance_id: meta}, + topology=topo, + ) + + async def _failing_load(_model_id: ModelId) -> ModelCard: + raise RuntimeError("Network error") + + monkeypatch.setattr( + mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_failing_load)}) + ) + + reconciler = MetaInstanceReconciler() + events = await reconciler.reconcile(state) + + placement_failed = [e for e in events if isinstance(e, MetaInstancePlacementFailed)] + assert len(placement_failed) == 1 + assert "Failed to load model card" in placement_failed[0].reason + assert meta.meta_instance_id == placement_failed[0].meta_instance_id + + +async def test_meta_instance_reconciler_model_load_error_skips_dedup( + monkeypatch: "pytest.MonkeyPatch", +): + """When ModelCard.load error matches existing placement_error, no duplicate event.""" + import exo.master.process_managers.meta_instance as mi_mod + + meta = _meta_instance(placement_error="Failed to load model card: Network error") + topo = _topology("node-a") + state = State( + meta_instances={meta.meta_instance_id: meta}, + topology=topo, + ) + + async def _failing_load(_model_id: ModelId) -> ModelCard: + raise RuntimeError("Network error") + + monkeypatch.setattr( + mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_failing_load)}) + ) + + reconciler = MetaInstanceReconciler() + events = await reconciler.reconcile(state) + + # Error matches existing placement_error, so no duplicate event emitted + assert len(events) == 0 + + +async def test_meta_instance_reconciler_continues_after_error( + monkeypatch: "pytest.MonkeyPatch", +): + """Reconciler should continue to next meta-instance after one fails to load.""" + import exo.master.process_managers.meta_instance as mi_mod + + meta_a = _meta_instance(model_id="org/model-a") + meta_b = _meta_instance(model_id="org/model-b") + topo = _topology("node-a") + state = State( + meta_instances={ + meta_a.meta_instance_id: meta_a, + meta_b.meta_instance_id: meta_b, + }, + topology=topo, + ) + + call_count = 0 + + async def _load_second_fails(model_id: ModelId) -> ModelCard: + nonlocal call_count + call_count += 1 + raise RuntimeError(f"Cannot load {model_id}") + + monkeypatch.setattr( + mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_load_second_fails)}) + ) + + reconciler = MetaInstanceReconciler() + events = await reconciler.reconcile(state) + + # Both meta-instances should have been attempted (not short-circuited) + assert call_count == 2 + # Both should have placement failed events + placement_failed = [e for e in events if isinstance(e, MetaInstancePlacementFailed)] + assert len(placement_failed) == 2 + + +# ============================================================================= +# 8. Cascade delete with task cancellation +# ============================================================================= + + +def test_cascade_delete_cancels_active_tasks(): + """Deleting a MetaInstance should cancel tasks on backing instances. + + Regression test: previously, cascade-deleting backing instances via + DeleteMetaInstance did not emit TaskStatusUpdated(Cancelled) for active + tasks, leaving orphaned task references in state. + """ + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + task_id = TaskId() + task = LoadModel(task_id=task_id, instance_id=iid, task_status=TaskStatus.Running) + + # Build state with meta-instance, backing instance, and active task + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + tasks={task_id: task}, + topology=_topology("node-a"), + ) + + # Simulate the cascade-delete event sequence produced by main.py: + # 1. MetaInstanceDeleted + # 2. TaskStatusUpdated(Cancelled) for active tasks + # 3. InstanceDeleted + idx = 0 + state = apply( + state, + IndexedEvent( + idx=idx, + event=MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id), + ), + ) + idx += 1 + state = apply( + state, + IndexedEvent( + idx=idx, + event=TaskStatusUpdated(task_id=task_id, task_status=TaskStatus.Cancelled), + ), + ) + idx += 1 + state = apply( + state, + IndexedEvent(idx=idx, event=InstanceDeleted(instance_id=iid)), + ) + + # Verify everything is cleaned up + assert len(state.meta_instances) == 0 + assert len(state.instances) == 0 + assert state.tasks[task_id].task_status == TaskStatus.Cancelled + + +def test_cascade_delete_skips_completed_tasks(): + """Cascade delete should only cancel Pending/Running tasks, not completed ones.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + + running_task_id = TaskId() + completed_task_id = TaskId() + running_task = LoadModel( + task_id=running_task_id, instance_id=iid, task_status=TaskStatus.Running + ) + completed_task = LoadModel( + task_id=completed_task_id, instance_id=iid, task_status=TaskStatus.Complete + ) + + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + tasks={running_task_id: running_task, completed_task_id: completed_task}, + topology=_topology("node-a"), + ) + + # Only the running task should be cancelled — we verify the logic pattern + # by checking which tasks are Pending or Running + active_tasks = [ + t + for t in state.tasks.values() + if t.instance_id == iid + and t.task_status in (TaskStatus.Pending, TaskStatus.Running) + ] + assert len(active_tasks) == 1 + assert active_tasks[0].task_id == running_task_id diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py index 245c4fd7..9c7ebada 100644 --- a/src/exo/master/tests/test_placement_utils.py +++ b/src/exo/master/tests/test_placement_utils.py @@ -3,10 +3,10 @@ import pytest from exo.master.placement_utils import ( allocate_layers_proportionally, filter_cycles_by_memory, + get_largest_cycles, get_mlx_jaccl_coordinators, get_shard_assignments, get_shard_assignments_for_pipeline_parallel, - get_smallest_cycles, ) from exo.master.tests.conftest import ( create_node_memory, @@ -143,7 +143,7 @@ def test_filter_multiple_cycles_by_memory(): } -def test_get_smallest_cycles(): +def test_get_largest_cycles(): # arrange node_a_id = NodeId() node_b_id = NodeId() @@ -175,12 +175,12 @@ def test_get_smallest_cycles(): cycles = [c for c in topology.get_cycles() if len(c) != 1] # ignore singletons # act - smallest_cycles = get_smallest_cycles(cycles) + largest_cycles = get_largest_cycles(cycles) # assert - assert len(smallest_cycles) == 1 - assert len(smallest_cycles[0]) == 2 - assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id} + assert len(largest_cycles) == 1 + assert len(largest_cycles[0]) == 3 + assert set(n for n in largest_cycles[0]) == {node_a_id, node_b_id, node_c_id} @pytest.mark.parametrize( diff --git a/src/exo/master/tests/test_reconcile.py b/src/exo/master/tests/test_reconcile.py new file mode 100644 index 00000000..e2d6e776 --- /dev/null +++ b/src/exo/master/tests/test_reconcile.py @@ -0,0 +1,742 @@ +from exo.master.process_managers.instance_health import InstanceHealthReconciler +from exo.master.reconcile import ( + find_unsatisfied_meta_instances, + instance_connections_healthy, + instance_runners_failed, + instance_satisfies_meta_instance, +) +from exo.shared.apply import apply +from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask +from exo.shared.topology import Topology +from exo.shared.types.common import Host, MetaInstanceId, NodeId +from exo.shared.types.events import ( + IndexedEvent, + InstanceCreated, + InstanceDeleted, + InstanceRetrying, + MetaInstanceCreated, + MetaInstanceDeleted, +) +from exo.shared.types.memory import Memory +from exo.shared.types.meta_instance import MetaInstance +from exo.shared.types.multiaddr import Multiaddr +from exo.shared.types.state import State +from exo.shared.types.topology import Connection, SocketConnection +from exo.shared.types.worker.instances import ( + InstanceId, + MlxRingInstance, +) +from exo.shared.types.worker.runners import ( + RunnerFailed, + RunnerId, + RunnerLoading, + RunnerReady, + RunnerShutdown, + ShardAssignments, +) +from exo.shared.types.worker.shards import PipelineShardMetadata + + +def _model_card(model_id: str = "test-org/test-model") -> ModelCard: + return ModelCard( + model_id=ModelId(model_id), + storage_size=Memory.from_kb(1000), + n_layers=10, + hidden_size=30, + supports_tensor=True, + tasks=[ModelTask.TextGeneration], + ) + + +def _topology(*node_ids: str, connect: bool = True) -> Topology: + """Build a topology with nodes connected in a bidirectional ring with unique IPs. + + Node at index ``i`` gets IP ``10.0.0.{i+1}``. Edges go in both directions + between consecutive nodes (including wrap-around). + """ + t = Topology() + nodes = [NodeId(n) for n in node_ids] + for n in nodes: + t.add_node(n) + if connect and len(nodes) > 1: + for i in range(len(nodes)): + j = (i + 1) % len(nodes) + t.add_connection( + Connection( + source=nodes[i], + sink=nodes[j], + edge=SocketConnection( + sink_multiaddr=Multiaddr( + address=f"/ip4/10.0.0.{j + 1}/tcp/50000" + ) + ), + ) + ) + t.add_connection( + Connection( + source=nodes[j], + sink=nodes[i], + edge=SocketConnection( + sink_multiaddr=Multiaddr( + address=f"/ip4/10.0.0.{i + 1}/tcp/50000" + ) + ), + ) + ) + return t + + +def _meta_instance( + model_id: str = "test-org/test-model", + *, + min_nodes: int = 1, + node_ids: list[NodeId] | None = None, + meta_instance_id: MetaInstanceId | None = None, +) -> MetaInstance: + return MetaInstance( + meta_instance_id=meta_instance_id or MetaInstanceId(), + model_id=ModelId(model_id), + min_nodes=min_nodes, + node_ids=node_ids, + ) + + +def _instance( + model_id: str = "test-org/test-model", + node_ids: list[str] | None = None, + instance_id: InstanceId | None = None, + meta_instance_id: MetaInstanceId | None = None, +) -> tuple[InstanceId, MlxRingInstance]: + """Create a test instance with hosts_by_node matching ``_topology()`` IPs.""" + iid = instance_id or InstanceId() + nodes = node_ids or ["node-a"] + n = len(nodes) + mc = _model_card(model_id) + ephemeral_port = 50000 + node_to_runner = {NodeId(nd): RunnerId() for nd in nodes} + runner_to_shard = { + runner_id: PipelineShardMetadata( + model_card=mc, + device_rank=i, + world_size=n, + start_layer=0, + end_layer=mc.n_layers, + n_layers=mc.n_layers, + ) + for i, runner_id in enumerate(node_to_runner.values()) + } + # Build hosts_by_node with IPs matching _topology() convention: + # node at index idx has IP 10.0.0.{idx+1} + hosts_by_node: dict[NodeId, list[Host]] = {} + for r, node_str in enumerate(nodes): + hosts: list[Host] = [] + for idx in range(n): + if idx == r: + hosts.append(Host(ip="0.0.0.0", port=ephemeral_port)) + elif n > 1 and idx in ((r - 1) % n, (r + 1) % n): + hosts.append(Host(ip=f"10.0.0.{idx + 1}", port=ephemeral_port)) + else: + hosts.append(Host(ip="198.51.100.1", port=0)) + hosts_by_node[NodeId(node_str)] = hosts + return iid, MlxRingInstance( + instance_id=iid, + shard_assignments=ShardAssignments( + model_id=ModelId(model_id), + runner_to_shard=runner_to_shard, + node_to_runner=node_to_runner, + ), + hosts_by_node=hosts_by_node, + ephemeral_port=ephemeral_port, + meta_instance_id=meta_instance_id, + ) + + +# --- instance_satisfies_meta_instance (pure constraint matching) --- + + +def test_satisfies_matching_model(): + meta = _meta_instance() + _, inst = _instance(node_ids=["node-a"]) + assert instance_satisfies_meta_instance(meta, inst) is True + + +def test_not_satisfies_wrong_model(): + meta = _meta_instance("test-org/model-a") + _, inst = _instance("test-org/model-b") + assert instance_satisfies_meta_instance(meta, inst) is False + + +def test_not_satisfies_missing_required_node(): + meta = _meta_instance(node_ids=[NodeId("node-c")]) + _, inst = _instance(node_ids=["node-a", "node-b"]) + assert instance_satisfies_meta_instance(meta, inst) is False + + +def test_not_satisfies_fewer_than_min_nodes(): + meta = _meta_instance(min_nodes=3) + _, inst = _instance(node_ids=["node-a", "node-b"]) + assert instance_satisfies_meta_instance(meta, inst) is False + + +def test_satisfies_with_node_ids_specified(): + meta = _meta_instance(node_ids=[NodeId("node-a"), NodeId("node-b")], min_nodes=2) + _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) + assert instance_satisfies_meta_instance(meta, inst) is True + + +# --- instance_connections_healthy --- + + +def test_healthy_single_node_present(): + _, inst = _instance(node_ids=["node-a"]) + topology = _topology("node-a") + assert instance_connections_healthy(inst, topology) is True + + +def test_unhealthy_single_node_missing(): + _, inst = _instance(node_ids=["node-a"]) + topology = Topology() # empty + assert instance_connections_healthy(inst, topology) is False + + +def test_healthy_two_node_ring(): + _, inst = _instance(node_ids=["node-a", "node-b"]) + topology = _topology("node-a", "node-b") + assert instance_connections_healthy(inst, topology) is True + + +def test_unhealthy_two_node_edge_removed(): + """Nodes present but edge removed — ring broken.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + topology = _topology("node-a", "node-b", connect=False) + assert instance_connections_healthy(inst, topology) is False + + +def test_unhealthy_two_node_ip_changed(): + """Edge exists but with a different IP than instance was configured with.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + # Build topology with different IPs than _instance() expects + topology = Topology() + topology.add_node(NodeId("node-a")) + topology.add_node(NodeId("node-b")) + topology.add_connection( + Connection( + source=NodeId("node-a"), + sink=NodeId("node-b"), + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/192.168.99.99/tcp/50000") + ), + ) + ) + topology.add_connection( + Connection( + source=NodeId("node-b"), + sink=NodeId("node-a"), + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/192.168.99.98/tcp/50000") + ), + ) + ) + assert instance_connections_healthy(inst, topology) is False + + +def test_healthy_three_node_ring(): + _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) + topology = _topology("node-a", "node-b", "node-c") + assert instance_connections_healthy(inst, topology) is True + + +def test_unhealthy_three_node_one_edge_removed(): + """Remove one edge from a three-node ring — instance unhealthy.""" + _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) + # Build topology with one direction of one edge missing + topology = Topology() + nodes = [NodeId("node-a"), NodeId("node-b"), NodeId("node-c")] + for n in nodes: + topology.add_node(n) + # Add all edges except node-a → node-b + topology.add_connection( + Connection( + source=nodes[1], + sink=nodes[0], + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/50000") + ), + ) + ) + topology.add_connection( + Connection( + source=nodes[1], + sink=nodes[2], + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.3/tcp/50000") + ), + ) + ) + topology.add_connection( + Connection( + source=nodes[2], + sink=nodes[1], + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.2/tcp/50000") + ), + ) + ) + topology.add_connection( + Connection( + source=nodes[2], + sink=nodes[0], + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/50000") + ), + ) + ) + topology.add_connection( + Connection( + source=nodes[0], + sink=nodes[2], + edge=SocketConnection( + sink_multiaddr=Multiaddr(address="/ip4/10.0.0.3/tcp/50000") + ), + ) + ) + # Missing: node-a → node-b (ip 10.0.0.2) + assert instance_connections_healthy(inst, topology) is False + + +def test_unhealthy_node_missing_from_topology(): + """Instance has a node that's not in the topology at all.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + topology = _topology("node-a") # node-b not present + assert instance_connections_healthy(inst, topology) is False + + +def test_healthy_extra_nodes_in_topology(): + """Extra nodes in topology don't affect instance health.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + topology = _topology("node-a", "node-b", "node-c") + assert instance_connections_healthy(inst, topology) is True + + +# --- find_unsatisfied_meta_instances --- + + +def test_unsatisfied_no_meta_instances(): + result = find_unsatisfied_meta_instances({}, {}, Topology()) + assert list(result) == [] + + +def test_unsatisfied_one_satisfied(): + meta = _meta_instance() + id_a, inst_a = _instance(meta_instance_id=meta.meta_instance_id) + topology = _topology("node-a") + result = find_unsatisfied_meta_instances( + {meta.meta_instance_id: meta}, + {id_a: inst_a}, + topology, + ) + assert list(result) == [] + + +def test_unsatisfied_one_not_satisfied(): + meta = _meta_instance("test-org/model-x") + id_a, inst_a = _instance("test-org/model-y") + topology = _topology("node-a") + result = find_unsatisfied_meta_instances( + {meta.meta_instance_id: meta}, {id_a: inst_a}, topology + ) + assert list(result) == [meta] + + +def test_unsatisfied_mix(): + meta_satisfied = _meta_instance("test-org/model-a") + meta_unsatisfied = _meta_instance("test-org/model-b") + id_a, inst_a = _instance( + "test-org/model-a", meta_instance_id=meta_satisfied.meta_instance_id + ) + topology = _topology("node-a") + result = find_unsatisfied_meta_instances( + { + meta_satisfied.meta_instance_id: meta_satisfied, + meta_unsatisfied.meta_instance_id: meta_unsatisfied, + }, + {id_a: inst_a}, + topology, + ) + assert list(result) == [meta_unsatisfied] + + +def test_unsatisfied_node_disconnect(): + meta = _meta_instance() + id_a, inst_a = _instance( + node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id + ) + topology = _topology("node-a") # node-b disconnected + result = find_unsatisfied_meta_instances( + {meta.meta_instance_id: meta}, + {id_a: inst_a}, + topology, + ) + assert list(result) == [meta] + + +def test_unsatisfied_edge_break(): + """Instance exists but its connections broke — meta-instance becomes unsatisfied.""" + meta = _meta_instance() + id_a, inst_a = _instance( + node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id + ) + topology = _topology("node-a", "node-b", connect=False) # nodes present, no edges + result = find_unsatisfied_meta_instances( + {meta.meta_instance_id: meta}, + {id_a: inst_a}, + topology, + ) + assert list(result) == [meta] + + +def test_unsatisfied_idempotent(): + meta = _meta_instance("test-org/model-x") + topology = _topology("node-a") + meta_instances = {meta.meta_instance_id: meta} + instances: dict[InstanceId, MlxRingInstance] = {} + result_1 = list( + find_unsatisfied_meta_instances(meta_instances, instances, topology) + ) + result_2 = list( + find_unsatisfied_meta_instances(meta_instances, instances, topology) + ) + assert result_1 == result_2 + + +def test_unsatisfied_exclusive_binding(): + """Two MetaInstances for the same model: one is bound via meta_instance_id, the other is unsatisfied.""" + meta_a = _meta_instance("test-org/model-x") + meta_b = _meta_instance("test-org/model-x") + id_inst, inst = _instance( + "test-org/model-x", meta_instance_id=meta_a.meta_instance_id + ) + topology = _topology("node-a") + result = find_unsatisfied_meta_instances( + { + meta_a.meta_instance_id: meta_a, + meta_b.meta_instance_id: meta_b, + }, + {id_inst: inst}, + topology, + ) + assert list(result) == [meta_b] + + +# --- apply handlers --- + + +def test_apply_meta_instance_created(): + state = State() + meta = _meta_instance() + event = MetaInstanceCreated(meta_instance=meta) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert meta.meta_instance_id in new_state.meta_instances + assert new_state.meta_instances[meta.meta_instance_id] == meta + + +def test_apply_meta_instance_deleted(): + meta = _meta_instance() + state = State(meta_instances={meta.meta_instance_id: meta}) + event = MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert meta.meta_instance_id not in new_state.meta_instances + + +def test_apply_meta_instance_deleted_clears_failure_info(): + meta = _meta_instance().model_copy( + update={"consecutive_failures": 2, "last_failure_error": "OOM"} + ) + state = State(meta_instances={meta.meta_instance_id: meta}) + event = MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert meta.meta_instance_id not in new_state.meta_instances + + +# --- instance_runners_failed --- + + +def test_runners_failed_all_failed(): + """All runners in RunnerFailed -> instance is failed.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + runners = { + rid: RunnerFailed(error_message="OOM") + for rid in inst.shard_assignments.node_to_runner.values() + } + is_failed, error = instance_runners_failed(inst, runners, {}) + assert is_failed is True + assert error is not None + assert "OOM" in error + + +def test_runners_failed_mixed_failed_shutdown(): + """One Failed + one Shutdown = failed.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + runners = { + runner_ids[0]: RunnerFailed(error_message="crash"), + runner_ids[1]: RunnerShutdown(), + } + is_failed, error = instance_runners_failed(inst, runners, {}) + assert is_failed is True + assert error is not None + assert "crash" in error + + +def test_runners_not_failed_all_shutdown(): + """All Shutdown (graceful) = not a failure.""" + _, inst = _instance(node_ids=["node-a"]) + runners = { + rid: RunnerShutdown() for rid in inst.shard_assignments.node_to_runner.values() + } + is_failed, _ = instance_runners_failed(inst, runners, {}) + assert is_failed is False + + +def test_runners_not_failed_still_active(): + """Some runners still active = not failed yet.""" + _, inst = _instance(node_ids=["node-a", "node-b"]) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + runners = { + runner_ids[0]: RunnerFailed(error_message="OOM"), + runner_ids[1]: RunnerLoading(), + } + is_failed, _ = instance_runners_failed(inst, runners, {}) + assert is_failed is False + + +def test_runners_not_failed_no_status(): + """Runner not yet reported = not failed.""" + _, inst = _instance(node_ids=["node-a"]) + is_failed, _ = instance_runners_failed(inst, {}, {}) + assert is_failed is False + + +def test_runners_not_failed_healthy(): + """Runners in Ready state = not failed.""" + _, inst = _instance(node_ids=["node-a"]) + runners = { + rid: RunnerReady() for rid in inst.shard_assignments.node_to_runner.values() + } + is_failed, _ = instance_runners_failed(inst, runners, {}) + assert is_failed is False + + +# --- failure tracking in apply_instance_deleted --- + + +def test_apply_instance_deleted_tracks_failure(): + """InstanceDeleted with failure_error increments meta instance failure count.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + event = InstanceDeleted(instance_id=iid, failure_error="Runner OOM") + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 1 + assert mi.last_failure_error == "Runner OOM" + + +def test_apply_instance_deleted_increments_failure(): + """Subsequent failures increment the counter.""" + meta = _meta_instance().model_copy( + update={"consecutive_failures": 2, "last_failure_error": "previous error"} + ) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + event = InstanceDeleted(instance_id=iid, failure_error="new error") + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 3 + assert mi.last_failure_error == "new error" + + +def test_apply_instance_deleted_no_failure_no_tracking(): + """InstanceDeleted without failure_error does not track.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + event = InstanceDeleted(instance_id=iid) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 0 + + +def test_apply_instance_deleted_orphan_no_tracking(): + """InstanceDeleted for orphan instance (no meta_instance_id) does not track.""" + iid, inst = _instance(node_ids=["node-a"]) + state = State(instances={iid: inst}) + event = InstanceDeleted(instance_id=iid, failure_error="crash") + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert len(new_state.meta_instances) == 0 + + +# --- InstanceRetrying --- + + +def test_apply_instance_retrying_removes_runners(): + """InstanceRetrying removes the instance's runners from state but keeps the instance.""" + meta = _meta_instance() + iid, inst = _instance( + node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id + ) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + runners = { + runner_ids[0]: RunnerFailed(error_message="OOM"), + runner_ids[1]: RunnerShutdown(), + } + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + runners=runners, + ) + event = InstanceRetrying( + instance_id=iid, + meta_instance_id=meta.meta_instance_id, + failure_error="OOM", + ) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + # Instance still exists + assert iid in new_state.instances + # Runners removed + assert runner_ids[0] not in new_state.runners + assert runner_ids[1] not in new_state.runners + + +def test_apply_instance_retrying_increments_failure(): + """InstanceRetrying increments consecutive_failures on the MetaInstance.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + event = InstanceRetrying( + instance_id=iid, + meta_instance_id=meta.meta_instance_id, + failure_error="crash", + ) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 1 + assert mi.last_failure_error == "crash" + + +def test_apply_instance_retrying_skips_missing_runners(): + """InstanceRetrying doesn't assert if runners haven't reported yet.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + # No runners in state at all + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + ) + event = InstanceRetrying( + instance_id=iid, + meta_instance_id=meta.meta_instance_id, + failure_error="crash", + ) + # Should not raise + new_state = apply(state, IndexedEvent(idx=0, event=event)) + assert iid in new_state.instances + + +def test_apply_instance_created_resets_failure_counter(): + """InstanceCreated resets consecutive_failures but preserves last_failure_error.""" + meta = _meta_instance().model_copy( + update={"consecutive_failures": 3, "last_failure_error": "old error"} + ) + _, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + state = State(meta_instances={meta.meta_instance_id: meta}) + event = InstanceCreated(instance=inst) + new_state = apply(state, IndexedEvent(idx=0, event=event)) + mi = new_state.meta_instances[meta.meta_instance_id] + assert mi.consecutive_failures == 0 + assert mi.last_failure_error == "old error" + assert mi.placement_error is None + + +# --- InstanceHealthReconciler retry-vs-delete --- + + +async def test_health_reconciler_retries_when_under_limit(): + """InstanceHealthReconciler emits InstanceRetrying when consecutive_failures < 3.""" + meta = _meta_instance() + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, + topology=_topology("node-a"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceRetrying) + assert events[0].instance_id == iid + assert events[0].meta_instance_id == meta.meta_instance_id + + +async def test_health_reconciler_deletes_when_limit_reached(): + """InstanceHealthReconciler emits InstanceDeleted when consecutive_failures >= 3.""" + meta = _meta_instance().model_copy(update={"consecutive_failures": 3}) + iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, + topology=_topology("node-a"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceDeleted) + + +async def test_health_reconciler_deletes_without_meta_instance(): + """Instances without a MetaInstance are deleted immediately on runner failure.""" + iid, inst = _instance(node_ids=["node-a"]) + runner_ids = list(inst.shard_assignments.node_to_runner.values()) + state = State( + instances={iid: inst}, + runners={runner_ids[0]: RunnerFailed(error_message="crash")}, + topology=_topology("node-a"), + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceDeleted) + + +async def test_health_reconciler_network_failure_always_deletes(): + """Network failure always triggers InstanceDeleted regardless of retry count.""" + meta = _meta_instance() + iid, inst = _instance( + node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id + ) + state = State( + meta_instances={meta.meta_instance_id: meta}, + instances={iid: inst}, + topology=_topology("node-a"), # node-b missing + ) + reconciler = InstanceHealthReconciler() + events = await reconciler.reconcile(state) + assert len(events) == 1 + assert isinstance(events[0], InstanceDeleted) + assert events[0].failure_error == "Network connection lost" diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index 94869dfe..f96c6b7f 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -4,7 +4,7 @@ from datetime import datetime from loguru import logger -from exo.shared.types.common import NodeId +from exo.shared.types.common import MetaInstanceId, NodeId from exo.shared.types.events import ( ChunkGenerated, Event, @@ -12,6 +12,12 @@ from exo.shared.types.events import ( InputChunkReceived, InstanceCreated, InstanceDeleted, + InstanceRetrying, + JacclSideChannelData, + JacclSideChannelGathered, + MetaInstanceCreated, + MetaInstanceDeleted, + MetaInstancePlacementFailed, NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, @@ -28,6 +34,7 @@ from exo.shared.types.events import ( TracesCollected, TracesMerged, ) +from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.profiling import ( NodeIdentity, NodeNetworkInfo, @@ -66,12 +73,22 @@ def event_apply(event: Event, state: State) -> State: | InputChunkReceived() | TracesCollected() | TracesMerged() + | JacclSideChannelData() + | JacclSideChannelGathered() ): # Pass-through events that don't modify state return state case InstanceCreated(): return apply_instance_created(event, state) case InstanceDeleted(): return apply_instance_deleted(event, state) + case InstanceRetrying(): + return apply_instance_retrying(event, state) + case MetaInstanceCreated(): + return apply_meta_instance_created(event, state) + case MetaInstanceDeleted(): + return apply_meta_instance_deleted(event, state) + case MetaInstancePlacementFailed(): + return apply_meta_instance_placement_failed(event, state) case NodeTimedOut(): return apply_node_timed_out(event, state) case NodeDownloadProgress(): @@ -174,20 +191,123 @@ def apply_task_failed(event: TaskFailed, state: State) -> State: return state.model_copy(update={"tasks": new_tasks}) +def _update_meta_instance( + state: State, mid: MetaInstanceId, **fields: object +) -> Mapping[MetaInstanceId, MetaInstance]: + mi = state.meta_instances[mid] + return {**state.meta_instances, mid: mi.model_copy(update=fields)} + + def apply_instance_created(event: InstanceCreated, state: State) -> State: instance = event.instance new_instances: Mapping[InstanceId, Instance] = { **state.instances, instance.instance_id: instance, } - return state.model_copy(update={"instances": new_instances}) + update: dict[str, object] = {"instances": new_instances} + # Reset failure tracking when a new instance is created for a meta-instance + if instance.meta_instance_id and instance.meta_instance_id in state.meta_instances: + mi = state.meta_instances[instance.meta_instance_id] + if mi.placement_error is not None or mi.consecutive_failures > 0: + update["meta_instances"] = _update_meta_instance( + state, + instance.meta_instance_id, + placement_error=None, + consecutive_failures=0, + ) + return state.model_copy(update=update) def apply_instance_deleted(event: InstanceDeleted, state: State) -> State: + deleted_instance = state.instances.get(event.instance_id) new_instances: Mapping[InstanceId, Instance] = { iid: inst for iid, inst in state.instances.items() if iid != event.instance_id } - return state.model_copy(update={"instances": new_instances}) + update: dict[str, object] = {"instances": new_instances} + + # Track failure on the MetaInstance itself + if ( + event.failure_error + and deleted_instance + and deleted_instance.meta_instance_id + and deleted_instance.meta_instance_id in state.meta_instances + ): + mid = deleted_instance.meta_instance_id + mi = state.meta_instances[mid] + update["meta_instances"] = { + **state.meta_instances, + mid: mi.model_copy( + update={ + "consecutive_failures": mi.consecutive_failures + 1, + "last_failure_error": event.failure_error, + } + ), + } + + return state.model_copy(update=update) + + +def apply_instance_retrying(event: InstanceRetrying, state: State) -> State: + """Runners failed but retry limit not reached — remove runners, keep instance.""" + instance = state.instances.get(event.instance_id) + if instance is None: + # Instance was already deleted (e.g. cascade from DeleteMetaInstance). + # The InstanceDeleted handler already incremented consecutive_failures + # on the MetaInstance, so skipping here avoids double-counting. + return state + + # Remove all runners belonging to this instance from state + runner_ids_to_remove = set(instance.shard_assignments.node_to_runner.values()) + new_runners: Mapping[RunnerId, RunnerStatus] = { + rid: rs for rid, rs in state.runners.items() if rid not in runner_ids_to_remove + } + + update: dict[str, object] = {"runners": new_runners} + + # Increment failure count on the MetaInstance + if event.meta_instance_id in state.meta_instances: + update["meta_instances"] = _update_meta_instance( + state, + event.meta_instance_id, + consecutive_failures=state.meta_instances[ + event.meta_instance_id + ].consecutive_failures + + 1, + last_failure_error=event.failure_error, + ) + + return state.model_copy(update=update) + + +def apply_meta_instance_created(event: MetaInstanceCreated, state: State) -> State: + new_meta: Mapping[MetaInstanceId, MetaInstance] = { + **state.meta_instances, + event.meta_instance.meta_instance_id: event.meta_instance, + } + return state.model_copy(update={"meta_instances": new_meta}) + + +def apply_meta_instance_deleted(event: MetaInstanceDeleted, state: State) -> State: + new_meta: Mapping[MetaInstanceId, MetaInstance] = { + mid: mi + for mid, mi in state.meta_instances.items() + if mid != event.meta_instance_id + } + return state.model_copy(update={"meta_instances": new_meta}) + + +def apply_meta_instance_placement_failed( + event: MetaInstancePlacementFailed, state: State +) -> State: + if event.meta_instance_id not in state.meta_instances: + return state + return state.model_copy( + update={ + "meta_instances": _update_meta_instance( + state, event.meta_instance_id, placement_error=event.reason + ) + } + ) def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State: diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 2756f0d4..2f0fafa8 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -6,7 +6,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field from exo.shared.models.model_cards import ModelCard, ModelId -from exo.shared.types.common import CommandId, NodeId +from exo.shared.types.common import CommandId, MetaInstanceId, NodeId from exo.shared.types.memory import Memory from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata @@ -262,6 +262,26 @@ class DeleteInstanceResponse(BaseModel): instance_id: InstanceId +class CreateMetaInstanceParams(BaseModel): + model_id: ModelId + sharding: Sharding = Sharding.Pipeline + instance_meta: InstanceMeta = InstanceMeta.MlxRing + min_nodes: int = 1 + node_ids: list[NodeId] | None = None + + +class CreateMetaInstanceResponse(BaseModel): + message: str + command_id: CommandId + meta_instance_id: MetaInstanceId + + +class DeleteMetaInstanceResponse(BaseModel): + message: str + command_id: CommandId + meta_instance_id: MetaInstanceId + + class AdvancedImageParams(BaseModel): seed: Annotated[int, Field(ge=0)] | None = None num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index 09c135aa..8697a6c2 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -6,7 +6,8 @@ from exo.shared.types.api import ( ImageGenerationTaskParams, ) from exo.shared.types.chunks import InputImageChunk -from exo.shared.types.common import CommandId, NodeId +from exo.shared.types.common import CommandId, MetaInstanceId, NodeId +from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata @@ -52,6 +53,14 @@ class TaskCancelled(BaseCommand): cancelled_command_id: CommandId +class CreateMetaInstance(BaseCommand): + meta_instance: MetaInstance + + +class DeleteMetaInstance(BaseCommand): + meta_instance_id: MetaInstanceId + + class TaskFinished(BaseCommand): finished_command_id: CommandId @@ -94,6 +103,8 @@ Command = ( | CreateInstance | DeleteInstance | TaskCancelled + | CreateMetaInstance + | DeleteMetaInstance | TaskFinished | SendInputChunk ) diff --git a/src/exo/shared/types/common.py b/src/exo/shared/types/common.py index 5db51cef..51806de2 100644 --- a/src/exo/shared/types/common.py +++ b/src/exo/shared/types/common.py @@ -42,6 +42,10 @@ class CommandId(Id): pass +class MetaInstanceId(Id): + """Identifier for a MetaInstance.""" + + class Host(CamelCaseModel): ip: str port: int diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 5cf93d0c..dd28d070 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -1,11 +1,14 @@ +import base64 +from collections.abc import Mapping from datetime import datetime -from typing import final +from typing import Annotated, final -from pydantic import Field +from pydantic import BeforeValidator, Field, PlainSerializer from exo.shared.topology import Connection from exo.shared.types.chunks import GenerationChunk, InputImageChunk -from exo.shared.types.common import CommandId, Id, NodeId, SessionId +from exo.shared.types.common import CommandId, Id, MetaInstanceId, NodeId, SessionId +from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId @@ -14,6 +17,28 @@ from exo.utils.info_gatherer.info_gatherer import GatheredInfo from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel +def _decode_base64_bytes(v: bytes | str) -> bytes: + if isinstance(v, bytes): + return v + return base64.b64decode(v) + + +def _encode_base64_bytes(v: bytes) -> str: + return base64.b64encode(v).decode("ascii") + + +Base64Bytes = Annotated[ + bytes, + BeforeValidator(_decode_base64_bytes), + PlainSerializer(_encode_base64_bytes, return_type=str), +] +"""bytes that serialize to/from base64 strings in JSON. + +Needed because TaggedModel's wrap validator converts JSON→Python validation +context, which breaks strict-mode bytes deserialization from JSON strings. +""" + + class EventId(Id): """ Newtype around `ID` @@ -66,6 +91,30 @@ class InstanceCreated(BaseEvent): class InstanceDeleted(BaseEvent): instance_id: InstanceId + failure_error: str | None = None + + +class MetaInstanceCreated(BaseEvent): + meta_instance: MetaInstance + + +class MetaInstanceDeleted(BaseEvent): + meta_instance_id: MetaInstanceId + + +@final +class MetaInstancePlacementFailed(BaseEvent): + meta_instance_id: MetaInstanceId + reason: str + + +@final +class InstanceRetrying(BaseEvent): + """Runners failed but retry count is below the limit — restart runners, keep instance.""" + + instance_id: InstanceId + meta_instance_id: MetaInstanceId + failure_error: str class RunnerStatusUpdated(BaseEvent): @@ -132,6 +181,25 @@ class TracesMerged(BaseEvent): traces: list[TraceEventData] +@final +class JacclSideChannelData(BaseEvent): + """A runner's local contribution to a JACCL SideChannel all_gather round.""" + + instance_id: InstanceId + runner_id: RunnerId + sequence: int + data: Base64Bytes + + +@final +class JacclSideChannelGathered(BaseEvent): + """Gathered result of a JACCL SideChannel all_gather round.""" + + instance_id: InstanceId + sequence: int + gathered_data: Mapping[RunnerId, Base64Bytes] + + Event = ( TestEvent | TaskCreated @@ -141,6 +209,10 @@ Event = ( | TaskAcknowledged | InstanceCreated | InstanceDeleted + | InstanceRetrying + | MetaInstanceCreated + | MetaInstanceDeleted + | MetaInstancePlacementFailed | RunnerStatusUpdated | RunnerDeleted | NodeTimedOut @@ -152,6 +224,8 @@ Event = ( | TopologyEdgeDeleted | TracesCollected | TracesMerged + | JacclSideChannelData + | JacclSideChannelGathered ) diff --git a/src/exo/shared/types/meta_instance.py b/src/exo/shared/types/meta_instance.py new file mode 100644 index 00000000..63052184 --- /dev/null +++ b/src/exo/shared/types/meta_instance.py @@ -0,0 +1,25 @@ +from typing import final + +from pydantic import Field + +from exo.shared.models.model_cards import ModelId +from exo.shared.types.common import MetaInstanceId, NodeId +from exo.shared.types.worker.instances import InstanceMeta +from exo.shared.types.worker.shards import Sharding +from exo.utils.pydantic_ext import FrozenModel + + +@final +class MetaInstance(FrozenModel): + """Declarative constraint: ensure an instance matching these parameters always exists.""" + + meta_instance_id: MetaInstanceId = Field(default_factory=MetaInstanceId) + model_id: ModelId + sharding: Sharding = Sharding.Pipeline + instance_meta: InstanceMeta = InstanceMeta.MlxRing + min_nodes: int = 1 + node_ids: list[NodeId] | None = None + # Failure tracking + placement_error: str | None = None + consecutive_failures: int = 0 + last_failure_error: str | None = None diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py index 7350cfb0..4ff1d4f7 100644 --- a/src/exo/shared/types/state.py +++ b/src/exo/shared/types/state.py @@ -6,7 +6,8 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator from pydantic.alias_generators import to_camel from exo.shared.topology import Topology, TopologySnapshot -from exo.shared.types.common import NodeId +from exo.shared.types.common import MetaInstanceId, NodeId +from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.profiling import ( DiskUsage, MemoryUsage, @@ -41,6 +42,7 @@ class State(CamelCaseModel): arbitrary_types_allowed=True, ) instances: Mapping[InstanceId, Instance] = {} + meta_instances: Mapping[MetaInstanceId, MetaInstance] = {} runners: Mapping[RunnerId, RunnerStatus] = {} downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {} tasks: Mapping[TaskId, Task] = {} diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py index 8d866456..cb88d401 100644 --- a/src/exo/shared/types/tasks.py +++ b/src/exo/shared/types/tasks.py @@ -61,7 +61,7 @@ class TextGeneration(BaseTask): # emitted by Master error_message: str | None = Field(default=None) -class CancelTask(BaseTask): +class CancelTask(BaseTask): # emitted by Worker when master cancels a task cancelled_task_id: TaskId runner_id: RunnerId diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py index cda11ffa..4254b998 100644 --- a/src/exo/shared/types/worker/instances.py +++ b/src/exo/shared/types/worker/instances.py @@ -2,7 +2,7 @@ from enum import Enum from pydantic import model_validator -from exo.shared.types.common import Host, Id, NodeId +from exo.shared.types.common import Host, Id, MetaInstanceId, NodeId from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel @@ -19,6 +19,7 @@ class InstanceMeta(str, Enum): class BaseInstance(TaggedModel): instance_id: InstanceId shard_assignments: ShardAssignments + meta_instance_id: MetaInstanceId | None = None def shard(self, runner_id: RunnerId) -> ShardMetadata | None: return self.shard_assignments.runner_to_shard.get(runner_id, None) diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py index 646ac8f6..ebf0165f 100644 --- a/src/exo/utils/channels.py +++ b/src/exo/utils/channels.py @@ -125,9 +125,7 @@ class MpSender[T]: self._state.buffer.put(item, block=True) async def send_async(self, item: T) -> None: - await to_thread.run_sync( - self.send, item, limiter=CapacityLimiter(1), abandon_on_cancel=True - ) + await to_thread.run_sync(self.send, item, limiter=CapacityLimiter(1)) def close(self) -> None: if not self._state.closed.is_set(): diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 3ed65ecc..670847eb 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -574,6 +574,11 @@ def mlx_cleanup( def mx_any(bool_: bool, group: Group | None) -> bool: + """Synchronize a boolean across all distributed nodes. + + Returns True if any node has bool_=True. Uses all_sum so every + node participates in the collective — preventing GPU deadlocks. + """ if group is None: return bool_ num_true = mx.distributed.all_sum( diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index af105652..6b2a9475 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -24,6 +24,7 @@ from exo.shared.types.events import ( ForwarderEvent, IndexedEvent, InputChunkReceived, + JacclSideChannelGathered, NodeGatheredInfo, TaskCreated, TaskStatusUpdated, @@ -33,7 +34,6 @@ from exo.shared.types.events import ( from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.state import State from exo.shared.types.tasks import ( - CancelTask, CreateRunner, DownloadModel, ImageEdits, @@ -159,6 +159,15 @@ class Worker: for idx, event in indexed_events: self.state = apply(self.state, IndexedEvent(idx=idx, event=event)) + # Dispatch JACCL gathered events to the relevant RunnerSupervisor + if isinstance(event, JacclSideChannelGathered): + for runner in self.runners.values(): + if ( + runner.bound_instance.instance.instance_id + == event.instance_id + ): + runner.notify_gathered(event) + # Buffer input image chunks for image editing if isinstance(event, InputChunkReceived): cmd_id = event.command_id @@ -225,22 +234,15 @@ class Worker: ) ) case Shutdown(runner_id=runner_id): - runner = self.runners.pop(runner_id) try: with fail_after(3): - await runner.start_task(task) + await self.runners.pop(runner_id).start_task(task) except TimeoutError: await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.TimedOut ) ) - finally: - runner.shutdown() - case CancelTask( - cancelled_task_id=cancelled_task_id, runner_id=runner_id - ): - await self.runners[runner_id].cancel_task(cancelled_task_id) case ImageEdits() if task.task_params.total_input_chunks > 0: # Assemble image from chunks and inject into task cmd_id = task.command_id @@ -278,18 +280,18 @@ class Worker: del self.input_chunk_buffer[cmd_id] if cmd_id in self.input_chunk_counts: del self.input_chunk_counts[cmd_id] - await self._start_runner_task(modified_task) + await self.runners[self._task_to_runner_id(task)].start_task( + modified_task + ) case task: - await self._start_runner_task(task) + await self.runners[self._task_to_runner_id(task)].start_task(task) def shutdown(self): self._tg.cancel_scope.cancel() - async def _start_runner_task(self, task: Task): - if (instance := self.state.instances.get(task.instance_id)) is not None: - await self.runners[ - instance.shard_assignments.node_to_runner[self.node_id] - ].start_task(task) + def _task_to_runner_id(self, task: Task): + instance = self.state.instances[task.instance_id] + return instance.shard_assignments.node_to_runner[self.node_id] async def _nack_request(self, since_idx: int) -> None: # We request all events after (and including) the missing index. diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index ce2eb4a9..4107d553 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -35,6 +35,7 @@ from exo.shared.types.worker.runners import ( RunnerLoading, RunnerReady, RunnerRunning, + RunnerShutdown, RunnerStatus, RunnerWarmingUp, ) @@ -56,7 +57,7 @@ def plan( return ( _cancel_tasks(runners, tasks) or _kill_runner(runners, all_runners, instances) - or _create_runner(node_id, runners, instances) + or _create_runner(node_id, runners, instances, all_runners) or _model_needs_download(node_id, runners, global_download_status) or _init_distributed_backend(runners, all_runners) or _load_model(runners, all_runners, global_download_status) @@ -75,6 +76,12 @@ def _kill_runner( if (instance_id := runner.bound_instance.instance.instance_id) not in instances: return Shutdown(instance_id=instance_id, runner_id=runner_id) + # Master removed our runner from state (retry signal) and process is dead + if runner_id not in all_runners and isinstance( + runner.status, (RunnerFailed, RunnerShutdown) + ): + return Shutdown(instance_id=instance_id, runner_id=runner_id) + for ( global_runner_id ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values(): @@ -92,6 +99,7 @@ def _create_runner( node_id: NodeId, runners: Mapping[RunnerId, RunnerSupervisor], instances: Mapping[InstanceId, Instance], + all_runners: Mapping[RunnerId, RunnerStatus], ) -> CreateRunner | None: for instance in instances.values(): runner_id = instance.shard_assignments.node_to_runner.get(node_id, None) @@ -101,6 +109,16 @@ def _create_runner( if runner_id in runners: continue + # Don't create while any peer runner is in a terminal state — wait for + # the master to emit InstanceRetrying which removes them from state. + has_terminal_peer = any( + isinstance(all_runners.get(peer_rid), (RunnerFailed, RunnerShutdown)) + for peer_rid in instance.shard_assignments.node_to_runner.values() + if peer_rid != runner_id + ) + if has_terminal_peer: + continue + shard = instance.shard(runner_id) assert shard is not None @@ -310,7 +328,8 @@ def _pending_tasks( def _cancel_tasks( runners: Mapping[RunnerId, RunnerSupervisor], tasks: Mapping[TaskId, Task], -) -> Task | None: +) -> CancelTask | None: + """Find a cancelled task that hasn't been sent to the runner yet.""" for task in tasks.values(): if task.task_status != TaskStatus.Cancelled: continue diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index ed420aab..69ef6c72 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/src/exo/worker/runner/bootstrap.py @@ -17,6 +17,7 @@ def entrypoint( task_receiver: MpReceiver[Task], cancel_receiver: MpReceiver[TaskId], _logger: "loguru.Logger", + pipe_fifo_paths: tuple[str, str] | None = None, ) -> None: fast_synch_override = os.environ.get("EXO_FAST_SYNCH") if fast_synch_override == "on" or ( @@ -30,6 +31,16 @@ def entrypoint( else: os.environ["MLX_METAL_FAST_SYNCH"] = "0" + # Open JACCL FIFOs by path and set env vars for C++ SideChannel. + # Named pipes (FIFOs) work across multiprocessing spawn (macOS default). + if pipe_fifo_paths is not None: + fifo_c2p, fifo_p2c = pipe_fifo_paths + # C++ reads gathered data from p2c (PIPE_IN), writes local data to c2p (PIPE_OUT) + pipe_in_fd = os.open(fifo_p2c, os.O_RDONLY) + pipe_out_fd = os.open(fifo_c2p, os.O_WRONLY) + os.environ["MLX_JACCL_PIPE_IN"] = str(pipe_in_fd) + os.environ["MLX_JACCL_PIPE_OUT"] = str(pipe_out_fd) + global logger logger = _logger @@ -56,7 +67,9 @@ def entrypoint( try: event_sender.close() task_receiver.close() + cancel_receiver.close() finally: event_sender.join() task_receiver.join() + cancel_receiver.join() logger.info("bye from the runner") diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index e55456d3..818bd9be 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -243,7 +243,7 @@ def main( assert inference_model assert tokenizer - t = time.monotonic() + t = time.perf_counter() toks = warmup_inference( model=inference_model, tokenizer=tokenizer, @@ -251,7 +251,7 @@ def main( ) logger.info(f"warmed up by generating {toks} tokens") check_for_cancel_every = min( - math.ceil(toks / min(time.monotonic() - t, 0.001)), 100 + math.ceil(toks / max(time.perf_counter() - t, 0.001)), 100 ) if group is not None: check_for_cancel_every = int( diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 5d39a881..519d7b07 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -1,6 +1,10 @@ import contextlib +import os import signal +import struct +import tempfile from dataclasses import dataclass, field +from functools import partial from multiprocessing import Process from typing import Self @@ -14,12 +18,14 @@ from loguru import logger from exo.shared.types.events import ( Event, + JacclSideChannelData, + JacclSideChannelGathered, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, ) from exo.shared.types.tasks import Task, TaskId, TaskStatus -from exo.shared.types.worker.instances import BoundInstance +from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance from exo.shared.types.worker.runners import ( RunnerConnecting, RunnerFailed, @@ -34,6 +40,26 @@ from exo.shared.types.worker.shards import ShardMetadata from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel from exo.worker.runner.bootstrap import entrypoint + +def _pipe_read_exact(fd: int, n: int) -> bytes | None: + """Read exactly n bytes from a file descriptor. Returns None on EOF.""" + data = b"" + while len(data) < n: + chunk = os.read(fd, n - len(data)) + if not chunk: + return None + data += chunk + return data + + +def _pipe_write_all(fd: int, data: bytes) -> None: + """Write all bytes to a file descriptor.""" + view = memoryview(data) + while view: + written = os.write(fd, view) + view = view[written:] + + PREFILL_TIMEOUT_SECONDS = 60 DECODE_TIMEOUT_SECONDS = 5 @@ -46,12 +72,21 @@ class RunnerSupervisor: initialize_timeout: float _ev_recv: MpReceiver[Event] _task_sender: MpSender[Task] - _event_sender: Sender[Event] _cancel_sender: MpSender[TaskId] + _event_sender: Sender[Event] + _pipe_read_fd: int | None = None # Python reads runner's pipe output + _pipe_write_fd: int | None = None # Python writes gathered data to runner + _child_pipe_fds: tuple[int, int] | None = None # fds to close after fork + _fifo_dir: str | None = None # Temp dir for FIFO files (for cleanup) + _fifo_c2p: str | None = None # FIFO path: C++ writes → Python reads + _fifo_p2c: str | None = None # FIFO path: Python writes → C++ reads status: RunnerStatus = field(default_factory=RunnerIdle, init=False) pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False) completed: set[TaskId] = field(default_factory=set, init=False) cancelled: set[TaskId] = field(default_factory=set, init=False) + _gathered_waiters: dict[ + int, tuple[anyio.Event, JacclSideChannelGathered | None] + ] = field(default_factory=dict, init=False) @classmethod def create( @@ -65,6 +100,23 @@ class RunnerSupervisor: task_sender, task_recv = mp_channel[Task]() cancel_sender, cancel_recv = mp_channel[TaskId]() + # For MlxJaccl instances, create named pipes (FIFOs) for SideChannel relay. + # Named pipes work across multiprocessing.Process spawn (macOS default). + # FIFO c2p: C++ writes local data → Python reads it + # FIFO p2c: Python writes gathered data → C++ reads it + fifo_dir: str | None = None + fifo_c2p: str | None = None + fifo_p2c: str | None = None + pipe_fifo_paths: tuple[str, str] | None = None + + if isinstance(bound_instance.instance, MlxJacclInstance): + fifo_dir = tempfile.mkdtemp(prefix="exo_jaccl_") + fifo_c2p = os.path.join(fifo_dir, "c2p") # C++ → Python + fifo_p2c = os.path.join(fifo_dir, "p2c") # Python → C++ + os.mkfifo(fifo_c2p) + os.mkfifo(fifo_p2c) + pipe_fifo_paths = (fifo_c2p, fifo_p2c) + runner_process = Process( target=entrypoint, args=( @@ -73,6 +125,7 @@ class RunnerSupervisor: task_recv, cancel_recv, logger, + pipe_fifo_paths, ), daemon=True, ) @@ -88,21 +141,54 @@ class RunnerSupervisor: _task_sender=task_sender, _cancel_sender=cancel_sender, _event_sender=event_sender, + _fifo_dir=fifo_dir, + _fifo_c2p=fifo_c2p, + _fifo_p2c=fifo_p2c, ) return self async def run(self): self.runner_process.start() - await self._forward_events() + + if self._fifo_c2p is not None and self._fifo_p2c is not None: + # Open FIFOs from parent side. These block until child opens the other end, + # so we run them in threads concurrently to avoid deadlock. + fifo_c2p = self._fifo_c2p + fifo_p2c = self._fifo_p2c + + async def open_read() -> None: + self._pipe_read_fd = await to_thread.run_sync( + partial(os.open, fifo_c2p, os.O_RDONLY) + ) + + async def open_write() -> None: + self._pipe_write_fd = await to_thread.run_sync( + partial(os.open, fifo_p2c, os.O_WRONLY) + ) + + async with anyio.create_task_group() as open_tg: + open_tg.start_soon(open_read) + open_tg.start_soon(open_write) + + logger.info( + f"JACCL pipe relay: FIFOs opened (read_fd={self._pipe_read_fd}, write_fd={self._pipe_write_fd})" + ) + + async with anyio.create_task_group() as tg: + tg.start_soon(self._pipe_relay) + tg.start_soon(self._forward_events) + else: + await self._forward_events() def shutdown(self): 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")) self._cancel_sender.close() + self._event_sender.close() + self._close_pipe_fds() self.runner_process.join(1) if not self.runner_process.is_alive(): logger.info("Runner process succesfully terminated") @@ -140,6 +226,7 @@ class RunnerSupervisor: await event.wait() async def cancel_task(self, task_id: TaskId): + """Send a cancellation signal to the runner process.""" if task_id in self.completed: logger.info(f"Unable to cancel {task_id} as it has been completed") return @@ -181,6 +268,110 @@ class RunnerSupervisor: for tid in self.pending: self.pending[tid].set() + def _close_pipe_fds(self) -> None: + if self._pipe_read_fd is not None: + with contextlib.suppress(OSError): + os.close(self._pipe_read_fd) + self._pipe_read_fd = None + if self._pipe_write_fd is not None: + with contextlib.suppress(OSError): + os.close(self._pipe_write_fd) + self._pipe_write_fd = None + if self._child_pipe_fds is not None: + for fd in self._child_pipe_fds: + with contextlib.suppress(OSError): + os.close(fd) + self._child_pipe_fds = None + # Clean up FIFO files + if self._fifo_c2p is not None: + with contextlib.suppress(OSError): + os.unlink(self._fifo_c2p) + self._fifo_c2p = None + if self._fifo_p2c is not None: + with contextlib.suppress(OSError): + os.unlink(self._fifo_p2c) + self._fifo_p2c = None + if self._fifo_dir is not None: + with contextlib.suppress(OSError): + os.rmdir(self._fifo_dir) + self._fifo_dir = None + + async def _pipe_relay(self) -> None: + """Relay JACCL SideChannel all_gather rounds between runner pipes and exo events.""" + assert self._pipe_read_fd is not None + assert self._pipe_write_fd is not None + read_fd = self._pipe_read_fd + write_fd = self._pipe_write_fd + sequence = 0 + + try: + while True: + # 1. Read local data from runner: [uint32 size][size bytes] + header = await to_thread.run_sync(partial(_pipe_read_exact, read_fd, 4)) + if header is None: + logger.info("JACCL pipe relay: runner closed pipe (EOF)") + break + data_size: int = struct.unpack(" None: + """Called by the worker when a JacclSideChannelGathered event arrives.""" + seq = event.sequence + if seq not in self._gathered_waiters: + logger.warning(f"JACCL: received gathered event for unknown sequence {seq}") + return + waiter, _ = self._gathered_waiters[seq] + self._gathered_waiters[seq] = (waiter, event) + waiter.set() + def __del__(self) -> None: if self.runner_process.is_alive(): logger.warning("RunnerSupervisor was not stopped cleanly.") 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..878aea99 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 @@ -1,9 +1,7 @@ # Check tasks are complete before runner is ever ready. -import unittest.mock from collections.abc import Iterable from typing import Callable -import mlx.core as mx import pytest import exo.worker.runner.runner as mlx_runner @@ -117,6 +115,12 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1)) monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin) monkeypatch.setattr(mlx_runner, "mx_any", make_nothin(False)) + + # Mock mx.distributed.all_gather so MockGroup doesn't hit real MLX C++ bindings. + def _mock_all_gather(x: object, **_kw: object) -> object: + return x + + monkeypatch.setattr(mlx_runner.mx.distributed, "all_gather", _mock_all_gather) # Mock apply_chat_template since we're using a fake tokenizer (integer 1). # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None. monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt")) @@ -178,16 +182,15 @@ def _run(tasks: Iterable[Task]): # this is some c++ nonsense task_receiver.close = nothin task_receiver.join = nothin - with unittest.mock.patch( - "exo.worker.runner.runner.mx.distributed.all_gather", - make_nothin(mx.array([1])), - ): - mlx_runner.main( - bound_instance, - event_sender, # pyright: ignore[reportArgumentType] - task_receiver, - cancel_receiver, - ) + cancel_receiver.close = nothin + cancel_receiver.join = nothin + + mlx_runner.main( + bound_instance, + event_sender, # pyright: ignore[reportArgumentType] + task_receiver, + cancel_receiver, + ) return event_sender.events From facf2d4d03165de6134162ff27f65eafd40a8798 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Tue, 17 Feb 2026 17:48:43 +0000 Subject: [PATCH 08/45] Use custom fork that resolves GPU locks (#1489) ## Motivation There is an issue on Macs that means that an explicit synchronization is necessary for memory to be updated from L1 cache. This means that GPU locks can occur when a spin wait does not see the updated timestamp. ## Changes Updated in my own personal fork. ## Why It Works https://github.com/ARM-software/acle/releases ## Test Plan ### Manual Testing Tested manually that no GPU locks occur (even with multiple simultaneous instances running) and that the performance differential is negligible (267 vs 269 tps on Llama 3.2 1B at an approx 10k context.) ------------------------------------------------------ I have seen a GPU lock, specifically when sending a particularly large chat completion while the model was loading. However, I have since been unable to reproduce and this may be something I did wrong. Please do create an issue and tag me if any GPU locks do occur. --------- Co-authored-by: Jake Hillion Co-authored-by: Claude Opus 4.6 --- README.md | 11 +++++++++-- flake.nix | 2 +- nix/mlx.nix | 10 +++++----- pyproject.toml | 3 ++- python/parts.nix | 32 +++++++++++++++++++++++++++++--- uv.lock | 40 ++++++++++++++++------------------------ 6 files changed, 62 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 58d41c37..ff1fe04b 100644 --- a/README.md +++ b/README.md @@ -72,16 +72,23 @@ 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 (after accepting the Cachix cache): + +```bash +nix run .#exo +``` + **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/flake.nix b/flake.nix index 9c2ca1ef..e90e0bd2 100644 --- a/flake.nix +++ b/flake.nix @@ -115,7 +115,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..8a40e11b 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.dev20260217+50487b41"; 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 = "50487b4141f3c951122655db3b83df5146c1fbeb"; + hash = "sha256-IL4a9vMX5nocgJU1WG4zE8hArHkHJtnh4sdYh3od5zU="; }; patches = [ diff --git a/pyproject.toml b/pyproject.toml index 5d8d79a5..02aa6071 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ 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", "tiktoken>=0.12.0", # required for kimi k2 tokenizer @@ -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 } diff --git a/python/parts.nix b/python/parts.nix index 46b4abdf..bac8ddab 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; diff --git a/uv.lock b/uv.lock index 627e6951..74687ad3 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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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,7 +416,7 @@ 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 = "msgspec", specifier = ">=0.19.0" }, @@ -1020,8 +1020,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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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 +1048,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 +1066,14 @@ cuda13 = [ { name = "mlx-cuda-13", marker = "sys_platform == 'linux'" }, ] +[[package]] +name = "mlx" +version = "0.30.7.dev20260217+50487b41" +source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" } +resolution-markers = [ + "sys_platform == 'darwin'", +] + [[package]] name = "mlx-cpu" version = "0.30.6" @@ -1102,7 +1104,7 @@ version = "0.30.6" 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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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'" }, @@ -1114,16 +1116,6 @@ 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" }, -] - [[package]] name = "more-itertools" version = "10.8.0" From 490d2e46ba469834215f318a358b37deb5e0e9a5 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Tue, 17 Feb 2026 09:52:49 -0800 Subject: [PATCH 09/45] feat: better onboarding UX for new users (#1479) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - **Auto-open dashboard** in browser on first launch (uses `~/.exo/.dashboard_opened` marker) - **Welcome overlay** with "Choose a Model" CTA button when no model instance is running - **Tutorial progress messages** during model download → loading → ready lifecycle stages - **Fix conversation sidebar** text contrast — bumped to white text, added active state background - **Simplify technical jargon** — sharding/instance type/min nodes hidden behind collapsible "Advanced Options" toggle; strategy display hidden behind debug mode - **Polished DMG installer** with drag-to-Applications layout, custom branded background, and AppleScript-configured window positioning ## Test plan - [ ] Launch exo for the first time (delete `~/.exo/.dashboard_opened` to simulate) — browser should auto-open - [ ] Verify welcome overlay appears on topology when no model is loaded - [ ] Launch a model and verify download/loading/ready messages appear in instance cards - [ ] Check conversation sidebar text is readable (white on dark, yellow when active) - [ ] Verify "Advanced Options" toggle hides/shows sharding controls - [ ] Build DMG with `packaging/dmg/create-dmg.sh` and verify drag-to-Applications layout 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 --- .github/workflows/build-app.yml | 5 +- app/EXO/EXO/ContentView.swift | 480 +-- app/EXO/EXO/EXOApp.swift | 16 +- app/EXO/EXO/ExoProcessController.swift | 20 + app/EXO/EXO/Views/FirstLaunchPopout.swift | 192 + app/EXO/EXO/Views/SettingsView.swift | 478 +++ .../EXO/Views/SettingsWindowController.swift | 47 + dashboard/src/app.css | 37 + dashboard/src/lib/components/ChatForm.svelte | 4 + .../src/lib/components/ChatMessages.svelte | 5 +- .../src/lib/components/ChatSidebar.svelte | 40 +- .../lib/components/ConnectionBanner.svelte | 20 + dashboard/src/lib/components/HeaderNav.svelte | 77 +- dashboard/src/lib/components/ModelCard.svelte | 26 + .../lib/components/ModelPickerModal.svelte | 67 +- .../src/lib/components/ToastContainer.svelte | 117 + dashboard/src/lib/stores/app.svelte.ts | 25 +- dashboard/src/lib/stores/toast.svelte.ts | 87 + dashboard/src/routes/+layout.svelte | 4 + dashboard/src/routes/+page.svelte | 3584 +++++++++++------ packaging/dmg/background.png | Bin 0 -> 7010 bytes packaging/dmg/create-dmg.sh | 112 + packaging/dmg/generate-background.py | 91 + src/exo/utils/banner.py | 30 + 24 files changed, 3974 insertions(+), 1590 deletions(-) create mode 100644 app/EXO/EXO/Views/FirstLaunchPopout.swift create mode 100644 app/EXO/EXO/Views/SettingsView.swift create mode 100644 app/EXO/EXO/Views/SettingsWindowController.swift create mode 100644 dashboard/src/lib/components/ConnectionBanner.svelte create mode 100644 dashboard/src/lib/components/ToastContainer.svelte create mode 100644 dashboard/src/lib/stores/toast.svelte.ts create mode 100644 packaging/dmg/background.png create mode 100755 packaging/dmg/create-dmg.sh create mode 100644 packaging/dmg/generate-background.py diff --git a/.github/workflows/build-app.yml b/.github/workflows/build-app.yml index d20bfc12..861f4697 100644 --- a/.github/workflows/build-app.yml +++ b/.github/workflows/build-app.yml @@ -303,11 +303,8 @@ jobs: SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}') /usr/bin/codesign --deep --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" EXO.app - mkdir -p dmg-root - cp -R EXO.app dmg-root/ - ln -s /Applications dmg-root/Applications DMG_NAME="EXO-${RELEASE_VERSION}.dmg" - hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME" + bash "$GITHUB_WORKSPACE/packaging/dmg/create-dmg.sh" EXO.app "$DMG_NAME" "EXO" /usr/bin/codesign --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" "$DMG_NAME" if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then diff --git a/app/EXO/EXO/ContentView.swift b/app/EXO/EXO/ContentView.swift index 688fba5b..ca09fe2f 100644 --- a/app/EXO/EXO/ContentView.swift +++ b/app/EXO/EXO/ContentView.swift @@ -15,18 +15,12 @@ struct ContentView: View { @EnvironmentObject private var localNetworkChecker: LocalNetworkChecker @EnvironmentObject private var updater: SparkleUpdater @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService + @EnvironmentObject private var settingsWindowController: SettingsWindowController @State private var focusedNode: NodeViewModel? @State private var deletingInstanceIDs: Set = [] @State private var showAllNodes = false @State private var showAllInstances = false - @State private var showAdvanced = false - @State private var showDebugInfo = false - @State private var bugReportInFlight = false - @State private var bugReportMessage: String? - @State private var uninstallInProgress = false - @State private var pendingNamespace: String = "" - @State private var pendingHFToken: String = "" - @State private var pendingEnableImageModels = false + @State private var baseURLCopied = false var body: some View { VStack(alignment: .leading, spacing: 12) { @@ -258,139 +252,79 @@ struct ContentView: View { VStack(alignment: .leading, spacing: 0) { if controller.status != .stopped { dashboardButton + baseURLRow Divider() .padding(.vertical, 8) } else { Divider() .padding(.vertical, 4) } - advancedSection - .padding(.bottom, 8) - controlButton(title: "Quit", tint: .secondary) { + HoverButton( + title: "Settings", + tint: .primary, + trailingSystemImage: "gear" + ) { + settingsWindowController.open( + controller: controller, + updater: updater, + networkStatusService: networkStatusService, + thunderboltBridgeService: thunderboltBridgeService, + stateService: stateService + ) + } + HoverButton( + title: "Check for Updates", + tint: .primary, + trailingSystemImage: "arrow.triangle.2.circlepath" + ) { + updater.checkForUpdates() + } + .padding(.bottom, 8) + HoverButton(title: "Quit", tint: .secondary) { controller.stop() NSApplication.shared.terminate(nil) } } } - private var advancedSection: some View { - VStack(alignment: .leading, spacing: 6) { - HStack { - Text("Advanced") - .font(.caption) - .foregroundColor(.secondary) - Spacer() - collapseButton(isExpanded: $showAdvanced) - } - .animation(nil, value: showAdvanced) - if showAdvanced { - VStack(alignment: .leading, spacing: 8) { - VStack(alignment: .leading, spacing: 4) { - Text("Cluster Namespace") - .font(.caption2) - .foregroundColor(.secondary) - HStack { - TextField("optional", text: $pendingNamespace) - .textFieldStyle(.roundedBorder) - .font(.caption2) - .onAppear { - pendingNamespace = controller.customNamespace - } - Button("Save & Restart") { - controller.customNamespace = pendingNamespace - if controller.status == .running || controller.status == .starting { - controller.restart() - } - } - .font(.caption2) - .disabled(pendingNamespace == controller.customNamespace) - } - } - VStack(alignment: .leading, spacing: 4) { - Text("HuggingFace Token") - .font(.caption2) - .foregroundColor(.secondary) - HStack { - SecureField("optional", text: $pendingHFToken) - .textFieldStyle(.roundedBorder) - .font(.caption2) - .onAppear { - pendingHFToken = controller.hfToken - } - Button("Save & Restart") { - controller.hfToken = pendingHFToken - if controller.status == .running || controller.status == .starting { - controller.restart() - } - } - .font(.caption2) - .disabled(pendingHFToken == controller.hfToken) - } - } - Divider() - HStack { - Toggle( - "Enable Image Models (experimental)", isOn: $pendingEnableImageModels - ) - .toggleStyle(.switch) - .font(.caption2) - .onAppear { - pendingEnableImageModels = controller.enableImageModels - } - - Spacer() - - Button("Save & Restart") { - controller.enableImageModels = pendingEnableImageModels - if controller.status == .running || controller.status == .starting { - controller.restart() - } - } - .font(.caption2) - .disabled(pendingEnableImageModels == controller.enableImageModels) - } - HoverButton(title: "Check for Updates", small: true) { - updater.checkForUpdates() - } - debugSection - HoverButton(title: "Uninstall", tint: .red, small: true) { - showUninstallConfirmationAlert() - } - .disabled(uninstallInProgress) - } - .transition(.opacity) - } - } - .animation(.easeInOut(duration: 0.25), value: showAdvanced) - } - - private func controlButton(title: String, tint: Color = .primary, action: @escaping () -> Void) - -> some View - { - HoverButton(title: title, tint: tint, trailingSystemImage: nil, action: action) - } - private var dashboardButton: some View { - Button { + HoverButton( + title: "Web Dashboard", + tint: .primary, + trailingSystemImage: "arrow.up.right" + ) { guard let url = URL(string: "http://localhost:52415/") else { return } NSWorkspace.shared.open(url) - } label: { - HStack { - Image(systemName: "arrow.up.right.square") - .imageScale(.small) - Text("Dashboard") - .fontWeight(.medium) - Spacer() - } - .padding(.vertical, 8) - .padding(.horizontal, 10) - .background( - RoundedRectangle(cornerRadius: 8, style: .continuous) - .fill(Color(red: 1.0, green: 0.87, blue: 0.0).opacity(0.2)) - ) } - .buttonStyle(.plain) - .padding(.bottom, 4) + } + + private var baseURLRow: some View { + HStack(spacing: 6) { + Image(systemName: "link") + .imageScale(.small) + .foregroundColor(.secondary) + Text("localhost:52415/v1") + .font(.system(.caption, design: .monospaced)) + .foregroundColor(.primary) + Spacer() + Button { + NSPasteboard.general.clearContents() + NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string) + baseURLCopied = true + DispatchQueue.main.asyncAfter(deadline: .now() + 2) { + baseURLCopied = false + } + } label: { + Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc") + .imageScale(.small) + .foregroundColor(baseURLCopied ? .green : .secondary) + .contentTransition(.symbolEffect(.replace)) + } + .buttonStyle(.plain) + .help("Copy API base URL") + } + .padding(.vertical, 4) + .padding(.horizontal, 8) } private func collapseButton(isExpanded: Binding) -> some View { @@ -445,207 +379,6 @@ struct ContentView: View { } } - private var thunderboltStatusText: String { - switch networkStatusService.status.thunderboltBridgeState { - case .some(.disabled): - return "Thunderbolt Bridge: Disabled" - case .some(.deleted): - return "Thunderbolt Bridge: Deleted" - case .some(.enabled): - return "Thunderbolt Bridge: Enabled" - case nil: - return "Thunderbolt Bridge: Unknown" - } - } - - private var thunderboltStatusColor: Color { - switch networkStatusService.status.thunderboltBridgeState { - case .some(.disabled), .some(.deleted): - return .green - case .some(.enabled): - return .red - case nil: - return .secondary - } - } - - /// Shows TB bridge status for all nodes from exo cluster state - private var clusterThunderboltBridgeView: some View { - let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:] - let localNodeId = stateService.localNodeId - let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] - - return VStack(alignment: .leading, spacing: 1) { - if bridgeStatuses.isEmpty { - Text("Cluster TB Bridge: No data") - .font(.caption2) - .foregroundColor(.secondary) - } else { - Text("Cluster TB Bridge Status:") - .font(.caption2) - .foregroundColor(.secondary) - ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in - if let status = bridgeStatuses[nodeId] { - let nodeName = - nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) - let isLocal = nodeId == localNodeId - let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" - let statusText = - !status.exists - ? "N/A" - : (status.enabled ? "Enabled" : "Disabled") - let color: Color = - !status.exists - ? .secondary - : (status.enabled ? .red : .green) - Text("\(prefix) \(statusText)") - .font(.caption2) - .foregroundColor(color) - } - } - } - } - } - - private var interfaceIpList: some View { - let statuses = networkStatusService.status.interfaceStatuses - return VStack(alignment: .leading, spacing: 1) { - Text("Interfaces (en0–en7):") - .font(.caption2) - .foregroundColor(.secondary) - if statuses.isEmpty { - Text(" Unknown") - .font(.caption2) - .foregroundColor(.secondary) - } else { - ForEach(statuses, id: \.interfaceName) { status in - let ipText = status.ipAddress ?? "No IP" - Text(" \(status.interfaceName): \(ipText)") - .font(.caption2) - .foregroundColor(status.ipAddress == nil ? .red : .green) - } - } - } - } - - private var debugSection: some View { - VStack(alignment: .leading, spacing: 4) { - HoverButton( - title: "Debug Info", - tint: .primary, - trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down", - small: true - ) { - showDebugInfo.toggle() - } - if showDebugInfo { - VStack(alignment: .leading, spacing: 4) { - Text("Version: \(buildTag)") - .font(.caption2) - .foregroundColor(.secondary) - Text("Commit: \(buildCommit)") - .font(.caption2) - .foregroundColor(.secondary) - Text(thunderboltStatusText) - .font(.caption2) - .foregroundColor(thunderboltStatusColor) - clusterThunderboltBridgeView - interfaceIpList - rdmaStatusView - sendBugReportButton - .padding(.top, 6) - } - .padding(.leading, 8) - .transition(.opacity) - } - } - .animation(.easeInOut(duration: 0.25), value: showDebugInfo) - } - - private var rdmaStatusView: some View { - let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:] - let localNodeId = stateService.localNodeId - let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] - let localDevices = networkStatusService.status.localRdmaDevices - let localPorts = networkStatusService.status.localRdmaActivePorts - - return VStack(alignment: .leading, spacing: 1) { - if rdmaStatuses.isEmpty { - Text("Cluster RDMA: No data") - .font(.caption2) - .foregroundColor(.secondary) - } else { - Text("Cluster RDMA Status:") - .font(.caption2) - .foregroundColor(.secondary) - ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in - if let status = rdmaStatuses[nodeId] { - let nodeName = - nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) - let isLocal = nodeId == localNodeId - let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" - let statusText = status.enabled ? "Enabled" : "Disabled" - let color: Color = status.enabled ? .green : .orange - Text("\(prefix) \(statusText)") - .font(.caption2) - .foregroundColor(color) - } - } - } - if !localDevices.isEmpty { - Text(" Local Devices: \(localDevices.joined(separator: ", "))") - .font(.caption2) - .foregroundColor(.secondary) - } - if !localPorts.isEmpty { - Text(" Local Active Ports:") - .font(.caption2) - .foregroundColor(.secondary) - ForEach(localPorts, id: \.device) { port in - Text(" \(port.device) port \(port.port): \(port.state)") - .font(.caption2) - .foregroundColor(.green) - } - } - } - } - - private var sendBugReportButton: some View { - VStack(alignment: .leading, spacing: 4) { - Button { - Task { - await sendBugReport() - } - } label: { - HStack { - if bugReportInFlight { - ProgressView() - .scaleEffect(0.6) - } - Text("Send Bug Report") - .font(.caption) - .fontWeight(.semibold) - Spacer() - } - .padding(.vertical, 6) - .padding(.horizontal, 8) - .background( - RoundedRectangle(cornerRadius: 6) - .fill(Color.accentColor.opacity(0.12)) - ) - } - .buttonStyle(.plain) - .disabled(bugReportInFlight) - - if let message = bugReportMessage { - Text(message) - .font(.caption2) - .foregroundColor(.secondary) - .fixedSize(horizontal: false, vertical: true) - } - } - } - private var processToggleBinding: Binding { Binding( get: { @@ -686,101 +419,6 @@ struct ContentView: View { ) } - private func sendBugReport() async { - bugReportInFlight = true - bugReportMessage = "Collecting logs..." - let service = BugReportService() - do { - let outcome = try await service.sendReport(isManual: true) - bugReportMessage = outcome.message - } catch { - bugReportMessage = error.localizedDescription - } - bugReportInFlight = false - } - - private func showUninstallConfirmationAlert() { - let alert = NSAlert() - alert.messageText = "Uninstall EXO" - alert.informativeText = """ - This will remove EXO and all its system components: - - • Network configuration daemon - • Launch at login registration - • EXO network location - - The app will be moved to Trash. - """ - alert.alertStyle = .warning - alert.addButton(withTitle: "Uninstall") - alert.addButton(withTitle: "Cancel") - - // Style the Uninstall button as destructive - if let uninstallButton = alert.buttons.first { - uninstallButton.hasDestructiveAction = true - } - - let response = alert.runModal() - if response == .alertFirstButtonReturn { - performUninstall() - } - } - - private func performUninstall() { - uninstallInProgress = true - - // Stop EXO process first - controller.cancelPendingLaunch() - controller.stop() - stateService.stopPolling() - - // Run the privileged uninstall on a background thread - // Using .utility QoS to avoid priority inversion with NSAppleScript's subprocess - DispatchQueue.global(qos: .utility).async { - do { - // Remove network setup daemon and components (requires admin privileges) - try NetworkSetupHelper.uninstall() - - DispatchQueue.main.async { - // Unregister from launch at login - LaunchAtLoginHelper.disable() - - // Move app to trash - self.moveAppToTrash() - - // Quit the app - DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { - NSApplication.shared.terminate(nil) - } - } - } catch { - DispatchQueue.main.async { - self.showErrorAlert(message: error.localizedDescription) - self.uninstallInProgress = false - } - } - } - } - - private func showErrorAlert(message: String) { - let alert = NSAlert() - alert.messageText = "Uninstall Failed" - alert.informativeText = message - alert.alertStyle = .critical - alert.addButton(withTitle: "OK") - alert.runModal() - } - - private func moveAppToTrash() { - guard let appURL = Bundle.main.bundleURL as URL? else { return } - do { - try FileManager.default.trashItem(at: appURL, resultingItemURL: nil) - } catch { - // If we can't trash the app, that's OK - user can do it manually - // The important system components have already been cleaned up - } - } - private var buildTag: String { Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown" } diff --git a/app/EXO/EXO/EXOApp.swift b/app/EXO/EXO/EXOApp.swift index 3aff58c3..bba26e24 100644 --- a/app/EXO/EXO/EXOApp.swift +++ b/app/EXO/EXO/EXOApp.swift @@ -21,7 +21,9 @@ struct EXOApp: App { @StateObject private var localNetworkChecker: LocalNetworkChecker @StateObject private var updater: SparkleUpdater @StateObject private var thunderboltBridgeService: ThunderboltBridgeService + @StateObject private var settingsWindowController: SettingsWindowController private let terminationObserver: TerminationObserver + private let firstLaunchPopout = FirstLaunchPopout() private let ciContext = CIContext(options: nil) init() { @@ -43,12 +45,13 @@ struct EXOApp: App { _updater = StateObject(wrappedValue: updater) let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service) _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge) + _settingsWindowController = StateObject(wrappedValue: SettingsWindowController()) enableLaunchAtLoginIfNeeded() // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops) NetworkSetupHelper.promptAndInstallIfNeeded() // Check local network access periodically (warning disappears when user grants permission) localNetwork.startPeriodicChecking(interval: 10) - controller.scheduleLaunch(after: 15) + controller.scheduleLaunch(after: 5) service.startPolling() networkStatus.startPolling() } @@ -62,8 +65,19 @@ struct EXOApp: App { .environmentObject(localNetworkChecker) .environmentObject(updater) .environmentObject(thunderboltBridgeService) + .environmentObject(settingsWindowController) } label: { menuBarIcon + .onReceive(controller.$isFirstLaunchReady) { ready in + if ready { + DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { + self.firstLaunchPopout.onComplete = { [weak controller] in + controller?.markOnboardingCompleted() + } + self.firstLaunchPopout.show() + } + } + } } .menuBarExtraStyle(.window) } diff --git a/app/EXO/EXO/ExoProcessController.swift b/app/EXO/EXO/ExoProcessController.swift index 7566674b..2350c118 100644 --- a/app/EXO/EXO/ExoProcessController.swift +++ b/app/EXO/EXO/ExoProcessController.swift @@ -5,6 +5,7 @@ import Foundation private let customNamespaceKey = "EXOCustomNamespace" private let hfTokenKey = "EXOHFToken" private let enableImageModelsKey = "EXOEnableImageModels" +private let onboardingCompletedKey = "EXOOnboardingCompleted" @MainActor final class ExoProcessController: ObservableObject { @@ -60,6 +61,9 @@ final class ExoProcessController: ObservableObject { } } + /// Fires once when EXO transitions to `.running` for the very first time (fresh install). + @Published private(set) var isFirstLaunchReady = false + private var process: Process? private var runtimeDirectoryURL: URL? private var pendingLaunchTask: Task? @@ -113,6 +117,11 @@ final class ExoProcessController: ObservableObject { try child.run() process = child status = .running + + // Show welcome popout if onboarding was never completed + if !UserDefaults.standard.bool(forKey: onboardingCompletedKey) { + isFirstLaunchReady = true + } } catch { process = nil status = .failed(message: "Launch error") @@ -164,6 +173,17 @@ final class ExoProcessController: ObservableObject { launch() } + /// Mark onboarding as completed (user interacted with the welcome popout). + func markOnboardingCompleted() { + UserDefaults.standard.set(true, forKey: onboardingCompletedKey) + } + + /// Reset onboarding so the welcome popout appears on next launch. + func resetOnboarding() { + UserDefaults.standard.removeObject(forKey: onboardingCompletedKey) + isFirstLaunchReady = false + } + func scheduleLaunch(after seconds: TimeInterval) { cancelPendingLaunch() let start = max(1, Int(ceil(seconds))) diff --git a/app/EXO/EXO/Views/FirstLaunchPopout.swift b/app/EXO/EXO/Views/FirstLaunchPopout.swift new file mode 100644 index 00000000..1a10b3ac --- /dev/null +++ b/app/EXO/EXO/Views/FirstLaunchPopout.swift @@ -0,0 +1,192 @@ +import AppKit +import SwiftUI + +/// A popover callout anchored to the menu bar icon on first launch, +/// pointing the user to the web dashboard with an arrow connecting to the icon. +@MainActor +final class FirstLaunchPopout { + private var popover: NSPopover? + private var countdownTask: Task? + private static let dashboardURL = "http://localhost:52415/" + + /// Called when the user completes onboarding (clicks Open Dashboard or dismisses). + var onComplete: (() -> Void)? + + func show() { + guard popover == nil else { return } + + // The status bar button may not exist yet on first launch; retry a few times. + showWithRetry(attemptsRemaining: 5) + } + + private func showWithRetry(attemptsRemaining: Int) { + guard attemptsRemaining > 0 else { return } + + guard let button = Self.findStatusItemButton() else { + DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { [weak self] in + self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1) + } + return + } + + let pop = NSPopover() + pop.behavior = .applicationDefined + pop.animates = true + pop.contentSize = NSSize(width: 280, height: 120) + pop.contentViewController = NSHostingController( + rootView: WelcomeCalloutView( + countdownDuration: 30, + onDismiss: { [weak self] in + self?.onComplete?() + self?.dismiss() + }, + onOpen: { [weak self] in + self?.openDashboard() + self?.onComplete?() + self?.dismiss() + } + ) + ) + + self.popover = pop + pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY) + + // Auto-open dashboard after 30s then dismiss + countdownTask = Task { + try? await Task.sleep(nanoseconds: 30_000_000_000) + if !Task.isCancelled { + openDashboard() + onComplete?() + dismiss() + } + } + } + + func dismiss() { + countdownTask?.cancel() + countdownTask = nil + UserDefaults.standard.set(true, forKey: "EXOOnboardingCompleted") + guard let pop = popover else { return } + popover = nil + pop.performClose(nil) + } + + private func openDashboard() { + guard let url = URL(string: Self.dashboardURL) else { return } + NSWorkspace.shared.open(url) + } + + /// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra. + /// Walks the view hierarchy to find the actual button rather than the content view. + private static func findStatusItemButton() -> NSView? { + for window in NSApp.windows { + let className = NSStringFromClass(type(of: window)) + if className.contains("NSStatusBarWindow") { + // Try to find the actual status bar button in the view hierarchy + if let content = window.contentView { + if let button = findButton(in: content) { + return button + } + return content + } + } + } + return nil + } + + /// Recursively searches the view hierarchy for an NSStatusBarButton. + private static func findButton(in view: NSView) -> NSView? { + let className = NSStringFromClass(type(of: view)) + if className.contains("StatusBarButton") { + return view + } + for subview in view.subviews { + if let found = findButton(in: subview) { + return found + } + } + return nil + } +} + +/// Minimal welcome callout — friendly pointer, not a wall of text. +/// Rendered inside the NSPopover which provides its own chrome and arrow. +private struct WelcomeCalloutView: View { + let countdownDuration: Int + let onDismiss: () -> Void + let onOpen: () -> Void + @State private var countdown: Int + @State private var timerTask: Task? + + init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) { + self.countdownDuration = countdownDuration + self.onDismiss = onDismiss + self.onOpen = onOpen + self._countdown = State(initialValue: countdownDuration) + } + + var body: some View { + VStack(alignment: .leading, spacing: 10) { + HStack(alignment: .top) { + Text("Welcome to EXO!") + .font(.system(.headline, design: .rounded)) + .fontWeight(.semibold) + .foregroundColor(.primary) + Spacer() + Button { + onDismiss() + } label: { + Image(systemName: "xmark.circle.fill") + .font(.system(size: 14)) + .foregroundStyle(.tertiary) + } + .buttonStyle(.plain) + } + + Text("Run your first model here:") + .font(.system(.subheadline, design: .default)) + .foregroundColor(.secondary) + + HStack { + Button { + onOpen() + } label: { + Label("Open Dashboard", systemImage: "arrow.up.right.square") + .font(.system(.caption, design: .default)) + .fontWeight(.medium) + } + .buttonStyle(.borderedProminent) + .tint(.accentColor) + .controlSize(.small) + + Spacer() + + if countdown > 0 { + Text("Auto-opens in \(countdown)s") + .font(.system(.caption2, design: .default)) + .foregroundColor(.secondary.opacity(0.6)) + .monospacedDigit() + } + } + } + .padding(14) + .onAppear { + startCountdown() + } + .onDisappear { + timerTask?.cancel() + timerTask = nil + } + } + + private func startCountdown() { + timerTask = Task { + while countdown > 0 { + try? await Task.sleep(nanoseconds: 1_000_000_000) + if !Task.isCancelled { + countdown -= 1 + } + } + } + } +} diff --git a/app/EXO/EXO/Views/SettingsView.swift b/app/EXO/EXO/Views/SettingsView.swift new file mode 100644 index 00000000..ed221c89 --- /dev/null +++ b/app/EXO/EXO/Views/SettingsView.swift @@ -0,0 +1,478 @@ +import AppKit +import SwiftUI + +/// Native macOS Settings window following Apple HIG. +/// Organized into General, Model, Advanced, and About sections. +struct SettingsView: View { + @EnvironmentObject private var controller: ExoProcessController + @EnvironmentObject private var updater: SparkleUpdater + @EnvironmentObject private var networkStatusService: NetworkStatusService + @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService + @EnvironmentObject private var stateService: ClusterStateService + + @State private var pendingNamespace: String = "" + @State private var pendingHFToken: String = "" + @State private var pendingEnableImageModels = false + @State private var needsRestart = false + @State private var bugReportInFlight = false + @State private var bugReportMessage: String? + @State private var uninstallInProgress = false + + var body: some View { + TabView { + generalTab + .tabItem { + Label("General", systemImage: "gear") + } + modelTab + .tabItem { + Label("Model", systemImage: "cube") + } + advancedTab + .tabItem { + Label("Advanced", systemImage: "wrench.and.screwdriver") + } + aboutTab + .tabItem { + Label("About", systemImage: "info.circle") + } + } + .frame(width: 450, height: 400) + .onAppear { + pendingNamespace = controller.customNamespace + pendingHFToken = controller.hfToken + pendingEnableImageModels = controller.enableImageModels + needsRestart = false + } + } + + // MARK: - General Tab + + private var generalTab: some View { + Form { + Section { + LabeledContent("Cluster Namespace") { + TextField("default", text: $pendingNamespace) + .textFieldStyle(.roundedBorder) + .frame(width: 200) + } + Text("Nodes with the same namespace form a cluster. Leave empty for default.") + .font(.caption) + .foregroundColor(.secondary) + } + + Section { + LabeledContent("HuggingFace Token") { + SecureField("optional", text: $pendingHFToken) + .textFieldStyle(.roundedBorder) + .frame(width: 200) + } + Text("Required for gated models. Get yours at huggingface.co/settings/tokens") + .font(.caption) + .foregroundColor(.secondary) + } + + Section { + HStack { + Spacer() + Button("Save & Restart") { + applyGeneralSettings() + } + .disabled(!hasGeneralChanges) + } + } + } + .formStyle(.grouped) + .padding() + } + + // MARK: - Model Tab + + private var modelTab: some View { + Form { + Section { + Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels) + Text("Allow text-to-image and image-to-image models in the model picker.") + .font(.caption) + .foregroundColor(.secondary) + } + + Section { + HStack { + Spacer() + Button("Save & Restart") { + applyModelSettings() + } + .disabled(!hasModelChanges) + } + } + } + .formStyle(.grouped) + .padding() + } + + // MARK: - Advanced Tab + + private var advancedTab: some View { + Form { + Section("Onboarding") { + HStack { + VStack(alignment: .leading) { + Text("Reset Onboarding") + Text("Opens the dashboard and resets the onboarding wizard.") + .font(.caption) + .foregroundColor(.secondary) + } + Spacer() + Button("Reset") { + guard let url = URL(string: "http://localhost:52415/?reset-onboarding") + else { return } + NSWorkspace.shared.open(url) + } + } + } + + Section("Debug Info") { + LabeledContent("Thunderbolt Bridge") { + Text(thunderboltStatusText) + .foregroundColor(thunderboltStatusColor) + } + + VStack(alignment: .leading, spacing: 2) { + clusterThunderboltBridgeView + } + + VStack(alignment: .leading, spacing: 2) { + interfaceIpList + } + + VStack(alignment: .leading, spacing: 2) { + rdmaStatusView + } + + sendBugReportButton + } + + Section("Danger Zone") { + Button(role: .destructive) { + showUninstallConfirmationAlert() + } label: { + HStack { + Text("Uninstall EXO") + Spacer() + Image(systemName: "trash") + .imageScale(.small) + } + } + .disabled(uninstallInProgress) + } + } + .formStyle(.grouped) + .padding() + } + + // MARK: - About Tab + + private var aboutTab: some View { + Form { + Section { + LabeledContent("Version") { + Text(buildTag) + .textSelection(.enabled) + } + LabeledContent("Commit") { + Text(buildCommit) + .font(.system(.body, design: .monospaced)) + .textSelection(.enabled) + } + } + + Section { + Button("Check for Updates") { + updater.checkForUpdates() + } + } + } + .formStyle(.grouped) + .padding() + } + + // MARK: - Debug Info Views (moved from ContentView) + + private var thunderboltStatusText: String { + switch networkStatusService.status.thunderboltBridgeState { + case .some(.disabled): + return "Disabled" + case .some(.deleted): + return "Deleted" + case .some(.enabled): + return "Enabled" + case nil: + return "Unknown" + } + } + + private var thunderboltStatusColor: Color { + switch networkStatusService.status.thunderboltBridgeState { + case .some(.disabled), .some(.deleted): + return .green + case .some(.enabled): + return .red + case nil: + return .secondary + } + } + + private var clusterThunderboltBridgeView: some View { + let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:] + let localNodeId = stateService.localNodeId + let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] + + return VStack(alignment: .leading, spacing: 1) { + if bridgeStatuses.isEmpty { + Text("Cluster TB Bridge: No data") + .font(.caption2) + .foregroundColor(.secondary) + } else { + Text("Cluster TB Bridge Status:") + .font(.caption2) + .foregroundColor(.secondary) + ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in + if let status = bridgeStatuses[nodeId] { + let nodeName = + nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) + let isLocal = nodeId == localNodeId + let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" + let statusText = + !status.exists + ? "N/A" + : (status.enabled ? "Enabled" : "Disabled") + let color: Color = + !status.exists + ? .secondary + : (status.enabled ? .red : .green) + Text("\(prefix) \(statusText)") + .font(.caption2) + .foregroundColor(color) + } + } + } + } + } + + private var interfaceIpList: some View { + let statuses = networkStatusService.status.interfaceStatuses + return VStack(alignment: .leading, spacing: 1) { + Text("Interfaces (en0–en7):") + .font(.caption2) + .foregroundColor(.secondary) + if statuses.isEmpty { + Text(" Unknown") + .font(.caption2) + .foregroundColor(.secondary) + } else { + ForEach(statuses, id: \.interfaceName) { status in + let ipText = status.ipAddress ?? "No IP" + Text(" \(status.interfaceName): \(ipText)") + .font(.caption2) + .foregroundColor(status.ipAddress == nil ? .red : .green) + } + } + } + } + + private var rdmaStatusView: some View { + let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:] + let localNodeId = stateService.localNodeId + let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] + let localDevices = networkStatusService.status.localRdmaDevices + let localPorts = networkStatusService.status.localRdmaActivePorts + + return VStack(alignment: .leading, spacing: 1) { + if rdmaStatuses.isEmpty { + Text("Cluster RDMA: No data") + .font(.caption2) + .foregroundColor(.secondary) + } else { + Text("Cluster RDMA Status:") + .font(.caption2) + .foregroundColor(.secondary) + ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in + if let status = rdmaStatuses[nodeId] { + let nodeName = + nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) + let isLocal = nodeId == localNodeId + let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" + let statusText = status.enabled ? "Enabled" : "Disabled" + let color: Color = status.enabled ? .green : .orange + Text("\(prefix) \(statusText)") + .font(.caption2) + .foregroundColor(color) + } + } + } + if !localDevices.isEmpty { + Text(" Local Devices: \(localDevices.joined(separator: ", "))") + .font(.caption2) + .foregroundColor(.secondary) + } + if !localPorts.isEmpty { + Text(" Local Active Ports:") + .font(.caption2) + .foregroundColor(.secondary) + ForEach(localPorts, id: \.device) { port in + Text(" \(port.device) port \(port.port): \(port.state)") + .font(.caption2) + .foregroundColor(.green) + } + } + } + } + + private var sendBugReportButton: some View { + VStack(alignment: .leading, spacing: 4) { + Button { + Task { + await sendBugReport() + } + } label: { + HStack { + if bugReportInFlight { + ProgressView() + .scaleEffect(0.6) + } + Text("Send Bug Report") + .font(.caption) + .fontWeight(.semibold) + Spacer() + } + } + .disabled(bugReportInFlight) + + if let message = bugReportMessage { + Text(message) + .font(.caption2) + .foregroundColor(.secondary) + .fixedSize(horizontal: false, vertical: true) + } + } + } + + // MARK: - Actions + + private func sendBugReport() async { + bugReportInFlight = true + bugReportMessage = "Collecting logs..." + let service = BugReportService() + do { + let outcome = try await service.sendReport(isManual: true) + bugReportMessage = outcome.message + } catch { + bugReportMessage = error.localizedDescription + } + bugReportInFlight = false + } + + private func showUninstallConfirmationAlert() { + let alert = NSAlert() + alert.messageText = "Uninstall EXO" + alert.informativeText = """ + This will remove EXO and all its system components: + + • Network configuration daemon + • Launch at login registration + • EXO network location + + The app will be moved to Trash. + """ + alert.alertStyle = .warning + alert.addButton(withTitle: "Uninstall") + alert.addButton(withTitle: "Cancel") + + if let uninstallButton = alert.buttons.first { + uninstallButton.hasDestructiveAction = true + } + + let response = alert.runModal() + if response == .alertFirstButtonReturn { + performUninstall() + } + } + + private func performUninstall() { + uninstallInProgress = true + + controller.cancelPendingLaunch() + controller.stop() + stateService.stopPolling() + + DispatchQueue.global(qos: .utility).async { + do { + try NetworkSetupHelper.uninstall() + + DispatchQueue.main.async { + LaunchAtLoginHelper.disable() + self.moveAppToTrash() + + DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { + NSApplication.shared.terminate(nil) + } + } + } catch { + DispatchQueue.main.async { + let errorAlert = NSAlert() + errorAlert.messageText = "Uninstall Failed" + errorAlert.informativeText = error.localizedDescription + errorAlert.alertStyle = .critical + errorAlert.addButton(withTitle: "OK") + errorAlert.runModal() + self.uninstallInProgress = false + } + } + } + } + + private func moveAppToTrash() { + guard let appURL = Bundle.main.bundleURL as URL? else { return } + do { + try FileManager.default.trashItem(at: appURL, resultingItemURL: nil) + } catch { + // If we can't trash the app, that's OK - user can do it manually + } + } + + // MARK: - Helpers + + private var hasGeneralChanges: Bool { + pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken + } + + private var hasModelChanges: Bool { + pendingEnableImageModels != controller.enableImageModels + } + + private func applyGeneralSettings() { + controller.customNamespace = pendingNamespace + controller.hfToken = pendingHFToken + restartIfRunning() + } + + private func applyModelSettings() { + controller.enableImageModels = pendingEnableImageModels + restartIfRunning() + } + + private func restartIfRunning() { + if controller.status == .running || controller.status == .starting { + controller.restart() + } + } + + private var buildTag: String { + Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown" + } + + private var buildCommit: String { + Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown" + } +} diff --git a/app/EXO/EXO/Views/SettingsWindowController.swift b/app/EXO/EXO/Views/SettingsWindowController.swift new file mode 100644 index 00000000..98517f92 --- /dev/null +++ b/app/EXO/EXO/Views/SettingsWindowController.swift @@ -0,0 +1,47 @@ +import AppKit +import SwiftUI + +/// Manages a standalone native macOS Settings window. +/// Ensures only one instance exists and brings it to front on repeated opens. +@MainActor +final class SettingsWindowController: ObservableObject { + private var window: NSWindow? + + func open( + controller: ExoProcessController, + updater: SparkleUpdater, + networkStatusService: NetworkStatusService, + thunderboltBridgeService: ThunderboltBridgeService, + stateService: ClusterStateService + ) { + if let existing = window, existing.isVisible { + existing.makeKeyAndOrderFront(nil) + NSApp.activate(ignoringOtherApps: true) + return + } + + let settingsView = SettingsView() + .environmentObject(controller) + .environmentObject(updater) + .environmentObject(networkStatusService) + .environmentObject(thunderboltBridgeService) + .environmentObject(stateService) + + let hostingView = NSHostingView(rootView: settingsView) + + let newWindow = NSWindow( + contentRect: NSRect(x: 0, y: 0, width: 450, height: 400), + styleMask: [.titled, .closable], + backing: .buffered, + defer: false + ) + newWindow.title = "EXO Settings" + newWindow.contentView = hostingView + newWindow.center() + newWindow.isReleasedWhenClosed = false + newWindow.makeKeyAndOrderFront(nil) + NSApp.activate(ignoringOtherApps: true) + + window = newWindow + } +} diff --git a/dashboard/src/app.css b/dashboard/src/app.css index fc532578..3e951163 100644 --- a/dashboard/src/app.css +++ b/dashboard/src/app.css @@ -202,6 +202,15 @@ filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5)); } +/* Onboarding step 2: connection line between devices */ +.onboarding-connection-line { + stroke: oklch(0.85 0.18 85 / 0.5); + stroke-width: 1.5px; + stroke-dasharray: 6, 6; + animation: flowAnimation 1s linear infinite; + filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4)); +} + .graph-link-active { stroke: oklch(0.85 0.18 85 / 0.8); stroke-width: 2px; @@ -320,3 +329,31 @@ input:focus, textarea:focus { transform: translate(400px, 400px); } } + +/* Respect reduced motion preference */ +@media (prefers-reduced-motion: reduce) { + .shooting-star, + .shooting-star::before { + animation: none !important; + opacity: 0 !important; + } + .graph-link { + animation: none; + } + .status-pulse { + animation: none; + } + .cursor-blink { + animation: none; + } + .onboarding-connection-line { + animation: none; + } + *, + *::before, + *::after { + transition-duration: 0.01ms !important; + animation-duration: 0.01ms !important; + animation-iteration-count: 1 !important; + } +} diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte index 2eb7aa40..1bc554e8 100644 --- a/dashboard/src/lib/components/ChatForm.svelte +++ b/dashboard/src/lib/components/ChatForm.svelte @@ -28,6 +28,7 @@ showModelSelector?: boolean; modelTasks?: Record; modelCapabilities?: Record; + onSend?: () => void; } let { @@ -38,6 +39,7 @@ showModelSelector = false, modelTasks = {}, modelCapabilities = {}, + onSend, }: Props = $props(); let message = $state(""); @@ -305,6 +307,8 @@ ); } + onSend?.(); + // Refocus the textarea after sending setTimeout(() => textareaRef?.focus(), 10); } diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte index ba5322a7..7c076459 100644 --- a/dashboard/src/lib/components/ChatMessages.svelte +++ b/dashboard/src/lib/components/ChatMessages.svelte @@ -802,8 +802,8 @@ > AWAITING INPUT

-

- ENTER A QUERY TO BEGIN +

+ Type a message below · Shift+Enter for newline

{/if} @@ -818,6 +818,7 @@ onclick={scrollToBottom} class="sticky bottom-4 left-1/2 -translate-x-1/2 w-10 h-10 rounded-full bg-exo-dark-gray/90 border border-exo-medium-gray/50 flex items-center justify-center text-exo-light-gray hover:text-exo-yellow hover:border-exo-yellow/50 transition-all shadow-lg cursor-pointer z-10" title="Scroll to bottom" + aria-label="Scroll to bottom of messages" >
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"} @@ -372,39 +372,37 @@ onkeydown={(e) => e.key === "Enter" && handleSelectConversation(conversation.id)} - class="group w-full flex items-center justify-between p-2 rounded mb-1 transition-all text-left cursor-pointer + class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer {activeId === conversation.id - ? 'bg-transparent border border-exo-yellow/30' - : 'hover:border-exo-yellow/20 border border-transparent'}" + ? 'bg-exo-yellow/5 border border-exo-yellow/30' + : 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}" >
{conversation.name}
-
+
{formatDate(conversation.updatedAt)}
-
+
{info.modelLabel}
-
- Strategy: {info.strategyLabel} -
{#if stats} -
- {#if stats.ttftMs}TTFT - {stats.ttftMs.toFixed( - 0, - )}ms{/if}{#if stats.ttftMs && stats.tps}•{/if}{#if stats.tps}{stats.tps.toFixed(1)} - tok/s{/if} +
+ {#if stats.ttftMs}TTFT + {stats.ttftMs.toFixed(0)}ms{/if}{#if stats.ttftMs && stats.tps}·{/if}{#if stats.tps}{stats.tps.toFixed(1)} + tok/s{/if}
{/if}
diff --git a/dashboard/src/lib/components/ConnectionBanner.svelte b/dashboard/src/lib/components/ConnectionBanner.svelte new file mode 100644 index 00000000..3339cf24 --- /dev/null +++ b/dashboard/src/lib/components/ConnectionBanner.svelte @@ -0,0 +1,20 @@ + + +{#if !connected} + +{/if} diff --git a/dashboard/src/lib/components/HeaderNav.svelte b/dashboard/src/lib/components/HeaderNav.svelte index 5bdcf1ba..d78bc809 100644 --- a/dashboard/src/lib/components/HeaderNav.svelte +++ b/dashboard/src/lib/components/HeaderNav.svelte @@ -6,6 +6,10 @@ export let showSidebarToggle = false; export let sidebarVisible = true; export let onToggleSidebar: (() => void) | null = null; + export let downloadProgress: { + count: number; + percentage: number; + } | null = null; function handleHome(): void { if (onHome) { @@ -35,11 +39,15 @@ onclick={handleToggleSidebar} class="p-2 rounded border border-exo-medium-gray/40 hover:border-exo-yellow/50 transition-colors cursor-pointer" title={sidebarVisible ? "Hide sidebar" : "Show sidebar"} + aria-label={sidebarVisible + ? "Hide conversation sidebar" + : "Show conversation sidebar"} + aria-pressed={sidebarVisible} > -
{#if showHome}
+ diff --git a/dashboard/src/lib/components/ModelCard.svelte b/dashboard/src/lib/components/ModelCard.svelte index 561c325b..b432b7a8 100644 --- a/dashboard/src/lib/components/ModelCard.svelte +++ b/dashboard/src/lib/components/ModelCard.svelte @@ -567,11 +567,17 @@
{sharding} {runtime === "MlxRing" ? "MLX Ring" @@ -581,6 +587,26 @@
+ + {#if isDownloading && progress} +
+
+ Downloading + {percentage.toFixed(1)}% · {formatSpeed(progress.speed)} + · {formatEta(progress.etaMs)} +
+
+
+
+
+ {/if} + {#if placementPreview().nodes.length > 0} {@const preview = placementPreview()} diff --git a/dashboard/src/lib/components/ModelPickerModal.svelte b/dashboard/src/lib/components/ModelPickerModal.svelte index 84a93ee7..cf21c727 100644 --- a/dashboard/src/lib/components/ModelPickerModal.svelte +++ b/dashboard/src/lib/components/ModelPickerModal.svelte @@ -512,6 +512,18 @@ ); }); + // Split filtered groups into recommended (fits_now) and others for visual separation + const recommendedGroups = $derived( + filteredGroups.filter((g) => + g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"), + ), + ); + const otherGroups = $derived( + filteredGroups.filter( + (g) => !g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"), + ), + ); + function toggleGroupExpanded(groupId: string) { const next = new Set(expandedGroups); if (next.has(groupId)) { @@ -840,7 +852,60 @@ {/if}
{:else} - {#each filteredGroups as group} + + {#if recommendedGroups.length > 0 && otherGroups.length > 0 && !searchQuery.trim()} +
+ + + + Recommended for your cluster + — fits in available memory +
+ {/if} + {#each recommendedGroups as group} + toggleGroupExpanded(group.id)} + onSelectModel={handleSelect} + {onToggleFavorite} + onShowInfo={(g) => (infoGroup = g)} + downloadStatusMap={getVariantDownloadMap(group)} + /> + {/each} + + {#if otherGroups.length > 0 && recommendedGroups.length > 0 && !searchQuery.trim()} +
+ Other models +
+ {/if} + {#each otherGroups as group} + import { toasts, dismissToast, type Toast } from "$lib/stores/toast.svelte"; + import { fly, fade } from "svelte/transition"; + import { flip } from "svelte/animate"; + + const items = $derived(toasts()); + + const typeStyles: Record< + Toast["type"], + { border: string; icon: string; iconColor: string } + > = { + success: { + border: "border-l-green-500", + icon: "M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z", + iconColor: "text-green-400", + }, + error: { + border: "border-l-red-500", + icon: "M12 9v3.75m9-.75a9 9 0 11-18 0 9 9 0 0118 0zm-9 3.75h.008v.008H12v-.008z", + iconColor: "text-red-400", + }, + warning: { + border: "border-l-yellow-500", + icon: "M12 9v3.75m-9.303 3.376c-.866 1.5.217 3.374 1.948 3.374h14.71c1.73 0 2.813-1.874 1.948-3.374L13.949 3.378c-.866-1.5-3.032-1.5-3.898 0L2.697 16.126z", + iconColor: "text-yellow-400", + }, + info: { + border: "border-l-blue-500", + icon: "M11.25 11.25l.041-.02a.75.75 0 011.063.852l-.708 2.836a.75.75 0 001.063.853l.041-.021M21 12a9 9 0 11-18 0 9 9 0 0118 0zm-9-3.75h.008v.008H12V8.25z", + iconColor: "text-blue-400", + }, + }; + + +{#if items.length > 0} +
+ {#each items as toast (toast.id)} + {@const style = typeStyles[toast.type]} + + {/each} +
+{/if} + + diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index ebb2d0df..5a5afc78 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -587,6 +587,12 @@ class AppStore { // Image editing state editingImage = $state(null); + /** True when the backend is reachable. */ + isConnected = $state(true); + /** Number of consecutive fetch failures. */ + private consecutiveFailures = 0; + private static readonly CONNECTION_LOST_THRESHOLD = 3; + private fetchInterval: ReturnType | null = null; private previewsInterval: ReturnType | null = null; private lastConversationPersistTs = 0; @@ -1290,7 +1296,19 @@ class AppStore { // Thunderbolt bridge status per node this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {}; this.lastUpdate = Date.now(); + // Connection recovered + if (!this.isConnected) { + this.isConnected = true; + } + this.consecutiveFailures = 0; } catch (error) { + this.consecutiveFailures++; + if ( + this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD && + this.isConnected + ) { + this.isConnected = false; + } console.error("Error fetching state:", error); } } @@ -1817,7 +1835,7 @@ class AppStore { assistantMessage.id, (msg) => { msg.content = - "Error: No model available. Please launch an instance first."; + "No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically."; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); @@ -2255,7 +2273,7 @@ class AppStore { const modelToUse = this.getModelForRequest(); if (!modelToUse) { throw new Error( - "No model selected and no running instances available. Please launch an instance first.", + "No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.", ); } @@ -3144,6 +3162,9 @@ export const setChatSidebarVisible = (visible: boolean) => appStore.setChatSidebarVisible(visible); export const refreshState = () => appStore.fetchState(); +// Connection status +export const isConnected = () => appStore.isConnected; + // Node identities (for OS version mismatch detection) export const nodeIdentities = () => appStore.nodeIdentities; diff --git a/dashboard/src/lib/stores/toast.svelte.ts b/dashboard/src/lib/stores/toast.svelte.ts new file mode 100644 index 00000000..e9e62019 --- /dev/null +++ b/dashboard/src/lib/stores/toast.svelte.ts @@ -0,0 +1,87 @@ +/** + * Toast notification store - Global notification system for the EXO dashboard. + * + * Usage: + * import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte"; + * addToast({ type: "success", message: "Model launched" }); + * addToast({ type: "error", message: "Connection lost", persistent: true }); + */ + +type ToastType = "success" | "error" | "warning" | "info"; + +export interface Toast { + id: string; + type: ToastType; + message: string; + /** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */ + duration: number; + createdAt: number; +} + +interface ToastInput { + type: ToastType; + message: string; + /** If true, toast stays until manually dismissed. Default: false. */ + persistent?: boolean; + /** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */ + duration?: number; +} + +const DEFAULT_DURATIONS: Record = { + success: 4000, + info: 4000, + warning: 6000, + error: 6000, +}; + +let toastList = $state([]); +const timers = new Map>(); + +function generateId(): string { + return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; +} + +export function addToast(input: ToastInput): string { + const id = generateId(); + const duration = input.persistent + ? 0 + : (input.duration ?? DEFAULT_DURATIONS[input.type]); + + const toast: Toast = { + id, + type: input.type, + message: input.message, + duration, + createdAt: Date.now(), + }; + + toastList = [...toastList, toast]; + + if (duration > 0) { + const timer = setTimeout(() => dismissToast(id), duration); + timers.set(id, timer); + } + + return id; +} + +export function dismissToast(id: string): void { + const timer = timers.get(id); + if (timer) { + clearTimeout(timer); + timers.delete(id); + } + toastList = toastList.filter((t) => t.id !== id); +} + +/** Dismiss all toasts matching a message (useful for dedup). */ +export function dismissByMessage(message: string): void { + const matching = toastList.filter((t) => t.message === message); + for (const t of matching) { + dismissToast(t.id); + } +} + +export function toasts(): Toast[] { + return toastList; +} diff --git a/dashboard/src/routes/+layout.svelte b/dashboard/src/routes/+layout.svelte index d249e9cf..295c9dc0 100644 --- a/dashboard/src/routes/+layout.svelte +++ b/dashboard/src/routes/+layout.svelte @@ -1,5 +1,7 @@ @@ -10,5 +12,7 @@
+ {@render children?.()} +
diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 5f7dcf04..20f196ee 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -36,6 +36,7 @@ createConversation, setSelectedChatModel, selectedChatModel, + sendMessage, debugMode, toggleDebugMode, topologyOnlyMode, @@ -52,9 +53,11 @@ type PlacementPreview, type MetaInstanceData, } from "$lib/stores/app.svelte"; + import { addToast } from "$lib/stores/toast.svelte"; import HeaderNav from "$lib/components/HeaderNav.svelte"; - import { fade, fly } from "svelte/transition"; - import { cubicInOut } from "svelte/easing"; + import { fade, fly, slide } from "svelte/transition"; + import { tweened } from "svelte/motion"; + import { cubicInOut, cubicOut } from "svelte/easing"; import { onMount } from "svelte"; const chatStarted = $derived(hasStartedChat()); @@ -140,6 +143,26 @@ const rdmaCtlData = $derived(nodeRdmaCtl()); const nodeFilter = $derived(previewNodeFilter()); + // Aggregate active download progress across all instances for header indicator + const activeDownloadSummary = $derived.by(() => { + let totalBytes = 0; + let downloadedBytes = 0; + let count = 0; + for (const [id, inst] of Object.entries(instanceData)) { + const status = getInstanceDownloadStatus(id, inst); + if (status.isDownloading && status.progress) { + count++; + totalBytes += status.progress.totalBytes || 0; + downloadedBytes += status.progress.downloadedBytes || 0; + } + } + if (count === 0) return null; + return { + count, + percentage: totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0, + }; + }); + // Detect macOS version mismatches across cluster nodes const macosVersionMismatch = $derived.by(() => { if (!identitiesData) return null; @@ -230,6 +253,171 @@ let mounted = $state(false); + // ── Onboarding wizard state ── + const ONBOARDING_COMPLETE_KEY = "exo-onboarding-complete"; + let onboardingStep = $state(0); // 0 = not in onboarding, 1-7 = wizard steps + let onboardingModelId = $state(null); // model selected during onboarding + const showOnboarding = $derived(onboardingStep > 0); + + // ── Step 2 animation state: "Add more devices, run bigger models" ── + let deviceAnimPhase = $state(0); // 0=waiting, 1=macbook, 2=studio joins, 3=connection+mid unlock, 4=big unlock + let showContinueStep2 = $state(false); + const studioX = tweened(540, { duration: 700, easing: cubicOut }); + const studioOpacity = tweened(0, { duration: 700, easing: cubicOut }); + + $effect(() => { + if (onboardingStep === 2) { + deviceAnimPhase = 0; + showContinueStep2 = false; + studioX.set(540, { duration: 0 }); + studioOpacity.set(0, { duration: 0 }); + + const t1 = setTimeout(() => { + deviceAnimPhase = 1; + }, 100); + const t2 = setTimeout(() => { + deviceAnimPhase = 2; + studioX.set(340); + studioOpacity.set(1); + }, 900); + const t3 = setTimeout(() => { + deviceAnimPhase = 3; + }, 1700); + const t4 = setTimeout(() => { + deviceAnimPhase = 4; + }, 2500); + const t5 = setTimeout(() => { + showContinueStep2 = true; + }, 3500); + + return () => { + clearTimeout(t1); + clearTimeout(t2); + clearTimeout(t3); + clearTimeout(t4); + clearTimeout(t5); + }; + } + }); + + // Recommended models for onboarding (sorted by fit, then size desc, limited to 6) + const onboardingModels = $derived.by(() => { + if (models.length === 0) return []; + return [...models] + .filter((m) => getModelMemoryFitStatus(m) !== "too_large") + .sort((a, b) => { + const aFit = hasEnoughMemory(a) ? 0 : 1; + const bFit = hasEnoughMemory(b) ? 0 : 1; + if (aFit !== bFit) return aFit - bFit; + return getModelSizeGB(b) - getModelSizeGB(a); + }) + .slice(0, 6); + }); + + // Track onboarding instance status for auto-advancing steps. + // Handles cached models: if no download is needed, skip step 5 entirely. + $effect(() => { + if (onboardingStep === 5 && instanceCount > 0) { + let anyDownloading = false; + let anyReady = false; + for (const [id, inst] of Object.entries(instanceData)) { + const status = getInstanceDownloadStatus(id, inst); + if (status.isDownloading) { + anyDownloading = true; + } + if ( + status.statusText === "READY" || + status.statusText === "LOADED" || + status.statusText === "RUNNING" + ) { + anyReady = true; + } + } + // Model already cached & ready — skip download AND loading steps + if (anyReady) { + onboardingStep = 7; + } else if (!anyDownloading) { + // Download finished (or was never needed) but not ready yet + onboardingStep = 6; + } + } + }); + + $effect(() => { + if (onboardingStep === 6 && instanceCount > 0) { + for (const [id, inst] of Object.entries(instanceData)) { + const status = getInstanceDownloadStatus(id, inst); + if ( + status.statusText === "READY" || + status.statusText === "LOADED" || + status.statusText === "RUNNING" + ) { + onboardingStep = 7; + break; + } + } + } + }); + + function completeOnboarding() { + onboardingStep = 0; + try { + localStorage.setItem(ONBOARDING_COMPLETE_KEY, "true"); + } catch { + // ignore + } + } + + let onboardingError = $state(null); + + async function onboardingLaunchModel(modelId: string) { + onboardingModelId = modelId; + onboardingError = null; + selectPreviewModel(modelId); + onboardingStep = 5; + // Launch via API + try { + const placementResponse = await fetch( + `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=1`, + ); + if (!placementResponse.ok) { + const errorText = await placementResponse.text(); + onboardingError = `Could not place model: ${errorText}`; + onboardingStep = 4; + return; + } + const placementData = await placementResponse.json(); + const response = await fetch("/instance", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ instance: placementData }), + }); + if (!response.ok) { + const errorText = await response.text(); + onboardingError = `Failed to launch: ${errorText}`; + onboardingStep = 4; + return; + } + setSelectedChatModel(modelId); + recordRecentLaunch(modelId); + } catch (error) { + onboardingError = `Network error: ${error}`; + onboardingStep = 4; + } + } + + // Helper to get onboarding download progress + const onboardingDownloadProgress = $derived.by(() => { + if (instanceCount === 0) return null; + for (const [id, inst] of Object.entries(instanceData)) { + const status = getInstanceDownloadStatus(id, inst); + if (status.isDownloading && status.progress) { + return status.progress; + } + } + return null; + }); + // Instance launch state let models = $state< Array<{ @@ -374,6 +562,9 @@ // Model picker modal state let isModelPickerOpen = $state(false); + // Advanced options toggle (hides technical jargon for new users) + let showAdvancedOptions = $state(false); + // Favorites state (reactive) const favoritesSet = $derived(getFavoritesSet()); @@ -703,6 +894,20 @@ onMount(() => { mounted = true; fetchModels(); + + // Handle reset-onboarding query parameter (triggered from native Settings) + const params = new URLSearchParams(window.location.search); + if (params.has("reset-onboarding")) { + localStorage.removeItem(ONBOARDING_COMPLETE_KEY); + window.history.replaceState({}, "", window.location.pathname); + onboardingStep = 1; + return; + } + + // Show onboarding wizard for first-time users + if (!localStorage.getItem(ONBOARDING_COMPLETE_KEY)) { + onboardingStep = 1; + } }); async function fetchModels() { @@ -782,7 +987,30 @@ ? Array.from(nodeFilter) : undefined; - const response = await fetch("/meta_instance", { + if (preview?.instance) { + // Use the instance from the preview + instanceData = preview.instance; + } else { + // Fallback: GET placement from API + const placementResponse = await fetch( + `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=${selectedMinNodes}`, + ); + + if (!placementResponse.ok) { + const errorText = await placementResponse.text(); + console.error("Failed to get placement:", errorText); + addToast({ + type: "error", + message: `Placement failed: ${errorText}`, + }); + return; + } + + instanceData = await placementResponse.json(); + } + + // POST the instance to create it + const response = await fetch("/instance", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ @@ -796,8 +1024,13 @@ if (!response.ok) { const errorText = await response.text(); - console.error("Failed to create meta instance:", errorText); + console.error("Failed to launch instance:", errorText); + addToast({ + type: "error", + message: `Failed to launch model: ${errorText}`, + }); } else { + addToast({ type: "success", message: `Model launched successfully` }); // Always auto-select the newly launched model so the user chats to what they just launched setSelectedChatModel(modelId); @@ -819,7 +1052,11 @@ setTimeout(scrollToBottom, 1000); } } catch (error) { - console.error("Error creating meta instance:", error); + console.error("Error launching instance:", error); + addToast({ + type: "error", + message: "Failed to launch model. Check console for details.", + }); } finally { launchingModelId = null; } @@ -1428,6 +1665,7 @@ if (!response.ok) { console.error("Failed to delete instance:", response.status); + addToast({ type: "error", message: "Failed to delete instance" }); } else if (wasSelected) { // If we deleted the currently selected model, switch to another available model // Find another instance that isn't the one we just deleted @@ -2416,146 +2654,979 @@ class="relative h-screen w-full flex flex-col bg-exo-dark-gray overflow-hidden" > -
+ {#if !showOnboarding} +
- -
-
-
-
-
- - {#if !topologyOnlyEnabled} - + +
+
+
+
+
{/if} - -
- - {#if !topologyOnlyEnabled && sidebarVisible} -
- -
- {/if} - - {#if topologyOnlyEnabled} - -
+ {#if showOnboarding} + + + +
+ {#if onboardingStep === 1} +
- - - {@render clusterWarnings()} - - - {#if tb5WithoutRdma && !tb5InfoDismissed} +
0} - class:top-4={tbBridgeCycles.length === 0} - role="status" + class="text-5xl font-mono font-bold text-exo-yellow tracking-wider mb-4" > + exo +
+

+ Welcome to exo +

+

+ Run AI models locally, across all your devices. +

+
+ +
+ {:else if onboardingStep === 2} + +
+
+

+ Add more devices, run bigger models +

+
+ + +
+ + + + + + + + + + + + + + + + + + + + + + + + {#if deviceAnimPhase >= 1} + + + + + + + + + + + + + MacBook Pro + + + 36 GB + + + {/if} + + + + + + + + + + + + + + + Mac Studio + + + 192 GB + + + + + {#if deviceAnimPhase >= 3} + + {/if} + + + {#if deviceAnimPhase >= 3} + + 228 GB combined + + {/if} + + + + + {#if deviceAnimPhase >= 3} + + + {/if} + + + {#if deviceAnimPhase >= 1} + + + + Qwen3 8B + + + 4 GB + + + {/if} + + + {#if deviceAnimPhase >= 1} + + = 3 + ? "rgba(255,215,0,0.06)" + : "rgba(0,0,0,0.03)"} + stroke={deviceAnimPhase >= 3 + ? "rgba(255,215,0,0.35)" + : "rgba(0,0,0,0.08)"} + stroke-width="1" + filter={deviceAnimPhase >= 3 + ? "url(#onb-gold-glow)" + : "none"} + style="transition: fill 500ms, stroke 500ms, filter 500ms;" + /> + {#if deviceAnimPhase < 3} + + + + + {:else} + + Qwen3 30B + + + 16 GB + + {/if} + + {/if} + + + {#if deviceAnimPhase >= 1} + + = 4 + ? "rgba(255,215,0,0.06)" + : "rgba(0,0,0,0.03)"} + stroke={deviceAnimPhase >= 4 + ? "rgba(255,215,0,0.35)" + : "rgba(0,0,0,0.08)"} + stroke-width="1" + filter={deviceAnimPhase >= 4 + ? "url(#onb-gold-glow)" + : "none"} + style="transition: fill 500ms, stroke 500ms, filter 500ms;" + /> + {#if deviceAnimPhase < 4} + + + + + {:else} + + Llama 72B + + + 36 GB + + {/if} + + {/if} + + + {#if deviceAnimPhase >= 1} + + = 4 + ? "rgba(255,215,0,0.08)" + : "rgba(0,0,0,0.03)"} + stroke={deviceAnimPhase >= 4 + ? "rgba(255,215,0,0.45)" + : "rgba(0,0,0,0.08)"} + stroke-width={deviceAnimPhase >= 4 ? "1.5" : "1"} + filter={deviceAnimPhase >= 4 + ? "url(#onb-gold-glow)" + : "none"} + style="transition: fill 700ms, stroke 700ms, filter 700ms, stroke-width 700ms;" + /> + {#if deviceAnimPhase < 4} + + + + + {:else} + + Llama 405B + + + 203 GB + + {/if} + + {/if} + + + {#if deviceAnimPhase >= 1} + + Models you can run + + {/if} + + +
+ + + {#if showContinueStep2} + -
+ {/if} - - +
+ {:else if onboardingStep === 3} + +
+
+

+ Your devices +

+

+ {nodeCount} device{nodeCount !== 1 ? "s" : ""} connected + {#if clusterTotalMemoryGB() > 0} + · {clusterTotalMemoryGB().toFixed(0)} GB total memory + {/if} +

+
+
+ +
+

+ Install exo on more devices on your network to combine their power — + they connect automatically. +

-
- {:else if !chatStarted} - -
- -
- + {:else if onboardingStep === 4} + +
+
+

+ Choose a model +

+

+ Pick a model to download and run locally. +

+
+ + {#if onboardingError} +
+ {onboardingError} +
+ {/if} + + {#if onboardingModels.length === 0} +
+
+ Loading models... +
+
+ {:else} +
+ {#each onboardingModels as model} + {@const sizeGB = getModelSizeGB(model)} + {@const fitsNow = hasEnoughMemory(model)} + {@const tags = modelTags()[model.id] || []} + + {/each} +
+ {/if} + + +
+ {:else if onboardingStep === 5} + +
+
+

+ Downloading +

+

+ {#if onboardingModelId} + {onboardingModelId} + {/if} +

+
+ + {#if onboardingDownloadProgress} +
+
+
+
+
+ {onboardingDownloadProgress.percentage.toFixed(1)}% + {formatBytes(onboardingDownloadProgress.downloadedBytes)} / + {formatBytes(onboardingDownloadProgress.totalBytes)} +
+
+ {formatSpeed(onboardingDownloadProgress.speed)} + ETA: {formatEta(onboardingDownloadProgress.etaMs)} +
+
+ {:else} +
+
+
+
+

+ Preparing download... +

+
+ {/if} + +

+ This may take a few minutes depending on your connection. +

+
+ {:else if onboardingStep === 6} + +
+
+

+ Loading into memory +

+

+ {#if onboardingModelId} + {onboardingModelId} + {/if} +

+
+ +
+
+
+ +

Almost ready...

+
+ {:else if onboardingStep === 7} + +
+
+ exo +
+ + + {#if onboardingModelId} +

+ {onboardingModelId.split("/").pop() ?? onboardingModelId} +

+ {/if} + + +
+ +
+ + +
+ {#each ["Write a poem about the ocean", "Explain quantum computing simply", "Help me debug my code", "Tell me a creative story"] as chip} + + {/each} +
+
+ {/if} +
+ + + {#if onboardingStep === 4} + m.id))} + canModelFit={(modelId) => { + const model = models.find((m) => m.id === modelId); + return model ? hasEnoughMemory(model) : false; + }} + getModelFitStatus={(modelId): ModelMemoryFitStatus => { + const model = models.find((m) => m.id === modelId); + return model ? getModelMemoryFitStatus(model) : "too_large"; + }} + onSelect={(modelId) => { + isModelPickerOpen = false; + onboardingLaunchModel(modelId); + }} + onClose={() => (isModelPickerOpen = false)} + onToggleFavorite={toggleFavorite} + onAddModel={addModelFromPicker} + onDeleteModel={deleteCustomModel} + totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} + usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} + {downloadsData} + topologyNodes={data?.nodes} + /> + {/if} + {:else} + + + + {#if !topologyOnlyEnabled} + + {/if} + + +
+ + {#if !topologyOnlyEnabled && sidebarVisible} + + {/if} + + {#if topologyOnlyEnabled} + +
+
- {#if tb5WithoutRdma && !tb5InfoDismissed}
0} class:top-4={tbBridgeCycles.length === 0} role="status" > -
- - - - - RDMA AVAILABLE - - -
- - - -
- {/if} - - - {#if isFilterActive()} - - {/if} -
- - -
-
- -
-
-
- - -
- -
- {:else} - -
- -
-
-
- -
-
- -
-
- -
+ +
- - - {#if minimized} -
+
+ {/if} + + {@render clusterWarnings()} + + + {#if tb5WithoutRdma && !tb5InfoDismissed} +
0} + class:top-4={tbBridgeCycles.length === 0} + role="status" + > +
+ + + + + RDMA AVAILABLE + + +
+ + + +
+ {/if} + + + {#if isFilterActive()} + + {/if} +
+ + +
+
+ {#if instanceCount === 0} + +
+

+ No model loaded yet. Select a model to get started. +

+
+ {/if} + - - {@render clusterWarningsCompact()}
- +
+
- + + - {/if} - - {/if} - + + {:else} + +
+ +
+
+
+ +
+
+ +
+
+ +
+
+
+ + + {#if minimized} + + {/if} +
+ {/if} + + {/if} - m.id))} - canModelFit={(modelId) => { - const model = models.find((m) => m.id === modelId); - return model ? hasEnoughMemory(model) : false; - }} - getModelFitStatus={(modelId): ModelMemoryFitStatus => { - const model = models.find((m) => m.id === modelId); - return model ? getModelMemoryFitStatus(model) : "too_large"; - }} - onSelect={handleModelPickerSelect} - onClose={() => (isModelPickerOpen = false)} - onToggleFavorite={toggleFavorite} - onAddModel={addModelFromPicker} - onDeleteModel={deleteCustomModel} - totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} - usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} - {downloadsData} - topologyNodes={data?.nodes} -/> +{#if !showOnboarding} + m.id))} + canModelFit={(modelId) => { + const model = models.find((m) => m.id === modelId); + return model ? hasEnoughMemory(model) : false; + }} + getModelFitStatus={(modelId): ModelMemoryFitStatus => { + const model = models.find((m) => m.id === modelId); + return model ? getModelMemoryFitStatus(model) : "too_large"; + }} + onSelect={handleModelPickerSelect} + onClose={() => (isModelPickerOpen = false)} + onToggleFavorite={toggleFavorite} + onAddModel={addModelFromPicker} + onDeleteModel={deleteCustomModel} + totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} + usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} + {downloadsData} + topologyNodes={data?.nodes} + /> +{/if} diff --git a/packaging/dmg/background.png b/packaging/dmg/background.png new file mode 100644 index 0000000000000000000000000000000000000000..5e56d8abd6e2f769097898c44a1331601d7e4cbc GIT binary patch literal 7010 zcmeHK`&U#|6#niEBQKR1R0NSgQ6@ul5?c!tmd(Oed zq4P&sO|SwmDriB#QecD%5Vcy2pqW*T{ULzkw4i|Zm+#Db_AoBgF=(u6Ns7^YAkxl# zYh6QFnOlmnM0t91evk5m?0M4dsCO#b;9E@J!?Ch+Kqy`(G z8JTnA_k>@@i-GSiCNvq3YiPE|Rt0Em{eSY>F9Ylg**#S*n|*hZ4$ya*NdaSS)0I_B zX*|S~j#6)9mVyMOoF~X1w9E5Lj|; z>I)|awa+r3Vk$$}+!7m)+f#536DdA5TrJXey&_v%h_gPbckP_jvqZ%QWQJ+rw=;v* zQeWQT;3VMHehR5HdaK1zi%R0LkoN!F>+xTIuCW!lN-3HZ2S~5*mizL@7<_La0d$#v zaGi348raO-J5p*CR}^a9=&3`-imLC3HI4^j(GlZQ$X6;LcczZp_5pe96nCxI_uUz9N@9r~D@7%eA$|k%n{* zE4z2}O6YK6T60XI`18%}{n{H&o}8xBLUov!PG5gUk~LL^YxYq!9Nwcd$dvEJ zV$&2ljP=!KPN=u_gQ@HOhUXl{F|8|U6I+9Cad?M*SN7$aAK-(ca4_!)A4@VS&+`dQOs zq5JvcQ%WGFXEzI6hm*o@D8G+Gvtk*?hkMbsAB8tTQSREVyW>WFyk_W literal 0 HcmV?d00001 diff --git a/packaging/dmg/create-dmg.sh b/packaging/dmg/create-dmg.sh new file mode 100755 index 00000000..cd68b2a7 --- /dev/null +++ b/packaging/dmg/create-dmg.sh @@ -0,0 +1,112 @@ +#!/usr/bin/env bash +# create-dmg.sh — Build a polished macOS DMG installer for EXO +# +# Usage: +# ./packaging/dmg/create-dmg.sh [volume-name] +# +# Example: +# ./packaging/dmg/create-dmg.sh output/EXO.app EXO-1.0.0.dmg "EXO" +# +# Creates a DMG with: +# - Custom background image with drag-to-Applications arrow +# - App icon on left, Applications alias on right +# - Proper window size and icon positioning +set -euo pipefail + +APP_PATH="${1:?Usage: create-dmg.sh [volume-name]}" +OUTPUT_DMG="${2:?Usage: create-dmg.sh [volume-name]}" +VOLUME_NAME="${3:-EXO}" + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +BACKGROUND_SCRIPT="${SCRIPT_DIR}/generate-background.py" +TEMP_DIR="$(mktemp -d)" +DMG_STAGING="${TEMP_DIR}/dmg-root" +TEMP_DMG="${TEMP_DIR}/temp.dmg" +BACKGROUND_PNG="${TEMP_DIR}/background.png" + +cleanup() { rm -rf "$TEMP_DIR"; } +trap cleanup EXIT + +echo "==> Creating DMG installer for ${VOLUME_NAME}" + +# ── Step 1: Generate background image ──────────────────────────────────────── +if command -v python3 &>/dev/null; then + python3 "$BACKGROUND_SCRIPT" "$BACKGROUND_PNG" + echo " Background image generated" +else + echo " Warning: python3 not found, skipping custom background" + BACKGROUND_PNG="" +fi + +# ── Step 2: Prepare staging directory ───────────────────────────────────────── +mkdir -p "$DMG_STAGING" +cp -R "$APP_PATH" "$DMG_STAGING/" +ln -s /Applications "$DMG_STAGING/Applications" + +# ── Step 3: Create writable DMG ────────────────────────────────────────────── +# Calculate required size (app size + 20MB headroom) +APP_SIZE_KB=$(du -sk "$APP_PATH" | cut -f1) +DMG_SIZE_KB=$((APP_SIZE_KB + 20480)) + +hdiutil create \ + -volname "$VOLUME_NAME" \ + -size "${DMG_SIZE_KB}k" \ + -fs HFS+ \ + -layout SPUD \ + "$TEMP_DMG" + +# ── Step 4: Mount and configure ────────────────────────────────────────────── +MOUNT_DIR=$(hdiutil attach "$TEMP_DMG" -readwrite -noverify | awk -F'\t' '/Apple_HFS/ {gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF); print $NF}') +echo " Mounted at: $MOUNT_DIR" + +# Copy contents +cp -R "$DMG_STAGING/"* "$MOUNT_DIR/" + +# Add background image +if [[ -n $BACKGROUND_PNG && -f $BACKGROUND_PNG ]]; then + mkdir -p "$MOUNT_DIR/.background" + cp "$BACKGROUND_PNG" "$MOUNT_DIR/.background/background.png" +fi + +# ── Step 5: Configure window appearance via AppleScript ────────────────────── +# Window: 800×400, app icon on left, Applications on right (matches Ollama layout) +# Background image is 1600×740 (2× retina for 800×400 logical window). +APP_NAME="$(basename "$APP_PATH")" + +osascript < DMG created: $OUTPUT_DMG" +echo " Size: $(du -h "$OUTPUT_DMG" | cut -f1)" diff --git a/packaging/dmg/generate-background.py b/packaging/dmg/generate-background.py new file mode 100644 index 00000000..ac3dc649 --- /dev/null +++ b/packaging/dmg/generate-background.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Generate the DMG background image with a centered drag-to-Applications arrow. + +The output is a 1600×740 retina PNG (2× for 800×400 logical window). +Icons are positioned at (200, 190) and (600, 190) in logical coordinates; +the arrow is drawn centered between them. + +Usage: + python3 generate-background.py [output.png] + +If no output path is given, overwrites the bundled background.png in-place. +""" + +from __future__ import annotations + +import math +import sys +from pathlib import Path + +from PIL import Image, ImageDraw + +# Retina dimensions (2× logical 800×400) +WIDTH = 1600 +HEIGHT = 740 + +# Icon positions in logical coords → retina coords +# App icon at (200, 190), Applications at (600, 190) +APP_X = 200 * 2 # 400 +APPS_X = 600 * 2 # 1200 +ICON_Y = 190 * 2 # 380 + +# Arrow drawn between icons, slightly above icon center +ARROW_START_X = APP_X + 160 # past the icon +ARROW_END_X = APPS_X - 160 # before the Applications icon +ARROW_Y = ICON_Y # same height as icons +ARROW_RISE = 120 # upward arc height + + +def draw_arrow(draw: ImageDraw.ImageDraw) -> None: + """Draw a hand-drawn-style curved arrow from app icon toward Applications.""" + color = (30, 30, 30) + line_width = 8 + + # Compute bezier curve points for a gentle upward arc + points: list[tuple[float, float]] = [] + steps = 80 + for i in range(steps + 1): + t = i / steps + # Quadratic bezier: start → control → end + cx = (ARROW_START_X + ARROW_END_X) / 2 + cy = ARROW_Y - ARROW_RISE + x = (1 - t) ** 2 * ARROW_START_X + 2 * (1 - t) * t * cx + t**2 * ARROW_END_X + y = (1 - t) ** 2 * ARROW_Y + 2 * (1 - t) * t * cy + t**2 * ARROW_Y + points.append((x, y)) + + # Draw the curve as connected line segments + for i in range(len(points) - 1): + draw.line([points[i], points[i + 1]], fill=color, width=line_width) + + # Arrowhead at the end + end_x, end_y = points[-1] + # Direction from second-to-last to last point + prev_x, prev_y = points[-3] + angle = math.atan2(end_y - prev_y, end_x - prev_x) + head_len = 36 + head_angle = math.radians(25) + + left_x = end_x - head_len * math.cos(angle - head_angle) + left_y = end_y - head_len * math.sin(angle - head_angle) + right_x = end_x - head_len * math.cos(angle + head_angle) + right_y = end_y - head_len * math.sin(angle + head_angle) + + draw.polygon( + [(end_x, end_y), (left_x, left_y), (right_x, right_y)], + fill=color, + ) + + +def generate_background(output_path: str) -> None: + """Generate a white DMG background with a centered arrow.""" + img = Image.new("RGBA", (WIDTH, HEIGHT), (255, 255, 255, 255)) + draw = ImageDraw.Draw(img) + draw_arrow(draw) + img.save(output_path, "PNG") + + +if __name__ == "__main__": + default_output = str(Path(__file__).parent / "background.png") + out = sys.argv[1] if len(sys.argv) >= 2 else default_output + generate_background(out) + print(f"Background image written to {out}") diff --git a/src/exo/utils/banner.py b/src/exo/utils/banner.py index ffdb5458..2742832e 100644 --- a/src/exo/utils/banner.py +++ b/src/exo/utils/banner.py @@ -1,8 +1,27 @@ +import logging +import os import sys +import webbrowser + +from exo.shared.constants import EXO_CONFIG_HOME + +logger = logging.getLogger(__name__) + +_FIRST_RUN_MARKER = EXO_CONFIG_HOME / ".dashboard_opened" + + +def _is_first_run() -> bool: + return not _FIRST_RUN_MARKER.exists() + + +def _mark_first_run_done() -> None: + _FIRST_RUN_MARKER.parent.mkdir(parents=True, exist_ok=True) + _FIRST_RUN_MARKER.touch() def print_startup_banner(port: int) -> None: dashboard_url = f"http://localhost:{port}" + first_run = _is_first_run() banner = f""" ╔═══════════════════════════════════════════════════════════════════════╗ ║ ║ @@ -30,3 +49,14 @@ def print_startup_banner(port: int) -> None: """ print(banner, file=sys.stderr) + + if first_run: + # Skip browser open when running inside the native macOS app — + # FirstLaunchPopout.swift handles the auto-open with a countdown. + if not os.environ.get("EXO_RUNTIME_DIR"): + try: + webbrowser.open(dashboard_url) + 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() From c8997217cfab22835d3bea407a0709d3f710318b Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Tue, 17 Feb 2026 17:56:26 +0000 Subject: [PATCH 10/45] Revert "feat: better onboarding UX for new users (#1479)" This reverts commit 490d2e46ba469834215f318a358b37deb5e0e9a5. --- .github/workflows/build-app.yml | 5 +- app/EXO/EXO/ContentView.swift | 476 ++- app/EXO/EXO/EXOApp.swift | 16 +- app/EXO/EXO/ExoProcessController.swift | 20 - app/EXO/EXO/Views/FirstLaunchPopout.swift | 192 - app/EXO/EXO/Views/SettingsView.swift | 478 --- .../EXO/Views/SettingsWindowController.swift | 47 - dashboard/src/app.css | 37 - dashboard/src/lib/components/ChatForm.svelte | 4 - .../src/lib/components/ChatMessages.svelte | 5 +- .../src/lib/components/ChatSidebar.svelte | 40 +- .../lib/components/ConnectionBanner.svelte | 20 - dashboard/src/lib/components/HeaderNav.svelte | 77 +- dashboard/src/lib/components/ModelCard.svelte | 26 - .../lib/components/ModelPickerModal.svelte | 67 +- .../src/lib/components/ToastContainer.svelte | 117 - dashboard/src/lib/stores/app.svelte.ts | 25 +- dashboard/src/lib/stores/toast.svelte.ts | 87 - dashboard/src/routes/+layout.svelte | 4 - dashboard/src/routes/+page.svelte | 3564 +++++------------ packaging/dmg/background.png | Bin 7010 -> 0 bytes packaging/dmg/create-dmg.sh | 112 - packaging/dmg/generate-background.py | 91 - src/exo/utils/banner.py | 30 - 24 files changed, 1578 insertions(+), 3962 deletions(-) delete mode 100644 app/EXO/EXO/Views/FirstLaunchPopout.swift delete mode 100644 app/EXO/EXO/Views/SettingsView.swift delete mode 100644 app/EXO/EXO/Views/SettingsWindowController.swift delete mode 100644 dashboard/src/lib/components/ConnectionBanner.svelte delete mode 100644 dashboard/src/lib/components/ToastContainer.svelte delete mode 100644 dashboard/src/lib/stores/toast.svelte.ts delete mode 100644 packaging/dmg/background.png delete mode 100755 packaging/dmg/create-dmg.sh delete mode 100644 packaging/dmg/generate-background.py diff --git a/.github/workflows/build-app.yml b/.github/workflows/build-app.yml index 861f4697..d20bfc12 100644 --- a/.github/workflows/build-app.yml +++ b/.github/workflows/build-app.yml @@ -303,8 +303,11 @@ jobs: SIGNING_IDENTITY=$(security find-identity -v -p codesigning "$BUILD_KEYCHAIN_PATH" | awk -F '"' '{print $2}') /usr/bin/codesign --deep --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" EXO.app + mkdir -p dmg-root + cp -R EXO.app dmg-root/ + ln -s /Applications dmg-root/Applications DMG_NAME="EXO-${RELEASE_VERSION}.dmg" - bash "$GITHUB_WORKSPACE/packaging/dmg/create-dmg.sh" EXO.app "$DMG_NAME" "EXO" + hdiutil create -volname "EXO" -srcfolder dmg-root -ov -format UDZO "$DMG_NAME" /usr/bin/codesign --force --timestamp --options runtime \ --sign "$SIGNING_IDENTITY" "$DMG_NAME" if [[ -n "$APPLE_NOTARIZATION_USERNAME" ]]; then diff --git a/app/EXO/EXO/ContentView.swift b/app/EXO/EXO/ContentView.swift index ca09fe2f..688fba5b 100644 --- a/app/EXO/EXO/ContentView.swift +++ b/app/EXO/EXO/ContentView.swift @@ -15,12 +15,18 @@ struct ContentView: View { @EnvironmentObject private var localNetworkChecker: LocalNetworkChecker @EnvironmentObject private var updater: SparkleUpdater @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService - @EnvironmentObject private var settingsWindowController: SettingsWindowController @State private var focusedNode: NodeViewModel? @State private var deletingInstanceIDs: Set = [] @State private var showAllNodes = false @State private var showAllInstances = false - @State private var baseURLCopied = false + @State private var showAdvanced = false + @State private var showDebugInfo = false + @State private var bugReportInFlight = false + @State private var bugReportMessage: String? + @State private var uninstallInProgress = false + @State private var pendingNamespace: String = "" + @State private var pendingHFToken: String = "" + @State private var pendingEnableImageModels = false var body: some View { VStack(alignment: .leading, spacing: 12) { @@ -252,79 +258,139 @@ struct ContentView: View { VStack(alignment: .leading, spacing: 0) { if controller.status != .stopped { dashboardButton - baseURLRow Divider() .padding(.vertical, 8) } else { Divider() .padding(.vertical, 4) } - HoverButton( - title: "Settings", - tint: .primary, - trailingSystemImage: "gear" - ) { - settingsWindowController.open( - controller: controller, - updater: updater, - networkStatusService: networkStatusService, - thunderboltBridgeService: thunderboltBridgeService, - stateService: stateService - ) - } - HoverButton( - title: "Check for Updates", - tint: .primary, - trailingSystemImage: "arrow.triangle.2.circlepath" - ) { - updater.checkForUpdates() - } - .padding(.bottom, 8) - HoverButton(title: "Quit", tint: .secondary) { + advancedSection + .padding(.bottom, 8) + controlButton(title: "Quit", tint: .secondary) { controller.stop() NSApplication.shared.terminate(nil) } } } - private var dashboardButton: some View { - HoverButton( - title: "Web Dashboard", - tint: .primary, - trailingSystemImage: "arrow.up.right" - ) { - guard let url = URL(string: "http://localhost:52415/") else { return } - NSWorkspace.shared.open(url) + private var advancedSection: some View { + VStack(alignment: .leading, spacing: 6) { + HStack { + Text("Advanced") + .font(.caption) + .foregroundColor(.secondary) + Spacer() + collapseButton(isExpanded: $showAdvanced) + } + .animation(nil, value: showAdvanced) + if showAdvanced { + VStack(alignment: .leading, spacing: 8) { + VStack(alignment: .leading, spacing: 4) { + Text("Cluster Namespace") + .font(.caption2) + .foregroundColor(.secondary) + HStack { + TextField("optional", text: $pendingNamespace) + .textFieldStyle(.roundedBorder) + .font(.caption2) + .onAppear { + pendingNamespace = controller.customNamespace + } + Button("Save & Restart") { + controller.customNamespace = pendingNamespace + if controller.status == .running || controller.status == .starting { + controller.restart() + } + } + .font(.caption2) + .disabled(pendingNamespace == controller.customNamespace) + } + } + VStack(alignment: .leading, spacing: 4) { + Text("HuggingFace Token") + .font(.caption2) + .foregroundColor(.secondary) + HStack { + SecureField("optional", text: $pendingHFToken) + .textFieldStyle(.roundedBorder) + .font(.caption2) + .onAppear { + pendingHFToken = controller.hfToken + } + Button("Save & Restart") { + controller.hfToken = pendingHFToken + if controller.status == .running || controller.status == .starting { + controller.restart() + } + } + .font(.caption2) + .disabled(pendingHFToken == controller.hfToken) + } + } + Divider() + HStack { + Toggle( + "Enable Image Models (experimental)", isOn: $pendingEnableImageModels + ) + .toggleStyle(.switch) + .font(.caption2) + .onAppear { + pendingEnableImageModels = controller.enableImageModels + } + + Spacer() + + Button("Save & Restart") { + controller.enableImageModels = pendingEnableImageModels + if controller.status == .running || controller.status == .starting { + controller.restart() + } + } + .font(.caption2) + .disabled(pendingEnableImageModels == controller.enableImageModels) + } + HoverButton(title: "Check for Updates", small: true) { + updater.checkForUpdates() + } + debugSection + HoverButton(title: "Uninstall", tint: .red, small: true) { + showUninstallConfirmationAlert() + } + .disabled(uninstallInProgress) + } + .transition(.opacity) + } } + .animation(.easeInOut(duration: 0.25), value: showAdvanced) } - private var baseURLRow: some View { - HStack(spacing: 6) { - Image(systemName: "link") - .imageScale(.small) - .foregroundColor(.secondary) - Text("localhost:52415/v1") - .font(.system(.caption, design: .monospaced)) - .foregroundColor(.primary) - Spacer() - Button { - NSPasteboard.general.clearContents() - NSPasteboard.general.setString("http://localhost:52415/v1", forType: .string) - baseURLCopied = true - DispatchQueue.main.asyncAfter(deadline: .now() + 2) { - baseURLCopied = false - } - } label: { - Image(systemName: baseURLCopied ? "checkmark" : "doc.on.doc") + private func controlButton(title: String, tint: Color = .primary, action: @escaping () -> Void) + -> some View + { + HoverButton(title: title, tint: tint, trailingSystemImage: nil, action: action) + } + + private var dashboardButton: some View { + Button { + guard let url = URL(string: "http://localhost:52415/") else { return } + NSWorkspace.shared.open(url) + } label: { + HStack { + Image(systemName: "arrow.up.right.square") .imageScale(.small) - .foregroundColor(baseURLCopied ? .green : .secondary) - .contentTransition(.symbolEffect(.replace)) + Text("Dashboard") + .fontWeight(.medium) + Spacer() } - .buttonStyle(.plain) - .help("Copy API base URL") + .padding(.vertical, 8) + .padding(.horizontal, 10) + .background( + RoundedRectangle(cornerRadius: 8, style: .continuous) + .fill(Color(red: 1.0, green: 0.87, blue: 0.0).opacity(0.2)) + ) } - .padding(.vertical, 4) - .padding(.horizontal, 8) + .buttonStyle(.plain) + .padding(.bottom, 4) } private func collapseButton(isExpanded: Binding) -> some View { @@ -379,6 +445,207 @@ struct ContentView: View { } } + private var thunderboltStatusText: String { + switch networkStatusService.status.thunderboltBridgeState { + case .some(.disabled): + return "Thunderbolt Bridge: Disabled" + case .some(.deleted): + return "Thunderbolt Bridge: Deleted" + case .some(.enabled): + return "Thunderbolt Bridge: Enabled" + case nil: + return "Thunderbolt Bridge: Unknown" + } + } + + private var thunderboltStatusColor: Color { + switch networkStatusService.status.thunderboltBridgeState { + case .some(.disabled), .some(.deleted): + return .green + case .some(.enabled): + return .red + case nil: + return .secondary + } + } + + /// Shows TB bridge status for all nodes from exo cluster state + private var clusterThunderboltBridgeView: some View { + let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:] + let localNodeId = stateService.localNodeId + let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] + + return VStack(alignment: .leading, spacing: 1) { + if bridgeStatuses.isEmpty { + Text("Cluster TB Bridge: No data") + .font(.caption2) + .foregroundColor(.secondary) + } else { + Text("Cluster TB Bridge Status:") + .font(.caption2) + .foregroundColor(.secondary) + ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in + if let status = bridgeStatuses[nodeId] { + let nodeName = + nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) + let isLocal = nodeId == localNodeId + let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" + let statusText = + !status.exists + ? "N/A" + : (status.enabled ? "Enabled" : "Disabled") + let color: Color = + !status.exists + ? .secondary + : (status.enabled ? .red : .green) + Text("\(prefix) \(statusText)") + .font(.caption2) + .foregroundColor(color) + } + } + } + } + } + + private var interfaceIpList: some View { + let statuses = networkStatusService.status.interfaceStatuses + return VStack(alignment: .leading, spacing: 1) { + Text("Interfaces (en0–en7):") + .font(.caption2) + .foregroundColor(.secondary) + if statuses.isEmpty { + Text(" Unknown") + .font(.caption2) + .foregroundColor(.secondary) + } else { + ForEach(statuses, id: \.interfaceName) { status in + let ipText = status.ipAddress ?? "No IP" + Text(" \(status.interfaceName): \(ipText)") + .font(.caption2) + .foregroundColor(status.ipAddress == nil ? .red : .green) + } + } + } + } + + private var debugSection: some View { + VStack(alignment: .leading, spacing: 4) { + HoverButton( + title: "Debug Info", + tint: .primary, + trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down", + small: true + ) { + showDebugInfo.toggle() + } + if showDebugInfo { + VStack(alignment: .leading, spacing: 4) { + Text("Version: \(buildTag)") + .font(.caption2) + .foregroundColor(.secondary) + Text("Commit: \(buildCommit)") + .font(.caption2) + .foregroundColor(.secondary) + Text(thunderboltStatusText) + .font(.caption2) + .foregroundColor(thunderboltStatusColor) + clusterThunderboltBridgeView + interfaceIpList + rdmaStatusView + sendBugReportButton + .padding(.top, 6) + } + .padding(.leading, 8) + .transition(.opacity) + } + } + .animation(.easeInOut(duration: 0.25), value: showDebugInfo) + } + + private var rdmaStatusView: some View { + let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:] + let localNodeId = stateService.localNodeId + let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] + let localDevices = networkStatusService.status.localRdmaDevices + let localPorts = networkStatusService.status.localRdmaActivePorts + + return VStack(alignment: .leading, spacing: 1) { + if rdmaStatuses.isEmpty { + Text("Cluster RDMA: No data") + .font(.caption2) + .foregroundColor(.secondary) + } else { + Text("Cluster RDMA Status:") + .font(.caption2) + .foregroundColor(.secondary) + ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in + if let status = rdmaStatuses[nodeId] { + let nodeName = + nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) + let isLocal = nodeId == localNodeId + let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" + let statusText = status.enabled ? "Enabled" : "Disabled" + let color: Color = status.enabled ? .green : .orange + Text("\(prefix) \(statusText)") + .font(.caption2) + .foregroundColor(color) + } + } + } + if !localDevices.isEmpty { + Text(" Local Devices: \(localDevices.joined(separator: ", "))") + .font(.caption2) + .foregroundColor(.secondary) + } + if !localPorts.isEmpty { + Text(" Local Active Ports:") + .font(.caption2) + .foregroundColor(.secondary) + ForEach(localPorts, id: \.device) { port in + Text(" \(port.device) port \(port.port): \(port.state)") + .font(.caption2) + .foregroundColor(.green) + } + } + } + } + + private var sendBugReportButton: some View { + VStack(alignment: .leading, spacing: 4) { + Button { + Task { + await sendBugReport() + } + } label: { + HStack { + if bugReportInFlight { + ProgressView() + .scaleEffect(0.6) + } + Text("Send Bug Report") + .font(.caption) + .fontWeight(.semibold) + Spacer() + } + .padding(.vertical, 6) + .padding(.horizontal, 8) + .background( + RoundedRectangle(cornerRadius: 6) + .fill(Color.accentColor.opacity(0.12)) + ) + } + .buttonStyle(.plain) + .disabled(bugReportInFlight) + + if let message = bugReportMessage { + Text(message) + .font(.caption2) + .foregroundColor(.secondary) + .fixedSize(horizontal: false, vertical: true) + } + } + } + private var processToggleBinding: Binding { Binding( get: { @@ -419,6 +686,101 @@ struct ContentView: View { ) } + private func sendBugReport() async { + bugReportInFlight = true + bugReportMessage = "Collecting logs..." + let service = BugReportService() + do { + let outcome = try await service.sendReport(isManual: true) + bugReportMessage = outcome.message + } catch { + bugReportMessage = error.localizedDescription + } + bugReportInFlight = false + } + + private func showUninstallConfirmationAlert() { + let alert = NSAlert() + alert.messageText = "Uninstall EXO" + alert.informativeText = """ + This will remove EXO and all its system components: + + • Network configuration daemon + • Launch at login registration + • EXO network location + + The app will be moved to Trash. + """ + alert.alertStyle = .warning + alert.addButton(withTitle: "Uninstall") + alert.addButton(withTitle: "Cancel") + + // Style the Uninstall button as destructive + if let uninstallButton = alert.buttons.first { + uninstallButton.hasDestructiveAction = true + } + + let response = alert.runModal() + if response == .alertFirstButtonReturn { + performUninstall() + } + } + + private func performUninstall() { + uninstallInProgress = true + + // Stop EXO process first + controller.cancelPendingLaunch() + controller.stop() + stateService.stopPolling() + + // Run the privileged uninstall on a background thread + // Using .utility QoS to avoid priority inversion with NSAppleScript's subprocess + DispatchQueue.global(qos: .utility).async { + do { + // Remove network setup daemon and components (requires admin privileges) + try NetworkSetupHelper.uninstall() + + DispatchQueue.main.async { + // Unregister from launch at login + LaunchAtLoginHelper.disable() + + // Move app to trash + self.moveAppToTrash() + + // Quit the app + DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { + NSApplication.shared.terminate(nil) + } + } + } catch { + DispatchQueue.main.async { + self.showErrorAlert(message: error.localizedDescription) + self.uninstallInProgress = false + } + } + } + } + + private func showErrorAlert(message: String) { + let alert = NSAlert() + alert.messageText = "Uninstall Failed" + alert.informativeText = message + alert.alertStyle = .critical + alert.addButton(withTitle: "OK") + alert.runModal() + } + + private func moveAppToTrash() { + guard let appURL = Bundle.main.bundleURL as URL? else { return } + do { + try FileManager.default.trashItem(at: appURL, resultingItemURL: nil) + } catch { + // If we can't trash the app, that's OK - user can do it manually + // The important system components have already been cleaned up + } + } + private var buildTag: String { Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown" } diff --git a/app/EXO/EXO/EXOApp.swift b/app/EXO/EXO/EXOApp.swift index bba26e24..3aff58c3 100644 --- a/app/EXO/EXO/EXOApp.swift +++ b/app/EXO/EXO/EXOApp.swift @@ -21,9 +21,7 @@ struct EXOApp: App { @StateObject private var localNetworkChecker: LocalNetworkChecker @StateObject private var updater: SparkleUpdater @StateObject private var thunderboltBridgeService: ThunderboltBridgeService - @StateObject private var settingsWindowController: SettingsWindowController private let terminationObserver: TerminationObserver - private let firstLaunchPopout = FirstLaunchPopout() private let ciContext = CIContext(options: nil) init() { @@ -45,13 +43,12 @@ struct EXOApp: App { _updater = StateObject(wrappedValue: updater) let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service) _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge) - _settingsWindowController = StateObject(wrappedValue: SettingsWindowController()) enableLaunchAtLoginIfNeeded() // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops) NetworkSetupHelper.promptAndInstallIfNeeded() // Check local network access periodically (warning disappears when user grants permission) localNetwork.startPeriodicChecking(interval: 10) - controller.scheduleLaunch(after: 5) + controller.scheduleLaunch(after: 15) service.startPolling() networkStatus.startPolling() } @@ -65,19 +62,8 @@ struct EXOApp: App { .environmentObject(localNetworkChecker) .environmentObject(updater) .environmentObject(thunderboltBridgeService) - .environmentObject(settingsWindowController) } label: { menuBarIcon - .onReceive(controller.$isFirstLaunchReady) { ready in - if ready { - DispatchQueue.main.asyncAfter(deadline: .now() + 1.0) { - self.firstLaunchPopout.onComplete = { [weak controller] in - controller?.markOnboardingCompleted() - } - self.firstLaunchPopout.show() - } - } - } } .menuBarExtraStyle(.window) } diff --git a/app/EXO/EXO/ExoProcessController.swift b/app/EXO/EXO/ExoProcessController.swift index 2350c118..7566674b 100644 --- a/app/EXO/EXO/ExoProcessController.swift +++ b/app/EXO/EXO/ExoProcessController.swift @@ -5,7 +5,6 @@ import Foundation private let customNamespaceKey = "EXOCustomNamespace" private let hfTokenKey = "EXOHFToken" private let enableImageModelsKey = "EXOEnableImageModels" -private let onboardingCompletedKey = "EXOOnboardingCompleted" @MainActor final class ExoProcessController: ObservableObject { @@ -61,9 +60,6 @@ final class ExoProcessController: ObservableObject { } } - /// Fires once when EXO transitions to `.running` for the very first time (fresh install). - @Published private(set) var isFirstLaunchReady = false - private var process: Process? private var runtimeDirectoryURL: URL? private var pendingLaunchTask: Task? @@ -117,11 +113,6 @@ final class ExoProcessController: ObservableObject { try child.run() process = child status = .running - - // Show welcome popout if onboarding was never completed - if !UserDefaults.standard.bool(forKey: onboardingCompletedKey) { - isFirstLaunchReady = true - } } catch { process = nil status = .failed(message: "Launch error") @@ -173,17 +164,6 @@ final class ExoProcessController: ObservableObject { launch() } - /// Mark onboarding as completed (user interacted with the welcome popout). - func markOnboardingCompleted() { - UserDefaults.standard.set(true, forKey: onboardingCompletedKey) - } - - /// Reset onboarding so the welcome popout appears on next launch. - func resetOnboarding() { - UserDefaults.standard.removeObject(forKey: onboardingCompletedKey) - isFirstLaunchReady = false - } - func scheduleLaunch(after seconds: TimeInterval) { cancelPendingLaunch() let start = max(1, Int(ceil(seconds))) diff --git a/app/EXO/EXO/Views/FirstLaunchPopout.swift b/app/EXO/EXO/Views/FirstLaunchPopout.swift deleted file mode 100644 index 1a10b3ac..00000000 --- a/app/EXO/EXO/Views/FirstLaunchPopout.swift +++ /dev/null @@ -1,192 +0,0 @@ -import AppKit -import SwiftUI - -/// A popover callout anchored to the menu bar icon on first launch, -/// pointing the user to the web dashboard with an arrow connecting to the icon. -@MainActor -final class FirstLaunchPopout { - private var popover: NSPopover? - private var countdownTask: Task? - private static let dashboardURL = "http://localhost:52415/" - - /// Called when the user completes onboarding (clicks Open Dashboard or dismisses). - var onComplete: (() -> Void)? - - func show() { - guard popover == nil else { return } - - // The status bar button may not exist yet on first launch; retry a few times. - showWithRetry(attemptsRemaining: 5) - } - - private func showWithRetry(attemptsRemaining: Int) { - guard attemptsRemaining > 0 else { return } - - guard let button = Self.findStatusItemButton() else { - DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { [weak self] in - self?.showWithRetry(attemptsRemaining: attemptsRemaining - 1) - } - return - } - - let pop = NSPopover() - pop.behavior = .applicationDefined - pop.animates = true - pop.contentSize = NSSize(width: 280, height: 120) - pop.contentViewController = NSHostingController( - rootView: WelcomeCalloutView( - countdownDuration: 30, - onDismiss: { [weak self] in - self?.onComplete?() - self?.dismiss() - }, - onOpen: { [weak self] in - self?.openDashboard() - self?.onComplete?() - self?.dismiss() - } - ) - ) - - self.popover = pop - pop.show(relativeTo: button.bounds, of: button, preferredEdge: .minY) - - // Auto-open dashboard after 30s then dismiss - countdownTask = Task { - try? await Task.sleep(nanoseconds: 30_000_000_000) - if !Task.isCancelled { - openDashboard() - onComplete?() - dismiss() - } - } - } - - func dismiss() { - countdownTask?.cancel() - countdownTask = nil - UserDefaults.standard.set(true, forKey: "EXOOnboardingCompleted") - guard let pop = popover else { return } - popover = nil - pop.performClose(nil) - } - - private func openDashboard() { - guard let url = URL(string: Self.dashboardURL) else { return } - NSWorkspace.shared.open(url) - } - - /// Finds the NSStatusBarButton created by SwiftUI's MenuBarExtra. - /// Walks the view hierarchy to find the actual button rather than the content view. - private static func findStatusItemButton() -> NSView? { - for window in NSApp.windows { - let className = NSStringFromClass(type(of: window)) - if className.contains("NSStatusBarWindow") { - // Try to find the actual status bar button in the view hierarchy - if let content = window.contentView { - if let button = findButton(in: content) { - return button - } - return content - } - } - } - return nil - } - - /// Recursively searches the view hierarchy for an NSStatusBarButton. - private static func findButton(in view: NSView) -> NSView? { - let className = NSStringFromClass(type(of: view)) - if className.contains("StatusBarButton") { - return view - } - for subview in view.subviews { - if let found = findButton(in: subview) { - return found - } - } - return nil - } -} - -/// Minimal welcome callout — friendly pointer, not a wall of text. -/// Rendered inside the NSPopover which provides its own chrome and arrow. -private struct WelcomeCalloutView: View { - let countdownDuration: Int - let onDismiss: () -> Void - let onOpen: () -> Void - @State private var countdown: Int - @State private var timerTask: Task? - - init(countdownDuration: Int, onDismiss: @escaping () -> Void, onOpen: @escaping () -> Void) { - self.countdownDuration = countdownDuration - self.onDismiss = onDismiss - self.onOpen = onOpen - self._countdown = State(initialValue: countdownDuration) - } - - var body: some View { - VStack(alignment: .leading, spacing: 10) { - HStack(alignment: .top) { - Text("Welcome to EXO!") - .font(.system(.headline, design: .rounded)) - .fontWeight(.semibold) - .foregroundColor(.primary) - Spacer() - Button { - onDismiss() - } label: { - Image(systemName: "xmark.circle.fill") - .font(.system(size: 14)) - .foregroundStyle(.tertiary) - } - .buttonStyle(.plain) - } - - Text("Run your first model here:") - .font(.system(.subheadline, design: .default)) - .foregroundColor(.secondary) - - HStack { - Button { - onOpen() - } label: { - Label("Open Dashboard", systemImage: "arrow.up.right.square") - .font(.system(.caption, design: .default)) - .fontWeight(.medium) - } - .buttonStyle(.borderedProminent) - .tint(.accentColor) - .controlSize(.small) - - Spacer() - - if countdown > 0 { - Text("Auto-opens in \(countdown)s") - .font(.system(.caption2, design: .default)) - .foregroundColor(.secondary.opacity(0.6)) - .monospacedDigit() - } - } - } - .padding(14) - .onAppear { - startCountdown() - } - .onDisappear { - timerTask?.cancel() - timerTask = nil - } - } - - private func startCountdown() { - timerTask = Task { - while countdown > 0 { - try? await Task.sleep(nanoseconds: 1_000_000_000) - if !Task.isCancelled { - countdown -= 1 - } - } - } - } -} diff --git a/app/EXO/EXO/Views/SettingsView.swift b/app/EXO/EXO/Views/SettingsView.swift deleted file mode 100644 index ed221c89..00000000 --- a/app/EXO/EXO/Views/SettingsView.swift +++ /dev/null @@ -1,478 +0,0 @@ -import AppKit -import SwiftUI - -/// Native macOS Settings window following Apple HIG. -/// Organized into General, Model, Advanced, and About sections. -struct SettingsView: View { - @EnvironmentObject private var controller: ExoProcessController - @EnvironmentObject private var updater: SparkleUpdater - @EnvironmentObject private var networkStatusService: NetworkStatusService - @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService - @EnvironmentObject private var stateService: ClusterStateService - - @State private var pendingNamespace: String = "" - @State private var pendingHFToken: String = "" - @State private var pendingEnableImageModels = false - @State private var needsRestart = false - @State private var bugReportInFlight = false - @State private var bugReportMessage: String? - @State private var uninstallInProgress = false - - var body: some View { - TabView { - generalTab - .tabItem { - Label("General", systemImage: "gear") - } - modelTab - .tabItem { - Label("Model", systemImage: "cube") - } - advancedTab - .tabItem { - Label("Advanced", systemImage: "wrench.and.screwdriver") - } - aboutTab - .tabItem { - Label("About", systemImage: "info.circle") - } - } - .frame(width: 450, height: 400) - .onAppear { - pendingNamespace = controller.customNamespace - pendingHFToken = controller.hfToken - pendingEnableImageModels = controller.enableImageModels - needsRestart = false - } - } - - // MARK: - General Tab - - private var generalTab: some View { - Form { - Section { - LabeledContent("Cluster Namespace") { - TextField("default", text: $pendingNamespace) - .textFieldStyle(.roundedBorder) - .frame(width: 200) - } - Text("Nodes with the same namespace form a cluster. Leave empty for default.") - .font(.caption) - .foregroundColor(.secondary) - } - - Section { - LabeledContent("HuggingFace Token") { - SecureField("optional", text: $pendingHFToken) - .textFieldStyle(.roundedBorder) - .frame(width: 200) - } - Text("Required for gated models. Get yours at huggingface.co/settings/tokens") - .font(.caption) - .foregroundColor(.secondary) - } - - Section { - HStack { - Spacer() - Button("Save & Restart") { - applyGeneralSettings() - } - .disabled(!hasGeneralChanges) - } - } - } - .formStyle(.grouped) - .padding() - } - - // MARK: - Model Tab - - private var modelTab: some View { - Form { - Section { - Toggle("Enable Image Models (experimental)", isOn: $pendingEnableImageModels) - Text("Allow text-to-image and image-to-image models in the model picker.") - .font(.caption) - .foregroundColor(.secondary) - } - - Section { - HStack { - Spacer() - Button("Save & Restart") { - applyModelSettings() - } - .disabled(!hasModelChanges) - } - } - } - .formStyle(.grouped) - .padding() - } - - // MARK: - Advanced Tab - - private var advancedTab: some View { - Form { - Section("Onboarding") { - HStack { - VStack(alignment: .leading) { - Text("Reset Onboarding") - Text("Opens the dashboard and resets the onboarding wizard.") - .font(.caption) - .foregroundColor(.secondary) - } - Spacer() - Button("Reset") { - guard let url = URL(string: "http://localhost:52415/?reset-onboarding") - else { return } - NSWorkspace.shared.open(url) - } - } - } - - Section("Debug Info") { - LabeledContent("Thunderbolt Bridge") { - Text(thunderboltStatusText) - .foregroundColor(thunderboltStatusColor) - } - - VStack(alignment: .leading, spacing: 2) { - clusterThunderboltBridgeView - } - - VStack(alignment: .leading, spacing: 2) { - interfaceIpList - } - - VStack(alignment: .leading, spacing: 2) { - rdmaStatusView - } - - sendBugReportButton - } - - Section("Danger Zone") { - Button(role: .destructive) { - showUninstallConfirmationAlert() - } label: { - HStack { - Text("Uninstall EXO") - Spacer() - Image(systemName: "trash") - .imageScale(.small) - } - } - .disabled(uninstallInProgress) - } - } - .formStyle(.grouped) - .padding() - } - - // MARK: - About Tab - - private var aboutTab: some View { - Form { - Section { - LabeledContent("Version") { - Text(buildTag) - .textSelection(.enabled) - } - LabeledContent("Commit") { - Text(buildCommit) - .font(.system(.body, design: .monospaced)) - .textSelection(.enabled) - } - } - - Section { - Button("Check for Updates") { - updater.checkForUpdates() - } - } - } - .formStyle(.grouped) - .padding() - } - - // MARK: - Debug Info Views (moved from ContentView) - - private var thunderboltStatusText: String { - switch networkStatusService.status.thunderboltBridgeState { - case .some(.disabled): - return "Disabled" - case .some(.deleted): - return "Deleted" - case .some(.enabled): - return "Enabled" - case nil: - return "Unknown" - } - } - - private var thunderboltStatusColor: Color { - switch networkStatusService.status.thunderboltBridgeState { - case .some(.disabled), .some(.deleted): - return .green - case .some(.enabled): - return .red - case nil: - return .secondary - } - } - - private var clusterThunderboltBridgeView: some View { - let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:] - let localNodeId = stateService.localNodeId - let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] - - return VStack(alignment: .leading, spacing: 1) { - if bridgeStatuses.isEmpty { - Text("Cluster TB Bridge: No data") - .font(.caption2) - .foregroundColor(.secondary) - } else { - Text("Cluster TB Bridge Status:") - .font(.caption2) - .foregroundColor(.secondary) - ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in - if let status = bridgeStatuses[nodeId] { - let nodeName = - nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) - let isLocal = nodeId == localNodeId - let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" - let statusText = - !status.exists - ? "N/A" - : (status.enabled ? "Enabled" : "Disabled") - let color: Color = - !status.exists - ? .secondary - : (status.enabled ? .red : .green) - Text("\(prefix) \(statusText)") - .font(.caption2) - .foregroundColor(color) - } - } - } - } - } - - private var interfaceIpList: some View { - let statuses = networkStatusService.status.interfaceStatuses - return VStack(alignment: .leading, spacing: 1) { - Text("Interfaces (en0–en7):") - .font(.caption2) - .foregroundColor(.secondary) - if statuses.isEmpty { - Text(" Unknown") - .font(.caption2) - .foregroundColor(.secondary) - } else { - ForEach(statuses, id: \.interfaceName) { status in - let ipText = status.ipAddress ?? "No IP" - Text(" \(status.interfaceName): \(ipText)") - .font(.caption2) - .foregroundColor(status.ipAddress == nil ? .red : .green) - } - } - } - } - - private var rdmaStatusView: some View { - let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:] - let localNodeId = stateService.localNodeId - let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:] - let localDevices = networkStatusService.status.localRdmaDevices - let localPorts = networkStatusService.status.localRdmaActivePorts - - return VStack(alignment: .leading, spacing: 1) { - if rdmaStatuses.isEmpty { - Text("Cluster RDMA: No data") - .font(.caption2) - .foregroundColor(.secondary) - } else { - Text("Cluster RDMA Status:") - .font(.caption2) - .foregroundColor(.secondary) - ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in - if let status = rdmaStatuses[nodeId] { - let nodeName = - nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8)) - let isLocal = nodeId == localNodeId - let prefix = isLocal ? " \(nodeName) (local):" : " \(nodeName):" - let statusText = status.enabled ? "Enabled" : "Disabled" - let color: Color = status.enabled ? .green : .orange - Text("\(prefix) \(statusText)") - .font(.caption2) - .foregroundColor(color) - } - } - } - if !localDevices.isEmpty { - Text(" Local Devices: \(localDevices.joined(separator: ", "))") - .font(.caption2) - .foregroundColor(.secondary) - } - if !localPorts.isEmpty { - Text(" Local Active Ports:") - .font(.caption2) - .foregroundColor(.secondary) - ForEach(localPorts, id: \.device) { port in - Text(" \(port.device) port \(port.port): \(port.state)") - .font(.caption2) - .foregroundColor(.green) - } - } - } - } - - private var sendBugReportButton: some View { - VStack(alignment: .leading, spacing: 4) { - Button { - Task { - await sendBugReport() - } - } label: { - HStack { - if bugReportInFlight { - ProgressView() - .scaleEffect(0.6) - } - Text("Send Bug Report") - .font(.caption) - .fontWeight(.semibold) - Spacer() - } - } - .disabled(bugReportInFlight) - - if let message = bugReportMessage { - Text(message) - .font(.caption2) - .foregroundColor(.secondary) - .fixedSize(horizontal: false, vertical: true) - } - } - } - - // MARK: - Actions - - private func sendBugReport() async { - bugReportInFlight = true - bugReportMessage = "Collecting logs..." - let service = BugReportService() - do { - let outcome = try await service.sendReport(isManual: true) - bugReportMessage = outcome.message - } catch { - bugReportMessage = error.localizedDescription - } - bugReportInFlight = false - } - - private func showUninstallConfirmationAlert() { - let alert = NSAlert() - alert.messageText = "Uninstall EXO" - alert.informativeText = """ - This will remove EXO and all its system components: - - • Network configuration daemon - • Launch at login registration - • EXO network location - - The app will be moved to Trash. - """ - alert.alertStyle = .warning - alert.addButton(withTitle: "Uninstall") - alert.addButton(withTitle: "Cancel") - - if let uninstallButton = alert.buttons.first { - uninstallButton.hasDestructiveAction = true - } - - let response = alert.runModal() - if response == .alertFirstButtonReturn { - performUninstall() - } - } - - private func performUninstall() { - uninstallInProgress = true - - controller.cancelPendingLaunch() - controller.stop() - stateService.stopPolling() - - DispatchQueue.global(qos: .utility).async { - do { - try NetworkSetupHelper.uninstall() - - DispatchQueue.main.async { - LaunchAtLoginHelper.disable() - self.moveAppToTrash() - - DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { - NSApplication.shared.terminate(nil) - } - } - } catch { - DispatchQueue.main.async { - let errorAlert = NSAlert() - errorAlert.messageText = "Uninstall Failed" - errorAlert.informativeText = error.localizedDescription - errorAlert.alertStyle = .critical - errorAlert.addButton(withTitle: "OK") - errorAlert.runModal() - self.uninstallInProgress = false - } - } - } - } - - private func moveAppToTrash() { - guard let appURL = Bundle.main.bundleURL as URL? else { return } - do { - try FileManager.default.trashItem(at: appURL, resultingItemURL: nil) - } catch { - // If we can't trash the app, that's OK - user can do it manually - } - } - - // MARK: - Helpers - - private var hasGeneralChanges: Bool { - pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken - } - - private var hasModelChanges: Bool { - pendingEnableImageModels != controller.enableImageModels - } - - private func applyGeneralSettings() { - controller.customNamespace = pendingNamespace - controller.hfToken = pendingHFToken - restartIfRunning() - } - - private func applyModelSettings() { - controller.enableImageModels = pendingEnableImageModels - restartIfRunning() - } - - private func restartIfRunning() { - if controller.status == .running || controller.status == .starting { - controller.restart() - } - } - - private var buildTag: String { - Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown" - } - - private var buildCommit: String { - Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown" - } -} diff --git a/app/EXO/EXO/Views/SettingsWindowController.swift b/app/EXO/EXO/Views/SettingsWindowController.swift deleted file mode 100644 index 98517f92..00000000 --- a/app/EXO/EXO/Views/SettingsWindowController.swift +++ /dev/null @@ -1,47 +0,0 @@ -import AppKit -import SwiftUI - -/// Manages a standalone native macOS Settings window. -/// Ensures only one instance exists and brings it to front on repeated opens. -@MainActor -final class SettingsWindowController: ObservableObject { - private var window: NSWindow? - - func open( - controller: ExoProcessController, - updater: SparkleUpdater, - networkStatusService: NetworkStatusService, - thunderboltBridgeService: ThunderboltBridgeService, - stateService: ClusterStateService - ) { - if let existing = window, existing.isVisible { - existing.makeKeyAndOrderFront(nil) - NSApp.activate(ignoringOtherApps: true) - return - } - - let settingsView = SettingsView() - .environmentObject(controller) - .environmentObject(updater) - .environmentObject(networkStatusService) - .environmentObject(thunderboltBridgeService) - .environmentObject(stateService) - - let hostingView = NSHostingView(rootView: settingsView) - - let newWindow = NSWindow( - contentRect: NSRect(x: 0, y: 0, width: 450, height: 400), - styleMask: [.titled, .closable], - backing: .buffered, - defer: false - ) - newWindow.title = "EXO Settings" - newWindow.contentView = hostingView - newWindow.center() - newWindow.isReleasedWhenClosed = false - newWindow.makeKeyAndOrderFront(nil) - NSApp.activate(ignoringOtherApps: true) - - window = newWindow - } -} diff --git a/dashboard/src/app.css b/dashboard/src/app.css index 3e951163..fc532578 100644 --- a/dashboard/src/app.css +++ b/dashboard/src/app.css @@ -202,15 +202,6 @@ filter: drop-shadow(0 0 3px oklch(0.85 0.18 85 / 0.5)); } -/* Onboarding step 2: connection line between devices */ -.onboarding-connection-line { - stroke: oklch(0.85 0.18 85 / 0.5); - stroke-width: 1.5px; - stroke-dasharray: 6, 6; - animation: flowAnimation 1s linear infinite; - filter: drop-shadow(0 0 4px oklch(0.85 0.18 85 / 0.4)); -} - .graph-link-active { stroke: oklch(0.85 0.18 85 / 0.8); stroke-width: 2px; @@ -329,31 +320,3 @@ input:focus, textarea:focus { transform: translate(400px, 400px); } } - -/* Respect reduced motion preference */ -@media (prefers-reduced-motion: reduce) { - .shooting-star, - .shooting-star::before { - animation: none !important; - opacity: 0 !important; - } - .graph-link { - animation: none; - } - .status-pulse { - animation: none; - } - .cursor-blink { - animation: none; - } - .onboarding-connection-line { - animation: none; - } - *, - *::before, - *::after { - transition-duration: 0.01ms !important; - animation-duration: 0.01ms !important; - animation-iteration-count: 1 !important; - } -} diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte index 1bc554e8..2eb7aa40 100644 --- a/dashboard/src/lib/components/ChatForm.svelte +++ b/dashboard/src/lib/components/ChatForm.svelte @@ -28,7 +28,6 @@ showModelSelector?: boolean; modelTasks?: Record; modelCapabilities?: Record; - onSend?: () => void; } let { @@ -39,7 +38,6 @@ showModelSelector = false, modelTasks = {}, modelCapabilities = {}, - onSend, }: Props = $props(); let message = $state(""); @@ -307,8 +305,6 @@ ); } - onSend?.(); - // Refocus the textarea after sending setTimeout(() => textareaRef?.focus(), 10); } diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte index 7c076459..ba5322a7 100644 --- a/dashboard/src/lib/components/ChatMessages.svelte +++ b/dashboard/src/lib/components/ChatMessages.svelte @@ -802,8 +802,8 @@ > AWAITING INPUT

-

- Type a message below · Shift+Enter for newline +

+ ENTER A QUERY TO BEGIN

{/if} @@ -818,7 +818,6 @@ onclick={scrollToBottom} class="sticky bottom-4 left-1/2 -translate-x-1/2 w-10 h-10 rounded-full bg-exo-dark-gray/90 border border-exo-medium-gray/50 flex items-center justify-center text-exo-light-gray hover:text-exo-yellow hover:border-exo-yellow/50 transition-all shadow-lg cursor-pointer z-10" title="Scroll to bottom" - aria-label="Scroll to bottom of messages" >
{searchQuery ? "SEARCH RESULTS" : "CONVERSATIONS"} @@ -372,37 +372,39 @@ onkeydown={(e) => e.key === "Enter" && handleSelectConversation(conversation.id)} - class="group w-full flex items-center justify-between p-2.5 rounded-lg mb-1 transition-all text-left cursor-pointer + class="group w-full flex items-center justify-between p-2 rounded mb-1 transition-all text-left cursor-pointer {activeId === conversation.id - ? 'bg-exo-yellow/5 border border-exo-yellow/30' - : 'hover:bg-white/[0.03] hover:border-white/10 border border-transparent'}" + ? 'bg-transparent border border-exo-yellow/30' + : 'hover:border-exo-yellow/20 border border-transparent'}" >
{conversation.name}
-
+
{formatDate(conversation.updatedAt)}
-
+
{info.modelLabel}
+
+ Strategy: {info.strategyLabel} +
{#if stats} -
- {#if stats.ttftMs}TTFT - {stats.ttftMs.toFixed(0)}ms{/if}{#if stats.ttftMs && stats.tps}·{/if}{#if stats.tps}{stats.tps.toFixed(1)} - tok/s{/if} +
+ {#if stats.ttftMs}TTFT + {stats.ttftMs.toFixed( + 0, + )}ms{/if}{#if stats.ttftMs && stats.tps}•{/if}{#if stats.tps}{stats.tps.toFixed(1)} + tok/s{/if}
{/if}
diff --git a/dashboard/src/lib/components/ConnectionBanner.svelte b/dashboard/src/lib/components/ConnectionBanner.svelte deleted file mode 100644 index 3339cf24..00000000 --- a/dashboard/src/lib/components/ConnectionBanner.svelte +++ /dev/null @@ -1,20 +0,0 @@ - - -{#if !connected} - -{/if} diff --git a/dashboard/src/lib/components/HeaderNav.svelte b/dashboard/src/lib/components/HeaderNav.svelte index d78bc809..5bdcf1ba 100644 --- a/dashboard/src/lib/components/HeaderNav.svelte +++ b/dashboard/src/lib/components/HeaderNav.svelte @@ -6,10 +6,6 @@ export let showSidebarToggle = false; export let sidebarVisible = true; export let onToggleSidebar: (() => void) | null = null; - export let downloadProgress: { - count: number; - percentage: number; - } | null = null; function handleHome(): void { if (onHome) { @@ -39,15 +35,11 @@ onclick={handleToggleSidebar} class="p-2 rounded border border-exo-medium-gray/40 hover:border-exo-yellow/50 transition-colors cursor-pointer" title={sidebarVisible ? "Hide sidebar" : "Show sidebar"} - aria-label={sidebarVisible - ? "Hide conversation sidebar" - : "Show conversation sidebar"} - aria-pressed={sidebarVisible} > -
- - - - -
- {downloadProgress.count} -
-
- {:else} - - - - - - {/if} + + + + + Downloads - +
diff --git a/dashboard/src/lib/components/ModelCard.svelte b/dashboard/src/lib/components/ModelCard.svelte index b432b7a8..561c325b 100644 --- a/dashboard/src/lib/components/ModelCard.svelte +++ b/dashboard/src/lib/components/ModelCard.svelte @@ -567,17 +567,11 @@
{sharding} {runtime === "MlxRing" ? "MLX Ring" @@ -587,26 +581,6 @@
- - {#if isDownloading && progress} -
-
- Downloading - {percentage.toFixed(1)}% · {formatSpeed(progress.speed)} - · {formatEta(progress.etaMs)} -
-
-
-
-
- {/if} - {#if placementPreview().nodes.length > 0} {@const preview = placementPreview()} diff --git a/dashboard/src/lib/components/ModelPickerModal.svelte b/dashboard/src/lib/components/ModelPickerModal.svelte index cf21c727..84a93ee7 100644 --- a/dashboard/src/lib/components/ModelPickerModal.svelte +++ b/dashboard/src/lib/components/ModelPickerModal.svelte @@ -512,18 +512,6 @@ ); }); - // Split filtered groups into recommended (fits_now) and others for visual separation - const recommendedGroups = $derived( - filteredGroups.filter((g) => - g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"), - ), - ); - const otherGroups = $derived( - filteredGroups.filter( - (g) => !g.variants.some((v) => getModelFitStatus(v.id) === "fits_now"), - ), - ); - function toggleGroupExpanded(groupId: string) { const next = new Set(expandedGroups); if (next.has(groupId)) { @@ -852,60 +840,7 @@ {/if}
{:else} - - {#if recommendedGroups.length > 0 && otherGroups.length > 0 && !searchQuery.trim()} -
- - - - Recommended for your cluster - — fits in available memory -
- {/if} - {#each recommendedGroups as group} - toggleGroupExpanded(group.id)} - onSelectModel={handleSelect} - {onToggleFavorite} - onShowInfo={(g) => (infoGroup = g)} - downloadStatusMap={getVariantDownloadMap(group)} - /> - {/each} - - {#if otherGroups.length > 0 && recommendedGroups.length > 0 && !searchQuery.trim()} -
- Other models -
- {/if} - {#each otherGroups as group} + {#each filteredGroups as group} - import { toasts, dismissToast, type Toast } from "$lib/stores/toast.svelte"; - import { fly, fade } from "svelte/transition"; - import { flip } from "svelte/animate"; - - const items = $derived(toasts()); - - const typeStyles: Record< - Toast["type"], - { border: string; icon: string; iconColor: string } - > = { - success: { - border: "border-l-green-500", - icon: "M9 12.75L11.25 15 15 9.75M21 12a9 9 0 11-18 0 9 9 0 0118 0z", - iconColor: "text-green-400", - }, - error: { - border: "border-l-red-500", - icon: "M12 9v3.75m9-.75a9 9 0 11-18 0 9 9 0 0118 0zm-9 3.75h.008v.008H12v-.008z", - iconColor: "text-red-400", - }, - warning: { - border: "border-l-yellow-500", - icon: "M12 9v3.75m-9.303 3.376c-.866 1.5.217 3.374 1.948 3.374h14.71c1.73 0 2.813-1.874 1.948-3.374L13.949 3.378c-.866-1.5-3.032-1.5-3.898 0L2.697 16.126z", - iconColor: "text-yellow-400", - }, - info: { - border: "border-l-blue-500", - icon: "M11.25 11.25l.041-.02a.75.75 0 011.063.852l-.708 2.836a.75.75 0 001.063.853l.041-.021M21 12a9 9 0 11-18 0 9 9 0 0118 0zm-9-3.75h.008v.008H12V8.25z", - iconColor: "text-blue-400", - }, - }; - - -{#if items.length > 0} -
- {#each items as toast (toast.id)} - {@const style = typeStyles[toast.type]} - - {/each} -
-{/if} - - diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 5a5afc78..ebb2d0df 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -587,12 +587,6 @@ class AppStore { // Image editing state editingImage = $state(null); - /** True when the backend is reachable. */ - isConnected = $state(true); - /** Number of consecutive fetch failures. */ - private consecutiveFailures = 0; - private static readonly CONNECTION_LOST_THRESHOLD = 3; - private fetchInterval: ReturnType | null = null; private previewsInterval: ReturnType | null = null; private lastConversationPersistTs = 0; @@ -1296,19 +1290,7 @@ class AppStore { // Thunderbolt bridge status per node this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {}; this.lastUpdate = Date.now(); - // Connection recovered - if (!this.isConnected) { - this.isConnected = true; - } - this.consecutiveFailures = 0; } catch (error) { - this.consecutiveFailures++; - if ( - this.consecutiveFailures >= AppStore.CONNECTION_LOST_THRESHOLD && - this.isConnected - ) { - this.isConnected = false; - } console.error("Error fetching state:", error); } } @@ -1835,7 +1817,7 @@ class AppStore { assistantMessage.id, (msg) => { msg.content = - "No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically."; + "Error: No model available. Please launch an instance first."; }, ); this.syncActiveMessagesIfNeeded(targetConversationId); @@ -2273,7 +2255,7 @@ class AppStore { const modelToUse = this.getModelForRequest(); if (!modelToUse) { throw new Error( - "No model is loaded yet. Select a model from the sidebar to get started — it will download and load automatically.", + "No model selected and no running instances available. Please launch an instance first.", ); } @@ -3162,9 +3144,6 @@ export const setChatSidebarVisible = (visible: boolean) => appStore.setChatSidebarVisible(visible); export const refreshState = () => appStore.fetchState(); -// Connection status -export const isConnected = () => appStore.isConnected; - // Node identities (for OS version mismatch detection) export const nodeIdentities = () => appStore.nodeIdentities; diff --git a/dashboard/src/lib/stores/toast.svelte.ts b/dashboard/src/lib/stores/toast.svelte.ts deleted file mode 100644 index e9e62019..00000000 --- a/dashboard/src/lib/stores/toast.svelte.ts +++ /dev/null @@ -1,87 +0,0 @@ -/** - * Toast notification store - Global notification system for the EXO dashboard. - * - * Usage: - * import { addToast, dismissToast, toasts } from "$lib/stores/toast.svelte"; - * addToast({ type: "success", message: "Model launched" }); - * addToast({ type: "error", message: "Connection lost", persistent: true }); - */ - -type ToastType = "success" | "error" | "warning" | "info"; - -export interface Toast { - id: string; - type: ToastType; - message: string; - /** Auto-dismiss after this many ms. 0 = persistent (must be dismissed manually). */ - duration: number; - createdAt: number; -} - -interface ToastInput { - type: ToastType; - message: string; - /** If true, toast stays until manually dismissed. Default: false. */ - persistent?: boolean; - /** Auto-dismiss duration in ms. Default: 4000 for success/info, 6000 for error/warning. */ - duration?: number; -} - -const DEFAULT_DURATIONS: Record = { - success: 4000, - info: 4000, - warning: 6000, - error: 6000, -}; - -let toastList = $state([]); -const timers = new Map>(); - -function generateId(): string { - return `toast-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; -} - -export function addToast(input: ToastInput): string { - const id = generateId(); - const duration = input.persistent - ? 0 - : (input.duration ?? DEFAULT_DURATIONS[input.type]); - - const toast: Toast = { - id, - type: input.type, - message: input.message, - duration, - createdAt: Date.now(), - }; - - toastList = [...toastList, toast]; - - if (duration > 0) { - const timer = setTimeout(() => dismissToast(id), duration); - timers.set(id, timer); - } - - return id; -} - -export function dismissToast(id: string): void { - const timer = timers.get(id); - if (timer) { - clearTimeout(timer); - timers.delete(id); - } - toastList = toastList.filter((t) => t.id !== id); -} - -/** Dismiss all toasts matching a message (useful for dedup). */ -export function dismissByMessage(message: string): void { - const matching = toastList.filter((t) => t.message === message); - for (const t of matching) { - dismissToast(t.id); - } -} - -export function toasts(): Toast[] { - return toastList; -} diff --git a/dashboard/src/routes/+layout.svelte b/dashboard/src/routes/+layout.svelte index 295c9dc0..d249e9cf 100644 --- a/dashboard/src/routes/+layout.svelte +++ b/dashboard/src/routes/+layout.svelte @@ -1,7 +1,5 @@ @@ -12,7 +10,5 @@
- {@render children?.()} -
diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 20f196ee..5f7dcf04 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -36,7 +36,6 @@ createConversation, setSelectedChatModel, selectedChatModel, - sendMessage, debugMode, toggleDebugMode, topologyOnlyMode, @@ -53,11 +52,9 @@ type PlacementPreview, type MetaInstanceData, } from "$lib/stores/app.svelte"; - import { addToast } from "$lib/stores/toast.svelte"; import HeaderNav from "$lib/components/HeaderNav.svelte"; - import { fade, fly, slide } from "svelte/transition"; - import { tweened } from "svelte/motion"; - import { cubicInOut, cubicOut } from "svelte/easing"; + import { fade, fly } from "svelte/transition"; + import { cubicInOut } from "svelte/easing"; import { onMount } from "svelte"; const chatStarted = $derived(hasStartedChat()); @@ -143,26 +140,6 @@ const rdmaCtlData = $derived(nodeRdmaCtl()); const nodeFilter = $derived(previewNodeFilter()); - // Aggregate active download progress across all instances for header indicator - const activeDownloadSummary = $derived.by(() => { - let totalBytes = 0; - let downloadedBytes = 0; - let count = 0; - for (const [id, inst] of Object.entries(instanceData)) { - const status = getInstanceDownloadStatus(id, inst); - if (status.isDownloading && status.progress) { - count++; - totalBytes += status.progress.totalBytes || 0; - downloadedBytes += status.progress.downloadedBytes || 0; - } - } - if (count === 0) return null; - return { - count, - percentage: totalBytes > 0 ? (downloadedBytes / totalBytes) * 100 : 0, - }; - }); - // Detect macOS version mismatches across cluster nodes const macosVersionMismatch = $derived.by(() => { if (!identitiesData) return null; @@ -253,171 +230,6 @@ let mounted = $state(false); - // ── Onboarding wizard state ── - const ONBOARDING_COMPLETE_KEY = "exo-onboarding-complete"; - let onboardingStep = $state(0); // 0 = not in onboarding, 1-7 = wizard steps - let onboardingModelId = $state(null); // model selected during onboarding - const showOnboarding = $derived(onboardingStep > 0); - - // ── Step 2 animation state: "Add more devices, run bigger models" ── - let deviceAnimPhase = $state(0); // 0=waiting, 1=macbook, 2=studio joins, 3=connection+mid unlock, 4=big unlock - let showContinueStep2 = $state(false); - const studioX = tweened(540, { duration: 700, easing: cubicOut }); - const studioOpacity = tweened(0, { duration: 700, easing: cubicOut }); - - $effect(() => { - if (onboardingStep === 2) { - deviceAnimPhase = 0; - showContinueStep2 = false; - studioX.set(540, { duration: 0 }); - studioOpacity.set(0, { duration: 0 }); - - const t1 = setTimeout(() => { - deviceAnimPhase = 1; - }, 100); - const t2 = setTimeout(() => { - deviceAnimPhase = 2; - studioX.set(340); - studioOpacity.set(1); - }, 900); - const t3 = setTimeout(() => { - deviceAnimPhase = 3; - }, 1700); - const t4 = setTimeout(() => { - deviceAnimPhase = 4; - }, 2500); - const t5 = setTimeout(() => { - showContinueStep2 = true; - }, 3500); - - return () => { - clearTimeout(t1); - clearTimeout(t2); - clearTimeout(t3); - clearTimeout(t4); - clearTimeout(t5); - }; - } - }); - - // Recommended models for onboarding (sorted by fit, then size desc, limited to 6) - const onboardingModels = $derived.by(() => { - if (models.length === 0) return []; - return [...models] - .filter((m) => getModelMemoryFitStatus(m) !== "too_large") - .sort((a, b) => { - const aFit = hasEnoughMemory(a) ? 0 : 1; - const bFit = hasEnoughMemory(b) ? 0 : 1; - if (aFit !== bFit) return aFit - bFit; - return getModelSizeGB(b) - getModelSizeGB(a); - }) - .slice(0, 6); - }); - - // Track onboarding instance status for auto-advancing steps. - // Handles cached models: if no download is needed, skip step 5 entirely. - $effect(() => { - if (onboardingStep === 5 && instanceCount > 0) { - let anyDownloading = false; - let anyReady = false; - for (const [id, inst] of Object.entries(instanceData)) { - const status = getInstanceDownloadStatus(id, inst); - if (status.isDownloading) { - anyDownloading = true; - } - if ( - status.statusText === "READY" || - status.statusText === "LOADED" || - status.statusText === "RUNNING" - ) { - anyReady = true; - } - } - // Model already cached & ready — skip download AND loading steps - if (anyReady) { - onboardingStep = 7; - } else if (!anyDownloading) { - // Download finished (or was never needed) but not ready yet - onboardingStep = 6; - } - } - }); - - $effect(() => { - if (onboardingStep === 6 && instanceCount > 0) { - for (const [id, inst] of Object.entries(instanceData)) { - const status = getInstanceDownloadStatus(id, inst); - if ( - status.statusText === "READY" || - status.statusText === "LOADED" || - status.statusText === "RUNNING" - ) { - onboardingStep = 7; - break; - } - } - } - }); - - function completeOnboarding() { - onboardingStep = 0; - try { - localStorage.setItem(ONBOARDING_COMPLETE_KEY, "true"); - } catch { - // ignore - } - } - - let onboardingError = $state(null); - - async function onboardingLaunchModel(modelId: string) { - onboardingModelId = modelId; - onboardingError = null; - selectPreviewModel(modelId); - onboardingStep = 5; - // Launch via API - try { - const placementResponse = await fetch( - `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=1`, - ); - if (!placementResponse.ok) { - const errorText = await placementResponse.text(); - onboardingError = `Could not place model: ${errorText}`; - onboardingStep = 4; - return; - } - const placementData = await placementResponse.json(); - const response = await fetch("/instance", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ instance: placementData }), - }); - if (!response.ok) { - const errorText = await response.text(); - onboardingError = `Failed to launch: ${errorText}`; - onboardingStep = 4; - return; - } - setSelectedChatModel(modelId); - recordRecentLaunch(modelId); - } catch (error) { - onboardingError = `Network error: ${error}`; - onboardingStep = 4; - } - } - - // Helper to get onboarding download progress - const onboardingDownloadProgress = $derived.by(() => { - if (instanceCount === 0) return null; - for (const [id, inst] of Object.entries(instanceData)) { - const status = getInstanceDownloadStatus(id, inst); - if (status.isDownloading && status.progress) { - return status.progress; - } - } - return null; - }); - // Instance launch state let models = $state< Array<{ @@ -562,9 +374,6 @@ // Model picker modal state let isModelPickerOpen = $state(false); - // Advanced options toggle (hides technical jargon for new users) - let showAdvancedOptions = $state(false); - // Favorites state (reactive) const favoritesSet = $derived(getFavoritesSet()); @@ -894,20 +703,6 @@ onMount(() => { mounted = true; fetchModels(); - - // Handle reset-onboarding query parameter (triggered from native Settings) - const params = new URLSearchParams(window.location.search); - if (params.has("reset-onboarding")) { - localStorage.removeItem(ONBOARDING_COMPLETE_KEY); - window.history.replaceState({}, "", window.location.pathname); - onboardingStep = 1; - return; - } - - // Show onboarding wizard for first-time users - if (!localStorage.getItem(ONBOARDING_COMPLETE_KEY)) { - onboardingStep = 1; - } }); async function fetchModels() { @@ -987,30 +782,7 @@ ? Array.from(nodeFilter) : undefined; - if (preview?.instance) { - // Use the instance from the preview - instanceData = preview.instance; - } else { - // Fallback: GET placement from API - const placementResponse = await fetch( - `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=${selectedMinNodes}`, - ); - - if (!placementResponse.ok) { - const errorText = await placementResponse.text(); - console.error("Failed to get placement:", errorText); - addToast({ - type: "error", - message: `Placement failed: ${errorText}`, - }); - return; - } - - instanceData = await placementResponse.json(); - } - - // POST the instance to create it - const response = await fetch("/instance", { + const response = await fetch("/meta_instance", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ @@ -1024,13 +796,8 @@ if (!response.ok) { const errorText = await response.text(); - console.error("Failed to launch instance:", errorText); - addToast({ - type: "error", - message: `Failed to launch model: ${errorText}`, - }); + console.error("Failed to create meta instance:", errorText); } else { - addToast({ type: "success", message: `Model launched successfully` }); // Always auto-select the newly launched model so the user chats to what they just launched setSelectedChatModel(modelId); @@ -1052,11 +819,7 @@ setTimeout(scrollToBottom, 1000); } } catch (error) { - console.error("Error launching instance:", error); - addToast({ - type: "error", - message: "Failed to launch model. Check console for details.", - }); + console.error("Error creating meta instance:", error); } finally { launchingModelId = null; } @@ -1665,7 +1428,6 @@ if (!response.ok) { console.error("Failed to delete instance:", response.status); - addToast({ type: "error", message: "Failed to delete instance" }); } else if (wasSelected) { // If we deleted the currently selected model, switch to another available model // Find another instance that isn't the one we just deleted @@ -2654,979 +2416,146 @@ class="relative h-screen w-full flex flex-col bg-exo-dark-gray overflow-hidden" > - {#if !showOnboarding} -
+
- -
-
-
-
-
+ +
+
+
+
+
+ + {#if !topologyOnlyEnabled} + {/if} - {#if showOnboarding} - - - -
- {#if onboardingStep === 1} - + +
+ + {#if !topologyOnlyEnabled && sidebarVisible} +
+ +
+ {/if} + + {#if topologyOnlyEnabled} + +
-
+ + + {@render clusterWarnings()} + + + {#if tb5WithoutRdma && !tb5InfoDismissed}
0} + class:top-4={tbBridgeCycles.length === 0} + role="status" > - exo -
-

- Welcome to exo -

-

- Run AI models locally, across all your devices. -

-
- -
- {:else if onboardingStep === 2} - -
-
-

- Add more devices, run bigger models -

-
- - -
- - - - - - - - - - - - - - - - - - - - - - - - {#if deviceAnimPhase >= 1} - - - - - - - - - - - - - MacBook Pro - - - 36 GB - - - {/if} - - - - - - - - - - - - - - - Mac Studio - - - 192 GB - - - - - {#if deviceAnimPhase >= 3} - - {/if} - - - {#if deviceAnimPhase >= 3} - - 228 GB combined - - {/if} - - - - - {#if deviceAnimPhase >= 3} - - - {/if} - - - {#if deviceAnimPhase >= 1} - - - - Qwen3 8B - - - 4 GB - - - {/if} - - - {#if deviceAnimPhase >= 1} - - = 3 - ? "rgba(255,215,0,0.06)" - : "rgba(0,0,0,0.03)"} - stroke={deviceAnimPhase >= 3 - ? "rgba(255,215,0,0.35)" - : "rgba(0,0,0,0.08)"} - stroke-width="1" - filter={deviceAnimPhase >= 3 - ? "url(#onb-gold-glow)" - : "none"} - style="transition: fill 500ms, stroke 500ms, filter 500ms;" - /> - {#if deviceAnimPhase < 3} - - - - - {:else} - - Qwen3 30B - - - 16 GB - - {/if} - - {/if} - - - {#if deviceAnimPhase >= 1} - - = 4 - ? "rgba(255,215,0,0.06)" - : "rgba(0,0,0,0.03)"} - stroke={deviceAnimPhase >= 4 - ? "rgba(255,215,0,0.35)" - : "rgba(0,0,0,0.08)"} - stroke-width="1" - filter={deviceAnimPhase >= 4 - ? "url(#onb-gold-glow)" - : "none"} - style="transition: fill 500ms, stroke 500ms, filter 500ms;" - /> - {#if deviceAnimPhase < 4} - - - - - {:else} - - Llama 72B - - - 36 GB - - {/if} - - {/if} - - - {#if deviceAnimPhase >= 1} - - = 4 - ? "rgba(255,215,0,0.08)" - : "rgba(0,0,0,0.03)"} - stroke={deviceAnimPhase >= 4 - ? "rgba(255,215,0,0.45)" - : "rgba(0,0,0,0.08)"} - stroke-width={deviceAnimPhase >= 4 ? "1.5" : "1"} - filter={deviceAnimPhase >= 4 - ? "url(#onb-gold-glow)" - : "none"} - style="transition: fill 700ms, stroke 700ms, filter 700ms, stroke-width 700ms;" - /> - {#if deviceAnimPhase < 4} - - - - - {:else} - - Llama 405B - - - 203 GB - - {/if} - - {/if} - - - {#if deviceAnimPhase >= 1} - - Models you can run - - {/if} - - -
- - - {#if showContinueStep2} - + + RDMA AVAILABLE + + +
{/if} -
- {:else if onboardingStep === 3} - -
-
-

- Your devices -

-

- {nodeCount} device{nodeCount !== 1 ? "s" : ""} connected - {#if clusterTotalMemoryGB() > 0} - · {clusterTotalMemoryGB().toFixed(0)} GB total memory - {/if} -

-
-
- -
-

- Install exo on more devices on your network to combine their power — - they connect automatically. -

+ +
- {:else if onboardingStep === 4} - -
-
-

- Choose a model -

-

- Pick a model to download and run locally. -

-
- - {#if onboardingError} -
- {onboardingError} -
- {/if} - - {#if onboardingModels.length === 0} -
-
- Loading models... -
-
- {:else} -
- {#each onboardingModels as model} - {@const sizeGB = getModelSizeGB(model)} - {@const fitsNow = hasEnoughMemory(model)} - {@const tags = modelTags()[model.id] || []} - - {/each} -
- {/if} - - -
- {:else if onboardingStep === 5} - -
-
-

- Downloading -

-

- {#if onboardingModelId} - {onboardingModelId} - {/if} -

-
- - {#if onboardingDownloadProgress} -
-
-
-
-
- {onboardingDownloadProgress.percentage.toFixed(1)}% - {formatBytes(onboardingDownloadProgress.downloadedBytes)} / - {formatBytes(onboardingDownloadProgress.totalBytes)} -
-
- {formatSpeed(onboardingDownloadProgress.speed)} - ETA: {formatEta(onboardingDownloadProgress.etaMs)} -
-
- {:else} -
-
-
-
-

- Preparing download... -

-
- {/if} - -

- This may take a few minutes depending on your connection. -

-
- {:else if onboardingStep === 6} - -
-
-

- Loading into memory -

-

- {#if onboardingModelId} - {onboardingModelId} - {/if} -

-
- -
-
-
- -

Almost ready...

-
- {:else if onboardingStep === 7} - -
- +
+ {:else if !chatStarted} + +
+ +
+
- exo -
- - - {#if onboardingModelId} -

- {onboardingModelId.split("/").pop() ?? onboardingModelId} -

- {/if} - - -
- -
- - -
- {#each ["Write a poem about the ocean", "Explain quantum computing simply", "Help me debug my code", "Tell me a creative story"] as chip} - - {/each} -
-
- {/if} -
- - - {#if onboardingStep === 4} - m.id))} - canModelFit={(modelId) => { - const model = models.find((m) => m.id === modelId); - return model ? hasEnoughMemory(model) : false; - }} - getModelFitStatus={(modelId): ModelMemoryFitStatus => { - const model = models.find((m) => m.id === modelId); - return model ? getModelMemoryFitStatus(model) : "too_large"; - }} - onSelect={(modelId) => { - isModelPickerOpen = false; - onboardingLaunchModel(modelId); - }} - onClose={() => (isModelPickerOpen = false)} - onToggleFavorite={toggleFavorite} - onAddModel={addModelFromPicker} - onDeleteModel={deleteCustomModel} - totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} - usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} - {downloadsData} - topologyNodes={data?.nodes} - /> - {/if} - {:else} - - - - {#if !topologyOnlyEnabled} - - {/if} - - -
- - {#if !topologyOnlyEnabled && sidebarVisible} - - {/if} - - {#if topologyOnlyEnabled} - -
-
+ {#if tb5WithoutRdma && !tb5InfoDismissed}
0} class:top-4={tbBridgeCycles.length === 0} role="status" > - - - - - RDMA AVAILABLE - - -
- {/if} - - - -
-
- {:else if !chatStarted} - -
- -
- -
- - - - - {#if !update} -
-
-
-

- Connecting to cluster… -

-
-
- {/if} - - - {#if instanceCount === 0 && update} -
-
-
-
- exo -
-

- {#if data && Object.keys(data.nodes).length > 1} - {Object.keys(data.nodes).length} devices connected. Choose - a model to start running AI across your cluster. - {:else if data && Object.keys(data.nodes).length === 1} - Your device is ready. Choose a model to start running - AI locally. - {:else} - Waiting for devices to connect… - {/if} -

-
- - - - -
- models download automatically - • - view downloads -
-
-
- {/if} - - {@render clusterWarnings()} - - - {#if tb5WithoutRdma && !tb5InfoDismissed} -
0} - class:top-4={tbBridgeCycles.length === 0} - role="status" - > -
+ RDMA AVAILABLE + + -
- - - +
- {/if} - - {#if isFilterActive()} -
+ {/if} + + + {#if isFilterActive()} + + {/if} +
+ + +
+
+ +
+
+
+ + + +
+ {:else} + +
+ +
+
+
+
- -
+
+ + + {#if minimized} + -
- {:else} - -
- -
-
-
- -
-
- -
-
- -
-
-
- - - {#if minimized} - - {/if} -
- {/if} - - {/if} + {/if} +
+ {/if} + -{#if !showOnboarding} - m.id))} - canModelFit={(modelId) => { - const model = models.find((m) => m.id === modelId); - return model ? hasEnoughMemory(model) : false; - }} - getModelFitStatus={(modelId): ModelMemoryFitStatus => { - const model = models.find((m) => m.id === modelId); - return model ? getModelMemoryFitStatus(model) : "too_large"; - }} - onSelect={handleModelPickerSelect} - onClose={() => (isModelPickerOpen = false)} - onToggleFavorite={toggleFavorite} - onAddModel={addModelFromPicker} - onDeleteModel={deleteCustomModel} - totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} - usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} - {downloadsData} - topologyNodes={data?.nodes} - /> -{/if} + m.id))} + canModelFit={(modelId) => { + const model = models.find((m) => m.id === modelId); + return model ? hasEnoughMemory(model) : false; + }} + getModelFitStatus={(modelId): ModelMemoryFitStatus => { + const model = models.find((m) => m.id === modelId); + return model ? getModelMemoryFitStatus(model) : "too_large"; + }} + onSelect={handleModelPickerSelect} + onClose={() => (isModelPickerOpen = false)} + onToggleFavorite={toggleFavorite} + onAddModel={addModelFromPicker} + onDeleteModel={deleteCustomModel} + totalMemoryGB={clusterMemory().total / (1024 * 1024 * 1024)} + usedMemoryGB={clusterMemory().used / (1024 * 1024 * 1024)} + {downloadsData} + topologyNodes={data?.nodes} +/> diff --git a/packaging/dmg/background.png b/packaging/dmg/background.png deleted file mode 100644 index 5e56d8abd6e2f769097898c44a1331601d7e4cbc..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 7010 zcmeHK`&U#|6#niEBQKR1R0NSgQ6@ul5?c!tmd(Oed zq4P&sO|SwmDriB#QecD%5Vcy2pqW*T{ULzkw4i|Zm+#Db_AoBgF=(u6Ns7^YAkxl# zYh6QFnOlmnM0t91evk5m?0M4dsCO#b;9E@J!?Ch+Kqy`(G z8JTnA_k>@@i-GSiCNvq3YiPE|Rt0Em{eSY>F9Ylg**#S*n|*hZ4$ya*NdaSS)0I_B zX*|S~j#6)9mVyMOoF~X1w9E5Lj|; z>I)|awa+r3Vk$$}+!7m)+f#536DdA5TrJXey&_v%h_gPbckP_jvqZ%QWQJ+rw=;v* zQeWQT;3VMHehR5HdaK1zi%R0LkoN!F>+xTIuCW!lN-3HZ2S~5*mizL@7<_La0d$#v zaGi348raO-J5p*CR}^a9=&3`-imLC3HI4^j(GlZQ$X6;LcczZp_5pe96nCxI_uUz9N@9r~D@7%eA$|k%n{* zE4z2}O6YK6T60XI`18%}{n{H&o}8xBLUov!PG5gUk~LL^YxYq!9Nwcd$dvEJ zV$&2ljP=!KPN=u_gQ@HOhUXl{F|8|U6I+9Cad?M*SN7$aAK-(ca4_!)A4@VS&+`dQOs zq5JvcQ%WGFXEzI6hm*o@D8G+Gvtk*?hkMbsAB8tTQSREVyW>WFyk_W diff --git a/packaging/dmg/create-dmg.sh b/packaging/dmg/create-dmg.sh deleted file mode 100755 index cd68b2a7..00000000 --- a/packaging/dmg/create-dmg.sh +++ /dev/null @@ -1,112 +0,0 @@ -#!/usr/bin/env bash -# create-dmg.sh — Build a polished macOS DMG installer for EXO -# -# Usage: -# ./packaging/dmg/create-dmg.sh [volume-name] -# -# Example: -# ./packaging/dmg/create-dmg.sh output/EXO.app EXO-1.0.0.dmg "EXO" -# -# Creates a DMG with: -# - Custom background image with drag-to-Applications arrow -# - App icon on left, Applications alias on right -# - Proper window size and icon positioning -set -euo pipefail - -APP_PATH="${1:?Usage: create-dmg.sh [volume-name]}" -OUTPUT_DMG="${2:?Usage: create-dmg.sh [volume-name]}" -VOLUME_NAME="${3:-EXO}" - -SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" -BACKGROUND_SCRIPT="${SCRIPT_DIR}/generate-background.py" -TEMP_DIR="$(mktemp -d)" -DMG_STAGING="${TEMP_DIR}/dmg-root" -TEMP_DMG="${TEMP_DIR}/temp.dmg" -BACKGROUND_PNG="${TEMP_DIR}/background.png" - -cleanup() { rm -rf "$TEMP_DIR"; } -trap cleanup EXIT - -echo "==> Creating DMG installer for ${VOLUME_NAME}" - -# ── Step 1: Generate background image ──────────────────────────────────────── -if command -v python3 &>/dev/null; then - python3 "$BACKGROUND_SCRIPT" "$BACKGROUND_PNG" - echo " Background image generated" -else - echo " Warning: python3 not found, skipping custom background" - BACKGROUND_PNG="" -fi - -# ── Step 2: Prepare staging directory ───────────────────────────────────────── -mkdir -p "$DMG_STAGING" -cp -R "$APP_PATH" "$DMG_STAGING/" -ln -s /Applications "$DMG_STAGING/Applications" - -# ── Step 3: Create writable DMG ────────────────────────────────────────────── -# Calculate required size (app size + 20MB headroom) -APP_SIZE_KB=$(du -sk "$APP_PATH" | cut -f1) -DMG_SIZE_KB=$((APP_SIZE_KB + 20480)) - -hdiutil create \ - -volname "$VOLUME_NAME" \ - -size "${DMG_SIZE_KB}k" \ - -fs HFS+ \ - -layout SPUD \ - "$TEMP_DMG" - -# ── Step 4: Mount and configure ────────────────────────────────────────────── -MOUNT_DIR=$(hdiutil attach "$TEMP_DMG" -readwrite -noverify | awk -F'\t' '/Apple_HFS/ {gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF); print $NF}') -echo " Mounted at: $MOUNT_DIR" - -# Copy contents -cp -R "$DMG_STAGING/"* "$MOUNT_DIR/" - -# Add background image -if [[ -n $BACKGROUND_PNG && -f $BACKGROUND_PNG ]]; then - mkdir -p "$MOUNT_DIR/.background" - cp "$BACKGROUND_PNG" "$MOUNT_DIR/.background/background.png" -fi - -# ── Step 5: Configure window appearance via AppleScript ────────────────────── -# Window: 800×400, app icon on left, Applications on right (matches Ollama layout) -# Background image is 1600×740 (2× retina for 800×400 logical window). -APP_NAME="$(basename "$APP_PATH")" - -osascript < DMG created: $OUTPUT_DMG" -echo " Size: $(du -h "$OUTPUT_DMG" | cut -f1)" diff --git a/packaging/dmg/generate-background.py b/packaging/dmg/generate-background.py deleted file mode 100644 index ac3dc649..00000000 --- a/packaging/dmg/generate-background.py +++ /dev/null @@ -1,91 +0,0 @@ -#!/usr/bin/env python3 -"""Generate the DMG background image with a centered drag-to-Applications arrow. - -The output is a 1600×740 retina PNG (2× for 800×400 logical window). -Icons are positioned at (200, 190) and (600, 190) in logical coordinates; -the arrow is drawn centered between them. - -Usage: - python3 generate-background.py [output.png] - -If no output path is given, overwrites the bundled background.png in-place. -""" - -from __future__ import annotations - -import math -import sys -from pathlib import Path - -from PIL import Image, ImageDraw - -# Retina dimensions (2× logical 800×400) -WIDTH = 1600 -HEIGHT = 740 - -# Icon positions in logical coords → retina coords -# App icon at (200, 190), Applications at (600, 190) -APP_X = 200 * 2 # 400 -APPS_X = 600 * 2 # 1200 -ICON_Y = 190 * 2 # 380 - -# Arrow drawn between icons, slightly above icon center -ARROW_START_X = APP_X + 160 # past the icon -ARROW_END_X = APPS_X - 160 # before the Applications icon -ARROW_Y = ICON_Y # same height as icons -ARROW_RISE = 120 # upward arc height - - -def draw_arrow(draw: ImageDraw.ImageDraw) -> None: - """Draw a hand-drawn-style curved arrow from app icon toward Applications.""" - color = (30, 30, 30) - line_width = 8 - - # Compute bezier curve points for a gentle upward arc - points: list[tuple[float, float]] = [] - steps = 80 - for i in range(steps + 1): - t = i / steps - # Quadratic bezier: start → control → end - cx = (ARROW_START_X + ARROW_END_X) / 2 - cy = ARROW_Y - ARROW_RISE - x = (1 - t) ** 2 * ARROW_START_X + 2 * (1 - t) * t * cx + t**2 * ARROW_END_X - y = (1 - t) ** 2 * ARROW_Y + 2 * (1 - t) * t * cy + t**2 * ARROW_Y - points.append((x, y)) - - # Draw the curve as connected line segments - for i in range(len(points) - 1): - draw.line([points[i], points[i + 1]], fill=color, width=line_width) - - # Arrowhead at the end - end_x, end_y = points[-1] - # Direction from second-to-last to last point - prev_x, prev_y = points[-3] - angle = math.atan2(end_y - prev_y, end_x - prev_x) - head_len = 36 - head_angle = math.radians(25) - - left_x = end_x - head_len * math.cos(angle - head_angle) - left_y = end_y - head_len * math.sin(angle - head_angle) - right_x = end_x - head_len * math.cos(angle + head_angle) - right_y = end_y - head_len * math.sin(angle + head_angle) - - draw.polygon( - [(end_x, end_y), (left_x, left_y), (right_x, right_y)], - fill=color, - ) - - -def generate_background(output_path: str) -> None: - """Generate a white DMG background with a centered arrow.""" - img = Image.new("RGBA", (WIDTH, HEIGHT), (255, 255, 255, 255)) - draw = ImageDraw.Draw(img) - draw_arrow(draw) - img.save(output_path, "PNG") - - -if __name__ == "__main__": - default_output = str(Path(__file__).parent / "background.png") - out = sys.argv[1] if len(sys.argv) >= 2 else default_output - generate_background(out) - print(f"Background image written to {out}") diff --git a/src/exo/utils/banner.py b/src/exo/utils/banner.py index 2742832e..ffdb5458 100644 --- a/src/exo/utils/banner.py +++ b/src/exo/utils/banner.py @@ -1,27 +1,8 @@ -import logging -import os import sys -import webbrowser - -from exo.shared.constants import EXO_CONFIG_HOME - -logger = logging.getLogger(__name__) - -_FIRST_RUN_MARKER = EXO_CONFIG_HOME / ".dashboard_opened" - - -def _is_first_run() -> bool: - return not _FIRST_RUN_MARKER.exists() - - -def _mark_first_run_done() -> None: - _FIRST_RUN_MARKER.parent.mkdir(parents=True, exist_ok=True) - _FIRST_RUN_MARKER.touch() def print_startup_banner(port: int) -> None: dashboard_url = f"http://localhost:{port}" - first_run = _is_first_run() banner = f""" ╔═══════════════════════════════════════════════════════════════════════╗ ║ ║ @@ -49,14 +30,3 @@ def print_startup_banner(port: int) -> None: """ print(banner, file=sys.stderr) - - if first_run: - # Skip browser open when running inside the native macOS app — - # FirstLaunchPopout.swift handles the auto-open with a countdown. - if not os.environ.get("EXO_RUNTIME_DIR"): - try: - webbrowser.open(dashboard_url) - 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() From eccc6298d10effe29b58362d5626da6feb45d518 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Tue, 17 Feb 2026 18:02:52 +0000 Subject: [PATCH 11/45] Revert "Add MetaInstance declarative layer (#1447)" This reverts commit a962a28afc7deb65ecb99d1f1cf83be1c5c10dee. --- .../src/lib/components/ChatSidebar.svelte | 6 +- dashboard/src/lib/components/ModelCard.svelte | 6 +- dashboard/src/lib/stores/app.svelte.ts | 27 +- dashboard/src/routes/+page.svelte | 774 +++-------------- src/exo/download/coordinator.py | 12 +- src/exo/main.py | 2 +- src/exo/master/api.py | 78 +- src/exo/master/main.py | 255 ++---- src/exo/master/placement.py | 20 +- src/exo/master/placement_utils.py | 6 +- src/exo/master/process_managers/__init__.py | 12 - .../process_managers/instance_health.py | 62 -- .../master/process_managers/meta_instance.py | 92 --- .../master/process_managers/node_timeout.py | 27 - src/exo/master/reconcile.py | 244 ------ .../tests/test_meta_instance_edge_cases.py | 778 ------------------ src/exo/master/tests/test_placement_utils.py | 12 +- src/exo/master/tests/test_reconcile.py | 742 ----------------- src/exo/shared/apply.py | 126 +-- src/exo/shared/types/api.py | 22 +- src/exo/shared/types/commands.py | 13 +- src/exo/shared/types/common.py | 4 - src/exo/shared/types/events.py | 80 +- src/exo/shared/types/meta_instance.py | 25 - src/exo/shared/types/state.py | 4 +- src/exo/shared/types/tasks.py | 2 +- src/exo/shared/types/worker/instances.py | 3 +- src/exo/utils/channels.py | 4 +- src/exo/worker/engines/mlx/utils_mlx.py | 5 - src/exo/worker/main.py | 34 +- src/exo/worker/plan.py | 23 +- src/exo/worker/runner/bootstrap.py | 13 - src/exo/worker/runner/runner.py | 4 +- src/exo/worker/runner/runner_supervisor.py | 199 +---- .../test_runner/test_event_ordering.py | 27 +- 35 files changed, 290 insertions(+), 3453 deletions(-) delete mode 100644 src/exo/master/process_managers/__init__.py delete mode 100644 src/exo/master/process_managers/instance_health.py delete mode 100644 src/exo/master/process_managers/meta_instance.py delete mode 100644 src/exo/master/process_managers/node_timeout.py delete mode 100644 src/exo/master/reconcile.py delete mode 100644 src/exo/master/tests/test_meta_instance_edge_cases.py delete mode 100644 src/exo/master/tests/test_reconcile.py delete mode 100644 src/exo/shared/types/meta_instance.py diff --git a/dashboard/src/lib/components/ChatSidebar.svelte b/dashboard/src/lib/components/ChatSidebar.svelte index 6a822ddf..b721b033 100644 --- a/dashboard/src/lib/components/ChatSidebar.svelte +++ b/dashboard/src/lib/components/ChatSidebar.svelte @@ -185,7 +185,11 @@ let instanceType: string | null = null; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; - else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; + else if ( + instanceTag === "MlxIbvInstance" || + instanceTag === "MlxJacclInstance" + ) + instanceType = "MLX RDMA"; let sharding: string | null = null; const inst = instance as { diff --git a/dashboard/src/lib/components/ModelCard.svelte b/dashboard/src/lib/components/ModelCard.svelte index 561c325b..9046d2a6 100644 --- a/dashboard/src/lib/components/ModelCard.svelte +++ b/dashboard/src/lib/components/ModelCard.svelte @@ -21,7 +21,7 @@ } | null; nodes?: Record; sharding?: "Pipeline" | "Tensor"; - runtime?: "MlxRing" | "MlxJaccl"; + runtime?: "MlxRing" | "MlxIbv" | "MlxJaccl"; onLaunch?: () => void; tags?: string[]; apiPreview?: PlacementPreview | null; @@ -348,7 +348,7 @@ // Debug mode state const isDebugMode = $derived(debugMode()); const topology = $derived(topologyData()); - const isRdma = $derived(runtime === "MlxJaccl"); + const isRdma = $derived(runtime === "MlxIbv" || runtime === "MlxJaccl"); // Get interface name for an IP from node data function getInterfaceForIp(nodeId: string, ip?: string): string | null { @@ -575,7 +575,7 @@ > {runtime === "MlxRing" ? "MLX Ring" - : runtime === "MlxJaccl" + : runtime === "MlxIbv" || runtime === "MlxJaccl" ? "MLX RDMA" : runtime} diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index ebb2d0df..e5dbf902 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -168,7 +168,7 @@ export interface ModelDownloadStatus { export interface PlacementPreview { model_id: string; sharding: "Pipeline" | "Tensor"; - instance_meta: "MlxRing" | "MlxJaccl"; + instance_meta: "MlxRing" | "MlxIbv" | "MlxJaccl"; instance: unknown | null; memory_delta_by_node: Record | null; error: string | null; @@ -219,6 +219,7 @@ interface RawStateResponse { string, { MlxRingInstance?: Instance; + MlxIbvInstance?: Instance; MlxJacclInstance?: Instance; } >; @@ -249,20 +250,6 @@ interface RawStateResponse { >; // Thunderbolt bridge cycles (nodes with bridge enabled forming loops) thunderboltBridgeCycles?: string[][]; - // MetaInstances (declarative instance constraints) - metaInstances?: Record; -} - -export interface MetaInstanceData { - metaInstanceId: string; - modelId: string; - sharding: string; - instanceMeta: string; - minNodes: number; - nodeIds: string[] | null; - placementError: string | null; - consecutiveFailures: number; - lastFailureError: string | null; } export interface MessageAttachment { @@ -550,7 +537,6 @@ class AppStore { previewNodeFilter = $state>(new Set()); lastUpdate = $state(null); nodeIdentities = $state>({}); - metaInstances = $state>({}); thunderboltBridgeCycles = $state([]); nodeThunderbolt = $state< Record< @@ -909,7 +895,11 @@ class AppStore { let instanceType: string | null = null; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; - else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; + else if ( + instanceTag === "MlxIbvInstance" || + instanceTag === "MlxJacclInstance" + ) + instanceType = "MLX RDMA"; let sharding: string | null = null; const inst = instance as { @@ -1283,8 +1273,6 @@ class AppStore { this.nodeThunderbolt = data.nodeThunderbolt ?? {}; // RDMA ctl status per node this.nodeRdmaCtl = data.nodeRdmaCtl ?? {}; - // MetaInstances - this.metaInstances = data.metaInstances ?? {}; // Thunderbolt bridge cycles this.thunderboltBridgeCycles = data.thunderboltBridgeCycles ?? []; // Thunderbolt bridge status per node @@ -3056,7 +3044,6 @@ export const tps = () => appStore.tps; export const totalTokens = () => appStore.totalTokens; export const topologyData = () => appStore.topologyData; export const instances = () => appStore.instances; -export const metaInstances = () => appStore.metaInstances; export const runners = () => appStore.runners; export const downloads = () => appStore.downloads; export const nodeDisk = () => appStore.nodeDisk; diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 5f7dcf04..2fdeb8ab 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -44,13 +44,11 @@ toggleChatSidebarVisible, nodeThunderbolt, nodeRdmaCtl, - metaInstances, thunderboltBridgeCycles, nodeThunderboltBridge, nodeIdentities, type DownloadProgress, type PlacementPreview, - type MetaInstanceData, } from "$lib/stores/app.svelte"; import HeaderNav from "$lib/components/HeaderNav.svelte"; import { fade, fly } from "svelte/transition"; @@ -70,70 +68,7 @@ const debugEnabled = $derived(debugMode()); const topologyOnlyEnabled = $derived(topologyOnlyMode()); const sidebarVisible = $derived(chatSidebarVisible()); - const metaInstancesData = $derived(metaInstances()); const tbBridgeCycles = $derived(thunderboltBridgeCycles()); - - // Get status for a MetaInstance that has no backing instance yet - function getMetaInstancePlacingStatus(metaInstanceId: string) { - const meta = metaInstancesData[metaInstanceId]; - const placementError = meta?.placementError; - const failures = meta?.consecutiveFailures ?? 0; - const lastError = meta?.lastFailureError; - - if (placementError) { - return { - statusText: "PLACEMENT FAILED", - statusClass: "failed", - isDownloading: false as const, - isFailed: true, - progress: null, - perNode: [] as Array<{ - nodeId: string; - nodeName: string; - progress: DownloadProgress; - }>, - perNodeStatus: [] as PerNodeRunnerStatus[], - errorMessage: placementError, - }; - } - - if (failures > 0) { - const retryPosition = ((failures - 1) % 3) + 1; - const isRecreated = failures % 3 === 0; - return { - statusText: isRecreated ? "PLACING" : `RETRYING (${retryPosition}/3)`, - statusClass: "starting", - isDownloading: false as const, - isFailed: false, - progress: null, - perNode: [] as Array<{ - nodeId: string; - nodeName: string; - progress: DownloadProgress; - }>, - perNodeStatus: [] as PerNodeRunnerStatus[], - errorMessage: isRecreated - ? `Instance re-created due to failure: ${lastError}` - : `Previous failure: ${lastError}`, - }; - } - - return { - statusText: "PLACING", - statusClass: "starting", - isDownloading: false as const, - isFailed: false, - progress: null, - perNode: [] as Array<{ - nodeId: string; - nodeName: string; - progress: DownloadProgress; - }>, - perNodeStatus: [] as PerNodeRunnerStatus[], - errorMessage: null, - }; - } - const tbBridgeData = $derived(nodeThunderboltBridge()); const identitiesData = $derived(nodeIdentities()); const tbIdentifiers = $derived(nodeThunderbolt()); @@ -179,17 +114,6 @@ }); let tb5InfoDismissed = $state(false); - // Detect [jaccl] RDMA driver errors from MetaInstance failure errors - const jacclError = $derived.by(() => { - for (const mi of Object.values(metaInstancesData)) { - if (mi.lastFailureError?.includes("[jaccl]")) { - return mi.lastFailureError; - } - } - return null; - }); - let jacclDismissedError = $state(null); - // Helper to get friendly node name from node ID function getNodeName(nodeId: string): string { const node = data?.nodes?.[nodeId]; @@ -300,7 +224,7 @@ return model.tasks.includes("ImageToImage"); } let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline"); - type InstanceMeta = "MlxRing" | "MlxJaccl"; + type InstanceMeta = "MlxRing" | "MlxIbv" | "MlxJaccl"; // Launch defaults persistence const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults"; @@ -557,7 +481,7 @@ const matchesSelectedRuntime = (runtime: InstanceMeta): boolean => selectedInstanceType === "MlxRing" ? runtime === "MlxRing" - : runtime === "MlxJaccl" || runtime === "MlxJaccl"; + : runtime === "MlxIbv" || runtime === "MlxJaccl"; // Helper to check if a model can be launched (has valid placement with >= minNodes) function canModelFit(modelId: string): boolean { @@ -773,30 +697,39 @@ launchingModelId = modelId; try { + // Use the specific preview if provided, otherwise fall back to filtered preview const preview = specificPreview ?? filteredPreview(); - // Extract node IDs from the preview the user is seeing - const previewNodeIds = preview?.memory_delta_by_node - ? Object.keys(preview.memory_delta_by_node) - : nodeFilter.size > 0 - ? Array.from(nodeFilter) - : undefined; + let instanceData: unknown; - const response = await fetch("/meta_instance", { + if (preview?.instance) { + // Use the instance from the preview + instanceData = preview.instance; + } else { + // Fallback: GET placement from API + const placementResponse = await fetch( + `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=${selectedMinNodes}`, + ); + + if (!placementResponse.ok) { + const errorText = await placementResponse.text(); + console.error("Failed to get placement:", errorText); + return; + } + + instanceData = await placementResponse.json(); + } + + // POST the instance to create it + const response = await fetch("/instance", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - model_id: modelId, - sharding: preview?.sharding ?? selectedSharding, - instance_meta: preview?.instance_meta ?? selectedInstanceType, - min_nodes: selectedMinNodes, - node_ids: previewNodeIds, - }), + body: JSON.stringify({ instance: instanceData }), }); if (!response.ok) { const errorText = await response.text(); - console.error("Failed to create meta instance:", errorText); + console.error("Failed to launch instance:", errorText); } else { // Always auto-select the newly launched model so the user chats to what they just launched setSelectedChatModel(modelId); @@ -819,7 +752,7 @@ setTimeout(scrollToBottom, 1000); } } catch (error) { - console.error("Error creating meta instance:", error); + console.error("Error launching instance:", error); } finally { launchingModelId = null; } @@ -1021,18 +954,15 @@ nodeName: string; progress: DownloadProgress; }>; - perNodeStatus: PerNodeRunnerStatus[]; } { if (!downloadsData || Object.keys(downloadsData).length === 0) { - const statusInfo = deriveInstanceStatus(instanceWrapped); return { isDownloading: false, - isFailed: statusInfo.statusText === "FAILED", - errorMessage: statusInfo.errorMessage, + isFailed: false, + errorMessage: null, progress: null, - statusText: statusInfo.statusText, + statusText: "RUNNING", perNode: [], - perNodeStatus: statusInfo.perNodeStatus, }; } @@ -1046,7 +976,6 @@ progress: null, statusText: "PREPARING", perNode: [], - perNodeStatus: [], }; } @@ -1115,7 +1044,6 @@ progress: null, statusText: "FAILED", perNode: [], - perNodeStatus: [], }; } } @@ -1156,11 +1084,10 @@ return { isDownloading: false, isFailed: statusInfo.statusText === "FAILED", - errorMessage: statusInfo.errorMessage, + errorMessage: null, progress: null, statusText: statusInfo.statusText, perNode: [], - perNodeStatus: statusInfo.perNodeStatus, }; } @@ -1184,223 +1111,92 @@ }, statusText: "DOWNLOADING", perNode, - perNodeStatus: [], }; } // Derive instance status from runners // Get color class for a status function getStatusColor(statusText: string): string { - if (statusText === "FAILED" || statusText === "PLACEMENT FAILED") - return "text-red-400"; - if (statusText.startsWith("RETRYING")) return "text-orange-400"; - if (statusText === "SHUTDOWN") return "text-gray-400"; - if (statusText === "DOWNLOADING") return "text-blue-400"; - if ( - statusText.startsWith("LOADING") || - statusText.startsWith("WARMING UP") || - statusText === "WAITING" || - statusText === "INITIALIZING" - ) - return "text-yellow-400"; - if (statusText === "RUNNING") return "text-teal-400"; - if (statusText === "READY" || statusText === "LOADED") - return "text-green-400"; - return "text-exo-light-gray"; - } - - const RUNNER_STATUS_MAP: Record = { - RunnerWaitingForInitialization: "WaitingForInitialization", - RunnerInitializingBackend: "InitializingBackend", - RunnerWaitingForModel: "WaitingForModel", - RunnerLoading: "Loading", - RunnerLoaded: "Loaded", - RunnerWarmingUp: "WarmingUp", - RunnerReady: "Ready", - RunnerRunning: "Running", - RunnerShutdown: "Shutdown", - RunnerFailed: "Failed", - }; - - // Friendly labels for display - const RUNNER_STATUS_DISPLAY: Record = { - WaitingForInitialization: "Initializing", - InitializingBackend: "Initializing", - WaitingForModel: "Waiting", - Loading: "Loading", - Loaded: "Loaded", - WarmingUp: "Warming Up", - Ready: "Ready", - Running: "Running", - Shutdown: "Shutdown", - Failed: "Failed", - }; - - interface PerNodeRunnerStatus { - nodeId: string; - nodeName: string; - status: string; // friendly display status + switch (statusText) { + case "FAILED": + return "text-red-400"; + case "SHUTDOWN": + return "text-gray-400"; + case "DOWNLOADING": + return "text-blue-400"; + case "LOADING": + case "WARMING UP": + case "WAITING": + case "INITIALIZING": + return "text-yellow-400"; + case "RUNNING": + return "text-teal-400"; + case "READY": + case "LOADED": + return "text-green-400"; + default: + return "text-exo-light-gray"; + } } function deriveInstanceStatus(instanceWrapped: unknown): { statusText: string; statusClass: string; - perNodeStatus: PerNodeRunnerStatus[]; - errorMessage: string | null; } { const [, instance] = getTagged(instanceWrapped); if (!instance || typeof instance !== "object") { - return { - statusText: "PREPARING", - statusClass: "inactive", - perNodeStatus: [], - errorMessage: null, - }; + return { statusText: "PREPARING", statusClass: "inactive" }; } const inst = instance as { - shardAssignments?: { - runnerToShard?: Record; - nodeToRunner?: Record; - }; + shardAssignments?: { runnerToShard?: Record }; }; - const nodeToRunner = inst.shardAssignments?.nodeToRunner || {}; const runnerIds = Object.keys(inst.shardAssignments?.runnerToShard || {}); - const totalNodes = runnerIds.length; - // Build per-node status and extract error messages from RunnerFailed - const perNodeStatus: PerNodeRunnerStatus[] = []; - const statuses: string[] = []; - const failedErrors: string[] = []; - for (const [nodeId, runnerId] of Object.entries(nodeToRunner)) { - const r = runnersData[runnerId]; - let status: string | null = null; - if (r) { - const [kind, runnerData] = getTagged(r); - status = kind ? RUNNER_STATUS_MAP[kind] || null : null; - // Extract error message from RunnerFailed - if ( - kind === "RunnerFailed" && - runnerData && - typeof runnerData === "object" - ) { - const rd = runnerData as { errorMessage?: string }; - if (rd.errorMessage) - failedErrors.push(`${getNodeName(nodeId)}: ${rd.errorMessage}`); - } - } - if (status) { - statuses.push(status); - perNodeStatus.push({ - nodeId, - nodeName: getNodeName(nodeId), - status: RUNNER_STATUS_DISPLAY[status] || status, - }); - } - } + const statuses = runnerIds + .map((rid) => { + const r = runnersData[rid]; + if (!r) return null; + const [kind] = getTagged(r); + const statusMap: Record = { + RunnerWaitingForInitialization: "WaitingForInitialization", + RunnerInitializingBackend: "InitializingBackend", + RunnerWaitingForModel: "WaitingForModel", + RunnerLoading: "Loading", + RunnerLoaded: "Loaded", + RunnerWarmingUp: "WarmingUp", + RunnerReady: "Ready", + RunnerRunning: "Running", + RunnerShutdown: "Shutdown", + RunnerFailed: "Failed", + }; + return kind ? statusMap[kind] || null : null; + }) + .filter((s): s is string => s !== null); const has = (s: string) => statuses.includes(s); - const count = (s: string) => statuses.filter((v) => v === s).length; if (statuses.length === 0) - return { - statusText: "PREPARING", - statusClass: "inactive", - perNodeStatus, - errorMessage: null, - }; - if (has("Failed")) - return { - statusText: "FAILED", - statusClass: "failed", - perNodeStatus, - errorMessage: failedErrors.length > 0 ? failedErrors.join("; ") : null, - }; + return { statusText: "PREPARING", statusClass: "inactive" }; + if (has("Failed")) return { statusText: "FAILED", statusClass: "failed" }; if (has("Shutdown")) - return { - statusText: "SHUTDOWN", - statusClass: "inactive", - perNodeStatus, - errorMessage: null, - }; - - // For loading/warming states, show node progress when multi-node - if (has("Loading")) { - const readyCount = count("Ready") + count("Running") + count("Loaded"); - const statusText = - totalNodes > 1 - ? `LOADING (${readyCount}/${totalNodes} nodes ready)` - : "LOADING"; - return { - statusText, - statusClass: "starting", - perNodeStatus, - errorMessage: null, - }; - } - if (has("WarmingUp")) { - const readyCount = count("Ready") + count("Running"); - const statusText = - totalNodes > 1 - ? `WARMING UP (${readyCount}/${totalNodes} nodes ready)` - : "WARMING UP"; - return { - statusText, - statusClass: "starting", - perNodeStatus, - errorMessage: null, - }; - } - + return { statusText: "SHUTDOWN", statusClass: "inactive" }; + if (has("Loading")) + return { statusText: "LOADING", statusClass: "starting" }; + if (has("WarmingUp")) + return { statusText: "WARMING UP", statusClass: "starting" }; if (has("Running")) - return { - statusText: "RUNNING", - statusClass: "running", - perNodeStatus, - errorMessage: null, - }; - if (has("Ready")) - return { - statusText: "READY", - statusClass: "loaded", - perNodeStatus, - errorMessage: null, - }; - if (has("Loaded")) - return { - statusText: "LOADED", - statusClass: "loaded", - perNodeStatus, - errorMessage: null, - }; + return { statusText: "RUNNING", statusClass: "running" }; + if (has("Ready")) return { statusText: "READY", statusClass: "loaded" }; + if (has("Loaded")) return { statusText: "LOADED", statusClass: "loaded" }; if (has("WaitingForModel")) - return { - statusText: "WAITING", - statusClass: "starting", - perNodeStatus, - errorMessage: null, - }; + return { statusText: "WAITING", statusClass: "starting" }; if (has("InitializingBackend")) - return { - statusText: "INITIALIZING", - statusClass: "starting", - perNodeStatus, - errorMessage: null, - }; + return { statusText: "INITIALIZING", statusClass: "starting" }; if (has("WaitingForInitialization")) - return { - statusText: "INITIALIZING", - statusClass: "starting", - perNodeStatus, - errorMessage: null, - }; + return { statusText: "INITIALIZING", statusClass: "starting" }; - return { - statusText: "RUNNING", - statusClass: "active", - perNodeStatus, - errorMessage: null, - }; + return { statusText: "RUNNING", statusClass: "active" }; } function getBytes(value: unknown): number { @@ -1459,75 +1255,6 @@ } } - async function deleteMetaInstance(metaInstanceId: string) { - const meta = metaInstancesData[metaInstanceId]; - const modelId = meta?.modelId ?? "unknown"; - if (!confirm(`Delete model ${modelId}?`)) return; - - const wasSelected = selectedChatModel() === modelId; - - try { - const response = await fetch(`/meta_instance/${metaInstanceId}`, { - method: "DELETE", - headers: { "Content-Type": "application/json" }, - }); - - if (!response.ok) { - console.error("Failed to delete meta instance:", response.status); - } else if (wasSelected) { - // Switch to another available model or clear selection - const remainingInstances = Object.entries(instanceData).filter( - ([id]) => id !== getBackingInstanceId(metaInstanceId), - ); - if (remainingInstances.length > 0) { - const [, lastInstance] = - remainingInstances[remainingInstances.length - 1]; - const newModelId = getInstanceModelId(lastInstance); - if ( - newModelId && - newModelId !== "Unknown" && - newModelId !== "Unknown Model" - ) { - setSelectedChatModel(newModelId); - } else { - setSelectedChatModel(""); - } - } else { - setSelectedChatModel(""); - } - } - } catch (error) { - console.error("Error deleting meta instance:", error); - } - } - - // Find the backing Instance ID for a MetaInstance by scanning instances - function getBackingInstanceId(metaInstanceId: string): string | null { - for (const [id, inst] of Object.entries(instanceData)) { - const [, inner] = getTagged(inst); - if ( - inner && - typeof inner === "object" && - (inner as Record).metaInstanceId === metaInstanceId - ) { - return id; - } - } - return null; - } - - // Get orphan Instance IDs (not backing any MetaInstance) - function getOrphanInstanceIds(): string[] { - return Object.keys(instanceData).filter((id) => { - const [, inner] = getTagged(instanceData[id]); - return ( - !inner || - typeof inner !== "object" || - !(inner as Record).metaInstanceId - ); - }); - } - // Helper to unwrap tagged unions like { MlxRingInstance: {...} } function getTagged(obj: unknown): [string | null, unknown] { if (!obj || typeof obj !== "object") return [null, null]; @@ -1568,7 +1295,11 @@ // Instance type from tag let instanceType = "Unknown"; if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring"; - else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA"; + else if ( + instanceTag === "MlxIbvInstance" || + instanceTag === "MlxJacclInstance" + ) + instanceType = "MLX RDMA"; const inst = instance as { shardAssignments?: { @@ -1916,51 +1647,7 @@ } const nodeCount = $derived(data ? Object.keys(data.nodes).length : 0); - const metaInstanceCount = $derived(Object.keys(metaInstancesData).length); - const orphanInstanceIds = $derived(getOrphanInstanceIds()); - const instanceCount = $derived(metaInstanceCount + orphanInstanceIds.length); - - // Unified display items: MetaInstances first, then orphan Instances - interface DisplayItem { - id: string; // MetaInstance ID or Instance ID (used as key and displayed) - modelId: string; - instance: unknown | null; // The backing/orphan instance (tagged union) or null if placing - instanceId: string | null; // The actual Instance ID (for topology hover) - isMetaInstance: boolean; - sharding: string | null; // From MetaInstance constraints (used when instance is null) - instanceMeta: string | null; // From MetaInstance constraints (used when instance is null) - } - - const unifiedDisplayItems = $derived.by((): DisplayItem[] => { - const items: DisplayItem[] = []; - // MetaInstances - for (const [metaId, meta] of Object.entries(metaInstancesData)) { - const backingId = getBackingInstanceId(metaId); - items.push({ - id: metaId, - modelId: meta.modelId, - instance: backingId ? instanceData[backingId] : null, - instanceId: backingId, - isMetaInstance: true, - sharding: meta.sharding, - instanceMeta: meta.instanceMeta, - }); - } - // Orphan Instances - for (const orphanId of getOrphanInstanceIds()) { - const inst = instanceData[orphanId]; - items.push({ - id: orphanId, - modelId: getInstanceModelId(inst), - instance: inst, - instanceId: orphanId, - isMetaInstance: false, - sharding: null, - instanceMeta: null, - }); - } - return items; - }); + const instanceCount = $derived(Object.keys(instanceData).length); // Helper to get the number of nodes in a placement preview function getPreviewNodeCount(preview: PlacementPreview): number { @@ -2078,71 +1765,8 @@ {#snippet clusterWarnings()} - {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (jacclError && jacclError !== jacclDismissedError)} + {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)}
- {#if jacclError && jacclError !== jacclDismissedError} - - {/if} - {#if tbBridgeCycles.length > 0} {@const cycle = tbBridgeCycles[0]} {@const serviceName = getTbBridgeServiceName(cycle)} @@ -2311,29 +1935,8 @@ {/snippet} {#snippet clusterWarningsCompact()} - {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed) || (jacclError && jacclError !== jacclDismissedError)} + {#if tbBridgeCycles.length > 0 || macosVersionMismatch || (tb5WithoutRdma && !tb5InfoDismissed)}
- {#if jacclError && jacclError !== jacclDismissedError} -
- - - - JACCL ERROR -
- {/if} {#if tbBridgeCycles.length > 0}
- {#each unifiedDisplayItems as item (item.id)} - {@const id = item.id} - {@const instance = item.instance} - {@const downloadInfo = instance - ? getInstanceDownloadStatus(item.instanceId ?? id, instance) - : getMetaInstancePlacingStatus(id)} - {@const metaData = item.isMetaInstance - ? metaInstancesData[id] - : null} - {@const retryError = - metaData?.lastFailureError && !downloadInfo.isFailed - ? metaData.consecutiveFailures > 0 - ? `(${((metaData.consecutiveFailures - 1) % 3) + 1}/3) ${metaData.lastFailureError}` - : metaData.lastFailureError - : null} + {#each Object.entries(instanceData) as [id, instance]} + {@const downloadInfo = getInstanceDownloadStatus( + id, + instance, + )} {@const statusText = downloadInfo.statusText} {@const isDownloading = downloadInfo.isDownloading} - {@const isFailed = - statusText === "FAILED" || - statusText === "PLACEMENT FAILED"} + {@const isFailed = statusText === "FAILED"} {@const isLoading = - statusText.startsWith("LOADING") || - statusText.startsWith("WARMING UP") || - statusText === "WAITING" || - statusText === "PLACING" || - statusText.startsWith("RETRYING")} + statusText === "LOADING" || + statusText === "WARMING UP" || + statusText === "WAITING"} {@const isReady = statusText === "READY" || statusText === "LOADED"} {@const isRunning = statusText === "RUNNING"} - {@const instanceModelId = item.modelId} - {@const instanceInfo = instance - ? getInstanceInfo(instance) - : { - instanceType: - item.instanceMeta === "MlxRing" - ? "MLX Ring" - : item.instanceMeta === "MlxJaccl" - ? "MLX RDMA" - : "Unknown", - sharding: item.sharding ?? "Unknown", - nodeNames: [] as string[], - nodeIds: [] as string[], - nodeCount: 0, - }} - {@const instanceConnections = instance - ? getInstanceConnections(instance) - : []} + {@const instanceModelId = getInstanceModelId(instance)} + {@const instanceInfo = getInstanceInfo(instance)} + {@const instanceConnections = + getInstanceConnections(instance)}
- (hoveredInstanceId = item.instanceId ?? id)} + onmouseenter={() => (hoveredInstanceId = id)} onmouseleave={() => (hoveredInstanceId = null)} onclick={() => { if ( @@ -2864,10 +2438,7 @@ >
@@ -3337,21 +2884,21 @@
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index f661c878..db13ccef 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -314,17 +314,7 @@ class DownloadCoordinator: ), ) elif progress.status in ["in_progress", "not_started"]: - if ( - progress.downloaded_bytes.in_bytes - >= progress.total_bytes.in_bytes - > 0 - ): - status = DownloadCompleted( - node_id=self.node_id, - shard_metadata=progress.shard, - total_bytes=progress.total_bytes, - ) - elif progress.downloaded_bytes_this_session.in_bytes == 0: + if progress.downloaded_bytes_this_session.in_bytes == 0: status = DownloadPending( node_id=self.node_id, shard_metadata=progress.shard, diff --git a/src/exo/main.py b/src/exo/main.py index 8f0c5a41..1d358975 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -254,7 +254,7 @@ def main(): target = min(max(soft, 65535), hard) resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard)) - mp.set_start_method("spawn", force=True) + mp.set_start_method("spawn") # TODO: Refactor the current verbosity system logger_setup(EXO_LOG, args.verbosity) logger.info("Starting EXO") diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 91c74f41..b8476334 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -71,11 +71,8 @@ from exo.shared.types.api import ( ChatCompletionResponse, CreateInstanceParams, CreateInstanceResponse, - CreateMetaInstanceParams, - CreateMetaInstanceResponse, DeleteDownloadResponse, DeleteInstanceResponse, - DeleteMetaInstanceResponse, ErrorInfo, ErrorResponse, FinishReason, @@ -118,10 +115,8 @@ from exo.shared.types.claude_api import ( from exo.shared.types.commands import ( Command, CreateInstance, - CreateMetaInstance, DeleteDownload, DeleteInstance, - DeleteMetaInstance, DownloadCommand, ForwarderCommand, ForwarderDownloadCommand, @@ -134,7 +129,7 @@ from exo.shared.types.commands import ( TaskFinished, TextGeneration, ) -from exo.shared.types.common import CommandId, Id, MetaInstanceId, NodeId, SessionId +from exo.shared.types.common import CommandId, Id, NodeId, SessionId from exo.shared.types.events import ( ChunkGenerated, Event, @@ -143,7 +138,6 @@ from exo.shared.types.events import ( TracesMerged, ) from exo.shared.types.memory import Memory -from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.openai_responses import ( ResponsesRequest, ResponsesResponse, @@ -282,9 +276,6 @@ class API: self.app.get("/instance/previews")(self.get_placement_previews) self.app.get("/instance/{instance_id}")(self.get_instance) self.app.delete("/instance/{instance_id}")(self.delete_instance) - self.app.get("/meta_instances")(self.list_meta_instances) - self.app.post("/meta_instance")(self.create_meta_instance) - self.app.delete("/meta_instance/{meta_instance_id}")(self.delete_meta_instance) self.app.get("/models")(self.get_models) self.app.get("/v1/models")(self.get_models) self.app.post("/models/add")(self.add_custom_model) @@ -314,27 +305,12 @@ class API: self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw) async def place_instance(self, payload: PlaceInstanceParams): - model_card = await ModelCard.load(payload.model_id) command = PlaceInstance( - model_card=model_card, + model_card=await ModelCard.load(payload.model_id), sharding=payload.sharding, instance_meta=payload.instance_meta, min_nodes=payload.min_nodes, ) - - # Validate placement before sending — fail fast with a clear error - # instead of silently dropping the command in the master. - try: - get_instance_placements( - command, - topology=self.state.topology, - current_instances=self.state.instances, - node_memory=self.state.node_memory, - node_network=self.state.node_network, - ) - except ValueError as exc: - raise HTTPException(status_code=400, detail=str(exc)) from exc - await self._send(command) return CreateInstanceResponse( @@ -546,44 +522,6 @@ class API: instance_id=instance_id, ) - def list_meta_instances(self) -> dict[MetaInstanceId, MetaInstance]: - return dict(self.state.meta_instances) - - async def create_meta_instance( - self, payload: CreateMetaInstanceParams - ) -> CreateMetaInstanceResponse: - meta_instance = MetaInstance( - model_id=payload.model_id, - sharding=payload.sharding, - instance_meta=payload.instance_meta, - min_nodes=payload.min_nodes, - node_ids=payload.node_ids, - ) - command = CreateMetaInstance(meta_instance=meta_instance) - await self._send(command) - return CreateMetaInstanceResponse( - message="Command received.", - command_id=command.command_id, - meta_instance_id=meta_instance.meta_instance_id, - ) - - async def delete_meta_instance( - self, meta_instance_id: MetaInstanceId - ) -> DeleteMetaInstanceResponse: - meta = self.state.meta_instances.get(meta_instance_id) - if not meta: - raise HTTPException(status_code=404, detail="MetaInstance not found") - - # Command processor handles cascade-deleting backing instances - command = DeleteMetaInstance(meta_instance_id=meta_instance_id) - await self._send(command) - - return DeleteMetaInstanceResponse( - message="Command received.", - command_id=command.command_id, - meta_instance_id=meta_instance_id, - ) - async def _token_chunk_stream( self, command_id: CommandId ) -> AsyncGenerator[ErrorChunk | ToolCallChunk | TokenChunk, None]: @@ -603,10 +541,10 @@ class API: break except anyio.get_cancelled_exc_class(): - cancel_command = TaskCancelled(cancelled_command_id=command_id) + command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( - ForwarderCommand(origin=self.node_id, command=cancel_command) + ForwarderCommand(origin=self.node_id, command=command) ) raise finally: @@ -946,10 +884,10 @@ class API: del image_metadata[key] except anyio.get_cancelled_exc_class(): - cancel_command = TaskCancelled(cancelled_command_id=command_id) + command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( - ForwarderCommand(origin=self.node_id, command=cancel_command) + ForwarderCommand(origin=self.node_id, command=command) ) raise finally: @@ -1032,10 +970,10 @@ class API: return (images, stats if capture_stats else None) except anyio.get_cancelled_exc_class(): - cancel_command = TaskCancelled(cancelled_command_id=command_id) + command = TaskCancelled(cancelled_command_id=command_id) with anyio.CancelScope(shield=True): await self.command_sender.send( - ForwarderCommand(origin=self.node_id, command=cancel_command) + ForwarderCommand(origin=self.node_id, command=command) ) raise finally: diff --git a/src/exo/master/main.py b/src/exo/master/main.py index 405f6495..9c7cf578 100644 --- a/src/exo/master/main.py +++ b/src/exo/master/main.py @@ -1,5 +1,4 @@ -from collections.abc import Sequence -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone import anyio from anyio.abc import TaskGroup @@ -13,22 +12,11 @@ from exo.master.placement import ( get_transition_events, place_instance, ) -from exo.master.process_managers import ProcessManager -from exo.master.process_managers.instance_health import InstanceHealthReconciler -from exo.master.process_managers.meta_instance import MetaInstanceReconciler -from exo.master.process_managers.node_timeout import NodeTimeoutReconciler -from exo.master.reconcile import ( - find_unsatisfied_meta_instances, - try_place_for_meta_instance, -) from exo.shared.apply import apply from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED -from exo.shared.models.model_cards import ModelCard from exo.shared.types.commands import ( CreateInstance, - CreateMetaInstance, DeleteInstance, - DeleteMetaInstance, ForwarderCommand, ForwarderDownloadCommand, ImageEdits, @@ -48,12 +36,8 @@ from exo.shared.types.events import ( IndexedEvent, InputChunkReceived, InstanceDeleted, - JacclSideChannelData, - JacclSideChannelGathered, - MetaInstanceCreated, - MetaInstanceDeleted, - MetaInstancePlacementFailed, NodeGatheredInfo, + NodeTimedOut, TaskCreated, TaskDeleted, TaskStatusUpdated, @@ -76,8 +60,7 @@ from exo.shared.types.tasks import ( TextGeneration as TextGenerationTask, ) from exo.shared.types.worker.instances import InstanceId -from exo.shared.types.worker.runners import RunnerId -from exo.utils.channels import Receiver, Sender +from exo.utils.channels import Receiver, Sender, channel from exo.utils.event_buffer import MultiSourceBuffer @@ -101,16 +84,16 @@ class Master: self.local_event_receiver = local_event_receiver self.global_event_sender = global_event_sender self.download_command_sender = download_command_sender + send, recv = channel[Event]() + self.event_sender: Sender[Event] = send + self._loopback_event_receiver: Receiver[Event] = recv + self._loopback_event_sender: Sender[ForwarderEvent] = ( + local_event_receiver.clone_sender() + ) self._multi_buffer = MultiSourceBuffer[NodeId, Event]() self._event_log = DiskEventLog(EXO_EVENT_LOG_DIR / "master") self._pending_traces: dict[TaskId, dict[int, list[TraceEventData]]] = {} self._expected_ranks: dict[TaskId, set[int]] = {} - self._jaccl_pending: dict[InstanceId, dict[int, dict[RunnerId, bytes]]] = {} - self._process_managers: Sequence[ProcessManager] = [ - InstanceHealthReconciler(), - NodeTimeoutReconciler(), - MetaInstanceReconciler(), - ] async def run(self): logger.info("Starting Master") @@ -119,12 +102,15 @@ class Master: async with self._tg as tg: tg.start_soon(self._event_processor) tg.start_soon(self._command_processor) - tg.start_soon(self._reconcile) + tg.start_soon(self._loopback_processor) + tg.start_soon(self._plan) finally: self._event_log.close() self.global_event_sender.close() self.local_event_receiver.close() self.command_receiver.close() + self._loopback_event_sender.close() + self._loopback_event_receiver.close() async def shutdown(self): logger.info("Stopping Master") @@ -306,86 +292,6 @@ class Master: ) ) generated_events.extend(transition_events) - case CreateMetaInstance(): - logger.info( - f"Creating MetaInstance for {command.meta_instance.model_id}" - f" (min_nodes={command.meta_instance.min_nodes}," - f" sharding={command.meta_instance.sharding})" - ) - # Apply immediately so self.state is fresh across - # the await below and the reconciler won't race. - await self._apply_and_broadcast( - MetaInstanceCreated(meta_instance=command.meta_instance) - ) - # Immediate placement attempt for responsiveness - model_card = await ModelCard.load( - command.meta_instance.model_id - ) - # Re-check: reconciler may have satisfied it during the await - meta_id = command.meta_instance.meta_instance_id - still_unsatisfied = any( - m.meta_instance_id == meta_id - for m in find_unsatisfied_meta_instances( - self.state.meta_instances, - self.state.instances, - self.state.topology, - ) - ) - if still_unsatisfied: - result = try_place_for_meta_instance( - command.meta_instance, - model_card, - self.state.topology, - self.state.instances, - self.state.node_memory, - self.state.node_network, - self.state.tasks, - ) - generated_events.extend(result.events) - if result.error is not None: - generated_events.append( - MetaInstancePlacementFailed( - meta_instance_id=meta_id, - reason=result.error, - ) - ) - case DeleteMetaInstance(): - backing_count = sum( - 1 - for inst in self.state.instances.values() - if inst.meta_instance_id == command.meta_instance_id - ) - logger.info( - f"Deleting MetaInstance {command.meta_instance_id}" - f" (cascade-deleting {backing_count} backing instance(s))" - ) - generated_events.append( - MetaInstanceDeleted( - meta_instance_id=command.meta_instance_id - ) - ) - # Cascade-delete backing instances atomically, - # cancelling any active tasks first. - for iid, inst in self.state.instances.items(): - if inst.meta_instance_id == command.meta_instance_id: - for task in self.state.tasks.values(): - if ( - task.instance_id == iid - and task.task_status - in ( - TaskStatus.Pending, - TaskStatus.Running, - ) - ): - generated_events.append( - TaskStatusUpdated( - task_status=TaskStatus.Cancelled, - task_id=task.task_id, - ) - ) - generated_events.append( - InstanceDeleted(instance_id=iid) - ) case PlaceInstance(): placement = place_instance( command, @@ -417,19 +323,16 @@ class Master: ) case TaskCancelled(): if ( - command.cancelled_command_id - in self.command_task_mapping - ): + task_id := self.command_task_mapping.get( + command.cancelled_command_id + ) + ) is not None: generated_events.append( - TaskDeleted( - task_id=self.command_task_mapping[ - command.cancelled_command_id - ] + TaskStatusUpdated( + task_status=TaskStatus.Cancelled, + task_id=task_id, ) ) - del self.command_task_mapping[ - command.cancelled_command_id - ] case TaskFinished(): generated_events.append( TaskDeleted( @@ -438,10 +341,9 @@ class Master: ] ) ) - if command.finished_command_id in self.command_task_mapping: - del self.command_task_mapping[ - command.finished_command_id - ] + self.command_task_mapping.pop( + command.finished_command_id, None + ) case RequestEventLog(): # We should just be able to send everything, since other buffers will ignore old messages # rate limit to 1000 at a time @@ -452,32 +354,31 @@ class Master: ): await self._send_event(IndexedEvent(idx=i, event=event)) for event in generated_events: - await self._apply_and_broadcast(event) + await self.event_sender.send(event) except ValueError as e: logger.opt(exception=e).warning("Error in command processor") - async def _apply_and_broadcast(self, event: Event) -> None: - """Apply event to state, persist to disk, and broadcast to workers. - - State is updated synchronously (before any await), so callers can - rely on ``self.state`` reflecting this event immediately after the - call. Python's cooperative scheduling guarantees no interleaving - between the state read and write. - """ - logger.debug(f"Master indexing event: {str(event)[:100]}") - indexed = IndexedEvent(event=event, idx=len(self._event_log)) - self.state = apply(self.state, indexed) - event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage] - self._event_log.append(event) - await self._send_event(indexed) - - async def _reconcile(self) -> None: + # These plan loops are the cracks showing in our event sourcing architecture - more things could be commands + async def _plan(self) -> None: while True: - for pm in self._process_managers: - events = await pm.reconcile(self.state) - for event in events: - await self._apply_and_broadcast(event) - await anyio.sleep(1) + # kill broken instances + connected_node_ids = set(self.state.topology.list_nodes()) + for instance_id, instance in self.state.instances.items(): + for node_id in instance.shard_assignments.node_to_runner: + if node_id not in connected_node_ids: + await self.event_sender.send( + InstanceDeleted(instance_id=instance_id) + ) + break + + # time out dead nodes + for node_id, time in self.state.last_seen.items(): + now = datetime.now(tz=timezone.utc) + if now - time > timedelta(seconds=30): + logger.info(f"Manually removing node {node_id} due to inactivity") + await self.event_sender.send(NodeTimedOut(node_id=node_id)) + + await anyio.sleep(10) async def _event_processor(self) -> None: with self.local_event_receiver as local_events: @@ -495,15 +396,32 @@ class Master: await self._handle_traces_collected(event) continue - if isinstance(event, JacclSideChannelData): - await self._apply_and_broadcast(event) - await self._handle_jaccl_side_channel(event) - continue + logger.debug(f"Master indexing event: {str(event)[:100]}") + indexed = IndexedEvent(event=event, idx=len(self._event_log)) + self.state = apply(self.state, indexed) + event._master_time_stamp = datetime.now(tz=timezone.utc) # pyright: ignore[reportPrivateUsage] if isinstance(event, NodeGatheredInfo): event.when = str(datetime.now(tz=timezone.utc)) - await self._apply_and_broadcast(event) + self._event_log.append(event) + await self._send_event(indexed) + + async def _loopback_processor(self) -> None: + # this would ideally not be necessary. + # this is WAY less hacky than how I was working around this before + local_index = 0 + with self._loopback_event_receiver as events: + async for event in events: + await self._loopback_event_sender.send( + ForwarderEvent( + origin=NodeId(f"master_{self.node_id}"), + origin_idx=local_index, + session=self.session_id, + event=event, + ) + ) + local_index += 1 # This function is re-entrant, take care! async def _send_event(self, event: IndexedEvent): @@ -535,49 +453,10 @@ class Master: for trace_data in self._pending_traces[task_id].values(): all_trace_data.extend(trace_data) - await self._apply_and_broadcast( + await self.event_sender.send( TracesMerged(task_id=task_id, traces=all_trace_data) ) del self._pending_traces[task_id] if task_id in self._expected_ranks: del self._expected_ranks[task_id] - - async def _handle_jaccl_side_channel(self, event: JacclSideChannelData) -> None: - """Accumulate SideChannel contributions; when all runners for an instance - have submitted for the same sequence, emit JacclSideChannelGathered.""" - iid = event.instance_id - seq = event.sequence - - if iid not in self._jaccl_pending: - self._jaccl_pending[iid] = {} - if seq not in self._jaccl_pending[iid]: - self._jaccl_pending[iid][seq] = {} - self._jaccl_pending[iid][seq][event.runner_id] = event.data - - instance = self.state.instances.get(iid) - if instance is None: - logger.warning(f"JacclSideChannelData for unknown instance {iid}") - return - - expected_runners = set(instance.shard_assignments.runner_to_shard.keys()) - submitted = set(self._jaccl_pending[iid][seq].keys()) - - logger.info( - f"JACCL side channel: instance={iid} seq={seq} " - f"submitted={len(submitted)}/{len(expected_runners)}" - ) - - if submitted >= expected_runners: - gathered = dict(self._jaccl_pending[iid][seq]) - del self._jaccl_pending[iid][seq] - if not self._jaccl_pending[iid]: - del self._jaccl_pending[iid] - - await self._apply_and_broadcast( - JacclSideChannelGathered( - instance_id=iid, - sequence=seq, - gathered_data=gathered, - ) - ) diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index ab886c3e..cf31ca78 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -6,11 +6,11 @@ from typing import Sequence from exo.master.placement_utils import ( Cycle, filter_cycles_by_memory, - get_largest_cycles, get_mlx_jaccl_coordinators, get_mlx_jaccl_devices_matrix, get_mlx_ring_hosts_by_node, get_shard_assignments, + get_smallest_cycles, ) from exo.shared.models.model_cards import ModelId from exo.shared.topology import Topology @@ -106,27 +106,23 @@ def place_instance( "Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)" ) - largest_cycles = get_largest_cycles(cycles_with_sufficient_memory) + smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory) - largest_rdma_cycles = [ - cycle for cycle in largest_cycles if topology.is_rdma_cycle(cycle) + smallest_rdma_cycles = [ + cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle) ] - if command.instance_meta == InstanceMeta.MlxJaccl: - if not largest_rdma_cycles: - raise ValueError( - "Requested RDMA (MlxJaccl) but no RDMA-connected cycles available" - ) - largest_cycles = largest_rdma_cycles + if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []: + smallest_cycles = smallest_rdma_cycles cycles_with_leaf_nodes: list[Cycle] = [ cycle - for cycle in largest_cycles + for cycle in smallest_cycles if any(topology.node_is_leaf(node_id) for node_id in cycle) ] selected_cycle = max( - cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else largest_cycles, + cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles, key=lambda cycle: sum( (node_memory[node_id].ram_available for node_id in cycle), start=Memory(), diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index d47c40c0..b20a39cc 100644 --- a/src/exo/master/placement_utils.py +++ b/src/exo/master/placement_utils.py @@ -37,11 +37,11 @@ def filter_cycles_by_memory( return filtered_cycles -def get_largest_cycles( +def get_smallest_cycles( cycles: list[Cycle], ) -> list[Cycle]: - max_nodes = max(len(cycle) for cycle in cycles) - return [cycle for cycle in cycles if len(cycle) == max_nodes] + min_nodes = min(len(cycle) for cycle in cycles) + return [cycle for cycle in cycles if len(cycle) == min_nodes] def allocate_layers_proportionally( diff --git a/src/exo/master/process_managers/__init__.py b/src/exo/master/process_managers/__init__.py deleted file mode 100644 index f6a3ee0a..00000000 --- a/src/exo/master/process_managers/__init__.py +++ /dev/null @@ -1,12 +0,0 @@ -from collections.abc import Sequence -from typing import Protocol, runtime_checkable - -from exo.shared.types.events import Event -from exo.shared.types.state import State - - -@runtime_checkable -class ProcessManager(Protocol): - """A reconciliation step that examines state and returns corrective events.""" - - async def reconcile(self, state: State) -> Sequence[Event]: ... diff --git a/src/exo/master/process_managers/instance_health.py b/src/exo/master/process_managers/instance_health.py deleted file mode 100644 index f5f8e922..00000000 --- a/src/exo/master/process_managers/instance_health.py +++ /dev/null @@ -1,62 +0,0 @@ -from collections.abc import Sequence -from typing import final - -from loguru import logger - -from exo.master.reconcile import instance_connections_healthy, instance_runners_failed -from exo.shared.types.events import Event, InstanceDeleted, InstanceRetrying -from exo.shared.types.state import State - -MAX_INSTANCE_RETRIES = 3 - - -@final -class InstanceHealthReconciler: - """Delete instances whose network connections are broken or whose runners have all failed.""" - - async def reconcile(self, state: State) -> Sequence[Event]: - events: list[Event] = [] - for instance_id, instance in state.instances.items(): - if not instance_connections_healthy(instance, state.topology): - events.append( - InstanceDeleted( - instance_id=instance_id, - failure_error="Network connection lost", - ) - ) - continue - - is_failed, error_message = instance_runners_failed( - instance, state.runners, state.node_identities - ) - if is_failed: - # Retry within the same instance if backed by a MetaInstance - mid = instance.meta_instance_id - mi = state.meta_instances.get(mid) if mid else None - if mid and mi and mi.consecutive_failures < MAX_INSTANCE_RETRIES: - logger.info( - f"Instance {instance_id} failed (attempt" - f" {mi.consecutive_failures + 1}/{MAX_INSTANCE_RETRIES})," - f" retrying: {error_message}" - ) - events.append( - InstanceRetrying( - instance_id=instance_id, - meta_instance_id=mid, - failure_error=error_message or "Runner failed", - ) - ) - else: - if mid and mi: - logger.warning( - f"Instance {instance_id} exceeded retry limit" - f" ({MAX_INSTANCE_RETRIES}), deleting:" - f" {error_message}" - ) - events.append( - InstanceDeleted( - instance_id=instance_id, - failure_error=error_message, - ) - ) - return events diff --git a/src/exo/master/process_managers/meta_instance.py b/src/exo/master/process_managers/meta_instance.py deleted file mode 100644 index 93037ea8..00000000 --- a/src/exo/master/process_managers/meta_instance.py +++ /dev/null @@ -1,92 +0,0 @@ -from collections.abc import Sequence -from typing import final - -import anyio -from loguru import logger - -from exo.master.reconcile import ( - find_unsatisfied_meta_instances, - try_place_for_meta_instance, -) -from exo.shared.models.model_cards import ModelCard -from exo.shared.types.events import Event, InstanceCreated, MetaInstancePlacementFailed -from exo.shared.types.state import State -from exo.shared.types.worker.instances import Instance, InstanceId - -MODEL_CARD_LOAD_TIMEOUT_SECONDS = 10 - - -@final -class MetaInstanceReconciler: - """Place instances for unsatisfied MetaInstances.""" - - async def reconcile(self, state: State) -> Sequence[Event]: - all_events: list[Event] = [] - # Local copy for intermediate tracking — so placement of B - # sees A's instance and doesn't double-place on same resources. - current_instances: dict[InstanceId, Instance] = dict(state.instances) - - unsatisfied = find_unsatisfied_meta_instances( - state.meta_instances, - current_instances, - state.topology, - ) - for meta_instance in unsatisfied: - try: - with anyio.fail_after(MODEL_CARD_LOAD_TIMEOUT_SECONDS): - model_card = await ModelCard.load(meta_instance.model_id) - except TimeoutError: - logger.warning( - f"ModelCard.load timed out for {meta_instance.model_id}, skipping this cycle" - ) - continue - except Exception as exc: - logger.warning( - f"ModelCard.load failed for {meta_instance.model_id}: {exc}" - ) - error = f"Failed to load model card: {exc}" - if meta_instance.placement_error != error: - all_events.append( - MetaInstancePlacementFailed( - meta_instance_id=meta_instance.meta_instance_id, - reason=error, - ) - ) - continue - - result = try_place_for_meta_instance( - meta_instance, - model_card, - state.topology, - current_instances, - state.node_memory, - state.node_network, - state.tasks, - ) - # Update local instance map so next placement sees this one - for event in result.events: - if isinstance(event, InstanceCreated): - logger.info( - f"MetaInstance reconciler placed instance" - f" {event.instance.instance_id} for" - f" {meta_instance.model_id}" - ) - current_instances[event.instance.instance_id] = event.instance - all_events.extend(result.events) - - # Emit placement failure if error differs from what's already in state - if ( - result.error is not None - and meta_instance.placement_error != result.error - ): - logger.warning( - f"MetaInstance placement failed for" - f" {meta_instance.model_id}: {result.error}" - ) - all_events.append( - MetaInstancePlacementFailed( - meta_instance_id=meta_instance.meta_instance_id, - reason=result.error, - ) - ) - return all_events diff --git a/src/exo/master/process_managers/node_timeout.py b/src/exo/master/process_managers/node_timeout.py deleted file mode 100644 index 98045c25..00000000 --- a/src/exo/master/process_managers/node_timeout.py +++ /dev/null @@ -1,27 +0,0 @@ -from collections.abc import Sequence -from datetime import datetime, timedelta, timezone -from typing import final - -from loguru import logger - -from exo.shared.types.events import Event, NodeTimedOut -from exo.shared.types.state import State - -_DEFAULT_TIMEOUT = timedelta(seconds=30) - - -@final -class NodeTimeoutReconciler: - """Time out nodes that haven't been seen recently.""" - - def __init__(self, timeout: timedelta = _DEFAULT_TIMEOUT) -> None: - self.timeout = timeout - - async def reconcile(self, state: State) -> Sequence[Event]: - now = datetime.now(tz=timezone.utc) - events: list[Event] = [] - for node_id, last_seen in state.last_seen.items(): - if now - last_seen > self.timeout: - logger.info(f"Removing node {node_id} due to inactivity") - events.append(NodeTimedOut(node_id=node_id)) - return events diff --git a/src/exo/master/reconcile.py b/src/exo/master/reconcile.py deleted file mode 100644 index 4ca968b8..00000000 --- a/src/exo/master/reconcile.py +++ /dev/null @@ -1,244 +0,0 @@ -from collections.abc import Mapping, Sequence -from typing import NamedTuple - -from loguru import logger - -from exo.master.placement import get_transition_events, place_instance -from exo.shared.models.model_cards import ModelCard -from exo.shared.topology import Topology -from exo.shared.types.commands import PlaceInstance -from exo.shared.types.common import MetaInstanceId, NodeId -from exo.shared.types.events import Event -from exo.shared.types.meta_instance import MetaInstance -from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo -from exo.shared.types.tasks import Task, TaskId -from exo.shared.types.topology import RDMAConnection, SocketConnection -from exo.shared.types.worker.instances import ( - BaseInstance, - Instance, - InstanceId, - MlxJacclInstance, - MlxRingInstance, -) -from exo.shared.types.worker.runners import ( - RunnerFailed, - RunnerId, - RunnerShutdown, - RunnerStatus, -) - - -class PlacementResult(NamedTuple): - """Result of a placement attempt: events to apply and optional error reason.""" - - events: Sequence[Event] - error: str | None - - -def _get_ring_order(instance: BaseInstance) -> list[NodeId]: - """Reconstruct ring order from shard device_rank.""" - node_ranks: list[tuple[NodeId, int]] = [] - for node_id, runner_id in instance.shard_assignments.node_to_runner.items(): - shard = instance.shard_assignments.runner_to_shard[runner_id] - node_ranks.append((node_id, shard.device_rank)) - node_ranks.sort(key=lambda x: x[1]) - return [node_id for node_id, _ in node_ranks] - - -def _ring_connections_healthy(instance: MlxRingInstance, topology: Topology) -> bool: - """Check that the specific IPs used by a ring instance still exist in the topology.""" - ring = _get_ring_order(instance) - n = len(ring) - for node in ring: - hosts = instance.hosts_by_node[node] - for idx in range(n): - host = hosts[idx] - if host.ip in ("0.0.0.0", "198.51.100.1"): - continue # self or placeholder - # Real connection: node → ring[idx]. Check specific IP. - connections = topology.get_all_connections_between(node, ring[idx]) - if not any( - isinstance(c, SocketConnection) - and c.sink_multiaddr.ip_address == host.ip - for c in connections - ): - return False - return True - - -def _jaccl_connections_healthy(instance: MlxJacclInstance, topology: Topology) -> bool: - """Check that the specific RDMA interfaces used by a JACCL instance still exist.""" - ring = _get_ring_order(instance) - n = len(ring) - for i in range(n): - for j in range(n): - iface = instance.jaccl_devices[i][j] - if iface is None: - continue - connections = topology.get_all_connections_between(ring[i], ring[j]) - if not any( - isinstance(c, RDMAConnection) and c.source_rdma_iface == iface - for c in connections - ): - return False - return True - - -def instance_connections_healthy(instance: Instance, topology: Topology) -> bool: - """Check that an instance's nodes and specific connections are still in the topology.""" - instance_nodes = set(instance.shard_assignments.node_to_runner.keys()) - if not all(topology.contains_node(n) for n in instance_nodes): - return False - if len(instance_nodes) <= 1: - return True - match instance: - case MlxRingInstance(): - return _ring_connections_healthy(instance, topology) - case MlxJacclInstance(): - return _jaccl_connections_healthy(instance, topology) - - -def instance_runners_failed( - instance: Instance, - runners: Mapping[RunnerId, RunnerStatus], - node_identities: Mapping[NodeId, NodeIdentity], -) -> tuple[bool, str | None]: - """Check if an instance's runners have all reached terminal failure states. - - Returns ``(True, error_message)`` when ALL runners are terminal - (``RunnerFailed`` or ``RunnerShutdown``) and at least one is ``RunnerFailed``. - - Returns ``(False, None)`` when runners are still active, haven't reported - yet, or all gracefully shut down (no ``RunnerFailed``). - """ - instance_runner_ids = set(instance.shard_assignments.node_to_runner.values()) - - if not instance_runner_ids: - return False, None - - # Build reverse mapping: runner_id -> node_id - runner_to_node: dict[RunnerId, NodeId] = { - runner_id: node_id - for node_id, runner_id in instance.shard_assignments.node_to_runner.items() - } - - has_any_failed = False - error_messages: list[str] = [] - - for runner_id in instance_runner_ids: - status = runners.get(runner_id) - if status is None: - # Runner hasn't reported yet — instance is still starting - return False, None - if isinstance(status, RunnerFailed): - has_any_failed = True - if status.error_message: - node_id = runner_to_node.get(runner_id) - name = ( - node_identities[node_id].friendly_name - if node_id and node_id in node_identities - else node_id or "unknown" - ) - error_messages.append(f"{name}: {status.error_message}") - elif isinstance(status, RunnerShutdown): - pass # Terminal but not a failure indicator on its own - else: - # Runner is still active (connecting, loading, running, etc.) - return False, None - - if has_any_failed: - return True, "; ".join(error_messages) if error_messages else "Runner failed" - - # All runners are Shutdown but none Failed — graceful shutdown, not a failure - return False, None - - -def instance_satisfies_meta_instance( - meta_instance: MetaInstance, - instance: Instance, -) -> bool: - """Check if a single instance satisfies a meta-instance's constraints. - - This is a pure constraint check (model, min_nodes, node_ids). - Use ``instance_connections_healthy`` separately for topology health. - """ - if instance.shard_assignments.model_id != meta_instance.model_id: - return False - - instance_nodes = set(instance.shard_assignments.node_to_runner.keys()) - - if len(instance_nodes) < meta_instance.min_nodes: - return False - - return meta_instance.node_ids is None or set(meta_instance.node_ids).issubset( - instance_nodes - ) - - -def find_unsatisfied_meta_instances( - meta_instances: Mapping[MetaInstanceId, MetaInstance], - instances: Mapping[InstanceId, Instance], - topology: Topology, -) -> Sequence[MetaInstance]: - """Return meta-instances that have no healthy backing instance.""" - unsatisfied: list[MetaInstance] = [] - for meta_id, meta_instance in meta_instances.items(): - has_healthy_backing = any( - instance.meta_instance_id == meta_id - and instance_connections_healthy(instance, topology) - for instance in instances.values() - ) - if not has_healthy_backing: - unsatisfied.append(meta_instance) - return unsatisfied - - -def try_place_for_meta_instance( - meta_instance: MetaInstance, - model_card: ModelCard, - topology: Topology, - current_instances: Mapping[InstanceId, Instance], - node_memory: Mapping[NodeId, MemoryUsage], - node_network: Mapping[NodeId, NodeNetworkInfo], - tasks: Mapping[TaskId, Task], -) -> PlacementResult: - """Try to place an instance satisfying the meta-instance constraints. - - Returns a :class:`PlacementResult` with events on success, or an error - reason on failure. - """ - command = PlaceInstance( - model_card=model_card, - sharding=meta_instance.sharding, - instance_meta=meta_instance.instance_meta, - min_nodes=meta_instance.min_nodes, - ) - try: - target_instances = place_instance( - command, - topology, - current_instances, - node_memory, - node_network, - required_nodes=( - set(meta_instance.node_ids) if meta_instance.node_ids else None - ), - ) - # Tag the new instance with meta_instance_id - new_instance_ids = set(target_instances.keys()) - set(current_instances.keys()) - if new_instance_ids: - new_id = next(iter(new_instance_ids)) - target_instances[new_id] = target_instances[new_id].model_copy( - update={"meta_instance_id": meta_instance.meta_instance_id} - ) - return PlacementResult( - events=list( - get_transition_events(current_instances, target_instances, tasks) - ), - error=None, - ) - except ValueError as e: - logger.debug( - f"MetaInstance placement not possible for {meta_instance.model_id}: {e}" - ) - return PlacementResult(events=[], error=str(e)) diff --git a/src/exo/master/tests/test_meta_instance_edge_cases.py b/src/exo/master/tests/test_meta_instance_edge_cases.py deleted file mode 100644 index 26554834..00000000 --- a/src/exo/master/tests/test_meta_instance_edge_cases.py +++ /dev/null @@ -1,778 +0,0 @@ -"""Edge-case and regression tests for MetaInstance lifecycle, concurrent operations, and error handling.""" - -import pytest - -from exo.master.process_managers.instance_health import ( - MAX_INSTANCE_RETRIES, - InstanceHealthReconciler, -) -from exo.master.process_managers.meta_instance import MetaInstanceReconciler -from exo.master.reconcile import ( - find_unsatisfied_meta_instances, - instance_connections_healthy, - instance_runners_failed, - instance_satisfies_meta_instance, -) -from exo.shared.apply import apply -from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask -from exo.shared.topology import Topology -from exo.shared.types.common import Host, MetaInstanceId, NodeId -from exo.shared.types.events import ( - IndexedEvent, - InstanceCreated, - InstanceDeleted, - InstanceRetrying, - MetaInstanceCreated, - MetaInstanceDeleted, - MetaInstancePlacementFailed, - TaskStatusUpdated, -) -from exo.shared.types.memory import Memory -from exo.shared.types.meta_instance import MetaInstance -from exo.shared.types.multiaddr import Multiaddr -from exo.shared.types.profiling import NodeIdentity -from exo.shared.types.state import State -from exo.shared.types.tasks import LoadModel, TaskId, TaskStatus -from exo.shared.types.topology import Connection, SocketConnection -from exo.shared.types.worker.instances import ( - InstanceId, - MlxRingInstance, -) -from exo.shared.types.worker.runners import ( - RunnerFailed, - RunnerId, - RunnerReady, - ShardAssignments, -) -from exo.shared.types.worker.shards import PipelineShardMetadata - -# --- Helpers (copied from test_reconcile.py for independence) --- - - -def _model_card(model_id: str = "test-org/test-model") -> ModelCard: - return ModelCard( - model_id=ModelId(model_id), - storage_size=Memory.from_kb(1000), - n_layers=10, - hidden_size=30, - supports_tensor=True, - tasks=[ModelTask.TextGeneration], - ) - - -def _topology(*node_ids: str, connect: bool = True) -> Topology: - t = Topology() - nodes = [NodeId(n) for n in node_ids] - for n in nodes: - t.add_node(n) - if connect and len(nodes) > 1: - for i in range(len(nodes)): - j = (i + 1) % len(nodes) - t.add_connection( - Connection( - source=nodes[i], - sink=nodes[j], - edge=SocketConnection( - sink_multiaddr=Multiaddr( - address=f"/ip4/10.0.0.{j + 1}/tcp/50000" - ) - ), - ) - ) - t.add_connection( - Connection( - source=nodes[j], - sink=nodes[i], - edge=SocketConnection( - sink_multiaddr=Multiaddr( - address=f"/ip4/10.0.0.{i + 1}/tcp/50000" - ) - ), - ) - ) - return t - - -def _meta_instance( - model_id: str = "test-org/test-model", - *, - min_nodes: int = 1, - node_ids: list[NodeId] | None = None, - meta_instance_id: MetaInstanceId | None = None, - consecutive_failures: int = 0, - last_failure_error: str | None = None, - placement_error: str | None = None, -) -> MetaInstance: - return MetaInstance( - meta_instance_id=meta_instance_id or MetaInstanceId(), - model_id=ModelId(model_id), - min_nodes=min_nodes, - node_ids=node_ids, - consecutive_failures=consecutive_failures, - last_failure_error=last_failure_error, - placement_error=placement_error, - ) - - -def _instance( - model_id: str = "test-org/test-model", - node_ids: list[str] | None = None, - instance_id: InstanceId | None = None, - meta_instance_id: MetaInstanceId | None = None, -) -> tuple[InstanceId, MlxRingInstance]: - iid = instance_id or InstanceId() - nodes = node_ids or ["node-a"] - n = len(nodes) - mc = _model_card(model_id) - ephemeral_port = 50000 - node_to_runner = {NodeId(nd): RunnerId() for nd in nodes} - runner_to_shard = { - runner_id: PipelineShardMetadata( - model_card=mc, - device_rank=i, - world_size=n, - start_layer=0, - end_layer=mc.n_layers, - n_layers=mc.n_layers, - ) - for i, runner_id in enumerate(node_to_runner.values()) - } - hosts_by_node: dict[NodeId, list[Host]] = {} - for r, node_str in enumerate(nodes): - hosts: list[Host] = [] - for idx in range(n): - if idx == r: - hosts.append(Host(ip="0.0.0.0", port=ephemeral_port)) - elif n > 1 and idx in ((r - 1) % n, (r + 1) % n): - hosts.append(Host(ip=f"10.0.0.{idx + 1}", port=ephemeral_port)) - else: - hosts.append(Host(ip="198.51.100.1", port=0)) - hosts_by_node[NodeId(node_str)] = hosts - return iid, MlxRingInstance( - instance_id=iid, - shard_assignments=ShardAssignments( - model_id=ModelId(model_id), - runner_to_shard=runner_to_shard, - node_to_runner=node_to_runner, - ), - hosts_by_node=hosts_by_node, - ephemeral_port=ephemeral_port, - meta_instance_id=meta_instance_id, - ) - - -# ============================================================================= -# 1. MetaInstance lifecycle edge cases -# ============================================================================= - - -def test_meta_instance_model_is_frozen(): - """MetaInstance should be immutable (frozen model).""" - meta = _meta_instance() - try: - meta.model_id = ModelId("something-else") - raise AssertionError("Should have raised") - except Exception: - pass # Expected — frozen model - - -def test_meta_instance_created_then_deleted_roundtrip(): - """Create and delete a MetaInstance through apply — state should be clean.""" - state = State() - meta = _meta_instance() - state = apply( - state, IndexedEvent(idx=0, event=MetaInstanceCreated(meta_instance=meta)) - ) - assert meta.meta_instance_id in state.meta_instances - state = apply( - state, - IndexedEvent( - idx=1, event=MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id) - ), - ) - assert meta.meta_instance_id not in state.meta_instances - assert len(state.meta_instances) == 0 - - -def test_delete_nonexistent_meta_instance_is_safe(): - """Deleting a MetaInstance that doesn't exist should not crash.""" - state = State() - event = MetaInstanceDeleted(meta_instance_id=MetaInstanceId("nonexistent")) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert len(new_state.meta_instances) == 0 - - -def test_placement_failed_for_nonexistent_meta_instance_is_safe(): - """MetaInstancePlacementFailed for unknown ID should not crash.""" - state = State() - event = MetaInstancePlacementFailed( - meta_instance_id=MetaInstanceId("nonexistent"), - reason="test", - ) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert len(new_state.meta_instances) == 0 - - -def test_multiple_meta_instances_for_same_model(): - """Multiple MetaInstances for the same model are tracked independently.""" - state = State() - meta_a = _meta_instance("test-org/model-x") - meta_b = _meta_instance("test-org/model-x") - state = apply( - state, IndexedEvent(idx=0, event=MetaInstanceCreated(meta_instance=meta_a)) - ) - state = apply( - state, IndexedEvent(idx=1, event=MetaInstanceCreated(meta_instance=meta_b)) - ) - assert len(state.meta_instances) == 2 - assert meta_a.meta_instance_id in state.meta_instances - assert meta_b.meta_instance_id in state.meta_instances - - -# ============================================================================= -# 2. Retry logic edge cases -# ============================================================================= - - -def test_retry_counter_resets_on_successful_instance_creation(): - """When a new instance is created for a meta-instance, failures should reset.""" - meta = _meta_instance(consecutive_failures=2, last_failure_error="old") - _, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State(meta_instances={meta.meta_instance_id: meta}) - state = apply(state, IndexedEvent(idx=0, event=InstanceCreated(instance=inst))) - mi = state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 0 - # last_failure_error is preserved (for UI display) - assert mi.last_failure_error == "old" - - -async def test_retry_count_increments_through_full_cycle(): - """Walk through MAX_INSTANCE_RETRIES worth of retries, then verify delete.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - topology = _topology("node-a") - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - topology=topology, - ) - - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - for idx, i in enumerate(range(MAX_INSTANCE_RETRIES)): - # Simulate runners failing - state_with_runners = state.model_copy( - update={"runners": {runner_ids[0]: RunnerFailed(error_message=f"fail-{i}")}} - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state_with_runners) - assert len(events) == 1 - assert isinstance(events[0], InstanceRetrying), f"iteration {i}" - state = apply(state, IndexedEvent(idx=idx, event=events[0])) - - # After MAX_INSTANCE_RETRIES retries, failure counter should be at max - mi = state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == MAX_INSTANCE_RETRIES - - # Next failure should result in deletion - state_with_runners = state.model_copy( - update={"runners": {runner_ids[0]: RunnerFailed(error_message="final")}} - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state_with_runners) - assert len(events) == 1 - assert isinstance(events[0], InstanceDeleted) - - -async def test_health_reconciler_respects_exact_limit(): - """At exactly MAX_INSTANCE_RETRIES, reconciler should delete, not retry.""" - meta = _meta_instance(consecutive_failures=MAX_INSTANCE_RETRIES) - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, - topology=_topology("node-a"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceDeleted) - - -async def test_health_reconciler_at_limit_minus_one_retries(): - """At MAX_INSTANCE_RETRIES - 1, reconciler should still retry.""" - meta = _meta_instance(consecutive_failures=MAX_INSTANCE_RETRIES - 1) - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, - topology=_topology("node-a"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceRetrying) - - -# ============================================================================= -# 3. Error handling edge cases -# ============================================================================= - - -def test_runners_failed_with_empty_error_message(): - """RunnerFailed with empty error_message should still report as failed.""" - _, inst = _instance(node_ids=["node-a"]) - runners = { - rid: RunnerFailed(error_message="") - for rid in inst.shard_assignments.node_to_runner.values() - } - is_failed, error = instance_runners_failed(inst, runners, {}) - assert is_failed is True - # Empty error message means we get the fallback - assert error == "Runner failed" - - -def test_runners_failed_with_none_error_message(): - """RunnerFailed with None error_message should still report as failed.""" - _, inst = _instance(node_ids=["node-a"]) - runners = { - rid: RunnerFailed(error_message=None) - for rid in inst.shard_assignments.node_to_runner.values() - } - is_failed, error = instance_runners_failed(inst, runners, {}) - assert is_failed is True - assert error == "Runner failed" - - -def test_runners_failed_collects_all_error_messages(): - """With multiple failed runners, all error messages should be collected.""" - _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - runners = { - runner_ids[0]: RunnerFailed(error_message="OOM on GPU 0"), - runner_ids[1]: RunnerFailed(error_message="OOM on GPU 1"), - runner_ids[2]: RunnerFailed(error_message="OOM on GPU 2"), - } - is_failed, error = instance_runners_failed(inst, runners, {}) - assert is_failed is True - assert error is not None - assert "OOM on GPU 0" in error - assert "OOM on GPU 1" in error - assert "OOM on GPU 2" in error - - -def test_runners_failed_includes_friendly_name(): - """Error messages should include node friendly names when available.""" - _, inst = _instance(node_ids=["node-a"]) - node_id = NodeId("node-a") - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - runners = {runner_ids[0]: RunnerFailed(error_message="OOM")} - identities = {node_id: NodeIdentity(friendly_name="My Mac Studio")} - is_failed, error = instance_runners_failed(inst, runners, identities) - assert is_failed is True - assert error is not None - assert "My Mac Studio" in error - - -def test_instance_retrying_for_missing_instance_is_safe(): - """InstanceRetrying for an instance not in state should not crash. - - NOTE: When the instance is missing, the handler returns early WITHOUT - incrementing the MetaInstance failure counter. This means stale retry - events for already-deleted instances are silently dropped. This is - acceptable since the InstanceDeleted handler already increments failures. - """ - meta = _meta_instance() - state = State(meta_instances={meta.meta_instance_id: meta}) - event = InstanceRetrying( - instance_id=InstanceId("nonexistent"), - meta_instance_id=meta.meta_instance_id, - failure_error="crash", - ) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - # Does not crash, but failure count is NOT incremented (early return) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 0 - - -# ============================================================================= -# 4. Backward compatibility -# ============================================================================= - - -def test_instance_without_meta_instance_id_works(): - """Instances created without meta_instance_id should still function normally.""" - _, inst = _instance(node_ids=["node-a"]) - assert inst.meta_instance_id is None - topology = _topology("node-a") - assert instance_connections_healthy(inst, topology) is True - - -def test_instance_deleted_without_meta_does_not_affect_meta_instances(): - """Deleting an instance without meta_instance_id should not affect meta_instances.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"]) # no meta_instance_id - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - event = InstanceDeleted(instance_id=iid, failure_error="crash") - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 0 # unchanged - - -def test_satisfies_ignores_meta_instance_id_binding(): - """instance_satisfies_meta_instance checks constraints only, not binding.""" - meta = _meta_instance() - _, inst = _instance(node_ids=["node-a"]) # no meta_instance_id set - # Should match on constraints (model, min_nodes) regardless of binding - assert instance_satisfies_meta_instance(meta, inst) is True - - -def test_find_unsatisfied_uses_binding_not_constraints(): - """find_unsatisfied checks meta_instance_id binding, not just constraint matching.""" - meta = _meta_instance() - # Instance matches constraints but is NOT bound to this meta_instance - iid, inst = _instance(node_ids=["node-a"]) - topology = _topology("node-a") - result = find_unsatisfied_meta_instances( - {meta.meta_instance_id: meta}, {iid: inst}, topology - ) - # Should be unsatisfied because instance.meta_instance_id != meta.meta_instance_id - assert list(result) == [meta] - - -# ============================================================================= -# 5. Concurrent / multi-instance scenarios -# ============================================================================= - - -async def test_health_reconciler_handles_multiple_failing_instances(): - """Multiple instances failing simultaneously should each get their own event.""" - meta_a = _meta_instance() - meta_b = _meta_instance() - iid_a, inst_a = _instance( - node_ids=["node-a"], meta_instance_id=meta_a.meta_instance_id - ) - iid_b, inst_b = _instance( - node_ids=["node-b"], meta_instance_id=meta_b.meta_instance_id - ) - runner_ids_a = list(inst_a.shard_assignments.node_to_runner.values()) - runner_ids_b = list(inst_b.shard_assignments.node_to_runner.values()) - state = State( - meta_instances={ - meta_a.meta_instance_id: meta_a, - meta_b.meta_instance_id: meta_b, - }, - instances={iid_a: inst_a, iid_b: inst_b}, - runners={ - runner_ids_a[0]: RunnerFailed(error_message="OOM"), - runner_ids_b[0]: RunnerFailed(error_message="OOM"), - }, - topology=_topology("node-a", "node-b"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 2 - # Both should be InstanceRetrying since failures < MAX - assert all(isinstance(e, InstanceRetrying) for e in events) - instance_ids = {e.instance_id for e in events} # type: ignore[union-attr] - assert instance_ids == {iid_a, iid_b} - - -async def test_health_reconciler_mixed_healthy_and_failing(): - """Only failing instances should produce events; healthy ones should not.""" - meta_healthy = _meta_instance() - meta_failing = _meta_instance() - iid_h, inst_h = _instance( - node_ids=["node-a"], meta_instance_id=meta_healthy.meta_instance_id - ) - iid_f, inst_f = _instance( - node_ids=["node-b"], meta_instance_id=meta_failing.meta_instance_id - ) - runner_ids_h = list(inst_h.shard_assignments.node_to_runner.values()) - runner_ids_f = list(inst_f.shard_assignments.node_to_runner.values()) - state = State( - meta_instances={ - meta_healthy.meta_instance_id: meta_healthy, - meta_failing.meta_instance_id: meta_failing, - }, - instances={iid_h: inst_h, iid_f: inst_f}, - runners={ - runner_ids_h[0]: RunnerReady(), - runner_ids_f[0]: RunnerFailed(error_message="crash"), - }, - topology=_topology("node-a", "node-b"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceRetrying) - assert events[0].instance_id == iid_f - - -async def test_meta_instance_reconciler_empty_state(): - """MetaInstanceReconciler with no meta_instances should produce no events.""" - state = State() - reconciler = MetaInstanceReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 0 - - -# ============================================================================= -# 6. Placement error tracking -# ============================================================================= - - -def test_placement_failed_sets_error(): - """MetaInstancePlacementFailed should set placement_error on the MetaInstance.""" - meta = _meta_instance() - state = State(meta_instances={meta.meta_instance_id: meta}) - event = MetaInstancePlacementFailed( - meta_instance_id=meta.meta_instance_id, - reason="Not enough memory", - ) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.placement_error == "Not enough memory" - - -def test_instance_created_clears_placement_error(): - """InstanceCreated should clear placement_error on the MetaInstance.""" - meta = _meta_instance(placement_error="Not enough memory") - _, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State(meta_instances={meta.meta_instance_id: meta}) - state = apply(state, IndexedEvent(idx=0, event=InstanceCreated(instance=inst))) - mi = state.meta_instances[meta.meta_instance_id] - assert mi.placement_error is None - - -def test_placement_error_does_not_increment_failures(): - """Placement failures should only set placement_error, not increment consecutive_failures.""" - meta = _meta_instance() - state = State(meta_instances={meta.meta_instance_id: meta}) - event = MetaInstancePlacementFailed( - meta_instance_id=meta.meta_instance_id, - reason="No resources", - ) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 0 - assert mi.placement_error == "No resources" - - -# ============================================================================= -# 7. State serialization roundtrip -# ============================================================================= - - -def test_state_with_meta_instances_serializes(): - """State with meta_instances should serialize and deserialize correctly.""" - meta = _meta_instance(consecutive_failures=2, last_failure_error="test") - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - json_str = state.model_dump_json() - restored = State.model_validate_json(json_str) - assert meta.meta_instance_id in restored.meta_instances - mi = restored.meta_instances[meta.meta_instance_id] - assert mi.model_id == meta.model_id - assert mi.consecutive_failures == 2 - assert mi.last_failure_error == "test" - assert iid in restored.instances - assert restored.instances[iid].meta_instance_id == meta.meta_instance_id - - -# ============================================================================= -# 8. MetaInstanceReconciler error handling -# ============================================================================= - - -async def test_meta_instance_reconciler_model_load_error_emits_placement_failed( - monkeypatch: "pytest.MonkeyPatch", -): - """When ModelCard.load raises, reconciler emits MetaInstancePlacementFailed.""" - import exo.master.process_managers.meta_instance as mi_mod - - meta = _meta_instance() - topo = _topology("node-a") - state = State( - meta_instances={meta.meta_instance_id: meta}, - topology=topo, - ) - - async def _failing_load(_model_id: ModelId) -> ModelCard: - raise RuntimeError("Network error") - - monkeypatch.setattr( - mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_failing_load)}) - ) - - reconciler = MetaInstanceReconciler() - events = await reconciler.reconcile(state) - - placement_failed = [e for e in events if isinstance(e, MetaInstancePlacementFailed)] - assert len(placement_failed) == 1 - assert "Failed to load model card" in placement_failed[0].reason - assert meta.meta_instance_id == placement_failed[0].meta_instance_id - - -async def test_meta_instance_reconciler_model_load_error_skips_dedup( - monkeypatch: "pytest.MonkeyPatch", -): - """When ModelCard.load error matches existing placement_error, no duplicate event.""" - import exo.master.process_managers.meta_instance as mi_mod - - meta = _meta_instance(placement_error="Failed to load model card: Network error") - topo = _topology("node-a") - state = State( - meta_instances={meta.meta_instance_id: meta}, - topology=topo, - ) - - async def _failing_load(_model_id: ModelId) -> ModelCard: - raise RuntimeError("Network error") - - monkeypatch.setattr( - mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_failing_load)}) - ) - - reconciler = MetaInstanceReconciler() - events = await reconciler.reconcile(state) - - # Error matches existing placement_error, so no duplicate event emitted - assert len(events) == 0 - - -async def test_meta_instance_reconciler_continues_after_error( - monkeypatch: "pytest.MonkeyPatch", -): - """Reconciler should continue to next meta-instance after one fails to load.""" - import exo.master.process_managers.meta_instance as mi_mod - - meta_a = _meta_instance(model_id="org/model-a") - meta_b = _meta_instance(model_id="org/model-b") - topo = _topology("node-a") - state = State( - meta_instances={ - meta_a.meta_instance_id: meta_a, - meta_b.meta_instance_id: meta_b, - }, - topology=topo, - ) - - call_count = 0 - - async def _load_second_fails(model_id: ModelId) -> ModelCard: - nonlocal call_count - call_count += 1 - raise RuntimeError(f"Cannot load {model_id}") - - monkeypatch.setattr( - mi_mod, "ModelCard", type("MC", (), {"load": staticmethod(_load_second_fails)}) - ) - - reconciler = MetaInstanceReconciler() - events = await reconciler.reconcile(state) - - # Both meta-instances should have been attempted (not short-circuited) - assert call_count == 2 - # Both should have placement failed events - placement_failed = [e for e in events if isinstance(e, MetaInstancePlacementFailed)] - assert len(placement_failed) == 2 - - -# ============================================================================= -# 8. Cascade delete with task cancellation -# ============================================================================= - - -def test_cascade_delete_cancels_active_tasks(): - """Deleting a MetaInstance should cancel tasks on backing instances. - - Regression test: previously, cascade-deleting backing instances via - DeleteMetaInstance did not emit TaskStatusUpdated(Cancelled) for active - tasks, leaving orphaned task references in state. - """ - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - task_id = TaskId() - task = LoadModel(task_id=task_id, instance_id=iid, task_status=TaskStatus.Running) - - # Build state with meta-instance, backing instance, and active task - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - tasks={task_id: task}, - topology=_topology("node-a"), - ) - - # Simulate the cascade-delete event sequence produced by main.py: - # 1. MetaInstanceDeleted - # 2. TaskStatusUpdated(Cancelled) for active tasks - # 3. InstanceDeleted - idx = 0 - state = apply( - state, - IndexedEvent( - idx=idx, - event=MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id), - ), - ) - idx += 1 - state = apply( - state, - IndexedEvent( - idx=idx, - event=TaskStatusUpdated(task_id=task_id, task_status=TaskStatus.Cancelled), - ), - ) - idx += 1 - state = apply( - state, - IndexedEvent(idx=idx, event=InstanceDeleted(instance_id=iid)), - ) - - # Verify everything is cleaned up - assert len(state.meta_instances) == 0 - assert len(state.instances) == 0 - assert state.tasks[task_id].task_status == TaskStatus.Cancelled - - -def test_cascade_delete_skips_completed_tasks(): - """Cascade delete should only cancel Pending/Running tasks, not completed ones.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - - running_task_id = TaskId() - completed_task_id = TaskId() - running_task = LoadModel( - task_id=running_task_id, instance_id=iid, task_status=TaskStatus.Running - ) - completed_task = LoadModel( - task_id=completed_task_id, instance_id=iid, task_status=TaskStatus.Complete - ) - - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - tasks={running_task_id: running_task, completed_task_id: completed_task}, - topology=_topology("node-a"), - ) - - # Only the running task should be cancelled — we verify the logic pattern - # by checking which tasks are Pending or Running - active_tasks = [ - t - for t in state.tasks.values() - if t.instance_id == iid - and t.task_status in (TaskStatus.Pending, TaskStatus.Running) - ] - assert len(active_tasks) == 1 - assert active_tasks[0].task_id == running_task_id diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py index 9c7ebada..245c4fd7 100644 --- a/src/exo/master/tests/test_placement_utils.py +++ b/src/exo/master/tests/test_placement_utils.py @@ -3,10 +3,10 @@ import pytest from exo.master.placement_utils import ( allocate_layers_proportionally, filter_cycles_by_memory, - get_largest_cycles, get_mlx_jaccl_coordinators, get_shard_assignments, get_shard_assignments_for_pipeline_parallel, + get_smallest_cycles, ) from exo.master.tests.conftest import ( create_node_memory, @@ -143,7 +143,7 @@ def test_filter_multiple_cycles_by_memory(): } -def test_get_largest_cycles(): +def test_get_smallest_cycles(): # arrange node_a_id = NodeId() node_b_id = NodeId() @@ -175,12 +175,12 @@ def test_get_largest_cycles(): cycles = [c for c in topology.get_cycles() if len(c) != 1] # ignore singletons # act - largest_cycles = get_largest_cycles(cycles) + smallest_cycles = get_smallest_cycles(cycles) # assert - assert len(largest_cycles) == 1 - assert len(largest_cycles[0]) == 3 - assert set(n for n in largest_cycles[0]) == {node_a_id, node_b_id, node_c_id} + assert len(smallest_cycles) == 1 + assert len(smallest_cycles[0]) == 2 + assert set(n for n in smallest_cycles[0]) == {node_a_id, node_b_id} @pytest.mark.parametrize( diff --git a/src/exo/master/tests/test_reconcile.py b/src/exo/master/tests/test_reconcile.py deleted file mode 100644 index e2d6e776..00000000 --- a/src/exo/master/tests/test_reconcile.py +++ /dev/null @@ -1,742 +0,0 @@ -from exo.master.process_managers.instance_health import InstanceHealthReconciler -from exo.master.reconcile import ( - find_unsatisfied_meta_instances, - instance_connections_healthy, - instance_runners_failed, - instance_satisfies_meta_instance, -) -from exo.shared.apply import apply -from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask -from exo.shared.topology import Topology -from exo.shared.types.common import Host, MetaInstanceId, NodeId -from exo.shared.types.events import ( - IndexedEvent, - InstanceCreated, - InstanceDeleted, - InstanceRetrying, - MetaInstanceCreated, - MetaInstanceDeleted, -) -from exo.shared.types.memory import Memory -from exo.shared.types.meta_instance import MetaInstance -from exo.shared.types.multiaddr import Multiaddr -from exo.shared.types.state import State -from exo.shared.types.topology import Connection, SocketConnection -from exo.shared.types.worker.instances import ( - InstanceId, - MlxRingInstance, -) -from exo.shared.types.worker.runners import ( - RunnerFailed, - RunnerId, - RunnerLoading, - RunnerReady, - RunnerShutdown, - ShardAssignments, -) -from exo.shared.types.worker.shards import PipelineShardMetadata - - -def _model_card(model_id: str = "test-org/test-model") -> ModelCard: - return ModelCard( - model_id=ModelId(model_id), - storage_size=Memory.from_kb(1000), - n_layers=10, - hidden_size=30, - supports_tensor=True, - tasks=[ModelTask.TextGeneration], - ) - - -def _topology(*node_ids: str, connect: bool = True) -> Topology: - """Build a topology with nodes connected in a bidirectional ring with unique IPs. - - Node at index ``i`` gets IP ``10.0.0.{i+1}``. Edges go in both directions - between consecutive nodes (including wrap-around). - """ - t = Topology() - nodes = [NodeId(n) for n in node_ids] - for n in nodes: - t.add_node(n) - if connect and len(nodes) > 1: - for i in range(len(nodes)): - j = (i + 1) % len(nodes) - t.add_connection( - Connection( - source=nodes[i], - sink=nodes[j], - edge=SocketConnection( - sink_multiaddr=Multiaddr( - address=f"/ip4/10.0.0.{j + 1}/tcp/50000" - ) - ), - ) - ) - t.add_connection( - Connection( - source=nodes[j], - sink=nodes[i], - edge=SocketConnection( - sink_multiaddr=Multiaddr( - address=f"/ip4/10.0.0.{i + 1}/tcp/50000" - ) - ), - ) - ) - return t - - -def _meta_instance( - model_id: str = "test-org/test-model", - *, - min_nodes: int = 1, - node_ids: list[NodeId] | None = None, - meta_instance_id: MetaInstanceId | None = None, -) -> MetaInstance: - return MetaInstance( - meta_instance_id=meta_instance_id or MetaInstanceId(), - model_id=ModelId(model_id), - min_nodes=min_nodes, - node_ids=node_ids, - ) - - -def _instance( - model_id: str = "test-org/test-model", - node_ids: list[str] | None = None, - instance_id: InstanceId | None = None, - meta_instance_id: MetaInstanceId | None = None, -) -> tuple[InstanceId, MlxRingInstance]: - """Create a test instance with hosts_by_node matching ``_topology()`` IPs.""" - iid = instance_id or InstanceId() - nodes = node_ids or ["node-a"] - n = len(nodes) - mc = _model_card(model_id) - ephemeral_port = 50000 - node_to_runner = {NodeId(nd): RunnerId() for nd in nodes} - runner_to_shard = { - runner_id: PipelineShardMetadata( - model_card=mc, - device_rank=i, - world_size=n, - start_layer=0, - end_layer=mc.n_layers, - n_layers=mc.n_layers, - ) - for i, runner_id in enumerate(node_to_runner.values()) - } - # Build hosts_by_node with IPs matching _topology() convention: - # node at index idx has IP 10.0.0.{idx+1} - hosts_by_node: dict[NodeId, list[Host]] = {} - for r, node_str in enumerate(nodes): - hosts: list[Host] = [] - for idx in range(n): - if idx == r: - hosts.append(Host(ip="0.0.0.0", port=ephemeral_port)) - elif n > 1 and idx in ((r - 1) % n, (r + 1) % n): - hosts.append(Host(ip=f"10.0.0.{idx + 1}", port=ephemeral_port)) - else: - hosts.append(Host(ip="198.51.100.1", port=0)) - hosts_by_node[NodeId(node_str)] = hosts - return iid, MlxRingInstance( - instance_id=iid, - shard_assignments=ShardAssignments( - model_id=ModelId(model_id), - runner_to_shard=runner_to_shard, - node_to_runner=node_to_runner, - ), - hosts_by_node=hosts_by_node, - ephemeral_port=ephemeral_port, - meta_instance_id=meta_instance_id, - ) - - -# --- instance_satisfies_meta_instance (pure constraint matching) --- - - -def test_satisfies_matching_model(): - meta = _meta_instance() - _, inst = _instance(node_ids=["node-a"]) - assert instance_satisfies_meta_instance(meta, inst) is True - - -def test_not_satisfies_wrong_model(): - meta = _meta_instance("test-org/model-a") - _, inst = _instance("test-org/model-b") - assert instance_satisfies_meta_instance(meta, inst) is False - - -def test_not_satisfies_missing_required_node(): - meta = _meta_instance(node_ids=[NodeId("node-c")]) - _, inst = _instance(node_ids=["node-a", "node-b"]) - assert instance_satisfies_meta_instance(meta, inst) is False - - -def test_not_satisfies_fewer_than_min_nodes(): - meta = _meta_instance(min_nodes=3) - _, inst = _instance(node_ids=["node-a", "node-b"]) - assert instance_satisfies_meta_instance(meta, inst) is False - - -def test_satisfies_with_node_ids_specified(): - meta = _meta_instance(node_ids=[NodeId("node-a"), NodeId("node-b")], min_nodes=2) - _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) - assert instance_satisfies_meta_instance(meta, inst) is True - - -# --- instance_connections_healthy --- - - -def test_healthy_single_node_present(): - _, inst = _instance(node_ids=["node-a"]) - topology = _topology("node-a") - assert instance_connections_healthy(inst, topology) is True - - -def test_unhealthy_single_node_missing(): - _, inst = _instance(node_ids=["node-a"]) - topology = Topology() # empty - assert instance_connections_healthy(inst, topology) is False - - -def test_healthy_two_node_ring(): - _, inst = _instance(node_ids=["node-a", "node-b"]) - topology = _topology("node-a", "node-b") - assert instance_connections_healthy(inst, topology) is True - - -def test_unhealthy_two_node_edge_removed(): - """Nodes present but edge removed — ring broken.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - topology = _topology("node-a", "node-b", connect=False) - assert instance_connections_healthy(inst, topology) is False - - -def test_unhealthy_two_node_ip_changed(): - """Edge exists but with a different IP than instance was configured with.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - # Build topology with different IPs than _instance() expects - topology = Topology() - topology.add_node(NodeId("node-a")) - topology.add_node(NodeId("node-b")) - topology.add_connection( - Connection( - source=NodeId("node-a"), - sink=NodeId("node-b"), - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/192.168.99.99/tcp/50000") - ), - ) - ) - topology.add_connection( - Connection( - source=NodeId("node-b"), - sink=NodeId("node-a"), - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/192.168.99.98/tcp/50000") - ), - ) - ) - assert instance_connections_healthy(inst, topology) is False - - -def test_healthy_three_node_ring(): - _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) - topology = _topology("node-a", "node-b", "node-c") - assert instance_connections_healthy(inst, topology) is True - - -def test_unhealthy_three_node_one_edge_removed(): - """Remove one edge from a three-node ring — instance unhealthy.""" - _, inst = _instance(node_ids=["node-a", "node-b", "node-c"]) - # Build topology with one direction of one edge missing - topology = Topology() - nodes = [NodeId("node-a"), NodeId("node-b"), NodeId("node-c")] - for n in nodes: - topology.add_node(n) - # Add all edges except node-a → node-b - topology.add_connection( - Connection( - source=nodes[1], - sink=nodes[0], - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/50000") - ), - ) - ) - topology.add_connection( - Connection( - source=nodes[1], - sink=nodes[2], - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/10.0.0.3/tcp/50000") - ), - ) - ) - topology.add_connection( - Connection( - source=nodes[2], - sink=nodes[1], - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/10.0.0.2/tcp/50000") - ), - ) - ) - topology.add_connection( - Connection( - source=nodes[2], - sink=nodes[0], - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/50000") - ), - ) - ) - topology.add_connection( - Connection( - source=nodes[0], - sink=nodes[2], - edge=SocketConnection( - sink_multiaddr=Multiaddr(address="/ip4/10.0.0.3/tcp/50000") - ), - ) - ) - # Missing: node-a → node-b (ip 10.0.0.2) - assert instance_connections_healthy(inst, topology) is False - - -def test_unhealthy_node_missing_from_topology(): - """Instance has a node that's not in the topology at all.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - topology = _topology("node-a") # node-b not present - assert instance_connections_healthy(inst, topology) is False - - -def test_healthy_extra_nodes_in_topology(): - """Extra nodes in topology don't affect instance health.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - topology = _topology("node-a", "node-b", "node-c") - assert instance_connections_healthy(inst, topology) is True - - -# --- find_unsatisfied_meta_instances --- - - -def test_unsatisfied_no_meta_instances(): - result = find_unsatisfied_meta_instances({}, {}, Topology()) - assert list(result) == [] - - -def test_unsatisfied_one_satisfied(): - meta = _meta_instance() - id_a, inst_a = _instance(meta_instance_id=meta.meta_instance_id) - topology = _topology("node-a") - result = find_unsatisfied_meta_instances( - {meta.meta_instance_id: meta}, - {id_a: inst_a}, - topology, - ) - assert list(result) == [] - - -def test_unsatisfied_one_not_satisfied(): - meta = _meta_instance("test-org/model-x") - id_a, inst_a = _instance("test-org/model-y") - topology = _topology("node-a") - result = find_unsatisfied_meta_instances( - {meta.meta_instance_id: meta}, {id_a: inst_a}, topology - ) - assert list(result) == [meta] - - -def test_unsatisfied_mix(): - meta_satisfied = _meta_instance("test-org/model-a") - meta_unsatisfied = _meta_instance("test-org/model-b") - id_a, inst_a = _instance( - "test-org/model-a", meta_instance_id=meta_satisfied.meta_instance_id - ) - topology = _topology("node-a") - result = find_unsatisfied_meta_instances( - { - meta_satisfied.meta_instance_id: meta_satisfied, - meta_unsatisfied.meta_instance_id: meta_unsatisfied, - }, - {id_a: inst_a}, - topology, - ) - assert list(result) == [meta_unsatisfied] - - -def test_unsatisfied_node_disconnect(): - meta = _meta_instance() - id_a, inst_a = _instance( - node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id - ) - topology = _topology("node-a") # node-b disconnected - result = find_unsatisfied_meta_instances( - {meta.meta_instance_id: meta}, - {id_a: inst_a}, - topology, - ) - assert list(result) == [meta] - - -def test_unsatisfied_edge_break(): - """Instance exists but its connections broke — meta-instance becomes unsatisfied.""" - meta = _meta_instance() - id_a, inst_a = _instance( - node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id - ) - topology = _topology("node-a", "node-b", connect=False) # nodes present, no edges - result = find_unsatisfied_meta_instances( - {meta.meta_instance_id: meta}, - {id_a: inst_a}, - topology, - ) - assert list(result) == [meta] - - -def test_unsatisfied_idempotent(): - meta = _meta_instance("test-org/model-x") - topology = _topology("node-a") - meta_instances = {meta.meta_instance_id: meta} - instances: dict[InstanceId, MlxRingInstance] = {} - result_1 = list( - find_unsatisfied_meta_instances(meta_instances, instances, topology) - ) - result_2 = list( - find_unsatisfied_meta_instances(meta_instances, instances, topology) - ) - assert result_1 == result_2 - - -def test_unsatisfied_exclusive_binding(): - """Two MetaInstances for the same model: one is bound via meta_instance_id, the other is unsatisfied.""" - meta_a = _meta_instance("test-org/model-x") - meta_b = _meta_instance("test-org/model-x") - id_inst, inst = _instance( - "test-org/model-x", meta_instance_id=meta_a.meta_instance_id - ) - topology = _topology("node-a") - result = find_unsatisfied_meta_instances( - { - meta_a.meta_instance_id: meta_a, - meta_b.meta_instance_id: meta_b, - }, - {id_inst: inst}, - topology, - ) - assert list(result) == [meta_b] - - -# --- apply handlers --- - - -def test_apply_meta_instance_created(): - state = State() - meta = _meta_instance() - event = MetaInstanceCreated(meta_instance=meta) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert meta.meta_instance_id in new_state.meta_instances - assert new_state.meta_instances[meta.meta_instance_id] == meta - - -def test_apply_meta_instance_deleted(): - meta = _meta_instance() - state = State(meta_instances={meta.meta_instance_id: meta}) - event = MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert meta.meta_instance_id not in new_state.meta_instances - - -def test_apply_meta_instance_deleted_clears_failure_info(): - meta = _meta_instance().model_copy( - update={"consecutive_failures": 2, "last_failure_error": "OOM"} - ) - state = State(meta_instances={meta.meta_instance_id: meta}) - event = MetaInstanceDeleted(meta_instance_id=meta.meta_instance_id) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert meta.meta_instance_id not in new_state.meta_instances - - -# --- instance_runners_failed --- - - -def test_runners_failed_all_failed(): - """All runners in RunnerFailed -> instance is failed.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - runners = { - rid: RunnerFailed(error_message="OOM") - for rid in inst.shard_assignments.node_to_runner.values() - } - is_failed, error = instance_runners_failed(inst, runners, {}) - assert is_failed is True - assert error is not None - assert "OOM" in error - - -def test_runners_failed_mixed_failed_shutdown(): - """One Failed + one Shutdown = failed.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - runners = { - runner_ids[0]: RunnerFailed(error_message="crash"), - runner_ids[1]: RunnerShutdown(), - } - is_failed, error = instance_runners_failed(inst, runners, {}) - assert is_failed is True - assert error is not None - assert "crash" in error - - -def test_runners_not_failed_all_shutdown(): - """All Shutdown (graceful) = not a failure.""" - _, inst = _instance(node_ids=["node-a"]) - runners = { - rid: RunnerShutdown() for rid in inst.shard_assignments.node_to_runner.values() - } - is_failed, _ = instance_runners_failed(inst, runners, {}) - assert is_failed is False - - -def test_runners_not_failed_still_active(): - """Some runners still active = not failed yet.""" - _, inst = _instance(node_ids=["node-a", "node-b"]) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - runners = { - runner_ids[0]: RunnerFailed(error_message="OOM"), - runner_ids[1]: RunnerLoading(), - } - is_failed, _ = instance_runners_failed(inst, runners, {}) - assert is_failed is False - - -def test_runners_not_failed_no_status(): - """Runner not yet reported = not failed.""" - _, inst = _instance(node_ids=["node-a"]) - is_failed, _ = instance_runners_failed(inst, {}, {}) - assert is_failed is False - - -def test_runners_not_failed_healthy(): - """Runners in Ready state = not failed.""" - _, inst = _instance(node_ids=["node-a"]) - runners = { - rid: RunnerReady() for rid in inst.shard_assignments.node_to_runner.values() - } - is_failed, _ = instance_runners_failed(inst, runners, {}) - assert is_failed is False - - -# --- failure tracking in apply_instance_deleted --- - - -def test_apply_instance_deleted_tracks_failure(): - """InstanceDeleted with failure_error increments meta instance failure count.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - event = InstanceDeleted(instance_id=iid, failure_error="Runner OOM") - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 1 - assert mi.last_failure_error == "Runner OOM" - - -def test_apply_instance_deleted_increments_failure(): - """Subsequent failures increment the counter.""" - meta = _meta_instance().model_copy( - update={"consecutive_failures": 2, "last_failure_error": "previous error"} - ) - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - event = InstanceDeleted(instance_id=iid, failure_error="new error") - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 3 - assert mi.last_failure_error == "new error" - - -def test_apply_instance_deleted_no_failure_no_tracking(): - """InstanceDeleted without failure_error does not track.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - event = InstanceDeleted(instance_id=iid) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 0 - - -def test_apply_instance_deleted_orphan_no_tracking(): - """InstanceDeleted for orphan instance (no meta_instance_id) does not track.""" - iid, inst = _instance(node_ids=["node-a"]) - state = State(instances={iid: inst}) - event = InstanceDeleted(instance_id=iid, failure_error="crash") - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert len(new_state.meta_instances) == 0 - - -# --- InstanceRetrying --- - - -def test_apply_instance_retrying_removes_runners(): - """InstanceRetrying removes the instance's runners from state but keeps the instance.""" - meta = _meta_instance() - iid, inst = _instance( - node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id - ) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - runners = { - runner_ids[0]: RunnerFailed(error_message="OOM"), - runner_ids[1]: RunnerShutdown(), - } - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - runners=runners, - ) - event = InstanceRetrying( - instance_id=iid, - meta_instance_id=meta.meta_instance_id, - failure_error="OOM", - ) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - # Instance still exists - assert iid in new_state.instances - # Runners removed - assert runner_ids[0] not in new_state.runners - assert runner_ids[1] not in new_state.runners - - -def test_apply_instance_retrying_increments_failure(): - """InstanceRetrying increments consecutive_failures on the MetaInstance.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - event = InstanceRetrying( - instance_id=iid, - meta_instance_id=meta.meta_instance_id, - failure_error="crash", - ) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 1 - assert mi.last_failure_error == "crash" - - -def test_apply_instance_retrying_skips_missing_runners(): - """InstanceRetrying doesn't assert if runners haven't reported yet.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - # No runners in state at all - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - ) - event = InstanceRetrying( - instance_id=iid, - meta_instance_id=meta.meta_instance_id, - failure_error="crash", - ) - # Should not raise - new_state = apply(state, IndexedEvent(idx=0, event=event)) - assert iid in new_state.instances - - -def test_apply_instance_created_resets_failure_counter(): - """InstanceCreated resets consecutive_failures but preserves last_failure_error.""" - meta = _meta_instance().model_copy( - update={"consecutive_failures": 3, "last_failure_error": "old error"} - ) - _, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - state = State(meta_instances={meta.meta_instance_id: meta}) - event = InstanceCreated(instance=inst) - new_state = apply(state, IndexedEvent(idx=0, event=event)) - mi = new_state.meta_instances[meta.meta_instance_id] - assert mi.consecutive_failures == 0 - assert mi.last_failure_error == "old error" - assert mi.placement_error is None - - -# --- InstanceHealthReconciler retry-vs-delete --- - - -async def test_health_reconciler_retries_when_under_limit(): - """InstanceHealthReconciler emits InstanceRetrying when consecutive_failures < 3.""" - meta = _meta_instance() - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, - topology=_topology("node-a"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceRetrying) - assert events[0].instance_id == iid - assert events[0].meta_instance_id == meta.meta_instance_id - - -async def test_health_reconciler_deletes_when_limit_reached(): - """InstanceHealthReconciler emits InstanceDeleted when consecutive_failures >= 3.""" - meta = _meta_instance().model_copy(update={"consecutive_failures": 3}) - iid, inst = _instance(node_ids=["node-a"], meta_instance_id=meta.meta_instance_id) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - runners={runner_ids[0]: RunnerFailed(error_message="OOM")}, - topology=_topology("node-a"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceDeleted) - - -async def test_health_reconciler_deletes_without_meta_instance(): - """Instances without a MetaInstance are deleted immediately on runner failure.""" - iid, inst = _instance(node_ids=["node-a"]) - runner_ids = list(inst.shard_assignments.node_to_runner.values()) - state = State( - instances={iid: inst}, - runners={runner_ids[0]: RunnerFailed(error_message="crash")}, - topology=_topology("node-a"), - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceDeleted) - - -async def test_health_reconciler_network_failure_always_deletes(): - """Network failure always triggers InstanceDeleted regardless of retry count.""" - meta = _meta_instance() - iid, inst = _instance( - node_ids=["node-a", "node-b"], meta_instance_id=meta.meta_instance_id - ) - state = State( - meta_instances={meta.meta_instance_id: meta}, - instances={iid: inst}, - topology=_topology("node-a"), # node-b missing - ) - reconciler = InstanceHealthReconciler() - events = await reconciler.reconcile(state) - assert len(events) == 1 - assert isinstance(events[0], InstanceDeleted) - assert events[0].failure_error == "Network connection lost" diff --git a/src/exo/shared/apply.py b/src/exo/shared/apply.py index f96c6b7f..94869dfe 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -4,7 +4,7 @@ from datetime import datetime from loguru import logger -from exo.shared.types.common import MetaInstanceId, NodeId +from exo.shared.types.common import NodeId from exo.shared.types.events import ( ChunkGenerated, Event, @@ -12,12 +12,6 @@ from exo.shared.types.events import ( InputChunkReceived, InstanceCreated, InstanceDeleted, - InstanceRetrying, - JacclSideChannelData, - JacclSideChannelGathered, - MetaInstanceCreated, - MetaInstanceDeleted, - MetaInstancePlacementFailed, NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, @@ -34,7 +28,6 @@ from exo.shared.types.events import ( TracesCollected, TracesMerged, ) -from exo.shared.types.meta_instance import MetaInstance from exo.shared.types.profiling import ( NodeIdentity, NodeNetworkInfo, @@ -73,22 +66,12 @@ def event_apply(event: Event, state: State) -> State: | InputChunkReceived() | TracesCollected() | TracesMerged() - | JacclSideChannelData() - | JacclSideChannelGathered() ): # Pass-through events that don't modify state return state case InstanceCreated(): return apply_instance_created(event, state) case InstanceDeleted(): return apply_instance_deleted(event, state) - case InstanceRetrying(): - return apply_instance_retrying(event, state) - case MetaInstanceCreated(): - return apply_meta_instance_created(event, state) - case MetaInstanceDeleted(): - return apply_meta_instance_deleted(event, state) - case MetaInstancePlacementFailed(): - return apply_meta_instance_placement_failed(event, state) case NodeTimedOut(): return apply_node_timed_out(event, state) case NodeDownloadProgress(): @@ -191,123 +174,20 @@ def apply_task_failed(event: TaskFailed, state: State) -> State: return state.model_copy(update={"tasks": new_tasks}) -def _update_meta_instance( - state: State, mid: MetaInstanceId, **fields: object -) -> Mapping[MetaInstanceId, MetaInstance]: - mi = state.meta_instances[mid] - return {**state.meta_instances, mid: mi.model_copy(update=fields)} - - def apply_instance_created(event: InstanceCreated, state: State) -> State: instance = event.instance new_instances: Mapping[InstanceId, Instance] = { **state.instances, instance.instance_id: instance, } - update: dict[str, object] = {"instances": new_instances} - # Reset failure tracking when a new instance is created for a meta-instance - if instance.meta_instance_id and instance.meta_instance_id in state.meta_instances: - mi = state.meta_instances[instance.meta_instance_id] - if mi.placement_error is not None or mi.consecutive_failures > 0: - update["meta_instances"] = _update_meta_instance( - state, - instance.meta_instance_id, - placement_error=None, - consecutive_failures=0, - ) - return state.model_copy(update=update) + return state.model_copy(update={"instances": new_instances}) def apply_instance_deleted(event: InstanceDeleted, state: State) -> State: - deleted_instance = state.instances.get(event.instance_id) new_instances: Mapping[InstanceId, Instance] = { iid: inst for iid, inst in state.instances.items() if iid != event.instance_id } - update: dict[str, object] = {"instances": new_instances} - - # Track failure on the MetaInstance itself - if ( - event.failure_error - and deleted_instance - and deleted_instance.meta_instance_id - and deleted_instance.meta_instance_id in state.meta_instances - ): - mid = deleted_instance.meta_instance_id - mi = state.meta_instances[mid] - update["meta_instances"] = { - **state.meta_instances, - mid: mi.model_copy( - update={ - "consecutive_failures": mi.consecutive_failures + 1, - "last_failure_error": event.failure_error, - } - ), - } - - return state.model_copy(update=update) - - -def apply_instance_retrying(event: InstanceRetrying, state: State) -> State: - """Runners failed but retry limit not reached — remove runners, keep instance.""" - instance = state.instances.get(event.instance_id) - if instance is None: - # Instance was already deleted (e.g. cascade from DeleteMetaInstance). - # The InstanceDeleted handler already incremented consecutive_failures - # on the MetaInstance, so skipping here avoids double-counting. - return state - - # Remove all runners belonging to this instance from state - runner_ids_to_remove = set(instance.shard_assignments.node_to_runner.values()) - new_runners: Mapping[RunnerId, RunnerStatus] = { - rid: rs for rid, rs in state.runners.items() if rid not in runner_ids_to_remove - } - - update: dict[str, object] = {"runners": new_runners} - - # Increment failure count on the MetaInstance - if event.meta_instance_id in state.meta_instances: - update["meta_instances"] = _update_meta_instance( - state, - event.meta_instance_id, - consecutive_failures=state.meta_instances[ - event.meta_instance_id - ].consecutive_failures - + 1, - last_failure_error=event.failure_error, - ) - - return state.model_copy(update=update) - - -def apply_meta_instance_created(event: MetaInstanceCreated, state: State) -> State: - new_meta: Mapping[MetaInstanceId, MetaInstance] = { - **state.meta_instances, - event.meta_instance.meta_instance_id: event.meta_instance, - } - return state.model_copy(update={"meta_instances": new_meta}) - - -def apply_meta_instance_deleted(event: MetaInstanceDeleted, state: State) -> State: - new_meta: Mapping[MetaInstanceId, MetaInstance] = { - mid: mi - for mid, mi in state.meta_instances.items() - if mid != event.meta_instance_id - } - return state.model_copy(update={"meta_instances": new_meta}) - - -def apply_meta_instance_placement_failed( - event: MetaInstancePlacementFailed, state: State -) -> State: - if event.meta_instance_id not in state.meta_instances: - return state - return state.model_copy( - update={ - "meta_instances": _update_meta_instance( - state, event.meta_instance_id, placement_error=event.reason - ) - } - ) + return state.model_copy(update={"instances": new_instances}) def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State: diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 2f0fafa8..2756f0d4 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -6,7 +6,7 @@ from uuid import uuid4 from pydantic import BaseModel, Field from exo.shared.models.model_cards import ModelCard, ModelId -from exo.shared.types.common import CommandId, MetaInstanceId, NodeId +from exo.shared.types.common import CommandId, NodeId from exo.shared.types.memory import Memory from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata @@ -262,26 +262,6 @@ class DeleteInstanceResponse(BaseModel): instance_id: InstanceId -class CreateMetaInstanceParams(BaseModel): - model_id: ModelId - sharding: Sharding = Sharding.Pipeline - instance_meta: InstanceMeta = InstanceMeta.MlxRing - min_nodes: int = 1 - node_ids: list[NodeId] | None = None - - -class CreateMetaInstanceResponse(BaseModel): - message: str - command_id: CommandId - meta_instance_id: MetaInstanceId - - -class DeleteMetaInstanceResponse(BaseModel): - message: str - command_id: CommandId - meta_instance_id: MetaInstanceId - - class AdvancedImageParams(BaseModel): seed: Annotated[int, Field(ge=0)] | None = None num_inference_steps: Annotated[int, Field(ge=1, le=100)] | None = None diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py index 8697a6c2..09c135aa 100644 --- a/src/exo/shared/types/commands.py +++ b/src/exo/shared/types/commands.py @@ -6,8 +6,7 @@ from exo.shared.types.api import ( ImageGenerationTaskParams, ) from exo.shared.types.chunks import InputImageChunk -from exo.shared.types.common import CommandId, MetaInstanceId, NodeId -from exo.shared.types.meta_instance import MetaInstance +from exo.shared.types.common import CommandId, NodeId from exo.shared.types.text_generation import TextGenerationTaskParams from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta from exo.shared.types.worker.shards import Sharding, ShardMetadata @@ -53,14 +52,6 @@ class TaskCancelled(BaseCommand): cancelled_command_id: CommandId -class CreateMetaInstance(BaseCommand): - meta_instance: MetaInstance - - -class DeleteMetaInstance(BaseCommand): - meta_instance_id: MetaInstanceId - - class TaskFinished(BaseCommand): finished_command_id: CommandId @@ -103,8 +94,6 @@ Command = ( | CreateInstance | DeleteInstance | TaskCancelled - | CreateMetaInstance - | DeleteMetaInstance | TaskFinished | SendInputChunk ) diff --git a/src/exo/shared/types/common.py b/src/exo/shared/types/common.py index 51806de2..5db51cef 100644 --- a/src/exo/shared/types/common.py +++ b/src/exo/shared/types/common.py @@ -42,10 +42,6 @@ class CommandId(Id): pass -class MetaInstanceId(Id): - """Identifier for a MetaInstance.""" - - class Host(CamelCaseModel): ip: str port: int diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index dd28d070..5cf93d0c 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -1,14 +1,11 @@ -import base64 -from collections.abc import Mapping from datetime import datetime -from typing import Annotated, final +from typing import final -from pydantic import BeforeValidator, Field, PlainSerializer +from pydantic import Field from exo.shared.topology import Connection from exo.shared.types.chunks import GenerationChunk, InputImageChunk -from exo.shared.types.common import CommandId, Id, MetaInstanceId, NodeId, SessionId -from exo.shared.types.meta_instance import MetaInstance +from exo.shared.types.common import CommandId, Id, NodeId, SessionId from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId @@ -17,28 +14,6 @@ from exo.utils.info_gatherer.info_gatherer import GatheredInfo from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel -def _decode_base64_bytes(v: bytes | str) -> bytes: - if isinstance(v, bytes): - return v - return base64.b64decode(v) - - -def _encode_base64_bytes(v: bytes) -> str: - return base64.b64encode(v).decode("ascii") - - -Base64Bytes = Annotated[ - bytes, - BeforeValidator(_decode_base64_bytes), - PlainSerializer(_encode_base64_bytes, return_type=str), -] -"""bytes that serialize to/from base64 strings in JSON. - -Needed because TaggedModel's wrap validator converts JSON→Python validation -context, which breaks strict-mode bytes deserialization from JSON strings. -""" - - class EventId(Id): """ Newtype around `ID` @@ -91,30 +66,6 @@ class InstanceCreated(BaseEvent): class InstanceDeleted(BaseEvent): instance_id: InstanceId - failure_error: str | None = None - - -class MetaInstanceCreated(BaseEvent): - meta_instance: MetaInstance - - -class MetaInstanceDeleted(BaseEvent): - meta_instance_id: MetaInstanceId - - -@final -class MetaInstancePlacementFailed(BaseEvent): - meta_instance_id: MetaInstanceId - reason: str - - -@final -class InstanceRetrying(BaseEvent): - """Runners failed but retry count is below the limit — restart runners, keep instance.""" - - instance_id: InstanceId - meta_instance_id: MetaInstanceId - failure_error: str class RunnerStatusUpdated(BaseEvent): @@ -181,25 +132,6 @@ class TracesMerged(BaseEvent): traces: list[TraceEventData] -@final -class JacclSideChannelData(BaseEvent): - """A runner's local contribution to a JACCL SideChannel all_gather round.""" - - instance_id: InstanceId - runner_id: RunnerId - sequence: int - data: Base64Bytes - - -@final -class JacclSideChannelGathered(BaseEvent): - """Gathered result of a JACCL SideChannel all_gather round.""" - - instance_id: InstanceId - sequence: int - gathered_data: Mapping[RunnerId, Base64Bytes] - - Event = ( TestEvent | TaskCreated @@ -209,10 +141,6 @@ Event = ( | TaskAcknowledged | InstanceCreated | InstanceDeleted - | InstanceRetrying - | MetaInstanceCreated - | MetaInstanceDeleted - | MetaInstancePlacementFailed | RunnerStatusUpdated | RunnerDeleted | NodeTimedOut @@ -224,8 +152,6 @@ Event = ( | TopologyEdgeDeleted | TracesCollected | TracesMerged - | JacclSideChannelData - | JacclSideChannelGathered ) diff --git a/src/exo/shared/types/meta_instance.py b/src/exo/shared/types/meta_instance.py deleted file mode 100644 index 63052184..00000000 --- a/src/exo/shared/types/meta_instance.py +++ /dev/null @@ -1,25 +0,0 @@ -from typing import final - -from pydantic import Field - -from exo.shared.models.model_cards import ModelId -from exo.shared.types.common import MetaInstanceId, NodeId -from exo.shared.types.worker.instances import InstanceMeta -from exo.shared.types.worker.shards import Sharding -from exo.utils.pydantic_ext import FrozenModel - - -@final -class MetaInstance(FrozenModel): - """Declarative constraint: ensure an instance matching these parameters always exists.""" - - meta_instance_id: MetaInstanceId = Field(default_factory=MetaInstanceId) - model_id: ModelId - sharding: Sharding = Sharding.Pipeline - instance_meta: InstanceMeta = InstanceMeta.MlxRing - min_nodes: int = 1 - node_ids: list[NodeId] | None = None - # Failure tracking - placement_error: str | None = None - consecutive_failures: int = 0 - last_failure_error: str | None = None diff --git a/src/exo/shared/types/state.py b/src/exo/shared/types/state.py index 4ff1d4f7..7350cfb0 100644 --- a/src/exo/shared/types/state.py +++ b/src/exo/shared/types/state.py @@ -6,8 +6,7 @@ from pydantic import ConfigDict, Field, field_serializer, field_validator from pydantic.alias_generators import to_camel from exo.shared.topology import Topology, TopologySnapshot -from exo.shared.types.common import MetaInstanceId, NodeId -from exo.shared.types.meta_instance import MetaInstance +from exo.shared.types.common import NodeId from exo.shared.types.profiling import ( DiskUsage, MemoryUsage, @@ -42,7 +41,6 @@ class State(CamelCaseModel): arbitrary_types_allowed=True, ) instances: Mapping[InstanceId, Instance] = {} - meta_instances: Mapping[MetaInstanceId, MetaInstance] = {} runners: Mapping[RunnerId, RunnerStatus] = {} downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {} tasks: Mapping[TaskId, Task] = {} diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py index cb88d401..8d866456 100644 --- a/src/exo/shared/types/tasks.py +++ b/src/exo/shared/types/tasks.py @@ -61,7 +61,7 @@ class TextGeneration(BaseTask): # emitted by Master error_message: str | None = Field(default=None) -class CancelTask(BaseTask): # emitted by Worker when master cancels a task +class CancelTask(BaseTask): cancelled_task_id: TaskId runner_id: RunnerId diff --git a/src/exo/shared/types/worker/instances.py b/src/exo/shared/types/worker/instances.py index 4254b998..cda11ffa 100644 --- a/src/exo/shared/types/worker/instances.py +++ b/src/exo/shared/types/worker/instances.py @@ -2,7 +2,7 @@ from enum import Enum from pydantic import model_validator -from exo.shared.types.common import Host, Id, MetaInstanceId, NodeId +from exo.shared.types.common import Host, Id, NodeId from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel @@ -19,7 +19,6 @@ class InstanceMeta(str, Enum): class BaseInstance(TaggedModel): instance_id: InstanceId shard_assignments: ShardAssignments - meta_instance_id: MetaInstanceId | None = None def shard(self, runner_id: RunnerId) -> ShardMetadata | None: return self.shard_assignments.runner_to_shard.get(runner_id, None) diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py index ebf0165f..646ac8f6 100644 --- a/src/exo/utils/channels.py +++ b/src/exo/utils/channels.py @@ -125,7 +125,9 @@ class MpSender[T]: self._state.buffer.put(item, block=True) async def send_async(self, item: T) -> None: - await to_thread.run_sync(self.send, item, limiter=CapacityLimiter(1)) + await to_thread.run_sync( + self.send, item, limiter=CapacityLimiter(1), abandon_on_cancel=True + ) def close(self) -> None: if not self._state.closed.is_set(): diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 670847eb..3ed65ecc 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -574,11 +574,6 @@ def mlx_cleanup( def mx_any(bool_: bool, group: Group | None) -> bool: - """Synchronize a boolean across all distributed nodes. - - Returns True if any node has bool_=True. Uses all_sum so every - node participates in the collective — preventing GPU deadlocks. - """ if group is None: return bool_ num_true = mx.distributed.all_sum( diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 6b2a9475..af105652 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -24,7 +24,6 @@ from exo.shared.types.events import ( ForwarderEvent, IndexedEvent, InputChunkReceived, - JacclSideChannelGathered, NodeGatheredInfo, TaskCreated, TaskStatusUpdated, @@ -34,6 +33,7 @@ from exo.shared.types.events import ( from exo.shared.types.multiaddr import Multiaddr from exo.shared.types.state import State from exo.shared.types.tasks import ( + CancelTask, CreateRunner, DownloadModel, ImageEdits, @@ -159,15 +159,6 @@ class Worker: for idx, event in indexed_events: self.state = apply(self.state, IndexedEvent(idx=idx, event=event)) - # Dispatch JACCL gathered events to the relevant RunnerSupervisor - if isinstance(event, JacclSideChannelGathered): - for runner in self.runners.values(): - if ( - runner.bound_instance.instance.instance_id - == event.instance_id - ): - runner.notify_gathered(event) - # Buffer input image chunks for image editing if isinstance(event, InputChunkReceived): cmd_id = event.command_id @@ -234,15 +225,22 @@ class Worker: ) ) case Shutdown(runner_id=runner_id): + runner = self.runners.pop(runner_id) try: with fail_after(3): - await self.runners.pop(runner_id).start_task(task) + await runner.start_task(task) except TimeoutError: await self.event_sender.send( TaskStatusUpdated( task_id=task.task_id, task_status=TaskStatus.TimedOut ) ) + finally: + runner.shutdown() + case CancelTask( + cancelled_task_id=cancelled_task_id, runner_id=runner_id + ): + await self.runners[runner_id].cancel_task(cancelled_task_id) case ImageEdits() if task.task_params.total_input_chunks > 0: # Assemble image from chunks and inject into task cmd_id = task.command_id @@ -280,18 +278,18 @@ class Worker: del self.input_chunk_buffer[cmd_id] if cmd_id in self.input_chunk_counts: del self.input_chunk_counts[cmd_id] - await self.runners[self._task_to_runner_id(task)].start_task( - modified_task - ) + await self._start_runner_task(modified_task) case task: - await self.runners[self._task_to_runner_id(task)].start_task(task) + await self._start_runner_task(task) def shutdown(self): self._tg.cancel_scope.cancel() - def _task_to_runner_id(self, task: Task): - instance = self.state.instances[task.instance_id] - return instance.shard_assignments.node_to_runner[self.node_id] + async def _start_runner_task(self, task: Task): + if (instance := self.state.instances.get(task.instance_id)) is not None: + await self.runners[ + instance.shard_assignments.node_to_runner[self.node_id] + ].start_task(task) async def _nack_request(self, since_idx: int) -> None: # We request all events after (and including) the missing index. diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 4107d553..ce2eb4a9 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -35,7 +35,6 @@ from exo.shared.types.worker.runners import ( RunnerLoading, RunnerReady, RunnerRunning, - RunnerShutdown, RunnerStatus, RunnerWarmingUp, ) @@ -57,7 +56,7 @@ def plan( return ( _cancel_tasks(runners, tasks) or _kill_runner(runners, all_runners, instances) - or _create_runner(node_id, runners, instances, all_runners) + or _create_runner(node_id, runners, instances) or _model_needs_download(node_id, runners, global_download_status) or _init_distributed_backend(runners, all_runners) or _load_model(runners, all_runners, global_download_status) @@ -76,12 +75,6 @@ def _kill_runner( if (instance_id := runner.bound_instance.instance.instance_id) not in instances: return Shutdown(instance_id=instance_id, runner_id=runner_id) - # Master removed our runner from state (retry signal) and process is dead - if runner_id not in all_runners and isinstance( - runner.status, (RunnerFailed, RunnerShutdown) - ): - return Shutdown(instance_id=instance_id, runner_id=runner_id) - for ( global_runner_id ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values(): @@ -99,7 +92,6 @@ def _create_runner( node_id: NodeId, runners: Mapping[RunnerId, RunnerSupervisor], instances: Mapping[InstanceId, Instance], - all_runners: Mapping[RunnerId, RunnerStatus], ) -> CreateRunner | None: for instance in instances.values(): runner_id = instance.shard_assignments.node_to_runner.get(node_id, None) @@ -109,16 +101,6 @@ def _create_runner( if runner_id in runners: continue - # Don't create while any peer runner is in a terminal state — wait for - # the master to emit InstanceRetrying which removes them from state. - has_terminal_peer = any( - isinstance(all_runners.get(peer_rid), (RunnerFailed, RunnerShutdown)) - for peer_rid in instance.shard_assignments.node_to_runner.values() - if peer_rid != runner_id - ) - if has_terminal_peer: - continue - shard = instance.shard(runner_id) assert shard is not None @@ -328,8 +310,7 @@ def _pending_tasks( def _cancel_tasks( runners: Mapping[RunnerId, RunnerSupervisor], tasks: Mapping[TaskId, Task], -) -> CancelTask | None: - """Find a cancelled task that hasn't been sent to the runner yet.""" +) -> Task | None: for task in tasks.values(): if task.task_status != TaskStatus.Cancelled: continue diff --git a/src/exo/worker/runner/bootstrap.py b/src/exo/worker/runner/bootstrap.py index 69ef6c72..ed420aab 100644 --- a/src/exo/worker/runner/bootstrap.py +++ b/src/exo/worker/runner/bootstrap.py @@ -17,7 +17,6 @@ def entrypoint( task_receiver: MpReceiver[Task], cancel_receiver: MpReceiver[TaskId], _logger: "loguru.Logger", - pipe_fifo_paths: tuple[str, str] | None = None, ) -> None: fast_synch_override = os.environ.get("EXO_FAST_SYNCH") if fast_synch_override == "on" or ( @@ -31,16 +30,6 @@ def entrypoint( else: os.environ["MLX_METAL_FAST_SYNCH"] = "0" - # Open JACCL FIFOs by path and set env vars for C++ SideChannel. - # Named pipes (FIFOs) work across multiprocessing spawn (macOS default). - if pipe_fifo_paths is not None: - fifo_c2p, fifo_p2c = pipe_fifo_paths - # C++ reads gathered data from p2c (PIPE_IN), writes local data to c2p (PIPE_OUT) - pipe_in_fd = os.open(fifo_p2c, os.O_RDONLY) - pipe_out_fd = os.open(fifo_c2p, os.O_WRONLY) - os.environ["MLX_JACCL_PIPE_IN"] = str(pipe_in_fd) - os.environ["MLX_JACCL_PIPE_OUT"] = str(pipe_out_fd) - global logger logger = _logger @@ -67,9 +56,7 @@ def entrypoint( try: event_sender.close() task_receiver.close() - cancel_receiver.close() finally: event_sender.join() task_receiver.join() - cancel_receiver.join() logger.info("bye from the runner") diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 818bd9be..e55456d3 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -243,7 +243,7 @@ def main( assert inference_model assert tokenizer - t = time.perf_counter() + t = time.monotonic() toks = warmup_inference( model=inference_model, tokenizer=tokenizer, @@ -251,7 +251,7 @@ def main( ) logger.info(f"warmed up by generating {toks} tokens") check_for_cancel_every = min( - math.ceil(toks / max(time.perf_counter() - t, 0.001)), 100 + math.ceil(toks / min(time.monotonic() - t, 0.001)), 100 ) if group is not None: check_for_cancel_every = int( diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 519d7b07..5d39a881 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -1,10 +1,6 @@ import contextlib -import os import signal -import struct -import tempfile from dataclasses import dataclass, field -from functools import partial from multiprocessing import Process from typing import Self @@ -18,14 +14,12 @@ from loguru import logger from exo.shared.types.events import ( Event, - JacclSideChannelData, - JacclSideChannelGathered, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, ) from exo.shared.types.tasks import Task, TaskId, TaskStatus -from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance +from exo.shared.types.worker.instances import BoundInstance from exo.shared.types.worker.runners import ( RunnerConnecting, RunnerFailed, @@ -40,26 +34,6 @@ from exo.shared.types.worker.shards import ShardMetadata from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel from exo.worker.runner.bootstrap import entrypoint - -def _pipe_read_exact(fd: int, n: int) -> bytes | None: - """Read exactly n bytes from a file descriptor. Returns None on EOF.""" - data = b"" - while len(data) < n: - chunk = os.read(fd, n - len(data)) - if not chunk: - return None - data += chunk - return data - - -def _pipe_write_all(fd: int, data: bytes) -> None: - """Write all bytes to a file descriptor.""" - view = memoryview(data) - while view: - written = os.write(fd, view) - view = view[written:] - - PREFILL_TIMEOUT_SECONDS = 60 DECODE_TIMEOUT_SECONDS = 5 @@ -72,21 +46,12 @@ class RunnerSupervisor: initialize_timeout: float _ev_recv: MpReceiver[Event] _task_sender: MpSender[Task] - _cancel_sender: MpSender[TaskId] _event_sender: Sender[Event] - _pipe_read_fd: int | None = None # Python reads runner's pipe output - _pipe_write_fd: int | None = None # Python writes gathered data to runner - _child_pipe_fds: tuple[int, int] | None = None # fds to close after fork - _fifo_dir: str | None = None # Temp dir for FIFO files (for cleanup) - _fifo_c2p: str | None = None # FIFO path: C++ writes → Python reads - _fifo_p2c: str | None = None # FIFO path: Python writes → C++ reads + _cancel_sender: MpSender[TaskId] status: RunnerStatus = field(default_factory=RunnerIdle, init=False) pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False) completed: set[TaskId] = field(default_factory=set, init=False) cancelled: set[TaskId] = field(default_factory=set, init=False) - _gathered_waiters: dict[ - int, tuple[anyio.Event, JacclSideChannelGathered | None] - ] = field(default_factory=dict, init=False) @classmethod def create( @@ -100,23 +65,6 @@ class RunnerSupervisor: task_sender, task_recv = mp_channel[Task]() cancel_sender, cancel_recv = mp_channel[TaskId]() - # For MlxJaccl instances, create named pipes (FIFOs) for SideChannel relay. - # Named pipes work across multiprocessing.Process spawn (macOS default). - # FIFO c2p: C++ writes local data → Python reads it - # FIFO p2c: Python writes gathered data → C++ reads it - fifo_dir: str | None = None - fifo_c2p: str | None = None - fifo_p2c: str | None = None - pipe_fifo_paths: tuple[str, str] | None = None - - if isinstance(bound_instance.instance, MlxJacclInstance): - fifo_dir = tempfile.mkdtemp(prefix="exo_jaccl_") - fifo_c2p = os.path.join(fifo_dir, "c2p") # C++ → Python - fifo_p2c = os.path.join(fifo_dir, "p2c") # Python → C++ - os.mkfifo(fifo_c2p) - os.mkfifo(fifo_p2c) - pipe_fifo_paths = (fifo_c2p, fifo_p2c) - runner_process = Process( target=entrypoint, args=( @@ -125,7 +73,6 @@ class RunnerSupervisor: task_recv, cancel_recv, logger, - pipe_fifo_paths, ), daemon=True, ) @@ -141,54 +88,21 @@ class RunnerSupervisor: _task_sender=task_sender, _cancel_sender=cancel_sender, _event_sender=event_sender, - _fifo_dir=fifo_dir, - _fifo_c2p=fifo_c2p, - _fifo_p2c=fifo_p2c, ) return self async def run(self): self.runner_process.start() - - if self._fifo_c2p is not None and self._fifo_p2c is not None: - # Open FIFOs from parent side. These block until child opens the other end, - # so we run them in threads concurrently to avoid deadlock. - fifo_c2p = self._fifo_c2p - fifo_p2c = self._fifo_p2c - - async def open_read() -> None: - self._pipe_read_fd = await to_thread.run_sync( - partial(os.open, fifo_c2p, os.O_RDONLY) - ) - - async def open_write() -> None: - self._pipe_write_fd = await to_thread.run_sync( - partial(os.open, fifo_p2c, os.O_WRONLY) - ) - - async with anyio.create_task_group() as open_tg: - open_tg.start_soon(open_read) - open_tg.start_soon(open_write) - - logger.info( - f"JACCL pipe relay: FIFOs opened (read_fd={self._pipe_read_fd}, write_fd={self._pipe_write_fd})" - ) - - async with anyio.create_task_group() as tg: - tg.start_soon(self._pipe_relay) - tg.start_soon(self._forward_events) - else: - await self._forward_events() + await self._forward_events() def shutdown(self): 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")) self._cancel_sender.close() - self._event_sender.close() - self._close_pipe_fds() self.runner_process.join(1) if not self.runner_process.is_alive(): logger.info("Runner process succesfully terminated") @@ -226,7 +140,6 @@ class RunnerSupervisor: await event.wait() async def cancel_task(self, task_id: TaskId): - """Send a cancellation signal to the runner process.""" if task_id in self.completed: logger.info(f"Unable to cancel {task_id} as it has been completed") return @@ -268,110 +181,6 @@ class RunnerSupervisor: for tid in self.pending: self.pending[tid].set() - def _close_pipe_fds(self) -> None: - if self._pipe_read_fd is not None: - with contextlib.suppress(OSError): - os.close(self._pipe_read_fd) - self._pipe_read_fd = None - if self._pipe_write_fd is not None: - with contextlib.suppress(OSError): - os.close(self._pipe_write_fd) - self._pipe_write_fd = None - if self._child_pipe_fds is not None: - for fd in self._child_pipe_fds: - with contextlib.suppress(OSError): - os.close(fd) - self._child_pipe_fds = None - # Clean up FIFO files - if self._fifo_c2p is not None: - with contextlib.suppress(OSError): - os.unlink(self._fifo_c2p) - self._fifo_c2p = None - if self._fifo_p2c is not None: - with contextlib.suppress(OSError): - os.unlink(self._fifo_p2c) - self._fifo_p2c = None - if self._fifo_dir is not None: - with contextlib.suppress(OSError): - os.rmdir(self._fifo_dir) - self._fifo_dir = None - - async def _pipe_relay(self) -> None: - """Relay JACCL SideChannel all_gather rounds between runner pipes and exo events.""" - assert self._pipe_read_fd is not None - assert self._pipe_write_fd is not None - read_fd = self._pipe_read_fd - write_fd = self._pipe_write_fd - sequence = 0 - - try: - while True: - # 1. Read local data from runner: [uint32 size][size bytes] - header = await to_thread.run_sync(partial(_pipe_read_exact, read_fd, 4)) - if header is None: - logger.info("JACCL pipe relay: runner closed pipe (EOF)") - break - data_size: int = struct.unpack(" None: - """Called by the worker when a JacclSideChannelGathered event arrives.""" - seq = event.sequence - if seq not in self._gathered_waiters: - logger.warning(f"JACCL: received gathered event for unknown sequence {seq}") - return - waiter, _ = self._gathered_waiters[seq] - self._gathered_waiters[seq] = (waiter, event) - waiter.set() - def __del__(self) -> None: if self.runner_process.is_alive(): logger.warning("RunnerSupervisor was not stopped cleanly.") 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 878aea99..38a0a921 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 @@ -1,7 +1,9 @@ # Check tasks are complete before runner is ever ready. +import unittest.mock from collections.abc import Iterable from typing import Callable +import mlx.core as mx import pytest import exo.worker.runner.runner as mlx_runner @@ -115,12 +117,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1)) monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin) monkeypatch.setattr(mlx_runner, "mx_any", make_nothin(False)) - - # Mock mx.distributed.all_gather so MockGroup doesn't hit real MLX C++ bindings. - def _mock_all_gather(x: object, **_kw: object) -> object: - return x - - monkeypatch.setattr(mlx_runner.mx.distributed, "all_gather", _mock_all_gather) # Mock apply_chat_template since we're using a fake tokenizer (integer 1). # Returns a prompt without thinking tag so detect_thinking_prompt_suffix returns None. monkeypatch.setattr(mlx_runner, "apply_chat_template", make_nothin("test prompt")) @@ -182,15 +178,16 @@ def _run(tasks: Iterable[Task]): # this is some c++ nonsense task_receiver.close = nothin task_receiver.join = nothin - cancel_receiver.close = nothin - cancel_receiver.join = nothin - - mlx_runner.main( - bound_instance, - event_sender, # pyright: ignore[reportArgumentType] - task_receiver, - cancel_receiver, - ) + with unittest.mock.patch( + "exo.worker.runner.runner.mx.distributed.all_gather", + make_nothin(mx.array([1])), + ): + mlx_runner.main( + bound_instance, + event_sender, # pyright: ignore[reportArgumentType] + task_receiver, + cancel_receiver, + ) return event_sender.events From 83af8c63fab15172d963a5abb9e4d004dc2201d7 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Tue, 17 Feb 2026 18:18:54 +0000 Subject: [PATCH 12/45] Revert "Use custom fork that resolves GPU locks" (#1502) Reverts exo-explore/exo#1489 Goddammit Claude... --- README.md | 11 ++--------- flake.nix | 2 +- nix/mlx.nix | 10 +++++----- pyproject.toml | 3 +-- python/parts.nix | 32 +++----------------------------- uv.lock | 40 ++++++++++++++++++++++++---------------- 6 files changed, 36 insertions(+), 62 deletions(-) diff --git a/README.md b/README.md index ff1fe04b..58d41c37 100644 --- a/README.md +++ b/README.md @@ -72,23 +72,16 @@ There are two ways to run exo: ### Run from Source (macOS) -If you have [Nix](https://nixos.org/) installed, you can skip most of the steps below and run exo directly (after accepting the Cachix cache): - -```bash -nix run .#exo -``` - **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/flake.nix b/flake.nix index e90e0bd2..9c2ca1ef 100644 --- a/flake.nix +++ b/flake.nix @@ -115,7 +115,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" && p.source ? git) uvLock.package); + mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx") uvLock.package); uvLockMlxVersion = mlxPackage.version; in { diff --git a/nix/mlx.nix b/nix/mlx.nix index 8a40e11b..f29217b8 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.7.dev20260217+50487b41"; in + version = let v = "0.30.6"; in assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix."; v; pyproject = true; src = fetchFromGitHub { - owner = "rltakashige"; - repo = "mlx-jaccl-fix-small-recv"; - rev = "50487b4141f3c951122655db3b83df5146c1fbeb"; - hash = "sha256-IL4a9vMX5nocgJU1WG4zE8hArHkHJtnh4sdYh3od5zU="; + owner = "ml-explore"; + repo = "mlx"; + tag = "v${version}"; + hash = "sha256-avD5EGhwgmPdXLAyQSqTO6AXk/W3ziH+f6AetjK3Sdo="; }; patches = [ diff --git a/pyproject.toml b/pyproject.toml index 02aa6071..5d8d79a5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "loguru>=0.7.3", "exo_pyo3_bindings", # rust bindings "anyio==4.11.0", - "mlx; sys_platform == 'darwin'", + "mlx==0.30.6; sys_platform == 'darwin'", "mlx[cpu]==0.30.6; sys_platform == 'linux'", "mlx-lm==0.30.6", "tiktoken>=0.12.0", # required for kimi k2 tokenizer @@ -64,7 +64,6 @@ 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 } diff --git a/python/parts.nix b/python/parts.nix index bac8ddab..46b4abdf 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -58,21 +58,6 @@ 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; } @@ -89,25 +74,14 @@ linuxOverlay ] ); - # 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; - }; + exoVenv = pythonSet.mkVirtualEnv "exo-env" workspace.deps.default; # 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; diff --git a/uv.lock b/uv.lock index 74687ad3..627e6951 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", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.30.7.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "mlx", extra = ["cpu"], marker = "sys_platform == 'linux'" }, { 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,7 +416,7 @@ 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'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" }, + { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.6" }, { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" }, { name = "mlx-lm", specifier = "==0.30.6" }, { name = "msgspec", specifier = ">=0.19.0" }, @@ -1020,8 +1020,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", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.30.7.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, + { name = "mlx", extra = ["cuda13"], marker = "sys_platform == 'linux'" }, { 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,12 +1048,18 @@ wheels = [ name = "mlx" version = "0.30.6" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "sys_platform == 'linux'", +dependencies = [ + { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, ] 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" }, ] @@ -1066,14 +1072,6 @@ cuda13 = [ { name = "mlx-cuda-13", marker = "sys_platform == 'linux'" }, ] -[[package]] -name = "mlx" -version = "0.30.7.dev20260217+50487b41" -source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" } -resolution-markers = [ - "sys_platform == 'darwin'", -] - [[package]] name = "mlx-cpu" version = "0.30.6" @@ -1104,7 +1102,7 @@ version = "0.30.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", version = "0.30.7.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", 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'" }, @@ -1116,6 +1114,16 @@ 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" }, +] + [[package]] name = "more-itertools" version = "10.8.0" From f2be92921167acf519990d14a2e7ff4eea233db3 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Tue, 17 Feb 2026 22:00:52 +0000 Subject: [PATCH 13/45] Leo/address rdma gpu locks 2 (#1515) Same as #1489 . Had to revert and redo thanks to Claude. --------- Co-authored-by: Jake Hillion Co-authored-by: Claude Opus 4.6 --- README.md | 11 +++++++++-- flake.nix | 2 +- nix/mlx.nix | 10 +++++----- pyproject.toml | 3 ++- python/parts.nix | 32 +++++++++++++++++++++++++++++--- uv.lock | 40 ++++++++++++++++------------------------ 6 files changed, 62 insertions(+), 36 deletions(-) diff --git a/README.md b/README.md index 58d41c37..ff1fe04b 100644 --- a/README.md +++ b/README.md @@ -72,16 +72,23 @@ 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 (after accepting the Cachix cache): + +```bash +nix run .#exo +``` + **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/flake.nix b/flake.nix index 9c2ca1ef..e90e0bd2 100644 --- a/flake.nix +++ b/flake.nix @@ -115,7 +115,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..8a40e11b 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.dev20260217+50487b41"; 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 = "50487b4141f3c951122655db3b83df5146c1fbeb"; + hash = "sha256-IL4a9vMX5nocgJU1WG4zE8hArHkHJtnh4sdYh3od5zU="; }; patches = [ diff --git a/pyproject.toml b/pyproject.toml index 5d8d79a5..02aa6071 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ 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", "tiktoken>=0.12.0", # required for kimi k2 tokenizer @@ -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 } diff --git a/python/parts.nix b/python/parts.nix index 46b4abdf..bac8ddab 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; diff --git a/uv.lock b/uv.lock index 627e6951..74687ad3 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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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,7 +416,7 @@ 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 = "msgspec", specifier = ">=0.19.0" }, @@ -1020,8 +1020,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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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 +1048,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 +1066,14 @@ cuda13 = [ { name = "mlx-cuda-13", marker = "sys_platform == 'linux'" }, ] +[[package]] +name = "mlx" +version = "0.30.7.dev20260217+50487b41" +source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" } +resolution-markers = [ + "sys_platform == 'darwin'", +] + [[package]] name = "mlx-cpu" version = "0.30.6" @@ -1102,7 +1104,7 @@ version = "0.30.6" 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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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'" }, @@ -1114,16 +1116,6 @@ 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" }, -] - [[package]] name = "more-itertools" version = "10.8.0" From 3addeadea8c8a0f83ec3f7d00e47e4e69b327861 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Wed, 18 Feb 2026 03:14:23 -0800 Subject: [PATCH 14/45] Update mlx-lm to 0.30.7 (#1520) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Bumps `mlx-lm` from 0.30.6 to 0.30.7 in `pyproject.toml` and `uv.lock` ## Test plan - [x] `uv lock` resolves successfully - [x] `basedpyright` — no new errors (63 pre-existing in unrelated `test_tool_call_tracker.py`) - [x] `ruff check` — all checks passed - [x] `nix fmt` — no formatting changes - [x] `pytest` — 188 passed, 1 skipped 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 --- pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 02aa6071..c4bfa550 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,7 @@ dependencies = [ "anyio==4.11.0", "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", diff --git a/uv.lock b/uv.lock index 74687ad3..59e5ab69 100644 --- a/uv.lock +++ b/uv.lock @@ -418,7 +418,7 @@ requires-dist = [ { name = "mflux", specifier = "==0.15.5" }, { 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" }, @@ -1100,7 +1100,7 @@ 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'" }, @@ -1111,9 +1111,9 @@ dependencies = [ { 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" }, + { 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]] From 8f01523ddbf86155e88cc25907690f1ddf416825 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Wed, 18 Feb 2026 11:43:27 +0000 Subject: [PATCH 15/45] remove dead code (#1496) --- Cargo.lock | 123 --------- Cargo.toml | 28 -- rust/exo_pyo3_bindings/Cargo.toml | 25 +- rust/exo_pyo3_bindings/src/allow_threading.rs | 6 +- rust/exo_pyo3_bindings/src/examples/mod.rs | 240 ------------------ rust/exo_pyo3_bindings/src/lib.rs | 29 +-- rust/exo_pyo3_bindings/src/networking.rs | 7 +- rust/networking/Cargo.toml | 9 +- rust/networking/examples/chatroom_manual.rs | 2 +- rust/networking/src/discovery.rs | 1 - rust/networking/src/keep_alive.rs | 44 ---- rust/networking/src/lib.rs | 20 -- 12 files changed, 14 insertions(+), 520 deletions(-) delete mode 100644 rust/exo_pyo3_bindings/src/examples/mod.rs delete mode 100644 rust/networking/src/keep_alive.rs diff --git a/Cargo.lock b/Cargo.lock index a45bfe9d..089f5b65 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", "libp2p", "log", "networking", - "once_cell", "pin-project", "pyo3", "pyo3-async-runtimes", "pyo3-log", "pyo3-stub-gen", - "thiserror 2.0.17", - "thread_local", "tokio", "util", ] @@ -1640,17 +1584,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 +1762,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 +2751,13 @@ name = "networking" version = "0.0.1" dependencies = [ "delegate", - "derive_more", "either", "extend", "futures", "futures-timer", - "impl-trait-for-tuples", "keccak-const", "libp2p", "log", - "thiserror 2.0.17", "tokio", "tracing-subscriber", "util", @@ -2918,17 +2842,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 +3192,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 +3640,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 +4504,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..4c0b8721 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,49 +26,21 @@ 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-timer = "3.0" # Data structures either = "1.15" -ordered-float = "5.0" -ahash = "0.8" # Tracing/logging log = "0.4" diff --git a/rust/exo_pyo3_bindings/Cargo.toml b/rust/exo_pyo3_bindings/Cargo.toml index 12803ab4..e97b9148 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 +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,8 +45,6 @@ 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 @@ -54,24 +52,11 @@ tokio = { workspace = true, features = ["full", "tracing"] } futures = { 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"] } diff --git a/rust/exo_pyo3_bindings/src/allow_threading.rs b/rust/exo_pyo3_bindings/src/allow_threading.rs index 3106e535..18a426aa 100644 --- a/rust/exo_pyo3_bindings/src/allow_threading.rs +++ b/rust/exo_pyo3_bindings/src/allow_threading.rs @@ -6,7 +6,7 @@ use pyo3::marker::Ungil; use pyo3::prelude::*; use std::{ future::Future, - pin::{Pin, pin}, + pin::Pin, task::{Context, Poll}, }; @@ -33,8 +33,6 @@ where 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/lib.rs b/rust/exo_pyo3_bindings/src/lib.rs index 4f591b8c..25f2865e 100644 --- a/rust/exo_pyo3_bindings/src/lib.rs +++ b/rust/exo_pyo3_bindings/src/lib.rs @@ -17,7 +17,6 @@ extern crate core; mod allow_threading; -mod examples; pub(crate) mod networking; pub(crate) mod pylibp2p; @@ -25,7 +24,6 @@ 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::{Bound, PyResult, pyclass, pymodule}; use pyo3_stub_gen::define_stub_info_gatherer; @@ -36,14 +34,10 @@ pub(crate) mod r#const { /// 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 @@ -51,7 +45,6 @@ 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 +55,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 +91,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 +168,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. diff --git a/rust/exo_pyo3_bindings/src/networking.rs b/rust/exo_pyo3_bindings/src/networking.rs index 01024cf9..03a0dcb6 100644 --- a/rust/exo_pyo3_bindings/src/networking.rs +++ b/rust/exo_pyo3_bindings/src/networking.rs @@ -11,9 +11,9 @@ use crate::ext::{ResultExt as _, TokioMpscReceiverExt as _, TokioMpscSenderExt a 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] @@ -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"); @@ -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/networking/Cargo.toml b/rust/networking/Cargo.toml index 47d61f41..fd5f1b1f 100644 --- a/rust/networking/Cargo.toml +++ b/rust/networking/Cargo.toml @@ -19,8 +19,6 @@ 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"] } @@ -29,11 +27,6 @@ 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_manual.rs b/rust/networking/examples/chatroom_manual.rs index 5d92ac86..e9f10a1f 100644 --- a/rust/networking/examples/chatroom_manual.rs +++ b/rust/networking/examples/chatroom_manual.rs @@ -24,8 +24,8 @@ use libp2p::{ swarm::{NetworkBehaviour, SwarmEvent}, tcp, yamux, }; +use std::error::Error; use std::time::Duration; -use std::{error::Error, hash::Hash}; use tokio::{io, io::AsyncBufReadExt, select}; use tracing_subscriber::EnvFilter; diff --git a/rust/networking/src/discovery.rs b/rust/networking/src/discovery.rs index b9a4052c..581f2200 100644 --- a/rust/networking/src/discovery.rs +++ b/rust/networking/src/discovery.rs @@ -1,5 +1,4 @@ use crate::ext::MultiaddrExt; -use crate::keep_alive; use delegate::delegate; use either::Either; use futures::FutureExt; 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 {} -} From 5cbd6377a2120d70d6e78d29916f83c84afa7190 Mon Sep 17 00:00:00 2001 From: Evan Date: Wed, 18 Feb 2026 12:48:42 +0000 Subject: [PATCH 16/45] prioritize official model cards over custom model cards our old model card search path would override official model cards with custom model cards - our packaged model cards should always be the default here --- src/exo/shared/models/model_cards.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index f9079854..99d1ed2e 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 From 48b8f863954a8f079e4bd89297dc064affb1d9b1 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Wed, 18 Feb 2026 14:04:06 +0000 Subject: [PATCH 17/45] Add support for GLM 5 (#1526) ## Motivation Add GLM 5 support in favor of #1513 ## Changes ## Why It Works ## Test Plan ### Manual Testing ### Automated Testing --- .mlx_typings/mlx_lm/models/glm_moe_dsa.pyi | 46 +++++++++++++++++++ nix/mlx.nix | 6 +-- .../mlx-community--GLM-5-8bit.toml | 12 +++++ .../mlx-community--GLM-5-MXFP4-Q8.toml | 12 +++++ .../mlx-community--GLM-5-bf16.toml | 12 +++++ src/exo/shared/models/model_cards.py | 1 + src/exo/worker/engines/mlx/auto_parallel.py | 39 ++++++++++------ .../worker/engines/mlx/generator/generate.py | 14 +++--- src/exo/worker/engines/mlx/utils_mlx.py | 4 +- src/exo/worker/runner/runner_supervisor.py | 2 +- uv.lock | 10 ++-- 11 files changed, 127 insertions(+), 31 deletions(-) create mode 100644 .mlx_typings/mlx_lm/models/glm_moe_dsa.pyi create mode 100644 resources/inference_model_cards/mlx-community--GLM-5-8bit.toml create mode 100644 resources/inference_model_cards/mlx-community--GLM-5-MXFP4-Q8.toml create mode 100644 resources/inference_model_cards/mlx-community--GLM-5-bf16.toml 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/nix/mlx.nix b/nix/mlx.nix index 8a40e11b..4434d3cc 100644 --- a/nix/mlx.nix +++ b/nix/mlx.nix @@ -41,7 +41,7 @@ let mlx = stdenv.mkDerivation rec { pname = "mlx"; - version = let v = "0.30.7.dev20260217+50487b41"; 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; @@ -49,8 +49,8 @@ let src = fetchFromGitHub { owner = "rltakashige"; repo = "mlx-jaccl-fix-small-recv"; - rev = "50487b4141f3c951122655db3b83df5146c1fbeb"; - hash = "sha256-IL4a9vMX5nocgJU1WG4zE8hArHkHJtnh4sdYh3od5zU="; + rev = "1484197707f35186ad3bd614357c7c47fdf86ebc"; + hash = "sha256-FupCMoK/SF/ldfKuvMSAKECcOP8c+ANgkQlPZttDsLk="; }; patches = [ 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/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index 99d1ed2e..d02271e9 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -183,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/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/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index ffa2f5c0..a9fcd64b 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -57,6 +57,7 @@ def prefill( sampler: Callable[[mx.array], mx.array], prompt_tokens: mx.array, cache: KVCacheType, + group: mx.distributed.Group | None = None, ) -> tuple[float, int, list[CacheSnapshot]]: """Prefill the KV cache with prompt tokens. @@ -86,6 +87,9 @@ def prefill( 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( @@ -305,16 +309,9 @@ def mlx_generate( ) max_stop_len = max((len(s) for s in stop_sequences), default=0) - mx_barrier(group) - logger.info("Starting prefill") - # Prefill cache with all tokens except the last one prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( - model, - tokenizer, - sampler, - prompt_tokens[:-1], - caches, + model, tokenizer, sampler, prompt_tokens[:-1], caches, group ) cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None @@ -331,6 +328,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( diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 3ed65ecc..30b489f4 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -285,10 +285,12 @@ 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] return None diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 5d39a881..0458edac 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -191,7 +191,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: diff --git a/uv.lock b/uv.lock index 59e5ab69..1232004e 100644 --- a/uv.lock +++ b/uv.lock @@ -378,7 +378,7 @@ dependencies = [ { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mflux", marker = "sys_platform == 'darwin' or 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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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 = "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'" }, @@ -1021,7 +1021,7 @@ dependencies = [ { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "matplotlib", marker = "sys_platform == 'darwin' or 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.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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 = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1068,8 +1068,8 @@ cuda13 = [ [[package]] name = "mlx" -version = "0.30.7.dev20260217+50487b41" -source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" } +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'", ] @@ -1104,7 +1104,7 @@ version = "0.30.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", version = "0.30.7.dev20260217+50487b41", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#50487b4141f3c951122655db3b83df5146c1fbeb" }, 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'" }, From f54c80b1217701e92fbeb8a5d8992c527a3821f9 Mon Sep 17 00:00:00 2001 From: ciaranbor <81697641+ciaranbor@users.noreply.github.com> Date: Wed, 18 Feb 2026 16:05:39 +0000 Subject: [PATCH 18/45] Ciaran/image edit api (#1500) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation - Image editing previously ignored input image dimensions, always defaulting to 1024x1024 - Size dropdown was hidden in edit mode, giving users no control over output dimensions - Portrait/landscape presets used non-standard aspect ratios (1024x1365 / 1365x1024) ## Changes - Added "auto" size option that uses input image dimensions for edits, defaults to 1024x1024 for generation - Introduced ImageSize Literal type and normalize_image_size() validator (replaces raw str size fields) - Updated portrait/landscape presets to standard 1024x1536 / 1536x1024 - Made size selector visible in edit mode (previously hidden) - Default size changed from "1024x1024" to "auto" ## Why It Works - "auto" reads actual input image dimensions via PIL at generation time, so edits preserve the original aspect ratio - Pydantic field_validator on both ImageGenerationTaskParams and ImageEditsTaskParams normalizes None → "auto", keeping the API backward-compatible ## Test Plan ### Manual Testing - Verify image edits output at the input image's native resolution when size is "auto" - Verify size dropdown appears and works in both generate and edit modes --- .../lib/components/ImageParamsPanel.svelte | 169 +++++++++--------- dashboard/src/lib/stores/app.svelte.ts | 7 +- src/exo/master/api.py | 29 +-- src/exo/shared/types/api.py | 39 +++- src/exo/worker/engines/image/generate.py | 8 +- 5 files changed, 146 insertions(+), 106 deletions(-) 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/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index e5dbf902..1a6b1e3d 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -306,13 +306,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 +337,7 @@ export interface EditingImage { } const DEFAULT_IMAGE_PARAMS: ImageGenerationParams = { - size: "1024x1024", + size: "auto", quality: "medium", outputFormat: "png", numImages: 1, diff --git a/src/exo/master/api.py b/src/exo/master/api.py index b8476334..3c29b041 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -85,6 +85,7 @@ from exo.shared.types.api import ( ImageGenerationTaskParams, ImageListItem, ImageListResponse, + ImageSize, ModelList, ModelListModel, PlaceInstanceParams, @@ -100,6 +101,7 @@ from exo.shared.types.api import ( TraceRankStats, TraceResponse, TraceStatsResponse, + normalize_image_size, ) from exo.shared.types.chunks import ( ErrorChunk, @@ -751,9 +753,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 +1013,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 +1040,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 +1110,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 +1136,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 +1172,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 +1192,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, diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py index 2756f0d4..824c2c91 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 @@ -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/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py index a59e4eed..f5526c8f 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 From 2ebe6216b431643d7a7dc16f3914e8dba2db9cd6 Mon Sep 17 00:00:00 2001 From: vskiwi <141816715+vskiwi@users.noreply.github.com> Date: Wed, 18 Feb 2026 19:18:09 +0300 Subject: [PATCH 19/45] feat: add explicit --offline mode for air-gapped clusters (#1525) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation Closes #1510 There is currently no reliable way to run exo on an air-gapped or offline cluster where models are pre-staged on local disks. The two existing mechanisms — `--no-downloads` and `HF_HUB_OFFLINE=1` — each cover only a subset of the problem: 1. **`--no-downloads` blocks model loading**: When passed, `DownloadCoordinator` is not created. No `NodeDownloadProgress` events are ever emitted, so `_model_needs_download()` in `plan.py` perpetually returns `DownloadModel`, short-circuiting `_load_model()` and preventing the model from ever being loaded. 2. **`HF_HUB_OFFLINE=1` doesn't cover exo's aiohttp code**: exo's download pipeline primarily uses raw `aiohttp` for HTTP operations (file list fetching, file downloads, HEAD verification), not the `huggingface_hub` library. These calls will attempt connections and time out on air-gapped networks. 3. **`skip_internet` is not propagated to `download_file_with_retry()`**: Even when `internet_connection = False`, the `_download_file()` function still makes HTTP HEAD calls via `file_meta()` to verify local files and unconditionally attempts downloads for missing files. ## Changes ### `src/exo/main.py` - Add `--offline` flag to `Args` with env var detection (`EXO_OFFLINE=1`, `HF_HUB_OFFLINE=1`) - Pass `offline` to `DownloadCoordinator` at creation and re-creation (election loop) ### `src/exo/download/coordinator.py` - Add `offline: bool = False` field - In offline mode: set `internet_connection = False` immediately in `__post_init__`, skip `_test_internet_connection()` ping (avoids 3s timeout), skip `_check_internet_connection` periodic loop - In `_start_download()`: if model is not fully available locally, emit `DownloadFailed` with clear message instead of starting a download task ### `src/exo/download/download_utils.py` - Add `skip_internet: bool` parameter to `download_file_with_retry()` and `_download_file()` - When `skip_internet=True` in `_download_file()`: return local file immediately without HTTP HEAD verification; raise `FileNotFoundError` for missing files - Propagate `skip_internet` from `download_shard()` to `download_file_with_retry()` ### `src/exo/download/tests/test_offline_mode.py` (new) - 8 tests covering `_download_file`, `download_file_with_retry`, and `fetch_file_list_with_cache` in offline mode ## Why It Works Unlike `--no-downloads` which disables `DownloadCoordinator` entirely, `--offline` keeps the coordinator running in a restricted mode. The existing `_emit_existing_download_progress()` disk scanner still runs every 60 seconds, emitting `DownloadCompleted` events for pre-staged models. These events flow through the event-sourcing pipeline and populate `state.downloads`, which unblocks `_model_needs_download()` in `plan.py` — no changes to the planning logic required. ``` --offline flag → DownloadCoordinator (offline mode) → Skip 1.1.1.1 ping, internet_connection = False → _emit_existing_download_progress scans disk → Emits DownloadCompleted for pre-staged models → _model_needs_download sees DownloadCompleted → _load_model proceeds normally ``` ## Test Plan ### Automated Testing - `ruff check` — passes - 8 new tests in `test_offline_mode.py` — all pass - 11 existing download tests in `test_download_verification.py` — all pass (no regressions) ### Manual Testing 1. Pre-stage a model on disk (e.g., `~/.exo/models/mlx-community--Qwen3-0.6B-4bit/`) 2. Start exo with `--offline` (or `EXO_OFFLINE=1`) 3. Place an instance via API or dashboard 4. Verify: model loads into memory and inference works without any network calls ### Environment - macOS (Apple Silicon), multi-node cluster with Thunderbolt interconnect - Models pre-staged via rsync / NFS mount --- src/exo/download/coordinator.py | 27 ++- src/exo/download/download_utils.py | 13 +- src/exo/download/tests/test_offline_mode.py | 230 ++++++++++++++++++++ src/exo/main.py | 13 ++ 4 files changed, 279 insertions(+), 4 deletions(-) create mode 100644 src/exo/download/tests/test_offline_mode.py diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index db13ccef..899e4f14 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -47,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) @@ -62,6 +63,8 @@ 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: @@ -107,13 +110,17 @@ 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: @@ -202,6 +209,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) diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 7974d504..5691e5dd 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -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 @@ -814,6 +824,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/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 1d358975..ec203181 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -39,6 +39,7 @@ class Node: node_id: NodeId event_index_counter: Iterator[int] + offline: bool _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group) @classmethod @@ -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,6 +134,7 @@ class Node: api, node_id, event_index_counter, + args.offline, ) async def run(self): @@ -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", From 6c322ebb72c251ad0c8cd38c9009d5ae60ad7f07 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Wed, 18 Feb 2026 09:05:00 -0800 Subject: [PATCH 20/45] feat: only show thinking toggle for models that support it (#1497) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Adds `thinking_toggle` capability to 26 model cards that support toggling thinking mode on/off - GPT-OSS models (20b, 120b) excluded — they always think and don't support toggling - Dashboard UI updated to check for `thinking_toggle` capability before showing the toggle button ## Test plan - [x] `uv run basedpyright` — 0 errors - [x] `uv run ruff check` — all checks passed - [x] `nix fmt` — 0 files changed - [x] `uv run pytest` — 188 passed, 0 failed - [x] Security review passed (no secrets, eval/exec, innerHTML, or dep changes) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 --- dashboard/src/lib/components/ChatForm.svelte | 2 +- .../mlx-community--DeepSeek-V3.1-4bit.toml | 2 +- .../mlx-community--DeepSeek-V3.1-8bit.toml | 2 +- .../inference_model_cards/mlx-community--GLM-4.5-Air-8bit.toml | 2 +- .../inference_model_cards/mlx-community--GLM-4.5-Air-bf16.toml | 2 +- .../inference_model_cards/mlx-community--GLM-4.7-4bit.toml | 2 +- .../inference_model_cards/mlx-community--GLM-4.7-6bit.toml | 2 +- .../inference_model_cards/mlx-community--GLM-4.7-8bit-gs32.toml | 2 +- .../mlx-community--GLM-4.7-Flash-4bit.toml | 2 +- .../mlx-community--GLM-4.7-Flash-5bit.toml | 2 +- .../mlx-community--GLM-4.7-Flash-6bit.toml | 2 +- .../mlx-community--GLM-4.7-Flash-8bit.toml | 2 +- .../inference_model_cards/mlx-community--Kimi-K2-Thinking.toml | 2 +- resources/inference_model_cards/mlx-community--Kimi-K2.5.toml | 2 +- .../inference_model_cards/mlx-community--MiniMax-M2.1-3bit.toml | 2 +- .../inference_model_cards/mlx-community--MiniMax-M2.1-8bit.toml | 2 +- .../inference_model_cards/mlx-community--Qwen3-0.6B-4bit.toml | 2 +- .../inference_model_cards/mlx-community--Qwen3-0.6B-8bit.toml | 2 +- .../mlx-community--Qwen3-235B-A22B-Instruct-2507-4bit.toml | 2 +- .../mlx-community--Qwen3-235B-A22B-Instruct-2507-8bit.toml | 2 +- .../mlx-community--Qwen3-30B-A3B-4bit.toml | 2 +- .../mlx-community--Qwen3-30B-A3B-8bit.toml | 2 +- .../mlx-community--Qwen3-Next-80B-A3B-Thinking-4bit.toml | 2 +- .../mlx-community--Qwen3-Next-80B-A3B-Thinking-8bit.toml | 2 +- .../mlx-community--Step-3.5-Flash-4bit.toml | 2 +- .../mlx-community--Step-3.5-Flash-6bit.toml | 2 +- .../mlx-community--Step-3.5-Flash-8Bit.toml | 2 +- 27 files changed, 27 insertions(+), 27 deletions(-) diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte index 2eb7aa40..6ccba979 100644 --- a/dashboard/src/lib/components/ChatForm.svelte +++ b/dashboard/src/lib/components/ChatForm.svelte @@ -103,7 +103,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( 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--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--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 From c2f2111b887e9ff521d81746a9b93bff20ab2322 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Wed, 18 Feb 2026 20:29:18 +0000 Subject: [PATCH 21/45] Fix tool calling (#1529) ## Motivation GPT OSS tool calling issues. ## Changes Fixes those and adds a bunch of evals for tool calling. Fixes GLM5 prefix caching, where CacheList wasn't getting handled properly. Extracts a bunch of the setup functionality of exo bench to a harness that can be reused elsewhere, such as in the tool calling eval. ## Test Plan ### Automated Testing Let's run the evals for all models --- .mlx_typings/mlx/nn/layers/base.pyi | 2 +- .mlx_typings/mlx/utils.pyi | 19 +- bench/eval_tool_calls.py | 1046 +++++++++++++++++ bench/exo_bench.py | 481 +------- bench/harness.py | 327 ++++++ bench/pyproject.toml | 1 + bench/scenarios.toml | 240 ++++ python/parts.nix | 1 + src/exo/shared/types/mlx.py | 5 +- src/exo/worker/engines/mlx/cache.py | 51 +- .../worker/engines/mlx/generator/generate.py | 17 +- src/exo/worker/engines/mlx/utils_mlx.py | 2 + src/exo/worker/runner/runner.py | 7 +- src/exo/worker/runner/runner_supervisor.py | 2 +- .../tests/unittests/test_mlx/conftest.py | 9 +- .../test_mlx/test_kv_prefix_cache.py | 60 +- .../test_prefix_cache_architectures.py | 297 +++++ .../unittests/test_mlx/test_tokenizers.py | 10 +- .../test_runner/test_parse_gpt_oss.py | 162 +++ tests/eval_tool_calls.sh | 55 + tool_call_eval.py | 691 +++++++++++ uv.lock | 2 + 22 files changed, 2965 insertions(+), 522 deletions(-) create mode 100644 bench/eval_tool_calls.py create mode 100644 bench/harness.py create mode 100644 bench/scenarios.toml create mode 100644 src/exo/worker/tests/unittests/test_mlx/test_prefix_cache_architectures.py create mode 100644 src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py create mode 100755 tests/eval_tool_calls.sh create mode 100644 tool_call_eval.py 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/bench/eval_tool_calls.py b/bench/eval_tool_calls.py new file mode 100644 index 00000000..6cd84396 --- /dev/null +++ b/bench/eval_tool_calls.py @@ -0,0 +1,1046 @@ +# 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, + 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 + + +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, + ) + ) + + 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 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 + 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] + 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..7e19a7a6 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -4,26 +4,29 @@ 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, + 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 +106,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 +124,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 +215,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 +233,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 +242,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 +256,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 +290,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,16 +317,6 @@ def main() -> int: if args.dry_run: return 0 - logger.info("Planning phase: checking downloads...") - run_planning_phase( - client, - full_model_id, - selected[0], - args.danger_delete_downloads, - args.timeout, - settle_deadline, - ) - all_rows: list[dict[str, Any]] = [] for preview in selected: diff --git a/bench/harness.py b/bench/harness.py new file mode 100644 index 00000000..c8ae9318 --- /dev/null +++ b/bench/harness.py @@ -0,0 +1,327 @@ +# 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 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).", + ) 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..e258b5c6 --- /dev/null +++ b/bench/scenarios.toml @@ -0,0 +1,240 @@ +# 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" + +# -- 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"}' + +# -- 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/python/parts.nix b/python/parts.nix index bac8ddab..9ba703a0 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -158,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/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/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index c747ba4e..7669f1c1 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 = psutil.virtual_memory().total / (1024**3) + 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/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index a9fcd64b..7b78cbd9 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -48,7 +48,7 @@ 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 def prefill( @@ -57,7 +57,7 @@ def prefill( sampler: Callable[[mx.array], mx.array], prompt_tokens: mx.array, cache: KVCacheType, - group: mx.distributed.Group | None = None, + group: mx.distributed.Group | None, ) -> tuple[float, int, list[CacheSnapshot]]: """Prefill the KV cache with prompt tokens. @@ -133,7 +133,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." @@ -255,8 +255,8 @@ 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, ) -> Generator[GenerationResponse]: # Ensure that generation stats only contains peak memory for this generation mx.reset_peak_memory() @@ -436,9 +436,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 30b489f4..d9cebae9 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -292,6 +292,8 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None: 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 diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index e55456d3..c749159a 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -11,6 +11,7 @@ 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, @@ -588,7 +589,11 @@ 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 diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 0458edac..58ac778e 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -103,7 +103,7 @@ class RunnerSupervisor: self._event_sender.close() 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 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_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..0a7ba102 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py @@ -0,0 +1,162 @@ +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) + + # Should have thinking tags + content + tool call + text_parts = [r.text for r in results if isinstance(r, GenerationResponse)] + combined = "".join(text_parts) + assert "" in combined + assert "" in combined + assert "Let me think about this." in combined + + # 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/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/tool_call_eval.py b/tool_call_eval.py new file mode 100644 index 00000000..cced0bf2 --- /dev/null +++ b/tool_call_eval.py @@ -0,0 +1,691 @@ +#!/usr/bin/env python3 +"""Tool-calling eval for exo's OpenAI-compatible API. + +Tests whether models correctly: +- Trigger tool calls when appropriate +- Return valid JSON arguments matching function schemas +- Handle multi-turn tool use (call -> result -> final answer) +- Avoid calling tools when unnecessary + +Start exo with a model first, then run: + uv run python tool_call_eval.py --model + 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 json +import os +import sys +import time +from dataclasses import dataclass, field + +import httpx + +# --------------------------------------------------------------------------- +# Tool definitions +# --------------------------------------------------------------------------- + +WEATHER_TOOL = { + "type": "function", + "function": { + "name": "get_current_weather", + "description": "Get the current weather in a given location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City and state, e.g. San Francisco, CA", + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"], + "description": "Temperature unit", + }, + }, + "required": ["location"], + }, + }, +} + +CALCULATOR_TOOL = { + "type": "function", + "function": { + "name": "calculate", + "description": "Evaluate a mathematical expression and return the numeric result", + "parameters": { + "type": "object", + "properties": { + "expression": { + "type": "string", + "description": "The math expression to evaluate, e.g. '2 + 3 * 4'", + }, + }, + "required": ["expression"], + }, + }, +} + +SEARCH_TOOL = { + "type": "function", + "function": { + "name": "search_products", + "description": "Search for products in a catalog by query, category, and price", + "parameters": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Search query string", + }, + "category": { + "type": "string", + "enum": ["electronics", "clothing", "food", "books"], + "description": "Product category to filter by", + }, + "max_price": { + "type": "number", + "description": "Maximum price in USD", + }, + }, + "required": ["query"], + }, + }, +} + +ALL_TOOLS = [WEATHER_TOOL, CALCULATOR_TOOL, SEARCH_TOOL] + +# --------------------------------------------------------------------------- +# Scenarios +# --------------------------------------------------------------------------- + + +@dataclass +class Scenario: + name: str + description: str + messages: list[dict[str, object]] + tools: list[dict[str, object]] + expect_tool_call: bool + expected_function: str | None = None + required_arg_keys: list[str] | None = None + # For multi-turn: fake tool result to inject, then verify the follow-up. + tool_result: str | None = None + + +SCENARIOS = [ + # -- Should call a tool -------------------------------------------------- + Scenario( + name="weather_simple", + description="Basic weather query -> get_current_weather", + messages=[ + {"role": "user", "content": "What's the weather like in Tokyo right now?"} + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="get_current_weather", + required_arg_keys=["location"], + ), + Scenario( + name="calculator_simple", + description="Math question -> calculate", + messages=[ + { + "role": "user", + "content": "Use the calculator to compute 3847 * 926 + 17293", + } + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="calculate", + required_arg_keys=["expression"], + ), + Scenario( + name="search_with_filters", + description="Product search with category and price filter", + messages=[{"role": "user", "content": "Find me electronics under $50"}], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="search_products", + required_arg_keys=["query"], + ), + # -- Multi-turn: tool call then follow-up -------------------------------- + Scenario( + name="weather_multi_turn", + description="Weather query -> tool result -> natural language summary", + messages=[{"role": "user", "content": "What's the weather in Paris?"}], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="get_current_weather", + required_arg_keys=["location"], + tool_result=json.dumps( + { + "temperature": "18C", + "condition": "partly cloudy", + "humidity": "65%", + "wind": "12 km/h NW", + } + ), + ), + Scenario( + name="calculator_multi_turn", + description="Math query -> tool result -> model reports the answer", + messages=[ + { + "role": "user", + "content": "Use the calculator to compute 1847 * 263 + 5921", + } + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="calculate", + required_arg_keys=["expression"], + tool_result=json.dumps({"result": 491682}), + ), + Scenario( + name="search_multi_turn", + description="Search query -> tool result -> model summarizes products", + messages=[ + {"role": "user", "content": "Search for books about machine learning"} + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="search_products", + required_arg_keys=["query"], + tool_result=json.dumps( + { + "results": [ + { + "name": "Hands-On Machine Learning", + "price": 45.99, + "rating": 4.8, + }, + { + "name": "Deep Learning with Python", + "price": 39.99, + "rating": 4.6, + }, + ] + } + ), + ), + # -- Sequential tool calls: thinking + tool call, NO final answer ---------- + # This is the critical scenario for the Harmony recipient placement fix. + # + # When an assistant message has both thinking content and a tool_call, + # AND there is no subsequent final-answer assistant message, the Jinja + # template renders BOTH the analysis and the tool call: + # + # <|start|>assistant<|channel|>analysis<|message|>thinking...<|end|> + # <|start|>assistant to=functions.X<|channel|>commentary json<|message|>...<|call|> + # + # The two consecutive assistant messages have INCONSISTENT start patterns + # (one has <|channel|> immediately, the other has to= first). + # This confuses the model when it needs to generate its own tool call. + # + # The reformat fix makes both start with <|start|>assistant<|channel|>, + # only differing in the channel name (analysis vs commentary). + Scenario( + name="chained_tool_calls_same", + description="Thinking + weather(Tokyo) -> result -> model must call weather(London)", + messages=[ + {"role": "user", "content": "Compare the weather in Tokyo and London."}, + { + "role": "assistant", + "content": "I'll check both cities. Let me start with Tokyo.", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": json.dumps({"location": "Tokyo"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": json.dumps({"temperature": "25C", "condition": "sunny"}), + }, + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="get_current_weather", + required_arg_keys=["location"], + ), + Scenario( + name="chained_tool_calls_different", + description="Thinking + weather(Berlin) -> result -> model must call calculator", + messages=[ + { + "role": "user", + "content": "What's the weather in Berlin, and also use the calculator to compute 4819 * 37 + 291.", + }, + { + "role": "assistant", + "content": "I'll handle both. Let me check Berlin's weather first.", + "tool_calls": [ + { + "id": "call_2", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": json.dumps({"location": "Berlin"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_2", + "content": json.dumps({"temperature": "12C", "condition": "rainy"}), + }, + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="calculate", + required_arg_keys=["expression"], + ), + Scenario( + name="chained_tool_calls_three", + description="Two prior thinking+tool calls -> results -> model must make a third", + messages=[ + {"role": "user", "content": "Compare weather in Tokyo, Paris, and London."}, + { + "role": "assistant", + "content": "I'll check all three cities. Starting with Tokyo.", + "tool_calls": [ + { + "id": "call_3", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": json.dumps({"location": "Tokyo"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_3", + "content": json.dumps({"temperature": "25C", "condition": "sunny"}), + }, + { + "role": "assistant", + "content": "Got Tokyo. Now checking Paris.", + "tool_calls": [ + { + "id": "call_4", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": json.dumps({"location": "Paris"}), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_4", + "content": json.dumps({"temperature": "18C", "condition": "cloudy"}), + }, + ], + tools=ALL_TOOLS, + expect_tool_call=True, + expected_function="get_current_weather", + required_arg_keys=["location"], + ), + # -- Should NOT call a tool ---------------------------------------------- + Scenario( + name="no_tool_joke", + description="Joke request should NOT trigger any tool", + messages=[{"role": "user", "content": "Tell me a funny joke about cats."}], + tools=ALL_TOOLS, + expect_tool_call=False, + ), + Scenario( + name="no_tool_factual", + description="Factual question answerable from training data", + messages=[{"role": "user", "content": "What is the capital of Japan?"}], + tools=ALL_TOOLS, + expect_tool_call=False, + ), +] + +# --------------------------------------------------------------------------- +# Result tracking +# --------------------------------------------------------------------------- + + +@dataclass +class ScenarioResult: + name: 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 + + +# --------------------------------------------------------------------------- +# Evaluation helpers +# --------------------------------------------------------------------------- + + +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 e: + return False, f"Invalid JSON: {e}" + 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 call_api( + client: httpx.Client, + base_url: str, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]], + timeout: float, +) -> tuple[dict[str, object], float]: + """POST to /chat/completions, return (response_json, latency_ms).""" + url = f"{base_url.rstrip('/')}/chat/completions" + body: dict[str, object] = { + "model": model, + "messages": messages, + "tools": tools, + "temperature": 0.0, + "max_tokens": 4096, + } + t0 = time.monotonic() + resp = client.post(url, json=body, timeout=timeout) + latency = (time.monotonic() - t0) * 1000 + resp.raise_for_status() + return resp.json(), latency + + +# --------------------------------------------------------------------------- +# Scenario runner +# --------------------------------------------------------------------------- + + +def run_scenario( + client: httpx.Client, + base_url: str, + model: str, + scenario: Scenario, + timeout: float, + verbose: bool, +) -> list[ScenarioResult]: + results: list[ScenarioResult] = [] + + # --- Phase 1: initial request --- + try: + data, latency = call_api( + client, base_url, model, scenario.messages, scenario.tools, timeout + ) + except Exception as e: + results.append( + ScenarioResult( + name=scenario.name, + phase="tool_call", + passed=False, + error=f"API error: {e}", + ) + ) + return results + + if verbose: + print(f" response: {json.dumps(data, indent=2)}") + + choice = data["choices"][0] + finish_reason = choice.get("finish_reason") + message = choice.get("message", {}) + tool_calls = message.get("tool_calls") + content = message.get("content") + + checks: dict[str, bool] = {} + + if scenario.expect_tool_call: + checks["finish_reason_tool_calls"] = finish_reason == "tool_calls" + checks["has_tool_call"] = isinstance(tool_calls, list) and len(tool_calls) > 0 + + args_err: str | None = None + if checks["has_tool_call"]: + tc = tool_calls[0] + fn = tc.get("function", {}) + checks["correct_function"] = ( + scenario.expected_function is None + or fn.get("name") == scenario.expected_function + ) + if scenario.required_arg_keys: + ok, args_err = validate_args( + fn.get("arguments", ""), scenario.required_arg_keys + ) + checks["valid_arguments"] = ok + else: + checks["valid_arguments"] = True + 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"] = finish_reason == "stop" + checks["no_tool_call"] = tool_calls is None or len(tool_calls) == 0 + checks["has_content"] = isinstance(content, str) and len(content.strip()) > 0 + passed = all(checks.values()) + error = ( + None + if passed + else ( + f"finish_reason={finish_reason}, " + f"tool_calls={'yes' if tool_calls else 'no'}, " + f"content={'yes' if content else 'no'}" + ) + ) + + results.append( + ScenarioResult( + name=scenario.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 checks.get("has_tool_call"): + tc = tool_calls[0] + fn = tc.get("function", {}) + follow_up_messages: list[dict[str, object]] = list(scenario.messages) + [ + { + "role": "assistant", + "tool_calls": [ + { + "id": tc.get("id", "call_0"), + "type": "function", + "function": { + "name": fn.get("name", ""), + "arguments": fn.get("arguments", "{}"), + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": tc.get("id", "call_0"), + "content": scenario.tool_result, + }, + ] + + try: + data2, latency2 = call_api( + client, + base_url, + model, + follow_up_messages, + scenario.tools, + timeout, + ) + except Exception as e: + results.append( + ScenarioResult( + name=scenario.name, + phase="follow_up", + passed=False, + error=f"API error: {e}", + ) + ) + return results + + if verbose: + print(f" follow_up response: {json.dumps(data2, indent=2)}") + + choice2 = data2["choices"][0] + message2 = choice2.get("message", {}) + checks2: dict[str, bool] = {} + checks2["finish_reason_stop"] = choice2.get("finish_reason") == "stop" + tc2 = message2.get("tool_calls") + checks2["no_tool_call"] = tc2 is None or len(tc2) == 0 + c2 = message2.get("content") + checks2["has_content"] = isinstance(c2, str) and len(c2.strip()) > 0 + + passed2 = all(checks2.values()) + error2 = None + if not passed2: + error2 = ( + f"finish_reason={choice2.get('finish_reason')}, " + f"tool_calls={'yes' if tc2 else 'no'}, " + f"content={'yes' if c2 else 'no'}" + ) + results.append( + ScenarioResult( + name=scenario.name, + phase="follow_up", + passed=passed2, + checks=checks2, + error=error2, + latency_ms=latency2, + ) + ) + + return results + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + parser = argparse.ArgumentParser(description="Tool-calling eval for exo") + parser.add_argument("--model", required=True, help="Model ID to test") + parser.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost")) + parser.add_argument( + "--port", + type=int, + default=int(os.environ.get("EXO_PORT", "52415")), + ) + parser.add_argument( + "--timeout", type=float, default=120, help="Per-request timeout (seconds)" + ) + parser.add_argument( + "--repeat", type=int, default=1, help="Repeat each scenario N times" + ) + parser.add_argument( + "--scenarios", nargs="*", help="Run only these scenarios (by name)" + ) + parser.add_argument( + "--verbose", action="store_true", help="Print full API responses" + ) + args = parser.parse_args() + + scenarios = SCENARIOS + if args.scenarios: + scenarios = [s for s in SCENARIOS if s.name in args.scenarios] + if not scenarios: + print(f"No matching scenarios. Available: {[s.name for s in SCENARIOS]}") + sys.exit(1) + + base_url = f"http://{args.host}:{args.port}/v1" + total_runs = len(scenarios) * args.repeat + print(f"Model: {args.model}") + print(f"Endpoint: {base_url}") + print(f"Scenarios: {len(scenarios)} x {args.repeat} = {total_runs} runs") + print("=" * 64) + + all_results: list[ScenarioResult] = [] + + with httpx.Client() as client: + for run_idx in range(args.repeat): + if args.repeat > 1: + print(f"\n--- Run {run_idx + 1}/{args.repeat} ---") + + for scenario in scenarios: + print(f"\n {scenario.name}: {scenario.description}") + + results = run_scenario( + client, + base_url, + args.model, + scenario, + args.timeout, + args.verbose, + ) + all_results.extend(results) + + for r in results: + status = "PASS" if r.passed else "FAIL" + print(f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)") + for check_name, check_ok in r.checks.items(): + mark = "+" if check_ok else "-" + print(f" {mark} {check_name}") + if r.error: + print(f" ! {r.error}") + + # --- Summary --- + print(f"\n{'=' * 64}") + + 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}%)") + print(f"Tool call: {tc_passed}/{len(tool_call_results)} passed") + if follow_up_results: + print(f"Follow-up: {fu_passed}/{len(follow_up_results)} passed") + print(f"Avg latency: {avg_latency:.0f}ms") + + if passed < total: + print("\nFailed:") + for r in all_results: + if not r.passed: + print(f" - {r.name} [{r.phase}]: {r.error}") + + sys.exit(0 if passed == total else 1) + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index 1232004e..b14a1f69 100644 --- a/uv.lock +++ b/uv.lock @@ -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" }, From ce5a65d3b979eb5560be1fe975add239c8921e06 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Wed, 18 Feb 2026 13:11:13 -0800 Subject: [PATCH 22/45] Add MiniMax M2.5 model cards (#1514) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Adds model cards for MiniMax M2.5 in three quantizations: 4bit (~129 GB), 6bit (~186 GB), 8bit (~243 GB) - No code changes needed — `MiniMaxM2ForCausalLM` is already in the tensor parallel whitelist and `MiniMaxShardingStrategy` is already implemented in `auto_parallel.py` - Credit to @vskiwi for confirming MiniMax M2.5 works out of the box with existing code Closes #1480 ## Test plan - [x] `basedpyright` passes with 0 errors - [x] `ruff check` passes - [x] `pytest` passes (260 passed, 1 skipped) - [ ] Verify MiniMax M2.5 models appear in model selector on dashboard 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 Co-authored-by: rltakashige --- .../mlx-community--MiniMax-M2.5-4bit.toml | 12 ++++++++++++ .../mlx-community--MiniMax-M2.5-6bit.toml | 12 ++++++++++++ .../mlx-community--MiniMax-M2.5-8bit.toml | 12 ++++++++++++ 3 files changed, 36 insertions(+) create mode 100644 resources/inference_model_cards/mlx-community--MiniMax-M2.5-4bit.toml create mode 100644 resources/inference_model_cards/mlx-community--MiniMax-M2.5-6bit.toml create mode 100644 resources/inference_model_cards/mlx-community--MiniMax-M2.5-8bit.toml 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 From 315992549b9f92d7d7538e46fd045558e9ecff73 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Wed, 18 Feb 2026 13:59:02 -0800 Subject: [PATCH 23/45] fix: unblock MpReceiver.close() to prevent shutdown hang (#1511) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - `MpReceiver.close()` did not unblock threads stuck on `queue.get()` in `receive_async()`, causing abandoned threads (via `abandon_on_cancel=True`) to keep the Python process alive indefinitely after tests pass - This caused the `aarch64-darwin` CI jobs in PR #1462 to hang for ~6 hours until the GitHub Actions timeout killed them - Sends an `_MpEndOfStream` sentinel before closing the buffer, mirroring what `MpSender.close()` already does ## Test plan - [x] `uv run basedpyright` — 0 errors - [x] `uv run ruff check` — clean - [x] `nix fmt` — 0 changed - [x] `uv run pytest` — 188 passed, 1 skipped in 12s (no hang) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: rltakashige Co-authored-by: Ryuichi Leo Takashige --- src/exo/utils/channels.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py index 646ac8f6..c9336215 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 == @@ -204,6 +206,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 == From 24e99ce1970af89dfb004140c5bda222d9e9da13 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Wed, 18 Feb 2026 22:05:26 +0000 Subject: [PATCH 24/45] Cleanup mistakes (#1537) Oops --- bench/exo_bench.py | 17 +- tool_call_eval.py | 691 --------------------------------------------- 2 files changed, 16 insertions(+), 692 deletions(-) delete mode 100644 tool_call_eval.py diff --git a/bench/exo_bench.py b/bench/exo_bench.py index 7e19a7a6..9f6f5f02 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -1,5 +1,20 @@ +# 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 diff --git a/tool_call_eval.py b/tool_call_eval.py deleted file mode 100644 index cced0bf2..00000000 --- a/tool_call_eval.py +++ /dev/null @@ -1,691 +0,0 @@ -#!/usr/bin/env python3 -"""Tool-calling eval for exo's OpenAI-compatible API. - -Tests whether models correctly: -- Trigger tool calls when appropriate -- Return valid JSON arguments matching function schemas -- Handle multi-turn tool use (call -> result -> final answer) -- Avoid calling tools when unnecessary - -Start exo with a model first, then run: - uv run python tool_call_eval.py --model - 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 json -import os -import sys -import time -from dataclasses import dataclass, field - -import httpx - -# --------------------------------------------------------------------------- -# Tool definitions -# --------------------------------------------------------------------------- - -WEATHER_TOOL = { - "type": "function", - "function": { - "name": "get_current_weather", - "description": "Get the current weather in a given location", - "parameters": { - "type": "object", - "properties": { - "location": { - "type": "string", - "description": "City and state, e.g. San Francisco, CA", - }, - "unit": { - "type": "string", - "enum": ["celsius", "fahrenheit"], - "description": "Temperature unit", - }, - }, - "required": ["location"], - }, - }, -} - -CALCULATOR_TOOL = { - "type": "function", - "function": { - "name": "calculate", - "description": "Evaluate a mathematical expression and return the numeric result", - "parameters": { - "type": "object", - "properties": { - "expression": { - "type": "string", - "description": "The math expression to evaluate, e.g. '2 + 3 * 4'", - }, - }, - "required": ["expression"], - }, - }, -} - -SEARCH_TOOL = { - "type": "function", - "function": { - "name": "search_products", - "description": "Search for products in a catalog by query, category, and price", - "parameters": { - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "Search query string", - }, - "category": { - "type": "string", - "enum": ["electronics", "clothing", "food", "books"], - "description": "Product category to filter by", - }, - "max_price": { - "type": "number", - "description": "Maximum price in USD", - }, - }, - "required": ["query"], - }, - }, -} - -ALL_TOOLS = [WEATHER_TOOL, CALCULATOR_TOOL, SEARCH_TOOL] - -# --------------------------------------------------------------------------- -# Scenarios -# --------------------------------------------------------------------------- - - -@dataclass -class Scenario: - name: str - description: str - messages: list[dict[str, object]] - tools: list[dict[str, object]] - expect_tool_call: bool - expected_function: str | None = None - required_arg_keys: list[str] | None = None - # For multi-turn: fake tool result to inject, then verify the follow-up. - tool_result: str | None = None - - -SCENARIOS = [ - # -- Should call a tool -------------------------------------------------- - Scenario( - name="weather_simple", - description="Basic weather query -> get_current_weather", - messages=[ - {"role": "user", "content": "What's the weather like in Tokyo right now?"} - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="get_current_weather", - required_arg_keys=["location"], - ), - Scenario( - name="calculator_simple", - description="Math question -> calculate", - messages=[ - { - "role": "user", - "content": "Use the calculator to compute 3847 * 926 + 17293", - } - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="calculate", - required_arg_keys=["expression"], - ), - Scenario( - name="search_with_filters", - description="Product search with category and price filter", - messages=[{"role": "user", "content": "Find me electronics under $50"}], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="search_products", - required_arg_keys=["query"], - ), - # -- Multi-turn: tool call then follow-up -------------------------------- - Scenario( - name="weather_multi_turn", - description="Weather query -> tool result -> natural language summary", - messages=[{"role": "user", "content": "What's the weather in Paris?"}], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="get_current_weather", - required_arg_keys=["location"], - tool_result=json.dumps( - { - "temperature": "18C", - "condition": "partly cloudy", - "humidity": "65%", - "wind": "12 km/h NW", - } - ), - ), - Scenario( - name="calculator_multi_turn", - description="Math query -> tool result -> model reports the answer", - messages=[ - { - "role": "user", - "content": "Use the calculator to compute 1847 * 263 + 5921", - } - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="calculate", - required_arg_keys=["expression"], - tool_result=json.dumps({"result": 491682}), - ), - Scenario( - name="search_multi_turn", - description="Search query -> tool result -> model summarizes products", - messages=[ - {"role": "user", "content": "Search for books about machine learning"} - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="search_products", - required_arg_keys=["query"], - tool_result=json.dumps( - { - "results": [ - { - "name": "Hands-On Machine Learning", - "price": 45.99, - "rating": 4.8, - }, - { - "name": "Deep Learning with Python", - "price": 39.99, - "rating": 4.6, - }, - ] - } - ), - ), - # -- Sequential tool calls: thinking + tool call, NO final answer ---------- - # This is the critical scenario for the Harmony recipient placement fix. - # - # When an assistant message has both thinking content and a tool_call, - # AND there is no subsequent final-answer assistant message, the Jinja - # template renders BOTH the analysis and the tool call: - # - # <|start|>assistant<|channel|>analysis<|message|>thinking...<|end|> - # <|start|>assistant to=functions.X<|channel|>commentary json<|message|>...<|call|> - # - # The two consecutive assistant messages have INCONSISTENT start patterns - # (one has <|channel|> immediately, the other has to= first). - # This confuses the model when it needs to generate its own tool call. - # - # The reformat fix makes both start with <|start|>assistant<|channel|>, - # only differing in the channel name (analysis vs commentary). - Scenario( - name="chained_tool_calls_same", - description="Thinking + weather(Tokyo) -> result -> model must call weather(London)", - messages=[ - {"role": "user", "content": "Compare the weather in Tokyo and London."}, - { - "role": "assistant", - "content": "I'll check both cities. Let me start with Tokyo.", - "tool_calls": [ - { - "id": "call_1", - "type": "function", - "function": { - "name": "get_current_weather", - "arguments": json.dumps({"location": "Tokyo"}), - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_1", - "content": json.dumps({"temperature": "25C", "condition": "sunny"}), - }, - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="get_current_weather", - required_arg_keys=["location"], - ), - Scenario( - name="chained_tool_calls_different", - description="Thinking + weather(Berlin) -> result -> model must call calculator", - messages=[ - { - "role": "user", - "content": "What's the weather in Berlin, and also use the calculator to compute 4819 * 37 + 291.", - }, - { - "role": "assistant", - "content": "I'll handle both. Let me check Berlin's weather first.", - "tool_calls": [ - { - "id": "call_2", - "type": "function", - "function": { - "name": "get_current_weather", - "arguments": json.dumps({"location": "Berlin"}), - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_2", - "content": json.dumps({"temperature": "12C", "condition": "rainy"}), - }, - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="calculate", - required_arg_keys=["expression"], - ), - Scenario( - name="chained_tool_calls_three", - description="Two prior thinking+tool calls -> results -> model must make a third", - messages=[ - {"role": "user", "content": "Compare weather in Tokyo, Paris, and London."}, - { - "role": "assistant", - "content": "I'll check all three cities. Starting with Tokyo.", - "tool_calls": [ - { - "id": "call_3", - "type": "function", - "function": { - "name": "get_current_weather", - "arguments": json.dumps({"location": "Tokyo"}), - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_3", - "content": json.dumps({"temperature": "25C", "condition": "sunny"}), - }, - { - "role": "assistant", - "content": "Got Tokyo. Now checking Paris.", - "tool_calls": [ - { - "id": "call_4", - "type": "function", - "function": { - "name": "get_current_weather", - "arguments": json.dumps({"location": "Paris"}), - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": "call_4", - "content": json.dumps({"temperature": "18C", "condition": "cloudy"}), - }, - ], - tools=ALL_TOOLS, - expect_tool_call=True, - expected_function="get_current_weather", - required_arg_keys=["location"], - ), - # -- Should NOT call a tool ---------------------------------------------- - Scenario( - name="no_tool_joke", - description="Joke request should NOT trigger any tool", - messages=[{"role": "user", "content": "Tell me a funny joke about cats."}], - tools=ALL_TOOLS, - expect_tool_call=False, - ), - Scenario( - name="no_tool_factual", - description="Factual question answerable from training data", - messages=[{"role": "user", "content": "What is the capital of Japan?"}], - tools=ALL_TOOLS, - expect_tool_call=False, - ), -] - -# --------------------------------------------------------------------------- -# Result tracking -# --------------------------------------------------------------------------- - - -@dataclass -class ScenarioResult: - name: 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 - - -# --------------------------------------------------------------------------- -# Evaluation helpers -# --------------------------------------------------------------------------- - - -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 e: - return False, f"Invalid JSON: {e}" - 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 call_api( - client: httpx.Client, - base_url: str, - model: str, - messages: list[dict[str, object]], - tools: list[dict[str, object]], - timeout: float, -) -> tuple[dict[str, object], float]: - """POST to /chat/completions, return (response_json, latency_ms).""" - url = f"{base_url.rstrip('/')}/chat/completions" - body: dict[str, object] = { - "model": model, - "messages": messages, - "tools": tools, - "temperature": 0.0, - "max_tokens": 4096, - } - t0 = time.monotonic() - resp = client.post(url, json=body, timeout=timeout) - latency = (time.monotonic() - t0) * 1000 - resp.raise_for_status() - return resp.json(), latency - - -# --------------------------------------------------------------------------- -# Scenario runner -# --------------------------------------------------------------------------- - - -def run_scenario( - client: httpx.Client, - base_url: str, - model: str, - scenario: Scenario, - timeout: float, - verbose: bool, -) -> list[ScenarioResult]: - results: list[ScenarioResult] = [] - - # --- Phase 1: initial request --- - try: - data, latency = call_api( - client, base_url, model, scenario.messages, scenario.tools, timeout - ) - except Exception as e: - results.append( - ScenarioResult( - name=scenario.name, - phase="tool_call", - passed=False, - error=f"API error: {e}", - ) - ) - return results - - if verbose: - print(f" response: {json.dumps(data, indent=2)}") - - choice = data["choices"][0] - finish_reason = choice.get("finish_reason") - message = choice.get("message", {}) - tool_calls = message.get("tool_calls") - content = message.get("content") - - checks: dict[str, bool] = {} - - if scenario.expect_tool_call: - checks["finish_reason_tool_calls"] = finish_reason == "tool_calls" - checks["has_tool_call"] = isinstance(tool_calls, list) and len(tool_calls) > 0 - - args_err: str | None = None - if checks["has_tool_call"]: - tc = tool_calls[0] - fn = tc.get("function", {}) - checks["correct_function"] = ( - scenario.expected_function is None - or fn.get("name") == scenario.expected_function - ) - if scenario.required_arg_keys: - ok, args_err = validate_args( - fn.get("arguments", ""), scenario.required_arg_keys - ) - checks["valid_arguments"] = ok - else: - checks["valid_arguments"] = True - 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"] = finish_reason == "stop" - checks["no_tool_call"] = tool_calls is None or len(tool_calls) == 0 - checks["has_content"] = isinstance(content, str) and len(content.strip()) > 0 - passed = all(checks.values()) - error = ( - None - if passed - else ( - f"finish_reason={finish_reason}, " - f"tool_calls={'yes' if tool_calls else 'no'}, " - f"content={'yes' if content else 'no'}" - ) - ) - - results.append( - ScenarioResult( - name=scenario.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 checks.get("has_tool_call"): - tc = tool_calls[0] - fn = tc.get("function", {}) - follow_up_messages: list[dict[str, object]] = list(scenario.messages) + [ - { - "role": "assistant", - "tool_calls": [ - { - "id": tc.get("id", "call_0"), - "type": "function", - "function": { - "name": fn.get("name", ""), - "arguments": fn.get("arguments", "{}"), - }, - } - ], - }, - { - "role": "tool", - "tool_call_id": tc.get("id", "call_0"), - "content": scenario.tool_result, - }, - ] - - try: - data2, latency2 = call_api( - client, - base_url, - model, - follow_up_messages, - scenario.tools, - timeout, - ) - except Exception as e: - results.append( - ScenarioResult( - name=scenario.name, - phase="follow_up", - passed=False, - error=f"API error: {e}", - ) - ) - return results - - if verbose: - print(f" follow_up response: {json.dumps(data2, indent=2)}") - - choice2 = data2["choices"][0] - message2 = choice2.get("message", {}) - checks2: dict[str, bool] = {} - checks2["finish_reason_stop"] = choice2.get("finish_reason") == "stop" - tc2 = message2.get("tool_calls") - checks2["no_tool_call"] = tc2 is None or len(tc2) == 0 - c2 = message2.get("content") - checks2["has_content"] = isinstance(c2, str) and len(c2.strip()) > 0 - - passed2 = all(checks2.values()) - error2 = None - if not passed2: - error2 = ( - f"finish_reason={choice2.get('finish_reason')}, " - f"tool_calls={'yes' if tc2 else 'no'}, " - f"content={'yes' if c2 else 'no'}" - ) - results.append( - ScenarioResult( - name=scenario.name, - phase="follow_up", - passed=passed2, - checks=checks2, - error=error2, - latency_ms=latency2, - ) - ) - - return results - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -def main() -> None: - parser = argparse.ArgumentParser(description="Tool-calling eval for exo") - parser.add_argument("--model", required=True, help="Model ID to test") - parser.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost")) - parser.add_argument( - "--port", - type=int, - default=int(os.environ.get("EXO_PORT", "52415")), - ) - parser.add_argument( - "--timeout", type=float, default=120, help="Per-request timeout (seconds)" - ) - parser.add_argument( - "--repeat", type=int, default=1, help="Repeat each scenario N times" - ) - parser.add_argument( - "--scenarios", nargs="*", help="Run only these scenarios (by name)" - ) - parser.add_argument( - "--verbose", action="store_true", help="Print full API responses" - ) - args = parser.parse_args() - - scenarios = SCENARIOS - if args.scenarios: - scenarios = [s for s in SCENARIOS if s.name in args.scenarios] - if not scenarios: - print(f"No matching scenarios. Available: {[s.name for s in SCENARIOS]}") - sys.exit(1) - - base_url = f"http://{args.host}:{args.port}/v1" - total_runs = len(scenarios) * args.repeat - print(f"Model: {args.model}") - print(f"Endpoint: {base_url}") - print(f"Scenarios: {len(scenarios)} x {args.repeat} = {total_runs} runs") - print("=" * 64) - - all_results: list[ScenarioResult] = [] - - with httpx.Client() as client: - for run_idx in range(args.repeat): - if args.repeat > 1: - print(f"\n--- Run {run_idx + 1}/{args.repeat} ---") - - for scenario in scenarios: - print(f"\n {scenario.name}: {scenario.description}") - - results = run_scenario( - client, - base_url, - args.model, - scenario, - args.timeout, - args.verbose, - ) - all_results.extend(results) - - for r in results: - status = "PASS" if r.passed else "FAIL" - print(f" [{r.phase:>10}] {status} ({r.latency_ms:.0f}ms)") - for check_name, check_ok in r.checks.items(): - mark = "+" if check_ok else "-" - print(f" {mark} {check_name}") - if r.error: - print(f" ! {r.error}") - - # --- Summary --- - print(f"\n{'=' * 64}") - - 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}%)") - print(f"Tool call: {tc_passed}/{len(tool_call_results)} passed") - if follow_up_results: - print(f"Follow-up: {fu_passed}/{len(follow_up_results)} passed") - print(f"Avg latency: {avg_latency:.0f}ms") - - if passed < total: - print("\nFailed:") - for r in all_results: - if not r.passed: - print(f" - {r.name} [{r.phase}]: {r.error}") - - sys.exit(0 if passed == total else 1) - - -if __name__ == "__main__": - main() From 7cadca4f277aaab731bbe0b7729d7e76ae0abfcb Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Wed, 18 Feb 2026 14:10:07 -0800 Subject: [PATCH 25/45] Try multiple endpoints for internet connectivity check (#1516) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - `_test_internet_connection()` previously only tried `1.1.1.1:443`, which some ISPs/networks block, causing exo to incorrectly report no internet and fail downloads on startup - Now tries `1.1.1.1`, `8.8.8.8`, and `1.0.0.1` in sequence, succeeding if any endpoint responds - Returns early on first success for minimal latency in the common case Fixes #1425 ## Test plan - [ ] Verify downloads work on networks that block `1.1.1.1` - [ ] Verify existing behavior unchanged on networks where `1.1.1.1` works - [ ] Verify `internet_connection` is set to `False` only when all three endpoints fail 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 Co-authored-by: rltakashige --- src/exo/download/coordinator.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 899e4f14..30e45a08 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -123,14 +123,17 @@ class DownloadCoordinator: 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 From 19bc09550d55ca52344d53ac44a5c8089bf20219 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Wed, 18 Feb 2026 22:34:11 +0000 Subject: [PATCH 26/45] Add status=downloaded filter for model endpoint (#1539) ## Motivation https://github.com/exo-explore/exo/issues/1346#issuecomment-3831427905 ## Test Plan ### Manual Testing **Without filter** Screenshot 2026-02-18 at 22 26 22 **With filter** Screenshot 2026-02-18 at 22 26 45 --- src/exo/master/api.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/src/exo/master/api.py b/src/exo/master/api.py index 3c29b041..c85999a3 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -145,6 +145,7 @@ from exo.shared.types.openai_responses import ( 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 @@ -1292,8 +1293,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( @@ -1311,7 +1322,7 @@ class API: base_model=card.base_model, capabilities=card.capabilities, ) - for card in await get_model_cards() + for card in cards ] ) From 025ed9fd82fd28b823250e45f72ff677cf27884b Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Wed, 18 Feb 2026 19:18:25 -0800 Subject: [PATCH 27/45] feat: add prefill progress bar for long prompts (#1181) ## Motivation Users processing long prompts have no visibility into when token generation will start. This feature adds a progress bar showing prefill progress, giving users real-time feedback during prompt processing. ## Changes ### Backend - Added `PrefillProgress` event type with `command_id`, `processed_tokens`, `total_tokens` - Added `PrefillProgressResponse` type (though now using direct callback approach) - Wired `prompt_progress_callback` through MLX's `stream_generate()` - Progress events sent directly from callback for real-time updates (not batched) - API generates SSE named events: `event: prefill_progress\ndata: {...}` - Added `PrefillProgressData` dataclass and `StreamEvent` union type in API ### Dashboard - Added `PrefillProgress` interface to store - Updated SSE parsing to handle `event:` lines (named events) - Created `PrefillProgressBar.svelte` with animated progress bar - Shows "Processing prompt: X/Y tokens" with percentage - Progress bar disappears when first token arrives ## Why It Works MLX's `stream_generate()` accepts a `prompt_progress_callback(processed, total)` that's called after each prefill chunk. By sending events directly from this callback (rather than yielding from the generator), progress updates are sent in real-time during prefill. Using SSE named events (`event: prefill_progress`) maintains full OpenAI/Claude API compatibility - standard clients ignore named events they don't recognize, while the exo dashboard explicitly listens for them. ## Test Plan ### Manual Testing - Hardware: MacBook Pro M3 Max - Set `prefill_step_size=256` for more frequent updates - Tested with long prompts (pasted large documents) - Verified progress bar updates incrementally during prefill - Confirmed progress bar disappears when generation starts - Tested with curl - standard `data:` events still work normally Here is it working: https://github.com/user-attachments/assets/5cc6f075-c5b2-4a44-bb4d-9efb246bc5fe ### Automated Testing - Type checker passes (0 errors) - All 192 tests pass - Dashboard builds successfully ### API Compatibility - Named SSE events are ignored by OpenAI SDK clients - Regular token data uses standard `data: {...}` format - `[DONE]` sentinel works as expected --- **Note:** `prefill_step_size` is temporarily set to 256 for testing. Should be changed back to 2048 before merging for production performance. --------- Co-authored-by: Claude Opus 4.5 Co-authored-by: Evan Co-authored-by: Ryuichi Leo Takashige --- bench/eval_tool_calls.py | 42 +++++ bench/scenarios.toml | 66 +++++++ .../src/lib/components/ChatMessages.svelte | 7 +- .../components/HuggingFaceResultItem.svelte | 3 +- .../lib/components/PrefillProgressBar.svelte | 52 ++++++ dashboard/src/lib/stores/app.svelte.ts | 49 +++++ dashboard/src/routes/+page.svelte | 7 - src/exo/master/adapters/chat_completions.py | 176 ++++++++++-------- src/exo/master/adapters/claude.py | 21 ++- src/exo/master/adapters/responses.py | 21 ++- src/exo/master/api.py | 31 ++- src/exo/shared/apply.py | 2 + src/exo/shared/types/chunks.py | 11 +- src/exo/shared/types/events.py | 10 +- .../shared/types/worker/runner_response.py | 5 + .../worker/engines/mlx/generator/generate.py | 15 +- src/exo/worker/engines/mlx/utils_mlx.py | 65 +++++++ src/exo/worker/runner/runner.py | 24 +++ 18 files changed, 509 insertions(+), 98 deletions(-) create mode 100644 dashboard/src/lib/components/PrefillProgressBar.svelte diff --git a/bench/eval_tool_calls.py b/bench/eval_tool_calls.py index 6cd84396..2d55d6ca 100644 --- a/bench/eval_tool_calls.py +++ b/bench/eval_tool_calls.py @@ -38,6 +38,8 @@ class Scenario: 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]: @@ -105,6 +107,8 @@ def load_scenarios(path: Path) -> list[Scenario]: 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"), ) ) @@ -147,6 +151,35 @@ def validate_args(args_str: str, required_keys: list[str]) -> tuple[bool, str | 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, @@ -699,6 +732,15 @@ def run_scenario( 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 diff --git a/bench/scenarios.toml b/bench/scenarios.toml index e258b5c6..892a044a 100644 --- a/bench/scenarios.toml +++ b/bench/scenarios.toml @@ -39,6 +39,30 @@ description = "Product category to filter by" 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]] @@ -219,6 +243,48 @@ 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]] diff --git a/dashboard/src/lib/components/ChatMessages.svelte b/dashboard/src/lib/components/ChatMessages.svelte index ba5322a7..1b1d2d07 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/PrefillProgressBar.svelte b/dashboard/src/lib/components/PrefillProgressBar.svelte new file mode 100644 index 00000000..ea08d11d --- /dev/null +++ b/dashboard/src/lib/components/PrefillProgressBar.svelte @@ -0,0 +1,52 @@ + + +
+
+ Processing prompt + + {formatTokenCount(progress.processed)} / {formatTokenCount( + progress.total, + )} tokens + +
+
+
+
+
+ {percentage}% +
+
+ + diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 1a6b1e3d..3e9363c7 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -273,6 +273,11 @@ export interface TokenData { topLogprobs: TopLogprob[]; } +export interface PrefillProgress { + processed: number; + total: number; +} + export interface Message { id: string; role: "user" | "assistant" | "system"; @@ -520,6 +525,7 @@ 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); // Topology state topologyData = $state(null); @@ -2005,6 +2011,7 @@ class AppStore { reader: ReadableStreamDefaultReader, targetConversationId: string, onChunk: (parsed: T) => void, + onEvent?: Record void>, ): Promise { const decoder = new TextDecoder(); let buffer = ""; @@ -2025,6 +2032,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; @@ -2309,6 +2334,11 @@ 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; @@ -2371,8 +2401,26 @@ 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, + }; + }, + }, ); + // Clear prefill progress after stream ends + this.prefillProgress = null; + // Calculate final TPS if (firstTokenTime !== null && tokenCount > 1) { const totalGenerationTime = performance.now() - firstTokenTime; @@ -3043,6 +3091,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; diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 2fdeb8ab..03ac3750 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -932,13 +932,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, diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py index b86c5ec1..a8ca8500 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 @@ -123,67 +128,81 @@ 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 @@ -197,38 +216,43 @@ 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 + 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) diff --git a/src/exo/master/adapters/claude.py b/src/exo/master/adapters/claude.py index ee7c263a..d9d52496 100644 --- a/src/exo/master/adapters/claude.py +++ b/src/exo/master/adapters/claude.py @@ -5,7 +5,12 @@ 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, @@ -160,7 +165,9 @@ 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 @@ -172,6 +179,9 @@ async def collect_claude_response( 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 @@ -230,7 +240,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 @@ -256,6 +268,9 @@ async def generate_claude_stream( next_block_index = 1 # text block is 0, tool blocks start at 1 async for chunk in chunk_stream: + if isinstance(chunk, PrefillProgressChunk): + continue + if isinstance(chunk, ErrorChunk): # Close text block and bail break diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py index 5e059fee..e55160b6 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, @@ -121,7 +126,9 @@ 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 @@ -134,6 +141,9 @@ async def collect_responses_response( 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 @@ -189,7 +199,9 @@ 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}" @@ -243,6 +255,9 @@ async def generate_responses_stream( next_output_index = 1 # message item is at 0 async for chunk in chunk_stream: + if isinstance(chunk, PrefillProgressChunk): + continue + if isinstance(chunk, ErrorChunk): break diff --git a/src/exo/master/api.py b/src/exo/master/api.py index c85999a3..c3811072 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -107,6 +107,7 @@ from exo.shared.types.chunks import ( ErrorChunk, ImageChunk, InputImageChunk, + PrefillProgressChunk, TokenChunk, ToolCallChunk, ) @@ -137,6 +138,7 @@ from exo.shared.types.events import ( Event, ForwarderEvent, IndexedEvent, + PrefillProgress, TracesMerged, ) from exo.shared.types.memory import Memory @@ -221,7 +223,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] @@ -527,19 +530,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 @@ -566,6 +573,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, @@ -1446,6 +1456,21 @@ class API: except BrokenResourceError: self._text_generation_queues.pop(event.command_id, None) + elif isinstance(event, PrefillProgress): + if queue := self._text_generation_queues.get( + event.command_id, None + ): + try: + await queue.send( + PrefillProgressChunk( + model=event.model, + processed_tokens=event.processed_tokens, + total_tokens=event.total_tokens, + ) + ) + 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/shared/apply.py b/src/exo/shared/apply.py index 94869dfe..ee5b2229 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -15,6 +15,7 @@ from exo.shared.types.events import ( NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, + PrefillProgress, RunnerDeleted, RunnerStatusUpdated, TaskAcknowledged, @@ -64,6 +65,7 @@ def event_apply(event: Event, state: State) -> State: | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived() + | PrefillProgress() | TracesCollected() | TracesMerged() ): # Pass-through events that don't modify state diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py index 5fe9eb1c..b97eee5f 100644 --- a/src/exo/shared/types/chunks.py +++ b/src/exo/shared/types/chunks.py @@ -76,4 +76,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/events.py b/src/exo/shared/types/events.py index 5cf93d0c..83e162ab 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -5,7 +5,7 @@ from pydantic import Field from exo.shared.topology import Connection from exo.shared.types.chunks import GenerationChunk, InputImageChunk -from exo.shared.types.common import CommandId, Id, NodeId, SessionId +from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId @@ -102,6 +102,13 @@ class InputChunkReceived(BaseEvent): chunk: InputImageChunk +class PrefillProgress(BaseEvent): + command_id: CommandId + model: ModelId + processed_tokens: int + total_tokens: int + + class TopologyEdgeCreated(BaseEvent): conn: Connection @@ -148,6 +155,7 @@ Event = ( | NodeDownloadProgress | ChunkGenerated | InputChunkReceived + | PrefillProgress | TopologyEdgeCreated | TopologyEdgeDeleted | TracesCollected diff --git a/src/exo/shared/types/worker/runner_response.py b/src/exo/shared/types/worker/runner_response.py index 5f18bf5a..a66f762c 100644 --- a/src/exo/shared/types/worker/runner_response.py +++ b/src/exo/shared/types/worker/runner_response.py @@ -67,3 +67,8 @@ class ToolCallResponse(BaseRunnerResponse): class FinishedResponse(BaseRunnerResponse): pass + + +class PrefillProgressResponse(BaseRunnerResponse): + processed_tokens: int + total_tokens: int diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 7b78cbd9..42fda040 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -58,6 +58,7 @@ def prefill( 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. @@ -85,6 +86,9 @@ 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) @@ -99,7 +103,7 @@ def prefill( max_tokens=1, sampler=sampler, prompt_cache=cache, - prefill_step_size=8192, + prefill_step_size=4096, kv_group_size=KV_GROUP_SIZE, kv_bits=KV_BITS, prompt_progress_callback=progress_callback, @@ -257,6 +261,7 @@ def mlx_generate( prompt: str, 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() @@ -311,7 +316,13 @@ def mlx_generate( # Prefill cache with all tokens except the last one prefill_tps, prefill_tokens, ssm_snapshots_list = prefill( - model, tokenizer, sampler, prompt_tokens[:-1], caches, group + model, + tokenizer, + sampler, + prompt_tokens[:-1], + caches, + group, + on_prefill_progress, ) cache_snapshots: list[CacheSnapshot] | None = ssm_snapshots_list or None diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index d9cebae9..567b0f91 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 @@ -407,6 +408,56 @@ 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 apply_chat_template( tokenizer: TokenizerWrapper, task_params: TextGenerationTaskParams, @@ -453,14 +504,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 diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index c749159a..7966353e 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -26,6 +26,7 @@ from exo.shared.types.common import CommandId from exo.shared.types.events import ( ChunkGenerated, Event, + PrefillProgress, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, @@ -298,6 +299,18 @@ def main( assert tokenizer assert check_for_cancel_every + # Define callback to send prefill progress events directly + def on_prefill_progress(processed: int, total: int) -> None: + if device_rank == 0: + event_sender.send( + PrefillProgress( + command_id=command_id, + model=shard_metadata.model_card.model_id, + processed_tokens=processed, + total_tokens=total, + ) + ) + try: _check_for_debug_prompts(task_params) @@ -311,6 +324,7 @@ def main( task=task_params, prompt=prompt, kv_prefix_cache=kv_prefix_cache, + on_prefill_progress=on_prefill_progress, group=group, ) @@ -599,11 +613,21 @@ def parse_gpt_oss( 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( From 51021f6fc6dce3f2072796c6af215c622d02903e Mon Sep 17 00:00:00 2001 From: rltakashige Date: Thu, 19 Feb 2026 11:40:59 +0000 Subject: [PATCH 28/45] Add cancellation button and the ability to cancel during prefill (#1540) ## Motivation There's no way to easily use the cancellation features we added! Also, prefill can take ages so let's allow cancelling out of that. ## Changes Wiring up our existing functionality to easily cancel during generation (and adding stuff to do so during prefill) ## Test Plan ### Manual Testing Tested it works during both prefill and decode. ### Automated testing Needs testing to see if this causes a GPU timeout error on large prefill on large models in pipeline parallel. However, from manually testing GLM 5 pipeline ring on 2 nodes, and from reading the code, it does not seem like this will be the case. --- README.md | 9 +- dashboard/src/lib/components/ChatForm.svelte | 155 +++++++++--------- dashboard/src/lib/stores/app.svelte.ts | 33 +++- .../worker/engines/mlx/generator/generate.py | 37 +++-- src/exo/worker/runner/runner.py | 26 ++- 5 files changed, 161 insertions(+), 99 deletions(-) diff --git a/README.md b/README.md index ff1fe04b..785f6fba 100644 --- a/README.md +++ b/README.md @@ -72,12 +72,19 @@ 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 (after accepting the Cachix cache): +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) diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte index 6ccba979..aa116c2f 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"; @@ -653,86 +654,92 @@ style="min-height: 28px; max-height: 150px;" > - + + {:else} + + {/if}
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 3e9363c7..d13e7f2b 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -527,6 +527,9 @@ class AppStore { 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); instances = $state>({}); @@ -2281,6 +2284,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: { @@ -2297,6 +2303,7 @@ class AppStore { enable_thinking: enableThinking, }), }), + signal: abortController.signal, }); if (!response.ok) { @@ -2451,20 +2458,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 */ @@ -3109,6 +3127,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/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 42fda040..6e926205 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -51,6 +51,10 @@ generation_stream = mx.new_stream(mx.default_device()) _MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 +class PrefillCancelled(BaseException): + """Raised when prefill is cancelled via the progress callback.""" + + def prefill( model: Model, tokenizer: TokenizerWrapper, @@ -66,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: @@ -77,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 @@ -96,19 +101,23 @@ def 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=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 + 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) diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 7966353e..2dc340c3 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -82,7 +82,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, @@ -299,8 +303,16 @@ def main( assert tokenizer assert check_for_cancel_every - # Define callback to send prefill progress events directly - def on_prefill_progress(processed: int, total: int) -> None: + # 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( PrefillProgress( @@ -310,6 +322,12 @@ def main( 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) @@ -406,6 +424,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: From cacb456cb2d6b5d12f4a390bfef5e9299f9368e5 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Thu, 19 Feb 2026 12:55:31 +0000 Subject: [PATCH 29/45] remove nightly (#1538) we have no good need for rust nightly (nor futures, for that matter) --- Cargo.lock | 13 +- Cargo.toml | 3 +- flake.nix | 1 - rust/clippy.toml | 2 - rust/exo_pyo3_bindings/Cargo.toml | 6 +- rust/exo_pyo3_bindings/src/allow_threading.rs | 5 +- .../src/{pylibp2p => }/ident.rs | 0 rust/exo_pyo3_bindings/src/lib.rs | 28 +--- rust/exo_pyo3_bindings/src/networking.rs | 2 +- rust/exo_pyo3_bindings/src/pylibp2p/mod.rs | 8 -- .../src/pylibp2p/multiaddr.rs | 81 ----------- rust/networking/Cargo.toml | 2 +- rust/networking/examples/chatroom.rs | 12 +- rust/networking/examples/chatroom_manual.rs | 127 ------------------ rust/networking/src/discovery.rs | 4 +- rust/networking/src/swarm.rs | 2 +- rust/parts.nix | 5 +- rust/rust-toolchain.toml | 2 - 18 files changed, 33 insertions(+), 270 deletions(-) delete mode 100644 rust/clippy.toml rename rust/exo_pyo3_bindings/src/{pylibp2p => }/ident.rs (100%) delete mode 100644 rust/exo_pyo3_bindings/src/pylibp2p/mod.rs delete mode 100644 rust/exo_pyo3_bindings/src/pylibp2p/multiaddr.rs delete mode 100644 rust/networking/examples/chatroom_manual.rs delete mode 100644 rust/rust-toolchain.toml diff --git a/Cargo.lock b/Cargo.lock index 089f5b65..8914f5e7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -890,7 +890,7 @@ dependencies = [ "delegate", "env_logger", "extend", - "futures", + "futures-lite", "libp2p", "log", "networking", @@ -914,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" @@ -1022,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", ] @@ -2753,7 +2762,7 @@ dependencies = [ "delegate", "either", "extend", - "futures", + "futures-lite", "futures-timer", "keccak-const", "libp2p", diff --git a/Cargo.toml b/Cargo.toml index 4c0b8721..ffa9022c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,14 +29,13 @@ util = { path = "rust/util" } # Macro dependecies extend = "1.2" delegate = "0.13" -pin-project = "1" # Utility dependencies keccak-const = "0.2" # Async dependencies tokio = "1.46" -futures = "0.3" +futures-lite = "2.6.1" futures-timer = "3.0" # Data structures diff --git a/flake.nix b/flake.nix index e90e0bd2..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 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 e97b9148..77777055 100644 --- a/rust/exo_pyo3_bindings/Cargo.toml +++ b/rust/exo_pyo3_bindings/Cargo.toml @@ -27,7 +27,7 @@ networking = { workspace = true } # interop pyo3 = { version = "0.27.2", features = [ # "abi3-py313", # tells pyo3 (and maturin) to build using the stable ABI with minimum Python version 3.13 - "nightly", # enables better-supported GIL integration + # "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) @@ -45,11 +45,10 @@ pyo3-log = "0.13.2" # macro dependencies extend = { workspace = true } delegate = { workspace = true } -pin-project = { workspace = true } # async runtime tokio = { workspace = true, features = ["full", "tracing"] } -futures = { workspace = true } +futures-lite = { workspace = true } # utility dependencies util = { workspace = true } @@ -60,3 +59,4 @@ env_logger = "0.11" # Networking libp2p = { workspace = true, features = ["full"] } +pin-project = "1.1.10" diff --git a/rust/exo_pyo3_bindings/src/allow_threading.rs b/rust/exo_pyo3_bindings/src/allow_threading.rs index 18a426aa..142ff0c6 100644 --- a/rust/exo_pyo3_bindings/src/allow_threading.rs +++ b/rust/exo_pyo3_bindings/src/allow_threading.rs @@ -2,7 +2,6 @@ //! use pin_project::pin_project; -use pyo3::marker::Ungil; use pyo3::prelude::*; use std::{ future::Future, @@ -26,8 +25,8 @@ where impl Future for AllowThreads where - F: Future + Ungil, - F::Output: Ungil, + F: Future + Send, + F::Output: Send, { type Output = F::Output; diff --git a/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs b/rust/exo_pyo3_bindings/src/ident.rs similarity index 100% rename from rust/exo_pyo3_bindings/src/pylibp2p/ident.rs rename to rust/exo_pyo3_bindings/src/ident.rs diff --git a/rust/exo_pyo3_bindings/src/lib.rs b/rust/exo_pyo3_bindings/src/lib.rs index 25f2865e..45825b21 100644 --- a/rust/exo_pyo3_bindings/src/lib.rs +++ b/rust/exo_pyo3_bindings/src/lib.rs @@ -4,25 +4,12 @@ //! //! -// 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; -pub(crate) mod networking; -pub(crate) mod pylibp2p; +mod ident; +mod networking; +use crate::ident::ident_submodule; use crate::networking::networking_submodule; -use crate::pylibp2p::ident::ident_submodule; -use crate::pylibp2p::multiaddr::multiaddr_submodule; use pyo3::prelude::PyModule; use pyo3::{Bound, PyResult, pyclass, pymodule}; use pyo3_stub_gen::define_stub_info_gatherer; @@ -32,14 +19,6 @@ 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::marker::Tuple; - - pub trait SendFn = - Fn + Send + 'static; -} - /// Namespace for crate-wide extension traits/methods pub(crate) mod ext { use crate::allow_threading::AllowThreads; @@ -180,7 +159,6 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> { // work with maturin, where the types generate correctly, in the right folder, without // too many importing issues... ident_submodule(m)?; - multiaddr_submodule(m)?; 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 03a0dcb6..b864d876 100644 --- a/rust/exo_pyo3_bindings/src/networking.rs +++ b/rust/exo_pyo3_bindings/src/networking.rs @@ -8,8 +8,8 @@ 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, PyPeerId}; use crate::pyclass; -use crate::pylibp2p::ident::{PyKeypair, PyPeerId}; use libp2p::futures::StreamExt as _; use libp2p::gossipsub; use libp2p::gossipsub::{IdentTopic, Message, MessageId, PublishError}; 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 fd5f1b1f..58ca4e0a 100644 --- a/rust/networking/Cargo.toml +++ b/rust/networking/Cargo.toml @@ -22,7 +22,7 @@ delegate = { workspace = true } # async tokio = { workspace = true, features = ["full"] } -futures = { workspace = true } +futures-lite = { workspace = true } futures-timer = { workspace = true } # utility dependencies 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 e9f10a1f..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::error::Error; -use std::time::Duration; -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 581f2200..a6cbd0e0 100644 --- a/rust/networking/src/discovery.rs +++ b/rust/networking/src/discovery.rs @@ -1,7 +1,7 @@ use crate::ext::MultiaddrExt; 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}; @@ -362,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/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 From 2e29605194af486f80d9a90bb28524ffa3bbe4e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mustafa=20Alp=20Y=C4=B1lmaz?= <96022931+mustafalpyilmaz@users.noreply.github.com> Date: Thu, 19 Feb 2026 16:27:34 +0300 Subject: [PATCH 30/45] fix: finalize cancel tasks (#1498) # Cancel task finalization (main.py) After forwarding the cancel to the runner supervisor, emit TaskStatusUpdated(Complete) for the cancel task itself. This ensures the cancel task is properly removed from state.tasks. --- src/exo/worker/main.py | 5 +++++ 1 file changed, 5 insertions(+) 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 From aa3f106fb96679cffdb998aa5d1ebc1c8ffed863 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Thu, 19 Feb 2026 05:40:24 -0800 Subject: [PATCH 31/45] fix: import ResponsesStreamEvent and DRY up SSE formatting (#1499) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - `ResponsesStreamEvent` was defined in `openai_responses.py` as a union of all 11 streaming event types but never imported or used anywhere in the codebase - Import it in the responses adapter and add a `_format_sse(event: ResponsesStreamEvent) -> str` helper - Replace 13 hardcoded `f"event: {type}\ndata: {event.model_dump_json()}\n\n"` strings with `_format_sse()` calls ## Test plan - [x] `uv run basedpyright` — 0 errors - [x] `uv run ruff check` — all checks passed - [x] `nix fmt` — 0 files changed - [x] `uv run pytest` — 188 passed, 1 skipped 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 --- src/exo/master/adapters/responses.py | 32 +++++++++++++++++----------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/src/exo/master/adapters/responses.py b/src/exo/master/adapters/responses.py index e55160b6..90fa7732 100644 --- a/src/exo/master/adapters/responses.py +++ b/src/exo/master/adapters/responses.py @@ -31,6 +31,7 @@ from exo.shared.types.openai_responses import ( ResponseOutputText, ResponsesRequest, ResponsesResponse, + ResponsesStreamEvent, ResponseTextDeltaEvent, ResponseTextDoneEvent, ResponseUsage, @@ -38,6 +39,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): @@ -219,13 +225,13 @@ 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" + yield _format_sse(in_progress_event) # response.output_item.added initial_item = ResponseMessageItem( @@ -236,7 +242,7 @@ async def generate_responses_stream( 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" + yield _format_sse(item_added) # response.content_part.added initial_part = ResponseOutputText(text="") @@ -247,7 +253,7 @@ async def generate_responses_stream( content_index=0, part=initial_part, ) - yield f"event: response.content_part.added\ndata: {part_added.model_dump_json()}\n\n" + yield _format_sse(part_added) accumulated_text = "" function_call_items: list[ResponseFunctionCallItem] = [] @@ -281,7 +287,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( @@ -290,7 +296,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( @@ -300,7 +306,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( @@ -315,7 +321,7 @@ 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 @@ -331,7 +337,7 @@ async def generate_responses_stream( 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) # response.output_text.done text_done = ResponseTextDoneEvent( @@ -341,7 +347,7 @@ async def generate_responses_stream( 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) @@ -352,7 +358,7 @@ async def generate_responses_stream( 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( @@ -363,7 +369,7 @@ async def generate_responses_stream( item_done = ResponseOutputItemDoneEvent( sequence_number=next(seq), output_index=0, 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 @@ -388,4 +394,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) From 42e1e7322bc0fef65df441b11668ff6b0391c903 Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Thu, 19 Feb 2026 15:42:02 +0000 Subject: [PATCH 32/45] bench: restore --danger-delete-downloads planning phase (#1542) c2f2111b extracted shared utilities from exo_bench.py into harness.py but accidentally dropped the run_planning_phase function and --danger-delete-downloads CLI argument in the process. Restored run_planning_phase in harness.py (where its dependencies now live) and re-added the --danger-delete-downloads argument to add_common_instance_args. Re-wired the planning phase call in exo_bench.py's main() before the benchmark loop. --- bench/eval_tool_calls.py | 16 +++++ bench/exo_bench.py | 15 ++++ bench/harness.py | 150 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 181 insertions(+) diff --git a/bench/eval_tool_calls.py b/bench/eval_tool_calls.py index 2d55d6ca..13cb2537 100644 --- a/bench/eval_tool_calls.py +++ b/bench/eval_tool_calls.py @@ -20,6 +20,7 @@ from harness import ( 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, @@ -962,6 +963,21 @@ Examples: 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"]) diff --git a/bench/exo_bench.py b/bench/exo_bench.py index 9f6f5f02..f6fcd342 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -35,6 +35,7 @@ from harness import ( 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, @@ -332,6 +333,20 @@ 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, + full_model_id, + selected[0], + args.danger_delete_downloads, + args.timeout, + settle_deadline, + ) + all_rows: list[dict[str, Any]] = [] for preview in selected: diff --git a/bench/harness.py b/bench/harness.py index c8ae9318..58aa8435 100644 --- a/bench/harness.py +++ b/bench/harness.py @@ -282,6 +282,151 @@ def settle_and_fetch_placements( 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( @@ -325,3 +470,8 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None: 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.", + ) From 4c4c6ce99f1bb821adb5425591ab94a0dcfe095e Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Thu, 19 Feb 2026 17:19:31 +0000 Subject: [PATCH 33/45] simplify rust ident module this is partly dead code, partly narrowing the rust-python boundary in prep for future rewrites. no testing as this is all type safe refactoring. --- rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi | 115 +------------- rust/exo_pyo3_bindings/src/ident.rs | 146 ++---------------- rust/exo_pyo3_bindings/src/lib.rs | 5 +- rust/exo_pyo3_bindings/src/networking.rs | 8 +- src/exo/main.py | 2 +- src/exo/master/tests/test_master.py | 4 +- src/exo/routing/connection_message.py | 2 +- src/exo/routing/router.py | 6 +- .../shared/tests/test_node_id_persistence.py | 2 +- 9 files changed, 39 insertions(+), 251 deletions(-) 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/ident.rs b/rust/exo_pyo3_bindings/src/ident.rs index 3c27526a..55f40bc6 100644 --- a/rust/exo_pyo3_bindings/src/ident.rs +++ b/rust/exo_pyo3_bindings/src/ident.rs @@ -1,8 +1,6 @@ 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::types::{PyBytes, PyBytesMethods as _}; use pyo3::{Bound, PyResult, Python, pyclass, pymethods}; use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; @@ -18,142 +16,32 @@ pub struct PyKeypair(pub Keypair); impl PyKeypair { /// Generate a new Ed25519 keypair. #[staticmethod] - fn generate_ed25519() -> Self { + fn generate() -> Self { Self(Keypair::generate_ed25519()) } - /// Generate a new ECDSA keypair. + /// Construct an Ed25519 keypair from secret key bytes #[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 { + fn 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()?; + /// 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`. - 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() + /// 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() } } - -pub fn ident_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - - Ok(()) -} diff --git a/rust/exo_pyo3_bindings/src/lib.rs b/rust/exo_pyo3_bindings/src/lib.rs index 45825b21..8c1fb6b5 100644 --- a/rust/exo_pyo3_bindings/src/lib.rs +++ b/rust/exo_pyo3_bindings/src/lib.rs @@ -8,9 +8,10 @@ mod allow_threading; mod ident; mod networking; -use crate::ident::ident_submodule; +use crate::ident::PyKeypair; use crate::networking::networking_submodule; use pyo3::prelude::PyModule; +use pyo3::types::PyModuleMethods; use pyo3::{Bound, PyResult, pyclass, pymodule}; use pyo3_stub_gen::define_stub_info_gatherer; @@ -158,7 +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)?; + 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 b864d876..2fb1d78e 100644 --- a/rust/exo_pyo3_bindings/src/networking.rs +++ b/rust/exo_pyo3_bindings/src/networking.rs @@ -8,7 +8,7 @@ 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, PyPeerId}; +use crate::ident::PyKeypair; use crate::pyclass; use libp2p::futures::StreamExt as _; use libp2p::gossipsub; @@ -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)] @@ -251,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 { @@ -272,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 { diff --git a/src/exo/main.py b/src/exo/main.py index ec203181..27c78165 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -45,7 +45,7 @@ class Node: @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) 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/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/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: From ed001f2409ef671a9d398d4aecdf259c9957cfb3 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Thu, 19 Feb 2026 18:23:28 +0000 Subject: [PATCH 34/45] remove prefillprogress event (#1550) this should never have been a separate event, but i didnt quite communicate that well when this was merged. convert PrefillProgress to a chunk like the rest of the runner responses. tested with Llama-3.3-70B, prefill progress events still show up in the dashboard as usual --- src/exo/master/api.py | 17 ----------------- src/exo/shared/apply.py | 2 -- src/exo/shared/types/events.py | 10 +--------- src/exo/worker/runner/runner.py | 19 +++++++++++++------ 4 files changed, 14 insertions(+), 34 deletions(-) diff --git a/src/exo/master/api.py b/src/exo/master/api.py index c3811072..319a8164 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -138,7 +138,6 @@ from exo.shared.types.events import ( Event, ForwarderEvent, IndexedEvent, - PrefillProgress, TracesMerged, ) from exo.shared.types.memory import Memory @@ -1455,22 +1454,6 @@ class API: await queue.send(event.chunk) except BrokenResourceError: self._text_generation_queues.pop(event.command_id, None) - - elif isinstance(event, PrefillProgress): - if queue := self._text_generation_queues.get( - event.command_id, None - ): - try: - await queue.send( - PrefillProgressChunk( - model=event.model, - processed_tokens=event.processed_tokens, - total_tokens=event.total_tokens, - ) - ) - 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/shared/apply.py b/src/exo/shared/apply.py index ee5b2229..94869dfe 100644 --- a/src/exo/shared/apply.py +++ b/src/exo/shared/apply.py @@ -15,7 +15,6 @@ from exo.shared.types.events import ( NodeDownloadProgress, NodeGatheredInfo, NodeTimedOut, - PrefillProgress, RunnerDeleted, RunnerStatusUpdated, TaskAcknowledged, @@ -65,7 +64,6 @@ def event_apply(event: Event, state: State) -> State: | ChunkGenerated() | TaskAcknowledged() | InputChunkReceived() - | PrefillProgress() | TracesCollected() | TracesMerged() ): # Pass-through events that don't modify state diff --git a/src/exo/shared/types/events.py b/src/exo/shared/types/events.py index 83e162ab..5cf93d0c 100644 --- a/src/exo/shared/types/events.py +++ b/src/exo/shared/types/events.py @@ -5,7 +5,7 @@ from pydantic import Field from exo.shared.topology import Connection from exo.shared.types.chunks import GenerationChunk, InputImageChunk -from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId +from exo.shared.types.common import CommandId, Id, NodeId, SessionId from exo.shared.types.tasks import Task, TaskId, TaskStatus from exo.shared.types.worker.downloads import DownloadProgress from exo.shared.types.worker.instances import Instance, InstanceId @@ -102,13 +102,6 @@ class InputChunkReceived(BaseEvent): chunk: InputImageChunk -class PrefillProgress(BaseEvent): - command_id: CommandId - model: ModelId - processed_tokens: int - total_tokens: int - - class TopologyEdgeCreated(BaseEvent): conn: Connection @@ -155,7 +148,6 @@ Event = ( | NodeDownloadProgress | ChunkGenerated | InputChunkReceived - | PrefillProgress | TopologyEdgeCreated | TopologyEdgeDeleted | TracesCollected diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 2dc340c3..51950192 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -21,12 +21,17 @@ 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, Event, - PrefillProgress, RunnerStatusUpdated, TaskAcknowledged, TaskStatusUpdated, @@ -315,11 +320,13 @@ def main( ) -> None: if device_rank == 0: event_sender.send( - PrefillProgress( + ChunkGenerated( command_id=command_id, - model=shard_metadata.model_card.model_id, - processed_tokens=processed, - total_tokens=total, + chunk=PrefillProgressChunk( + model=shard_metadata.model_card.model_id, + processed_tokens=processed, + total_tokens=total, + ), ) ) cancelled_tasks.update(cancel_receiver.collect()) From 423ed0f07f706b0555547a0f3512ea0402221207 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Thu, 19 Feb 2026 18:29:34 +0000 Subject: [PATCH 35/45] Strip Claude headers to improve prefix cache hit rates (#1552) ## Motivation Our hits are really bad at the moment (0.2%). This PR makes it 98.5% on average. ## Changes Also adds an example for how to run Claude using Exo. ## Why It Works Claude sends some billing and session headers that change with each message. ## Test Plan ### Manual Testing Works in manual testing. --- src/exo/master/adapters/claude.py | 19 +++++++++++++++++++ tmp/config_examples/claude_code.sh | 8 ++++++++ 2 files changed, 27 insertions(+) create mode 100755 tmp/config_examples/claude_code.sh diff --git a/src/exo/master/adapters/claude.py b/src/exo/master/adapters/claude.py index d9d52496..d5246bbc 100644 --- a/src/exo/master/adapters/claude.py +++ b/src/exo/master/adapters/claude.py @@ -1,6 +1,7 @@ """Claude Messages API adapter for converting requests/responses.""" import json +import re from collections.abc import AsyncGenerator from typing import Any @@ -61,6 +62,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: @@ -73,6 +90,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 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 From 94b2ce69222a98a60693c9b884559a5b8d5a814a Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Thu, 19 Feb 2026 10:39:17 -0800 Subject: [PATCH 36/45] feat: Mac Studio en2 RDMA port warning v2 (#1551) Rebuilt from scratch (replaces PR #1543). Detects when Mac Studio uses RDMA over en2 (TB5 port next to Ethernet) which does not support RDMA. Shows dismissible warning banner with hover tooltip showing affected devices, SVG rear panel illustration, and fix instructions. 205 lines in +page.svelte. --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: rltakashige --- dashboard/src/routes/+page.svelte | 341 +++++++++++++++++++++++++++++- 1 file changed, 339 insertions(+), 2 deletions(-) diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 03ac3750..21e71774 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -114,6 +114,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]; @@ -1758,7 +1826,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]} @@ -1923,12 +1991,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} From cf648a53b84bd0bf50a3e61d7fcd0bbb9b9ac457 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Thu, 19 Feb 2026 18:44:49 +0000 Subject: [PATCH 37/45] Add thinking in thinking blocks, and fix DeepSeek interleaved tool calls (#1548) ## Motivation OpenCode shows tags and not thinking blocks as we aren't following the API specs properly. Claude was also getting horrible prefix cache hits because it sends headers. ## Changes Handle thinking tokens properly by placing them in think tags for each of the API endpoints. Also support DeepSeekV3.2 tool calling properly as a minor feature. Strips Claude headers at the API level. ## Test Plan ### Manual Testing Tested OpenCode manually Needs testing with Claude. ### Automated Testing All CI and tests passing - added a new e2e test for DeepSeekV32 tool parsing. --- dashboard/src/lib/stores/app.svelte.ts | 86 +- src/exo/master/adapters/chat_completions.py | 21 +- src/exo/master/adapters/claude.py | 114 ++- src/exo/master/adapters/responses.py | 237 ++++- src/exo/master/tests/test_claude_tool_use.py | 9 +- src/exo/shared/types/api.py | 2 +- src/exo/shared/types/chunks.py | 1 + src/exo/shared/types/claude_api.py | 35 +- src/exo/shared/types/openai_responses.py | 74 +- .../shared/types/worker/runner_response.py | 1 + src/exo/worker/engines/mlx/dsml_encoding.py | 72 ++ src/exo/worker/engines/mlx/utils_mlx.py | 27 +- src/exo/worker/runner/runner.py | 220 +++- .../unittests/test_runner/test_dsml_e2e.py | 967 ++++++++++++++++++ .../test_runner/test_event_ordering.py | 1 + .../test_runner/test_parse_gpt_oss.py | 23 +- 16 files changed, 1772 insertions(+), 118 deletions(-) create mode 100644 src/exo/worker/engines/mlx/dsml_encoding.py create mode 100644 src/exo/worker/tests/unittests/test_runner/test_dsml_e2e.py diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index d13e7f2b..3e0074e3 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -1652,11 +1652,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; @@ -1677,6 +1678,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; @@ -1695,7 +1697,11 @@ class AppStore { } } - if (delta) { + if (thinkingDelta) { + streamedThinking += thinkingDelta; + } + + if (delta || thinkingDelta) { if (firstTokenTime === null) { firstTokenTime = performance.now(); this.ttftMs = firstTokenTime - requestStartTime; @@ -1709,9 +1715,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; @@ -1723,7 +1734,7 @@ class AppStore { messageId, (m) => { m.content = displayContent; - m.thinking = thinkingContent || undefined; + m.thinking = combinedThinking || undefined; m.tokens = [...collectedTokens]; }, ); @@ -1735,11 +1746,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; @@ -1847,11 +1861,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; @@ -1872,6 +1887,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; @@ -1890,10 +1906,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) { @@ -1906,7 +1931,7 @@ class AppStore { assistantMessage.id, (msg) => { msg.content = displayContent; - msg.thinking = thinkingContent || undefined; + msg.thinking = combinedThinking || undefined; msg.tokens = [...collectedTokens]; }, ); @@ -1918,14 +1943,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]; }, ); @@ -2317,10 +2345,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; @@ -2348,6 +2377,7 @@ class AppStore { const choice = parsed.choices?.[0]; const tokenContent = choice?.delta?.content; + const thinkingContent = choice?.delta?.reasoning_content; // Collect logprobs data const logprobsContent = choice?.logprobs?.content; @@ -2366,7 +2396,11 @@ class AppStore { } } - if (tokenContent) { + if (thinkingContent) { + streamedThinking += thinkingContent; + } + + if (tokenContent || thinkingContent) { // Track first token for TTFT if (firstTokenTime === null) { firstTokenTime = performance.now(); @@ -2383,11 +2417,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) { @@ -2400,7 +2439,7 @@ class AppStore { assistantMessage.id, (msg) => { msg.content = displayContent; - msg.thinking = thinkingContent || undefined; + msg.thinking = combinedThinking || undefined; msg.tokens = [...collectedTokens]; }, ); @@ -2436,14 +2475,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) { diff --git a/src/exo/master/adapters/chat_completions.py b/src/exo/master/adapters/chat_completions.py index a8ca8500..9773806f 100644 --- a/src/exo/master/adapters/chat_completions.py +++ b/src/exo/master/adapters/chat_completions.py @@ -59,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"): @@ -111,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()), @@ -118,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, ) @@ -208,6 +217,7 @@ async def collect_chat_response( # 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 @@ -228,7 +238,10 @@ async def collect_chat_response( if model is None: model = chunk.model last_usage = chunk.usage or last_usage - text_parts.append(chunk.text) + 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( @@ -258,6 +271,7 @@ async def collect_chat_response( 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( @@ -270,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 d5246bbc..7a31f6fd 100644 --- a/src/exo/master/adapters/claude.py +++ b/src/exo/master/adapters/claude.py @@ -29,6 +29,8 @@ from exo.shared.types.claude_api import ( ClaudeStopReason, ClaudeTextBlock, ClaudeTextDelta, + ClaudeThinkingBlock, + ClaudeThinkingDelta, ClaudeToolResultBlock, ClaudeToolUseBlock, ClaudeUsage, @@ -104,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( { @@ -125,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"): @@ -132,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( @@ -145,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 @@ -162,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 @@ -175,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, @@ -192,6 +211,7 @@ async def collect_claude_response( # 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 @@ -219,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) @@ -228,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) @@ -275,16 +301,16 @@ 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): @@ -329,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) @@ -343,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/responses.py b/src/exo/master/adapters/responses.py index 90fa7732..dc726ba7 100644 --- a/src/exo/master/adapters/responses.py +++ b/src/exo/master/adapters/responses.py @@ -29,6 +29,12 @@ from exo.shared.types.openai_responses import ( ResponseOutputItemAddedEvent, ResponseOutputItemDoneEvent, ResponseOutputText, + ResponseReasoningItem, + ResponseReasoningSummaryPartAddedEvent, + ResponseReasoningSummaryPartDoneEvent, + ResponseReasoningSummaryText, + ResponseReasoningSummaryTextDeltaEvent, + ResponseReasoningSummaryTextDoneEvent, ResponsesRequest, ResponsesResponse, ResponsesStreamEvent, @@ -141,7 +147,9 @@ async def collect_responses_response( """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 @@ -168,6 +176,10 @@ async def collect_responses_response( ) continue + if chunk.is_thinking: + thinking_parts.append(chunk.text) + continue + accumulated_text += chunk.text if error_message is not None: @@ -182,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( @@ -212,6 +232,7 @@ async def generate_responses_stream( """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 @@ -233,32 +254,17 @@ async def generate_responses_stream( ) yield _format_sse(in_progress_event) - # 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 _format_sse(item_added) - - # 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 _format_sse(part_added) - 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): @@ -327,23 +333,184 @@ async def generate_responses_stream( 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 _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, ) @@ -354,7 +521,7 @@ async def generate_responses_stream( 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, ) @@ -367,7 +534,9 @@ 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 _format_sse(item_done) @@ -381,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, 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/shared/types/api.py b/src/exo/shared/types/api.py index 824c2c91..23ca9b7b 100644 --- a/src/exo/shared/types/api.py +++ b/src/exo/shared/types/api.py @@ -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 diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py index b97eee5f..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): 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/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/runner_response.py b/src/exo/shared/types/worker/runner_response.py index a66f762c..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): 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/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 567b0f91..01bfbe50 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -458,6 +458,19 @@ def _patch_lossy_chat_template(template: str) -> str | None: 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, @@ -469,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: @@ -497,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". diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index 51950192..f22ae5c8 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -7,6 +7,7 @@ from functools import cache from typing import 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] @@ -353,16 +354,22 @@ def main( 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 + mlx_generator, + 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 + ), ) - # GPT-OSS specific parsing to match other model formats. + # 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) @@ -414,6 +421,7 @@ def main( stats=response.stats, logprob=response.logprob, top_logprobs=response.top_logprobs, + is_thinking=response.is_thinking, ), ) ) @@ -675,44 +683,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_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. + """ + 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. + """ + 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. """ - 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. - """ - first = True + 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( 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 index 0a7ba102..080f0389 100644 --- 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 @@ -149,12 +149,23 @@ class TestParseGptOssThinkingThenToolCall: def test_thinking_then_tool_call(self): results = _collect(THINKING_THEN_TOOL_TOKENS) - # Should have thinking tags + content + tool call - text_parts = [r.text for r in results if isinstance(r, GenerationResponse)] - combined = "".join(text_parts) - assert "" in combined - assert "" in combined - assert "Let me think about this." in combined + # 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) From 7031901ae5164f10469d2455c8a4f07b5cc4b072 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Thu, 19 Feb 2026 20:51:17 +0000 Subject: [PATCH 38/45] Prevent common fatal crashes (#1555) ## Motivation Occasionally, memory does not get released when we shut down. There is no reason to delay deleting the model. Also handles can become None during shutdown, causing TypeErrors which are not handled and bringing down exo. Similarly, we were closing the event sender in the wrong place. Also let's not verify the SSL certificate for http connections to local peers, as this is failing sometimes and crashing. ## Test Plan ### Manual Testing No more crashes as described. --- src/exo/utils/channels.py | 8 +++++++- src/exo/utils/info_gatherer/net_profile.py | 2 +- src/exo/worker/runner/runner.py | 15 +++++++++------ src/exo/worker/runner/runner_supervisor.py | 17 +++++++++++------ 4 files changed, 28 insertions(+), 14 deletions(-) diff --git a/src/exo/utils/channels.py b/src/exo/utils/channels.py index c9336215..6ce4849d 100644 --- a/src/exo/utils/channels.py +++ b/src/exo/utils/channels.py @@ -192,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 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/runner/runner.py b/src/exo/worker/runner/runner.py index f22ae5c8..47293d8c 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -4,7 +4,7 @@ import resource import time from collections.abc import Generator from functools import cache -from typing import Literal +from typing import TYPE_CHECKING, Literal import mlx.core as mx from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model @@ -588,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 @@ -612,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 diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 58ac778e..2cee74fd 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -100,7 +100,6 @@ 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")) self._cancel_sender.close() self.runner_process.join(5) @@ -180,6 +179,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(): @@ -208,10 +208,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() From c45ff9ad4307341d0162d2ddbcdad50baefe1799 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Thu, 19 Feb 2026 21:15:33 +0000 Subject: [PATCH 39/45] memory tidy (#1558) add some pythonic extensions to memory, did a bunch of cleanup. --- dashboard/src/lib/stores/app.svelte.ts | 5 + dashboard/src/routes/+page.svelte | 11 +- dashboard/src/routes/downloads/+page.svelte | 16 +-- src/exo/download/coordinator.py | 8 +- src/exo/download/download_utils.py | 35 +++--- src/exo/download/shard_downloader.py | 6 +- src/exo/master/api.py | 2 +- src/exo/master/placement_utils.py | 13 +- src/exo/master/tests/test_placement.py | 6 +- .../test_apply/test_apply_node_download.py | 6 +- src/exo/shared/types/memory.py | 112 +++++++++++++++--- src/exo/shared/types/worker/downloads.py | 14 +-- src/exo/worker/engines/image/generate.py | 4 +- src/exo/worker/engines/mlx/cache.py | 2 +- src/exo/worker/engines/mlx/utils_mlx.py | 21 ++-- .../test_plan/test_download_and_loading.py | 24 +--- 16 files changed, 170 insertions(+), 115 deletions(-) diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 3e0074e3..03379cee 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 { diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte index 21e71774..76a3dbfd 100644 --- a/dashboard/src/routes/+page.svelte +++ b/dashboard/src/routes/+page.svelte @@ -858,10 +858,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; @@ -874,8 +872,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, @@ -1264,7 +1262,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; diff --git a/dashboard/src/routes/downloads/+page.svelte b/dashboard/src/routes/downloads/+page.svelte index 91719dc1..a57a18d3 100644 --- a/dashboard/src/routes/downloads/+page.svelte +++ b/dashboard/src/routes/downloads/+page.svelte @@ -74,7 +74,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; @@ -231,23 +230,14 @@ undefined; let cell: CellStatus; if (tag === "DownloadCompleted") { - const totalBytes = getBytes( - payload.total_bytes ?? payload.totalBytes, - ); + const totalBytes = getBytes(payload.total); cell = { kind: "completed", totalBytes, modelDirectory }; } else if (tag === "DownloadOngoing") { const rawProgress = payload.download_progress ?? payload.downloadProgress ?? {}; const prog = rawProgress as Record; - const totalBytes = getBytes( - prog.total_bytes ?? - prog.totalBytes ?? - payload.total_bytes ?? - payload.totalBytes, - ); - const downloadedBytes = getBytes( - prog.downloaded_bytes ?? prog.downloadedBytes, - ); + const totalBytes = getBytes(prog.total ?? payload.total); + const downloadedBytes = getBytes(prog.downloaded); const speed = (prog.speed as number) ?? 0; const etaMs = (prog.eta_ms as number) ?? (prog.etaMs as number) ?? 0; diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 30e45a08..f2b44495 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -80,7 +80,7 @@ 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 @@ -203,7 +203,7 @@ 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 @@ -332,13 +332,13 @@ 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, diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 5691e5dd..868bbbe9 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, @@ -578,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) ) @@ -609,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, 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/master/api.py b/src/exo/master/api.py index 319a8164..0f3d8711 100644 --- a/src/exo/master/api.py +++ b/src/exo/master/api.py @@ -1322,7 +1322,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), diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index b20a39cc..b80d70c0 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 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/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/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/worker/downloads.py b/src/exo/shared/types/worker/downloads.py index a29edcbf..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 @@ -34,7 +34,7 @@ class DownloadPending(BaseDownloadProgress): class DownloadCompleted(BaseDownloadProgress): - total_bytes: Memory + total: Memory class DownloadFailed(BaseDownloadProgress): @@ -86,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/worker/engines/image/generate.py b/src/exo/worker/engines/image/generate.py index f5526c8f..2f4c5a4e 100644 --- a/src/exo/worker/engines/image/generate.py +++ b/src/exo/worker/engines/image/generate.py @@ -166,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, @@ -175,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/cache.py b/src/exo/worker/engines/mlx/cache.py index 7669f1c1..ae6f76fa 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -22,7 +22,7 @@ from exo.worker.runner.bootstrap import logger # Fraction of device memory above which LRU eviction kicks in. # Smaller machines need more aggressive eviction. def _default_memory_threshold() -> float: - total_gb = psutil.virtual_memory().total / (1024**3) + total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb if total_gb >= 128: return 0.85 if total_gb >= 64: diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 01bfbe50..360a018b 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -232,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: @@ -642,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}.") 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 } From f662c129dd228a5de66bf2db759db8386b4a114e Mon Sep 17 00:00:00 2001 From: rltakashige Date: Thu, 19 Feb 2026 21:32:48 +0000 Subject: [PATCH 40/45] Prioritise tb for ring instances (#1556) ## Motivation TB has better bandwidth and latency than ethernet. We should prioritise TB5 where possible. This drastically improves distributed image generation performance. ## Test Plan ### Manual Testing Saw on the dashboard that TB (169.254) addresses were prioritised. Tested that image models scale much better. ### Automated Testing No regression on Kimi K2.5 --- src/exo/master/placement_utils.py | 35 ++++++++++++++++------ src/exo/worker/runner/runner_supervisor.py | 3 +- 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py index b80d70c0..eb78c8fa 100644 --- a/src/exo/master/placement_utils.py +++ b/src/exo/master/placement_utils.py @@ -341,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. @@ -353,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)) @@ -399,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( @@ -430,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/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py index 2cee74fd..e8a06a77 100644 --- a/src/exo/worker/runner/runner_supervisor.py +++ b/src/exo/worker/runner/runner_supervisor.py @@ -100,7 +100,8 @@ class RunnerSupervisor: logger.info("Runner supervisor shutting down") self._ev_recv.close() self._task_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(5) if not self.runner_process.is_alive(): From 3006c8ea4ee43addf90e59e71a44c34135afba05 Mon Sep 17 00:00:00 2001 From: rltakashige Date: Fri, 20 Feb 2026 11:46:24 +0000 Subject: [PATCH 41/45] Ensure coordinator is rank 0 (#1559) ## Motivation Coordinator can be a random rank. Let's just fix this to rank 0 as that's what we typically assume. ## Test Plan ### Manual Testing Works as normal on 2 nodes. Let's wait for a little more testing to merge this. --------- Co-authored-by: Evan --- src/exo/master/placement.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) 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, From a16ff2c0476627381a2eecc7f7e08a6ec3cee415 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mustafa=20Alp=20Y=C4=B1lmaz?= <96022931+mustafalpyilmaz@users.noreply.github.com> Date: Fri, 20 Feb 2026 14:57:55 +0300 Subject: [PATCH 42/45] fix: correct misleading docstring in seed_models (#1561) ## Summary - Fixed stale docstring in `seed_models()` that referenced `.cache/huggingface/hub` when the function actually moves models to `EXO_MODELS_DIR` (resolved via `ensure_models_dir()`) - The old docstring was misleading for AI coding agents analyzing the codebase, causing incorrect conclusions about model storage paths ## Changes `src/exo/download/download_utils.py`: Updated docstring from `"Move model in resources folder of app to .cache/huggingface/hub"` to `"Move models from resources folder to EXO_MODELS_DIR."` Co-authored-by: rltakashige --- src/exo/download/download_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 868bbbe9..a9da4cc5 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -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(): From addf73a1441590b8fa20084e08fd0446dca3a25b Mon Sep 17 00:00:00 2001 From: rltakashige Date: Fri, 20 Feb 2026 12:03:27 +0000 Subject: [PATCH 43/45] Add support for Ollama API (#1560) ## Motivation Ollama has a bunch of integrations, such as OpenWebUI, that are very handy. Let's support it :) ## Test Plan ### Manual Testing image --- src/exo/master/adapters/ollama.py | 456 +++++++++++++++++++++++++++++ src/exo/master/api.py | 193 ++++++++++++ src/exo/shared/types/ollama_api.py | 148 ++++++++++ 3 files changed, 797 insertions(+) create mode 100644 src/exo/master/adapters/ollama.py create mode 100644 src/exo/shared/types/ollama_api.py 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/api.py b/src/exo/master/api.py index 0f3d8711..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, @@ -141,6 +149,19 @@ 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, @@ -300,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) @@ -1293,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() 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] From bddad7e79cb505f895f2fa0d8e21892f65b6ac7e Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Fri, 20 Feb 2026 04:37:56 -0800 Subject: [PATCH 44/45] feat: show ETA on prefill progress bar (#1557) ## Summary - Show estimated time remaining during prefill (prompt processing phase) - Track prefill start time via performance.now() and extrapolate from observed token throughput - Display ~Xs remaining or ~Xm Ys remaining next to the percentage on the progress bar - Wait 200ms before showing ETA to ensure a stable sample window ## Changes **PrefillProgressBar.svelte**: Add etaText derived computation that calculates remaining time from (remainingTokens / tokensPerMs). Renders in a new flex row below the progress bar alongside the percentage. **app.svelte.ts**: Add startedAt: number field to PrefillProgress interface. Set on first prefill_progress SSE event, preserved across subsequent updates. ## Test plan - [ ] Start inference with a long prompt (10k+ tokens) on a multi-node cluster - [ ] Verify the progress bar shows ~Xs remaining after ~200ms of prefill - [ ] Verify the ETA decreases as prefill progresses - [ ] Verify short prefills (<200ms) dont flash a briefly-visible ETA - [ ] Verify ETA disappears when prefill completes and token generation begins Co-authored-by: Claude Opus 4.6 Co-authored-by: rltakashige --- .../lib/components/PrefillProgressBar.svelte | 22 +++++++++++++++++-- dashboard/src/lib/stores/app.svelte.ts | 3 +++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/dashboard/src/lib/components/PrefillProgressBar.svelte b/dashboard/src/lib/components/PrefillProgressBar.svelte index ea08d11d..c07be12c 100644 --- a/dashboard/src/lib/components/PrefillProgressBar.svelte +++ b/dashboard/src/lib/components/PrefillProgressBar.svelte @@ -14,6 +14,21 @@ : 0, ); + const etaText = $derived.by(() => { + if (progress.processed <= 0 || progress.total <= 0) return null; + const elapsedMs = performance.now() - progress.startedAt; + if (elapsedMs < 200) return null; // need a minimum sample window + const tokensPerMs = progress.processed / elapsedMs; + const remainingTokens = progress.total - progress.processed; + const remainingMs = remainingTokens / tokensPerMs; + const remainingSec = Math.ceil(remainingMs / 1000); + if (remainingSec <= 0) return null; + if (remainingSec < 60) return `~${remainingSec}s remaining`; + const mins = Math.floor(remainingSec / 60); + const secs = remainingSec % 60; + return `~${mins}m ${secs}s remaining`; + }); + function formatTokenCount(count: number | undefined): string { if (count == null) return "0"; if (count >= 1000) { @@ -40,8 +55,11 @@ style="width: {percentage}%" >
-
- {percentage}% +
+ {etaText ?? ""} + {percentage}%
diff --git a/dashboard/src/lib/stores/app.svelte.ts b/dashboard/src/lib/stores/app.svelte.ts index 03379cee..c10c8e0a 100644 --- a/dashboard/src/lib/stores/app.svelte.ts +++ b/dashboard/src/lib/stores/app.svelte.ts @@ -281,6 +281,8 @@ export interface TokenData { export interface PrefillProgress { processed: number; total: number; + /** Timestamp (performance.now()) when prefill started. */ + startedAt: number; } export interface Message { @@ -2464,6 +2466,7 @@ class AppStore { this.prefillProgress = { processed: inner.processed_tokens, total: inner.total_tokens, + startedAt: this.prefillProgress?.startedAt ?? performance.now(), }; }, }, From e32b649d2f5691f41a64c4ba872dd2b7cfe36f03 Mon Sep 17 00:00:00 2001 From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com> Date: Fri, 20 Feb 2026 05:12:22 -0800 Subject: [PATCH 45/45] fix: enable psutil fallback for memory monitoring when macmon is missing (#1478) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - On macOS, memory monitoring relied exclusively on `macmon` — the psutil fallback was explicitly disabled (`memory_poll_rate = None`) - When `macmon` is not installed (e.g., mac-mini-2 through mac-mini-4 in our cluster), **no memory data was reported**, causing nodes to show 0GB memory in the cluster state - This blocked the scheduler from placing shards on those nodes since it had no memory data to work with - Fix: when `macmon` is not found on Darwin, fall back to psutil-based memory polling (`memory_poll_rate = 1`) ## Root cause `InfoGatherer` has two memory monitoring paths: 1. `macmon` (Darwin-only): provides memory + GPU/CPU/power stats 2. `psutil` (non-Darwin fallback): provides memory via `MemoryUsage.from_psutil()` Line 378 disabled psutil on Darwin: `memory_poll_rate = None if IS_DARWIN else 1` Line 389 only starts macmon if the binary exists: `if shutil.which("macmon") is not None` If macmon is missing on Darwin, **neither path runs** — zero memory reported. ## Test plan - [ ] Verify `uv run basedpyright` passes (0 errors confirmed) - [ ] Verify `uv run ruff check` passes (confirmed) - [ ] Verify `uv run pytest src/exo/utils/info_gatherer/` passes (2/2 confirmed) - [ ] Deploy to cluster nodes without macmon and verify memory appears in `/state` 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: rltakashige --- src/exo/utils/info_gatherer/info_gatherer.py | 6 ++++++ 1 file changed, 6 insertions(+) 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)