← back to Exo
collection of fixes for Shanghai demo
57073f35c344719836df2047f747ffdafce212ba · 2025-08-15 15:21:51 +0100 · Gelu Vrabie
Co-authored-by: Matt Beton <matthew.beton@gmail.com>
Co-authored-by: Gelu Vrabie <gelu@exolabs.net>
Files touched
M engines/mlx/utils_mlx.pyA hosts.jsonM master/api.pyM master/placement.pyA mlx-lm-checkA nodes.jsonM run.shA run_remote.shA scp_repo.shA scripts_guide.txtM shared/models/model_cards.pyM shared/types/api.pyM shared/types/worker/commands_runner.pyM shared/types/worker/common.pyM uv.lockM worker/pyproject.tomlM worker/runner/communication.pyM worker/runner/runner_supervisor.pyM worker/runner/utils.pyM worker/tests/test_supervisor/test_supervisor_sad.pyM worker/worker.py
Diff
commit 57073f35c344719836df2047f747ffdafce212ba
Author: Gelu Vrabie <gelu.vrabie.univ@gmail.com>
Date: Fri Aug 15 15:21:51 2025 +0100
collection of fixes for Shanghai demo
Co-authored-by: Matt Beton <matthew.beton@gmail.com>
Co-authored-by: Gelu Vrabie <gelu@exolabs.net>
---
engines/mlx/utils_mlx.py | 16 +--
hosts.json | 1 +
master/api.py | 19 ++++
master/placement.py | 38 +++++--
mlx-lm-check | 1 +
nodes.json | 1 +
run.sh | 2 +-
run_remote.sh | 79 ++++++++++++++
scp_repo.sh | 65 ++++++++++++
scripts_guide.txt | 22 ++++
shared/models/model_cards.py | 2 +-
shared/types/api.py | 1 +
shared/types/worker/commands_runner.py | 3 +-
shared/types/worker/common.py | 5 +-
uv.lock | 26 +++--
worker/pyproject.toml | 5 +-
worker/runner/communication.py | 7 +-
worker/runner/runner_supervisor.py | 114 ++++++++++++---------
worker/runner/utils.py | 22 ++--
.../tests/test_supervisor/test_supervisor_sad.py | 4 +-
worker/worker.py | 1 -
21 files changed, 340 insertions(+), 94 deletions(-)
diff --git a/engines/mlx/utils_mlx.py b/engines/mlx/utils_mlx.py
index a409b5ca..a04f0222 100644
--- a/engines/mlx/utils_mlx.py
+++ b/engines/mlx/utils_mlx.py
@@ -70,7 +70,7 @@ def initialize_mlx(
mx.random.seed(42)
if len(hosts) > 1:
mlx_distributed_init(model_shard_meta.device_rank, hosts)
- sampler: Callable[[mx.array], mx.array] = make_sampler(temp=0.7) # type: ignore
+ sampler: Callable[[mx.array], mx.array] = make_sampler(temp=0.7)
model, tokenizer = shard_and_load(model_shard_meta)
@@ -107,14 +107,18 @@ async def apply_chat_template(
messages = chat_task_data.messages
messages_dicts = [msg.model_dump() for msg in messages]
- # Filter out None values, keeping only 'role' and 'content' keys
+ # Filter out None values, keeping relevant keys for the model
formatted_messages = []
for message in messages_dicts:
filtered_message: dict[str, Any] = {k: v for k, v in message.items() if v is not None} # type: ignore
- # Verify we have exactly the expected keys
- assert set(filtered_message.keys()) == {"role", "content"}, (
- f"Expected only 'role' and 'content' keys, got: {filtered_message.keys()}"
- )
+
+ # Verify we have required fields
+ if "role" not in filtered_message:
+ raise ValueError(f"Message missing 'role' field: {filtered_message}")
+ if "content" not in filtered_message and "thinking" not in filtered_message:
+ # If neither content nor thinking is present, skip this message
+ continue
+
formatted_messages.append(filtered_message) # type: ignore
messages_dicts = formatted_messages
diff --git a/hosts.json b/hosts.json
new file mode 100644
index 00000000..fdf160cf
--- /dev/null
+++ b/hosts.json
@@ -0,0 +1 @@
+["s13@169.254.249.73", "s14@169.254.69.217", "s15@169.254.165.26", "s16@169.254.29.77"]
\ No newline at end of file
diff --git a/master/api.py b/master/api.py
index 40c7af10..60250bae 100644
--- a/master/api.py
+++ b/master/api.py
@@ -193,6 +193,25 @@ class API:
"""Handle chat completions with proper streaming response."""
model_meta = await resolve_model_meta(payload.model)
payload.model = model_meta.model_id
+
+ # Preprocess messages for GPT-OSS harmony format if needed
+ if "gpt-oss" in payload.model.lower():
+ import re
+ for message in payload.messages:
+ if message.content and "<|channel|>" in message.content:
+ # Parse harmony format tags
+ thinking_pattern = r'<\|channel\|>(.*?)(?=<\|message\|>|$)'
+ content_pattern = r'<\|message\|>(.*?)(?=<\|end\|>|$)'
+
+ thinking_match = re.search(thinking_pattern, message.content, re.DOTALL)
+ content_match = re.search(content_pattern, message.content, re.DOTALL)
+
+ if content_match:
+ # Extract the actual content
+ message.content = content_match.group(1).strip()
+ if thinking_match:
+ # Store thinking in the thinking field
+ message.thinking = thinking_match.group(1).strip()
for instance in self.get_state().instances.values():
if instance.shard_assignments.model_id == payload.model:
diff --git a/master/placement.py b/master/placement.py
index 26268853..ed25cc2a 100644
--- a/master/placement.py
+++ b/master/placement.py
@@ -1,3 +1,4 @@
+import json
import random
from collections.abc import Mapping
from copy import deepcopy
@@ -11,7 +12,7 @@ from master.utils.placement_utils import (
get_smallest_cycles,
)
from shared.topology import Topology
-from shared.types.common import Host
+from shared.types.common import Host, NodeId
from shared.types.events import Event, InstanceCreated, InstanceDeleted
from shared.types.events.commands import CreateInstanceCommand, DeleteInstanceCommand
from shared.types.worker.common import InstanceId
@@ -21,6 +22,13 @@ from shared.types.worker.instances import Instance, InstanceStatus
def random_ephemeral_port() -> int:
return random.randint(49152, 65535)
+DEVICE_ORDERING: list[str] = []
+with open('nodes.json', ('r')) as f:
+ device_json: list[str] = json.load(f) # type: ignore
+ for device in device_json:
+ DEVICE_ORDERING.append(NodeId(device))
+assert len(DEVICE_ORDERING) == 4
+
@singledispatch
def get_instance_placements(
command: CreateInstanceCommand,
@@ -42,12 +50,30 @@ def get_instance_placements(
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
selected_cycle = None
+
+ has_thunderbolt_cycle = any([
+ topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
+ for cycle in smallest_cycles
+ ])
+ if has_thunderbolt_cycle:
+ smallest_cycles = [
+ cycle for cycle in smallest_cycles
+ if topology.get_subgraph_from_nodes(cycle).is_thunderbolt_cycle(cycle)
+ ]
+
+ nodes_01, nodes_23 = None, None
for cycle in smallest_cycles:
- cycle_graph: Topology = topology.get_subgraph_from_nodes(cycle)
- if cycle_graph.is_thunderbolt_cycle(cycle):
- selected_cycle = cycle
- break
- if selected_cycle is None:
+ cycle_ids = [x.node_id for x in cycle]
+ if nodes_01 is None and set(cycle_ids) == set(DEVICE_ORDERING[:2]):
+ nodes_01= cycle
+ if nodes_23 is None and set(cycle_ids) == set(DEVICE_ORDERING[2:]):
+ nodes_23= cycle
+
+ if nodes_01:
+ selected_cycle = nodes_01
+ elif nodes_23:
+ selected_cycle = nodes_23
+ else:
selected_cycle = max(smallest_cycles, key=lambda cycle: sum(node.node_profile.memory.ram_available for node in cycle if node.node_profile is not None))
shard_assignments = get_shard_assignments(command.model_meta, selected_cycle)
diff --git a/mlx-lm-check b/mlx-lm-check
new file mode 160000
index 00000000..d5bdab1a
--- /dev/null
+++ b/mlx-lm-check
@@ -0,0 +1 @@
+Subproject commit d5bdab1a22b053d75194ce4d225df9fc1635a400
diff --git a/nodes.json b/nodes.json
new file mode 100644
index 00000000..8c44494e
--- /dev/null
+++ b/nodes.json
@@ -0,0 +1 @@
+["9gG9JZ5YY1zLE5xVYA2L8DoCTxkYKfxrGi33stPqq1cb", "F4p3DefvhUk9fGfToJXteT7GL9JuF4qMbUCvUKeB7VPZ", "J7AAM7DiMfnvxNvA1AXUFfencsSfwp4Qi851Y7v9hP1M", "7BbDVE6oN35avU6xY7e75m3r3EjADNBTm2ZMZB83EsLf"]
\ No newline at end of file
diff --git a/run.sh b/run.sh
index 89ca175a..74e81181 100755
--- a/run.sh
+++ b/run.sh
@@ -32,7 +32,7 @@ if [ "$CLEAN" = true ]; then
fi
# Configure MLX
-./configure_mlx.sh
+# ./configure_mlx.sh
# First command (worker) - changes based on replica flag
if [ "$REPLICA" = true ]; then
diff --git a/run_remote.sh b/run_remote.sh
new file mode 100755
index 00000000..87ee2638
--- /dev/null
+++ b/run_remote.sh
@@ -0,0 +1,79 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+###############################################################################
+# Args & prerequisites
+###############################################################################
+if [[ $# -lt 1 || $# -gt 2 ]]; then
+ echo "Usage: $0 <PASSWORD> [hosts_file]" >&2 ; exit 1
+fi
+PASSWORD=$1
+HOSTS_FILE=${2:-hosts.json}
+
+for prog in jq sshpass; do
+ command -v "$prog" >/dev/null ||
+ { echo "Error: $prog not installed."; exit 1; }
+done
+
+###############################################################################
+# Load hosts.json (works on macOS Bash 3.2 and Bash 4+)
+###############################################################################
+if builtin command -v mapfile >/dev/null 2>&1; then
+ mapfile -t HOSTS < <(jq -r '.[]' "$HOSTS_FILE")
+else
+ HOSTS=()
+ while IFS= read -r h; do HOSTS+=("$h"); done < <(jq -r '.[]' "$HOSTS_FILE")
+fi
+[[ ${#HOSTS[@]} -gt 0 ]] || { echo "No hosts found in $HOSTS_FILE"; exit 1; }
+
+###############################################################################
+# Helper – run a remote command and capture rc/stderr/stdout
+###############################################################################
+ssh_opts=(-o StrictHostKeyChecking=no
+ -o NumberOfPasswordPrompts=1 # allow sshpass to answer exactly once
+ -o LogLevel=ERROR)
+
+run_remote () { # $1 host $2 command
+ local host=$1 cmd=$2 rc
+ if sshpass -p "$PASSWORD" ssh "${ssh_opts[@]}" "$host" "$cmd"; then
+ rc=0
+ else
+ rc=$?
+ fi
+ return $rc
+}
+
+###############################################################################
+# Phase 1 – kill exo everywhere (parallel)
+###############################################################################
+echo "=== Stage 1: killing exo on ${#HOSTS[@]} host(s) ==="
+fail=0
+for h in "${HOSTS[@]}"; do
+ (
+ run_remote "$h" 'pkill -f exo || true'
+ ) || fail=1 &
+done
+wait
+(( fail == 0 )) || { echo "❌ Some hosts could not be reached—check password or SSH access."; exit 1; }
+echo "✓ exo processes killed on all reachable hosts."
+
+###############################################################################
+# Phase 2 – start new exo processes (parallel, with sudo -S)
+###############################################################################
+echo "=== Stage 2: starting new exo processes ==="
+fail=0
+for i in "${!HOSTS[@]}"; do
+ h=${HOSTS[$i]}
+
+ # one liner that pre-caches sudo and then runs the script
+ if [[ $i -eq 0 ]]; then
+ remote_cmd="cd ~/exo && ./run.sh -c"
+ else
+ remote_cmd="cd ~/exo && ./run.sh -rc"
+ fi
+
+ ( run_remote "$h" "$remote_cmd" ) || fail=1 &
+done
+wait
+(( fail == 0 )) && echo "🎉 Deployment finished!" || \
+ { echo "⚠️ Some starts failed—see above."; exit 1; }
diff --git a/scp_repo.sh b/scp_repo.sh
new file mode 100755
index 00000000..a38f58ec
--- /dev/null
+++ b/scp_repo.sh
@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+# bulk_scp.sh — Sync a local repo to many hosts, respecting .gitignore and continuing even if
+# some hosts fail. Tested on macOS Bash 3.x.
+#
+# ------------ User-tunable variables ------------
+LOCAL_DIR="." # Local directory you want to send
+REMOTE_DIR="~/exo" # Destination directory on the remote machines
+HOSTS_FILE="hosts.json" # JSON array of hosts (["user@ip", ...])
+# ------------ End of user-tunable section -------
+
+set -uo pipefail # Treat unset vars as error; fail pipelines, but we handle exit codes ourselves
+
+if [ "$#" -ne 1 ]; then
+ echo "Usage: $0 <password>" >&2
+ exit 1
+fi
+PASSWORD="$1"
+
+# Dependency checks
+for cmd in sshpass jq rsync git; do
+ if ! command -v "$cmd" >/dev/null 2>&1; then
+ echo "Error: $cmd is required but not installed." >&2
+ exit 1
+ fi
+done
+
+# Verify hosts file exists
+if [ ! -f "$HOSTS_FILE" ]; then
+ echo "Error: Hosts file '$HOSTS_FILE' not found." >&2
+ exit 1
+fi
+
+# Build a temporary exclude file containing every Git‑ignored path
+EXCLUDE_FILE=$(mktemp)
+trap 'rm -f "$EXCLUDE_FILE"' EXIT
+
+if git -C "$LOCAL_DIR" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+ git -C "$LOCAL_DIR" ls-files -z -o -i --exclude-standard \
+ | tr '\0' '\n' > "$EXCLUDE_FILE"
+else
+ # Fallback: just use top‑level .gitignore if present
+ [ -f "$LOCAL_DIR/.gitignore" ] && cat "$LOCAL_DIR/.gitignore" > "$EXCLUDE_FILE"
+fi
+
+# Iterate over hosts — process substitution keeps stdin free for rsync/ssh
+while IFS= read -r TARGET || [ -n "$TARGET" ]; do
+ [ -z "$TARGET" ] && continue # skip blanks
+ echo "\n—— Syncing $LOCAL_DIR → $TARGET:$REMOTE_DIR ——"
+
+# # Ensure remote directory exists (ignore failure but report)
+# if ! sshpass -p "$PASSWORD" ssh -o StrictHostKeyChecking=no "$TARGET" "mkdir -p $REMOTE_DIR" </dev/null; then
+# echo "✗ Failed to create $REMOTE_DIR on $TARGET" >&2
+# continue # move on to next host
+# fi
+
+ # Rsync with checksums; redirect stdin so rsync/ssh can't eat host list
+ if sshpass -p "$PASSWORD" rsync -azc --delete --exclude-from="$EXCLUDE_FILE" \
+ -e "ssh -o StrictHostKeyChecking=no" \
+ "$LOCAL_DIR/" "$TARGET:$REMOTE_DIR/" </dev/null; then
+ echo "✓ Success: $TARGET"
+ else
+ echo "✗ Failed: $TARGET" >&2
+ fi
+
+done < <(jq -r '.[]' "$HOSTS_FILE")
diff --git a/scripts_guide.txt b/scripts_guide.txt
new file mode 100644
index 00000000..5e3d6bde
--- /dev/null
+++ b/scripts_guide.txt
@@ -0,0 +1,22 @@
+you have 2 scripts now added:
+ 1. scp_repo.sh that you call like "./scp_repo.sh {password}"
+where password is the password for the studios. call this from the
+root of the repo and it will send any differences in your local repo
+to the machines. this should only be needed when things changed
+ 2. run_remote.sh, also called like "./run_remote.sh {password}"
+which kills all running exo process and starts new ones with fresh dbs
+
+both of these use the file hosts.json which is a json list of strings
+of the form user@ip where you need to put the studios with their username
+and THUNDERBOLT ips (get these manually from the machines after all of
+them and your laptop are hooked up via tb5 and have ips on the thunderbolt
+bridge in settings>network). the order here doesn't matter EXCEPT for the
+first entry which will be the master. so the script runs ./run.sh -c on the
+first entry in that list and ./run.sh -rc on all the others
+
+
+separately, there is now a nodes.json which is also a list of strings but this
+time of the node ids of the machines (the uuid that gets generated in python
+and printed when the process starts etc). here you do need them in the exact
+order the machines are connected in via thunderbolt. this is used to prefer
+spawning models across machines 1-2 and then 3-4 in that order if doable
\ No newline at end of file
diff --git a/shared/models/model_cards.py b/shared/models/model_cards.py
index 8f2dcd2c..62165ee1 100644
--- a/shared/models/model_cards.py
+++ b/shared/models/model_cards.py
@@ -66,7 +66,7 @@ MODEL_CARDS: dict[str, ModelCard] = {
metadata=ModelMetadata(
model_id="mlx-community/DeepSeek-R1-0528-8bit",
pretty_name="DeepSeek R1 671B (8-bit)",
- storage_size_kilobytes=754706307,
+ storage_size_kilobytes=754998771712//1024,
n_layers=61,
),
),
diff --git a/shared/types/api.py b/shared/types/api.py
index cdf9913e..a166866d 100644
--- a/shared/types/api.py
+++ b/shared/types/api.py
@@ -28,6 +28,7 @@ class ModelList(BaseModel):
class ChatCompletionMessage(BaseModel):
role: Literal["system", "user", "assistant", "developer", "tool", "function"]
content: str | None = None
+ thinking: str | None = None # Added for GPT-OSS harmony format support
name: str | None = None
tool_calls: list[dict[str, Any]] | None = None
tool_call_id: str | None = None
diff --git a/shared/types/worker/commands_runner.py b/shared/types/worker/commands_runner.py
index 3ca0bf22..0cf2f89c 100644
--- a/shared/types/worker/commands_runner.py
+++ b/shared/types/worker/commands_runner.py
@@ -101,8 +101,7 @@ class ErrorResponse(BaseRunnerResponse[RunnerResponseType.ErrorResponse]):
)
error_type: str
error_message: str
- traceback: str | None = None
-
+ traceback: str
RunnerResponse = Annotated[
InitializedResponse | GenerationResponse | PrintResponse | FinishedResponse | ErrorResponse,
diff --git a/shared/types/worker/common.py b/shared/types/worker/common.py
index 7eb298c8..2d22785c 100644
--- a/shared/types/worker/common.py
+++ b/shared/types/worker/common.py
@@ -1,5 +1,4 @@
from enum import Enum
-from typing import Optional
from shared.types.common import ID
@@ -19,8 +18,8 @@ class NodeStatus(str, Enum):
class RunnerError(Exception):
"""Exception raised when the runner process encounters an error."""
- def __init__(self, error_type: str, error_message: str, traceback: Optional[str] = None):
+ def __init__(self, error_type: str, error_message: str, traceback: str):
self.error_type = error_type
self.error_message = error_message
self.traceback = traceback
- super().__init__(f"{error_type}: {error_message}")
\ No newline at end of file
+ super().__init__(f"{error_type}: {error_message}. Traceback: {traceback}")
\ No newline at end of file
diff --git a/uv.lock b/uv.lock
index 68365b4f..9a0ec757 100644
--- a/uv.lock
+++ b/uv.lock
@@ -382,15 +382,17 @@ dependencies = [
{ name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "psutil", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+ { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
[package.metadata]
requires-dist = [
{ name = "exo-shared", editable = "shared" },
{ name = "huggingface-hub", specifier = ">=0.33.4" },
- { name = "mlx", specifier = "==0.26.3" },
- { name = "mlx-lm", specifier = ">=0.25.3" },
+ { name = "mlx", specifier = ">=0.26.3" },
+ { name = "mlx-lm", git = "https://github.com/ml-explore/mlx-lm.git" },
{ name = "psutil", specifier = ">=7.0.0" },
+ { name = "transformers", specifier = ">=4.55.0" },
]
[[package]]
@@ -539,7 +541,7 @@ wheels = [
[[package]]
name = "huggingface-hub"
-version = "0.33.4"
+version = "0.34.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -551,9 +553,9 @@ dependencies = [
{ name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/4b/9e/9366b7349fc125dd68b9d384a0fea84d67b7497753fe92c71b67e13f47c4/huggingface_hub-0.33.4.tar.gz", hash = "sha256:6af13478deae120e765bfd92adad0ae1aec1ad8c439b46f23058ad5956cbca0a", size = 426674, upload-time = "2025-07-11T12:32:48.694Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/91/b4/e6b465eca5386b52cf23cb6df8644ad318a6b0e12b4b96a7e0be09cbfbcc/huggingface_hub-0.34.3.tar.gz", hash = "sha256:d58130fd5aa7408480681475491c0abd7e835442082fbc3ef4d45b6c39f83853", size = 456800, upload-time = "2025-07-29T08:38:53.885Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/46/7b/98daa50a2db034cab6cd23a3de04fa2358cb691593d28e9130203eb7a805/huggingface_hub-0.33.4-py3-none-any.whl", hash = "sha256:09f9f4e7ca62547c70f8b82767eefadd2667f4e116acba2e3e62a5a81815a7bb", size = 515339, upload-time = "2025-07-11T12:32:46.346Z" },
+ { url = "https://files.pythonhosted.org/packages/59/a8/4677014e771ed1591a87b63a2392ce6923baf807193deef302dcfde17542/huggingface_hub-0.34.3-py3-none-any.whl", hash = "sha256:5444550099e2d86e68b2898b09e85878fbd788fc2957b506c6a79ce060e39492", size = 558847, upload-time = "2025-07-29T08:38:51.904Z" },
]
[[package]]
@@ -676,8 +678,8 @@ wheels = [
[[package]]
name = "mlx-lm"
-version = "0.26.0"
-source = { registry = "https://pypi.org/simple" }
+version = "0.26.3"
+source = { git = "https://github.com/ml-explore/mlx-lm.git#d5bdab1a22b053d75194ce4d225df9fc1635a400" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -686,10 +688,6 @@ dependencies = [
{ name = "pyyaml", 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/8d/aa/a2f02e67736a2bf57acefb3a1a342005586f1be8d7b2fb37ca5f3d4f3049/mlx_lm-0.26.0.tar.gz", hash = "sha256:78980ad994baf976779cc1c34c0d55c1c6b63dffef4899d67fec240d0c443b52", size = 159064, upload-time = "2025-07-08T20:21:31.393Z" }
-wheels = [
- { url = "https://files.pythonhosted.org/packages/08/e7/d0e576397b61bf90a0bb27819443f723258acd8dd1207684fdef29243ce4/mlx_lm-0.26.0-py3-none-any.whl", hash = "sha256:b00294c26242cd50db4b6e3ec3a2baf1cfdf8ca49a5e6057dce14642fabe0d21", size = 217671, upload-time = "2025-07-08T20:21:29.448Z" },
-]
[[package]]
name = "multidict"
@@ -1172,7 +1170,7 @@ wheels = [
[[package]]
name = "transformers"
-version = "4.53.3"
+version = "4.55.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1186,9 +1184,9 @@ dependencies = [
{ name = "tokenizers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/f1/5c/49182918b58eaa0b4c954fd0e37c79fc299e5643e69d70089d0b0eb0cd9b/transformers-4.53.3.tar.gz", hash = "sha256:b2eda1a261de79b78b97f7888fe2005fc0c3fabf5dad33d52cc02983f9f675d8", size = 9197478, upload-time = "2025-07-22T07:30:51.51Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/27/5d/f7dc746eef83336a6b34197311fe0c1da0d1192f637c726c6a5cf0d83502/transformers-4.55.0.tar.gz", hash = "sha256:15aa138a05d07a15b30d191ea2c45e23061ebf9fcc928a1318e03fe2234f3ae1", size = 9569089, upload-time = "2025-08-05T16:13:48.997Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/41/b1/d7520cc5cb69c825599042eb3a7c986fa9baa8a8d2dea9acd78e152c81e2/transformers-4.53.3-py3-none-any.whl", hash = "sha256:5aba81c92095806b6baf12df35d756cf23b66c356975fb2a7fa9e536138d7c75", size = 10826382, upload-time = "2025-07-22T07:30:48.458Z" },
+ { url = "https://files.pythonhosted.org/packages/1c/93/bcb22fb52ed65084c0199270832aa4cdd4b41296d896f3e7ade188bccb68/transformers-4.55.0-py3-none-any.whl", hash = "sha256:29d9b8800e32a4a831bb16efb5f762f6a9742fef9fce5d693ed018d19b106490", size = 11267905, upload-time = "2025-08-05T16:13:34.814Z" },
]
[[package]]
diff --git a/worker/pyproject.toml b/worker/pyproject.toml
index ca38f5d6..1eefe599 100644
--- a/worker/pyproject.toml
+++ b/worker/pyproject.toml
@@ -7,9 +7,10 @@ requires-python = ">=3.13"
dependencies = [
"exo-shared",
"huggingface_hub>=0.33.4",
- "mlx==0.26.3",
- "mlx-lm>=0.25.3",
+ "mlx>=0.26.3",
+ "mlx-lm @ https://github.com/ml-explore/mlx-lm.git",
"psutil>=7.0.0",
+ "transformers>=4.55.0",
]
[build-system]
diff --git a/worker/runner/communication.py b/worker/runner/communication.py
index 83076607..5cde6a46 100644
--- a/worker/runner/communication.py
+++ b/worker/runner/communication.py
@@ -63,11 +63,12 @@ async def supervisor_read_response(
"proc.stdout should not be None when created with stdout=PIPE"
)
line_bytes: bytes = await asyncio.wait_for(proc.stdout.readline(), timeout=180)
- line: str = line_bytes.decode("utf-8").strip()
-
- if not line:
+ if not line_bytes:
+ # return None
raise EOFError("No more data to read when reading response from runner")
+ line: str = line_bytes.decode("utf-8").strip()
+
try:
return RunnerResponseTypeAdapter.validate_json(line)
except Exception as err:
diff --git a/worker/runner/runner_supervisor.py b/worker/runner/runner_supervisor.py
index d9945a4c..185889e5 100644
--- a/worker/runner/runner_supervisor.py
+++ b/worker/runner/runner_supervisor.py
@@ -1,5 +1,6 @@
import asyncio
import contextlib
+import time
import traceback
from collections.abc import AsyncGenerator
from logging import Logger
@@ -55,13 +56,13 @@ class RunnerSupervisor:
self.hosts: list[Host] = hosts
self.runner_process: asyncio.subprocess.Process = runner_process
self.running: bool = True
- self.stderr_task = asyncio.create_task(self._watch_stderr(logger))
+
+ self.stderr_queue = asyncio.Queue[tuple[float, str]]()
+ self.stderr_task = asyncio.create_task(self._watch_stderr(logger, self.stderr_queue))
self.running_task: asyncio.Task[None] = asyncio.create_task(
self._watch_runner()
)
self.logger = logger
- self.stderr_buffer: list[str] = [] # Accumulate stderr lines
- self.crash_detected: bool = False
self.returncode: int | None = None
self.stderr_outpu: str | None = None
@@ -78,6 +79,7 @@ class RunnerSupervisor:
The .create() classmethod pattern is used to ensure the constructor is asynchronous.
"""
cmd: list[str] = get_runner_command()
+
runner_process: asyncio.subprocess.Process = (
await asyncio.create_subprocess_exec(
*cmd,
@@ -87,6 +89,14 @@ class RunnerSupervisor:
)
)
logger.info(f'initializing mlx instance with {model_shard_meta=}')
+
+ self = cls(
+ model_shard_meta=model_shard_meta,
+ hosts=hosts,
+ runner_process=runner_process,
+ logger=logger,
+ )
+
await supervisor_write_message(
runner_process,
SetupMessage(
@@ -97,13 +107,19 @@ class RunnerSupervisor:
async def read_initialization_message() -> None:
while True:
- line: RunnerResponse | None = await supervisor_read_response(
- runner_process
- )
- if line is None:
- continue
- elif isinstance(line, PrintResponse):
- logger.info(line)
+ try:
+ line: RunnerResponse | None = await supervisor_read_response(
+ self.runner_process
+ )
+ if line is None:
+ continue
+ except EOFError:
+ if not self.runner_process.returncode:
+ continue
+ raise await self._raise_crashed() from EOFError
+
+ if isinstance(line, PrintResponse):
+ self.logger.info(f"runner printed: {line.text}")
continue
elif isinstance(line, ErrorResponse):
raise RunnerError(line.error_type, line.error_message, line.traceback or "")
@@ -117,12 +133,8 @@ class RunnerSupervisor:
if not initialize_timeout:
initialize_timeout = get_init_timeout(model_shard_meta)
await asyncio.wait_for(read_initialization_message(), timeout=initialize_timeout)
- return cls(
- model_shard_meta=model_shard_meta,
- hosts=hosts,
- runner_process=runner_process,
- logger=logger,
- )
+
+ return self
async def astop(self) -> None:
# Cancel the stderr monitoring task
@@ -189,45 +201,45 @@ class RunnerSupervisor:
async def _watch_runner(self) -> None:
returncode = await self.runner_process.wait()
self.running = False
+
if returncode != 0:
- self.crash_detected = True
self.returncode = returncode # Will be picked up by _watch_stderr too
- async def _watch_stderr(self, logger: Logger) -> None:
+ await self.astop()
+
+ async def _watch_stderr(self, logger: Logger, stderr_queue: asyncio.Queue[tuple[float, str]]) -> None:
assert self.runner_process.stderr is not None
while self.running:
try:
line_bytes = await self.runner_process.stderr.readline()
if not line_bytes:
- break # EOF
+ break
line = line_bytes.decode('utf-8').strip()
- self.stderr_buffer.append(line)
- logger.error(f"Runner stderr: {line}")
- # Detect common crash patterns (extend as needed, e.g., for OOM: "Killed" or "Out of memory")
- self.crash_detected = True
- self.stderr_output = "\n".join(self.stderr_buffer)
- logger.critical(f"Runner crash detected: {self.stderr_output}")
- # Don't raise here—let callers (e.g., stream_response) detect via healthy/returncode
+ await stderr_queue.put((time.time(), line))
+ logger.warning(f"Runner stderr read: {line}")
except Exception as e:
- logger.error(f"Error reading runner stderr: {e}")
+ logger.warning(f"Error reading runner stderr: {e}")
break
- # After EOF, inspect returncode for confirmation (Unix-like: negative == signal)
- returncode = self.runner_process.returncode
- if returncode is not None and returncode != 0:
- self.crash_detected = True
- self.returncode = returncode
- self.stderr_output = "\n".join(self.stderr_buffer)
+ async def _raise_crashed(self) -> Exception:
+ await self.astop()
- def _raise_if_crashed(self) -> None:
- if self.crash_detected:
- self.logger.error(f'Error {self.returncode}: {self.stderr_output}')
- raise RunnerError(
- error_type="RunnerCrash",
- error_message=self.stderr_output,
- traceback=traceback.format_exc(),
- )
+ # Accumulate all stderr messages from the queue
+ stderr_output = ''
+ while not self.stderr_queue.empty():
+ try:
+ timestamp, line = self.stderr_queue.get_nowait()
+ stderr_output += f"[{timestamp}] {line}\n"
+ except asyncio.QueueEmpty:
+ break
+
+ self.logger.error(f'Error {self.returncode}: {stderr_output}')
+ return RunnerError(
+ error_type="MLXCrash",
+ error_message=stderr_output,
+ traceback=traceback.format_exc(),
+ )
def __del__(self) -> None:
if self.running:
@@ -262,7 +274,7 @@ class RunnerSupervisor:
async def stream_response(
self,
task: Task,
- request_started_callback: Callable[..., CoroutineType[Any, Any, None]] | None = None, # fyi this is async now
+ request_started_callback: Callable[..., CoroutineType[Any, Any, None]] | None = None,
) -> AsyncGenerator[GenerationChunk]:
"""
Streams a chat request from the model.
@@ -282,9 +294,11 @@ class RunnerSupervisor:
# This is easy for now. If we need more reliability, the runner can have a new 'ready' message type.
if request_started_callback is not None:
await request_started_callback()
- prefil_timeout = get_prefil_timeout(self.model_shard_meta)
+ prefil_timeout = get_prefil_timeout(task, self.model_shard_meta)
token_timeout = get_token_generate_timeout(self.model_shard_meta)
timeout = prefil_timeout
+ self.logger.info(f'starting chat completion with timeout {timeout}')
+
while True:
try:
line: RunnerResponse | None = await asyncio.wait_for(supervisor_read_response(
@@ -292,13 +306,19 @@ class RunnerSupervisor:
), timeout=timeout)
if line is None:
continue
- except (asyncio.TimeoutError, EOFError) as e:
- self._raise_if_crashed()
+ except asyncio.TimeoutError as e:
+ self.logger.info(f'timed out from timeout duration {timeout} - {"prefil" if timeout == prefil_timeout else "decoding stage"}')
+ await self.astop()
raise RunnerError(
error_type=type(e).__name__,
error_message=str(e),
- traceback="",
+ traceback=traceback.format_exc(),
) from e
+ # TODO: change this to a return none instead of error coming from the supervisor_Read_respons3
+ except EOFError as e:
+ if not self.runner_process.returncode:
+ continue
+ raise await self._raise_crashed() from e
match line:
case GenerationResponse():
yield TokenChunk(
@@ -319,4 +339,4 @@ class RunnerSupervisor:
self.logger.info(f"runner printed: {line.text}")
case ErrorResponse():
await self.astop()
- raise RunnerError(line.error_type, line.error_message, line.traceback or "")
\ No newline at end of file
+ raise RunnerError(line.error_type, line.error_message, line.traceback)
\ No newline at end of file
diff --git a/worker/runner/utils.py b/worker/runner/utils.py
index e89199bb..a3579ca1 100644
--- a/worker/runner/utils.py
+++ b/worker/runner/utils.py
@@ -1,6 +1,7 @@
import sys
-from shared.constants import LB_DISK_GBPS, LB_MEMBW_GBPS, LB_TFLOPS
+from shared.constants import LB_DISK_GBPS, LB_MEMBW_GBPS
+from shared.types.tasks import Task
from shared.types.worker.shards import ShardMetadata
@@ -18,13 +19,20 @@ def get_init_timeout(model_shard_meta: ShardMetadata) -> float:
return weights_size_kb / kbps_read + 2.0
-def get_prefil_timeout(model_shard_meta: ShardMetadata) -> float:
- weights_size_gb = get_weights_size_kb(model_shard_meta) / (1024 * 1024)
-
- tokens = 1000 # constant for now - the prompt is only tokenized in the device...
- prompt_gflops = tokens * weights_size_gb * 2
+def get_prefil_timeout(task: Task, model_shard_meta: ShardMetadata) -> float:
+ def get_prompt_str(task: Task) -> str:
+ messages = [x.content for x in task.task_params.messages if x.content]
+ return ''.join(messages)
- return LB_TFLOPS / (1024 * prompt_gflops) * 3 + 10.0
+ # TODO: made this timeout very long
+ tokens = len(get_prompt_str(task)) // 3 + 3000 # constant for now - the prompt is only tokenized in the device...
+
+ # TODO: For now we just hack and assume we prefil at 10tok/s
+ return tokens * 0.1
+
+ # prompt_gflops = tokens * weights_size_gb * 2
+
+ # return LB_TFLOPS / (1024 * prompt_gflops) * 3 + 10.0
def get_token_generate_timeout(model_shard_meta: ShardMetadata) -> float:
weights_size_kb = get_weights_size_kb(model_shard_meta)
diff --git a/worker/tests/test_supervisor/test_supervisor_sad.py b/worker/tests/test_supervisor/test_supervisor_sad.py
index 450612c3..40863786 100644
--- a/worker/tests/test_supervisor/test_supervisor_sad.py
+++ b/worker/tests/test_supervisor/test_supervisor_sad.py
@@ -90,4 +90,6 @@ async def test_supervisor_inference_timeout(
task.task_params.messages[0].content = 'EXO RUNNER MUST TIMEOUT'
with pytest.raises(RunnerError):
async for _ in supervisor.stream_response(task):
- pass
\ No newline at end of file
+ pass
+
+ await asyncio.sleep(0.1)
\ No newline at end of file
diff --git a/worker/worker.py b/worker/worker.py
index 6ac3f47c..7b0c3969 100644
--- a/worker/worker.py
+++ b/worker/worker.py
@@ -348,7 +348,6 @@ class Worker:
## Operation Planner
async def execute_op(self, op: RunnerOp) -> AsyncGenerator[Event, None]:
- ## It would be great if we can get rid of this async for ... yield pattern.
match op.op_type:
case RunnerOpType.ASSIGN_RUNNER:
event_generator = self._execute_assign_op(op)
← 7e19804a Integrate flake parts
·
back to Exo
·
discovery fixed a2a37c0e →