#!/bin/bash
# Launch the 103M Stage B pilots once Stage A has completed on M3U96a. Usage: scripts/launch_103m.sh [steps]
set -e
cd "$(dirname "$0")/.."
RUN=$(ssh M3U96a 'ls -d ~/wp31/prisme-slm/experiments/runs/pretrain103m_seed01_*' | tail -1)
STATE=$(ssh M3U96a "grep -o '\"state\": \"[a-z_]*\"' $RUN/manifest.json")
echo "stage A run: $RUN ($STATE)"
[[ "$STATE" == *completed* ]] || { echo "Stage A not completed"; exit 1; }
REL=${RUN#*/prisme-slm/}
for c in editor_pilot_103m dense_pilot_103m gr_pilot_103m; do
  sed -i '' "s|experiments/runs/PRETRAIN103M_RUN/checkpoint|$REL/checkpoint|" configs/experiments/$c.yaml
done
# copy the backbone checkpoint to the other pilot nodes
for n in M3U96b M2U64; do
  ssh M3U96a "cd ~/wp31/prisme-slm && rsync -az --relative ./$REL/checkpoint $n:~/wp31/prisme-slm/" || rsync -az "M3U96a:$RUN/checkpoint" /tmp/ck103m/ && rsync -az --relative "/tmp/ck103m/." "$n:~/wp31/prisme-slm/$REL/"
done
for n in M3U96a M3U96b M2U64; do rsync -az --exclude .venv --exclude /data --exclude /experiments --exclude .git ./ $n:~/wp31/prisme-slm/; done
ssh M3U96a "cd ~/wp31/prisme-slm && nohup .venv/bin/python -m training.train_editor --config configs/experiments/editor_pilot_103m.yaml --device mps > experiments/logs/editor_pilot_103m_seed1.log 2>&1 < /dev/null & disown; echo launched editor"
ssh M3U96b "cd ~/wp31/prisme-slm && nohup .venv/bin/python -m training.train_editor --config configs/experiments/gr_pilot_103m.yaml --device mps > experiments/logs/gr_pilot_103m_seed1.log 2>&1 < /dev/null & disown; echo launched gr"
ssh M2U64 "cd ~/wp31/prisme-slm && nohup .venv/bin/python -m training.train_editor --config configs/experiments/dense_pilot_103m.yaml --device mps > experiments/logs/dense_pilot_103m_seed1.log 2>&1 < /dev/null & disown; echo launched dense"