From 5d2156016a704ec2e8eaee57a623a053b2a0b615 Mon Sep 17 00:00:00 2001 From: Bawan Wang Date: Thu, 30 Jul 2026 16:39:44 +0800 Subject: [PATCH 1/3] feat(networking): allow dialing peers over unicast via EXO_ZENOH_CONNECT Peer discovery relies solely on IPv6 link-local multicast to the group ff12::e0a1:de89. Many Wi-Fi access points isolate wireless clients and do not forward link-local multicast between them, so two nodes on the same subnet never discover each other even though unicast between them works fine. There is currently no fallback: --bootstrap-peers raises "Bootstrap peers has been temporarily removed". Multicast is only used to learn a peer's address; the actual link is a unicast TCP connection established by connect_peer(). This adds an optional EXO_ZENOH_CONNECT environment variable holding a comma-separated list of zenoh endpoints (e.g. "tcp/192.168.1.2:52414") that is injected into the session's connect/endpoints, letting an operator dial known peers directly and bypass discovery entirely. The variable is opt-in: when unset, configuration is unchanged. Only one side needs to set it, since the TCP link is bidirectional and scouting/gossip/multihop is already enabled. Verified on two WSL2 nodes over Wi-Fi where discovery had never worked: tcpdump on udp/52413 showed only locally originated multicast and no inbound packets from the peer. With EXO_ZENOH_CONNECT set on one node only, both nodes reported nodes: 2 with a populated connections map, and the link remained stable. Co-Authored-By: Claude Opus 5 (1M context) --- rust/networking/src/lib.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/rust/networking/src/lib.rs b/rust/networking/src/lib.rs index 171337d719..b95ddb4cd5 100644 --- a/rust/networking/src/lib.rs +++ b/rust/networking/src/lib.rs @@ -32,6 +32,20 @@ pub fn cfg(identity: &str, listen_port: u16) -> Result { cfg.insert_json5("scouting/multicast/enabled", "false")?; cfg.insert_json5("scouting/multicast/autoconnect", "[]")?; cfg.insert_json5("scouting/gossip/multihop", "true")?; + // Peer discovery relies on IPv6 link-local multicast, which some networks + // (notably Wi-Fi access points isolating wireless clients) do not forward. + // EXO_ZENOH_CONNECT lets the operator dial known peers over unicast instead, + // as a comma-separated list of zenoh endpoints, e.g. "tcp/192.168.1.2:52414". + if let Ok(peers) = std::env::var("EXO_ZENOH_CONNECT") { + let endpoints: Vec = peers + .split(',') + .filter(|peer| !peer.is_empty()) + .map(|peer| format!("\"{peer}\"")) + .collect(); + if !endpoints.is_empty() { + cfg.insert_json5("connect/endpoints", &format!("[{}]", endpoints.join(",")))?; + } + } cfg.insert_json5("adminspace/enabled", "true")?; //cfg.insert_json5("transport/link/tx/batch_size", "9216")?; cfg.insert_json5("transport/link/rx/buffer_size", "16777216")?; From 7a6d0389fd4f8ecf8aa8b9e7eaffcb2d3f31c83c Mon Sep 17 00:00:00 2001 From: Bawan Wang Date: Sat, 1 Aug 2026 00:07:22 +0800 Subject: [PATCH 2/3] feat(mlx): make the prefill step size configurable via EXO_PREFILL_STEP_SIZE MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prefill step size was hardcoded to 4096. On a small card the shard's weights leave little room: with a 14B 4-bit model split across two 8GB GPUs each runner holds ~4.15GB of weights and has only 2.5-3.4GB left, and a 4096-token prefill chunk exceeds that. The request then dies with auto_parallel.py:167 mx.eval(output) RuntimeError: cudaMallocAsync(&data, size, stream) failed: out of memory which kills the runner and tears the instance down, so a client that sends long prompts (an agent system prompt plus a repo map) never reaches decode. Reading the value from the environment lets an operator trade prefill throughput for headroom without patching the source on every node. Unset keeps upstream behaviour. Measured on that 2x8GB setup: lowering it is not by itself a fix — 512 still OOMs on a ~1000-token prompt, so the dominant term is elsewhere. The knob is still worth having on machines where the shard leaves more room. Co-Authored-By: Claude Opus 5 --- src/exo/worker/engines/mlx/generator/batch_generate.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 5c7394ddec..4a028911fb 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -1,4 +1,5 @@ import contextlib +import os import time import uuid from dataclasses import dataclass, field @@ -102,10 +103,14 @@ class ExoBatchGenerator: _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False) def __post_init__(self) -> None: + # A 4096-token prefill chunk needs more activation memory than an 8GB card + # has left once the shard's weights are resident, so long prompts (agent + # system prompt + repo map) OOM before reaching decode. Make it tunable. + prefill_step_size = int(os.getenv("EXO_PREFILL_STEP_SIZE", "4096")) self._mlx_gen = MlxBatchGenerator( model=self.model, stop_tokens=[[t] for t in eos_ids_from_tokenizer(self.tokenizer)], - prefill_step_size=4096, + prefill_step_size=prefill_step_size, ) self._step_count = 0 From 5f7047ee778b93ef4e3ff3d7affd7baccadc3fb5 Mon Sep 17 00:00:00 2001 From: Bawan Wang Date: Sat, 1 Aug 2026 00:07:35 +0800 Subject: [PATCH 3/3] chore(scripts): add host-aware launcher for the two-node WSL2 test cluster MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both nodes need a handful of environment variables to start exo at all, and the two sets genuinely differ, so keeping them in one file in git is what stops the machines drifting apart — a split we have now paid for twice, once with mlx-cpu vs mlx-cuda and once with mlx-cuda12 vs mlx-cuda13. Shared: LD_PRELOAD (anaconda ships libstdc++ 3.4.29, transformers loads it first and libmlx then fails on GLIBCXX_3.4.30) and OVERRIDE_MEMORY_MB (exo reports system RAM but CUDA runs out of VRAM, so placement over-assigns layers without it). MSI only: EXO_ZENOH_CONNECT, since the Wi-Fi AP drops IPv6 link-local multicast and one side dialling the other over unicast is enough. APU-TPNB04 only: CUDA_HOME and LD_LIBRARY_PATH, because its /usr/local/cuda is 12.6 — that cuda_fp8.h has no __nv_fp8_e8m0 so nvrtc cannot JIT mlx kernels, and its lib64 only carries libcublasLt.so.12. The script refuses to run on an unknown host rather than guessing. Co-Authored-By: Claude Opus 5 --- scripts/local/start_exo.sh | 65 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100755 scripts/local/start_exo.sh diff --git a/scripts/local/start_exo.sh b/scripts/local/start_exo.sh new file mode 100755 index 0000000000..a29e90e9ad --- /dev/null +++ b/scripts/local/start_exo.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Start exo for the F002 stage-5 two-node setup (MSI 5060 + APU-TPNB04 4060, WSL2). +# +# Host-aware on purpose: the two nodes need genuinely different settings, and +# keeping one script in git is what stops them drifting apart (we lost a day to +# a cuda12/cuda13 split and another to a CPU-vs-CUDA split). +# +# LD_PRELOAD both - anaconda ships libstdc++ 3.4.29; transformers loads it +# first and libmlx then fails on GLIBCXX_3.4.30 +# OVERRIDE_MEMORY_MB both - exo reports system RAM (profiling.py), but CUDA runs +# out of VRAM; keeps placement's budget honest +# CUDA_HOME APU - its /usr/local/cuda is 12.6, whose cuda_fp8.h lacks +# __nv_fp8_e8m0, so nvrtc fails to JIT mlx kernels +# LD_LIBRARY_PATH APU - same reason: system lib64 only has libcublasLt.so.12 +# EXO_ZENOH_CONNECT MSI - Wi-Fi APs drop IPv6 link-local multicast; one side +# dialling the other over unicast is enough +# +# Optional: EXO_PREFILL_STEP_SIZE (see the local patch in batch_generate.py). +# Leave it unset for upstream behaviour (4096). Lowering it does NOT fix the +# stage-18 prefill OOM — 512 still OOMs on a 1k-token prompt. +set -euo pipefail + +SESSION="${SESSION:-exo_test}" +LOG="${LOG:-/tmp/exo_run.log}" +REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" + +env_common=( + "LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6" + "OVERRIDE_MEMORY_MB=${OVERRIDE_MEMORY_MB:-7000}" +) + +case "$(hostname)" in + MSI) + node_env=("EXO_ZENOH_CONNECT=tcp/${PEER_IP:-10.156.19.41}:52414") + ;; + APU-TPNB04) + sp="$REPO/.venv/lib/python3.13/site-packages/nvidia" + node_env=( + "CUDA_HOME=$HOME/cuda-13.0" + "LD_LIBRARY_PATH=$sp/cu13/lib:$sp/cudnn/lib:$sp/nccl/lib" + ) + ;; + *) + echo "unknown host $(hostname); add its stanza before running" >&2 + exit 1 + ;; +esac + +[ -n "${EXO_PREFILL_STEP_SIZE:-}" ] && + node_env+=("EXO_PREFILL_STEP_SIZE=$EXO_PREFILL_STEP_SIZE") + +uv_bin="$(command -v uv || echo "$HOME/.local/bin/uv")" + +tmux kill-session -t "$SESSION" 2>/dev/null || true +sleep 2 +tmux new-session -d -s "$SESSION" \ + "cd '$REPO' && env ${env_common[*]} ${node_env[*]} '$uv_bin' run exo 2>&1 | tee -a '$LOG'" + +echo "started on $(hostname): ${node_env[*]}" +for _ in $(seq 1 60); do + curl -s -m 2 http://localhost:52415/state >/dev/null 2>&1 && { echo "API up"; exit 0; } + sleep 2 +done +echo "API did not come up within 120s; tmux capture-pane -p -t $SESSION" >&2 +exit 1