#!/bin/bash
# launch_r3.sh — RB4 runda 3: fork z ckpt 320k, config 1:1 jak control, rozni sie TYLKO danymi.
# Czeka az obecny trening na podzie sie skonczy (WAITPID) i az dane beda kompletne (manifest),
# weryfikuje krok bazy (ck["step"] == 320000), dopiero wtedy startuje trener + autopush + pusher.
#
# Uzycie: bash launch_r3.sh WORK DATA OUT BASE HFSUB NAME WAITPID
set -euo pipefail
WORK="$1"; DATA="$2"; OUT="$3"; BASE="$4"; HFSUB="$5"; NAME="$6"; WAITPID="$7"
cd "$WORK"
echo "[$(date -u +%H:%M:%SZ)] launch_r3: czekam na pid $WAITPID i $DATA/train.bin.json"
while kill -0 "$WAITPID" 2>/dev/null; do sleep 20; done
while [ ! -f "$DATA/train.bin.json" ]; do sleep 20; done
[ -e "$DATA/val.bin" ] || { echo "!! brak $DATA/val.bin"; exit 1; }

STEP=$(python3 -c "import torch,sys; print(int(torch.load(sys.argv[1], map_location='cpu')['step']))" "$BASE")
[ "$STEP" = "320000" ] || { echo "!! baza $BASE ma step $STEP, oczekiwano 320000"; exit 1; }
mkdir -p "$OUT"
[ -e "$OUT/ckpt.pt" ] && { echo "!! $OUT/ckpt.pt juz istnieje, nie nadpisuje"; exit 1; }
cp "$BASE" "$OUT/ckpt.pt"
echo "[$(date -u +%H:%M:%SZ)] baza OK (step $STEP), start treningu $NAME"

nohup python -u train_gpt_ref.py --data-dir "$DATA" --out-dir "$OUT" \
  --n-layer 14 --n-embd 576 --n-head 9 --block 1024 --vocab 12288 --norm rmsnorm --pos rope \
  --rope-theta 100000.0 --ffn swiglu --ffn-mult 2.667 --value-residual --qk-norm \
  --batch 32 --steps 400000 --lr 6e-4 --min-lr 6e-5 --warmup 2000 --wd 0.1 \
  --optimizer muon --muon-lr 0.02 --dtype uint16 --seed 1337 \
  --ckpt-every 20000 --eval-every 999999 --log-every 200 \
  --events-jsonl "$OUT/events.jsonl" --compile --resume > "$OUT/train.log" 2>&1 &
echo "TRAIN-PID=$!"
sleep 5
nohup python -u ckpt_autopush.py "$OUT/train.log" "$OUT/ckpt.pt" "$HFSUB" > "$OUT/autopush.log" 2>&1 &
echo "AUTOPUSH-PID=$!"
set -a; . /root/.fabryka.env; set +a
nohup python -u fabryka_push.py "$OUT/events.jsonl" "$NAME" gollem-v5 13107200000 > "$OUT/fabryka_push.log" 2>&1 &
echo "PUSHER-PID=$!"
echo "[$(date -u +%H:%M:%SZ)] launch_r3: wszystko uruchomione"