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")?; 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 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