#!/bin/sh
# Distribute embedding shards over the nodes (run from M3U96b after `embed_corpus.py --prepare`).
# Each node gets a contiguous shard range proportional to its measured throughput; results are rsynced back to M3U96b.
# Usage: sh embeddings/run_multinode.sh allenai/specter2_base
cd ~/wp44_uqo; MODEL=${1:-allenai/specter2_base}; TAG=$(echo $MODEL | sed 's#/#__#g')
NS=$(ls -d data/processed/embeddings/input/shard=* | wc -l | tr -d ' ')
# weights: M3U96b 3, M3U96a 3, M4M64a 2, M4M64c 2 (M3 Ultra ≈ 300 texts/s, M4 Max ≈ 200 texts/s measured/estimated)
A=$((NS*3/10)); B=$((NS*3/10)); C=$((NS*2/10)); D=$((NS - A - B - C))
echo "shards: $NS → M3U96b 0-$A, M3U96a $A-$((A+B)), M4M64a $((A+B))-$((A+B+C)), M4M64c $((A+B+C))-$NS"
for spec in "M3U96a:$A:$((A+B))" "M4M64a:$((A+B)):$((A+B+C))" "M4M64c:$((A+B+C)):$NS"; do
  node=${spec%%:*}; rest=${spec#*:}; from=${rest%%:*}; to=${rest#*:}
  ssh $node "mkdir -p ~/wp44_uqo/data/processed/embeddings/input"
  k=$from; while [ $k -lt $to ]; do rsync -aq data/processed/embeddings/input/shard=$k $node:~/wp44_uqo/data/processed/embeddings/input/ ; k=$((k+1)); done
  ssh $node "cd ~/wp44_uqo; . .venv/bin/activate; nohup python embeddings/embed_corpus.py --root ~/wp44_uqo --run --model $MODEL --shard-from $from --shard-to $to > logs/embed_$TAG.log 2>&1 &" && echo "$node started $from-$to"
done
. .venv/bin/activate; nohup python embeddings/embed_corpus.py --root ~/wp44_uqo --run --model $MODEL --shard-from 0 --shard-to $A > logs/embed_$TAG.log 2>&1 &
echo "M3U96b started 0-$A; collect with: sh embeddings/collect_multinode.sh $MODEL"