feat(networking): allow dialing peers over unicast via EXO_ZENOH_CONNECT

rust/networking/src/lib.rs

@@ -32,6 +32,20 @@ pub fn cfg(identity: &str, listen_port: u16) -> Result<zenoh::Config> {
     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<String> = 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")?;

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

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