#!/bin/zsh
# Run the SubMConv3d benchmark on every Apple Silicon box in the fleet.
# Ships code only (no weights); each box gets a small mlx+numpy venv under ~/lato-bench.
set -u
SRC=/Users/m3ultra/Documents/lato.2_mrp_mlx
OUT=$SRC/bench
SSH_OPTS=(-o ConnectTimeout=10 -o BatchMode=yes -o StrictHostKeyChecking=accept-new)
SSH_STR="ssh -o ConnectTimeout=10 -o BatchMode=yes -o StrictHostKeyChecking=accept-new"

# host:label   (m3ultra runs locally, already done)
HOSTS=(
  "m1max@100.92.78.24:m1max"
  "m2max@100.120.83.110:m2max"
  "m4pro@100.69.21.128:m4pro"
  "johnking@100.91.239.7:m1ultra"
)

hb(){ mkdir -p ~/.jobs; echo "$(date '+%F %T') | $1 | $2" > ~/.jobs/lato-fleet-bench.status; }
hb "starting" "${#HOSTS[@]} hosts"

for entry in $HOSTS; do
  host=${entry%%:*}; label=${entry##*:}
  echo "===== $label ($host)"
  hb "$label" "connecting"

  if ! ssh "${SSH_OPTS[@]}" "$host" 'true' 2>/dev/null; then
    echo "  SKIP unreachable"; continue
  fi

  # m1max is disk-critical (~20GB free per fleet notes) - refuse rather than fill it
  free_mb=$(ssh "${SSH_OPTS[@]}" "$host" "df -m / | tail -1 | awk '{print \$4}'" 2>/dev/null)
  if [ -n "$free_mb" ] && [ "$free_mb" -lt 3000 ]; then
    echo "  SKIP only ${free_mb}MB free - not installing a venv here"; continue
  fi

  ssh "${SSH_OPTS[@]}" "$host" 'mkdir -p ~/lato-bench' 2>/dev/null
  rsync -a -e "$SSH_STR" --delete \
    "$SRC/lato_mlx" "$SRC/bench" "$host:~/lato-bench/" 2>&1 | tail -1

  hb "$label" "installing venv"
  ssh "${SSH_OPTS[@]}" "$host" '
    cd ~/lato-bench || exit 1
    UV=$(command -v uv || echo /opt/homebrew/bin/uv)
    if [ ! -x "$UV" ]; then echo "  no uv on this host"; exit 3; fi
    [ -d .venv ] || "$UV" venv --python 3.12 .venv >/dev/null 2>&1
    VIRTUAL_ENV=.venv "$UV" pip install -q mlx numpy >/dev/null 2>&1
    .venv/bin/python -c "import mlx.core" 2>/dev/null || { echo "  mlx install failed"; exit 4; }
  ' || { echo "  SKIP setup failed (rc=$?)"; continue; }

  hb "$label" "benchmarking"
  ssh "${SSH_OPTS[@]}" "$host" \
    'cd ~/lato-bench && .venv/bin/python bench/bench_subm.py --out bench/result.json' 2>&1 | tail -8
  rsync -a -e "$SSH_STR" "$host:~/lato-bench/bench/result.json" "$OUT/$label.json" 2>/dev/null \
    && echo "  -> $OUT/$label.json" || echo "  (no result pulled)"
done

hb "DONE" "results in $OUT"
echo "===== done"