#!/bin/sh
# Distribute embedding shards over four nodes. Run from the LAPTOP (which can reach every node by alias); transfers go
# node-to-node over the LAN IPs with agent forwarding during this session; each node then runs its range under nohup.
# Usage: sh embeddings/run_multinode.sh allenai/specter2_base
MODEL=${1:-allenai/specter2_base}; TAG=$(echo $MODEL | sed 's#/#__#g'); SRC_IP=192.168.2.85
NS=$(ssh M3U96b 'ls -d ~/wp44_uqo/data/processed/embeddings/input/shard=* | wc -l' | tr -d ' ')
A=$((NS*3/10)); B=$((NS*3/10)); C=$((NS*2/10))
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 -A $node "mkdir -p ~/wp44_uqo/data/processed/embeddings/input ~/wp44_uqo/logs; k=$from; while [ \$k -lt $to ]; do rsync -aq -e 'ssh -o StrictHostKeyChecking=no' $SRC_IP:~/wp44_uqo/data/processed/embeddings/input/shard=\$k ~/wp44_uqo/data/processed/embeddings/input/; k=\$((k+1)); done; ls ~/wp44_uqo/data/processed/embeddings/input | wc -l"
  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
ssh M3U96b "cd ~/wp44_uqo; . .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"