#!/bin/bash
# One-time switch to the new loop job + Stage 1 extension to 200B tokens.
cd /fsx/ajay-stu/AR
source .venv/bin/activate

# 1) wait for the 4 tokenizing jobs, then check the files
while [ -n "$(squeue -u $USER -h -n ar-tok)" ]; do sleep 300; done
for g in 003 004 005 006; do
  [ -s /fsx/ajay-stu/ar-data/bin/fw$g/fineweb.bin ] || { echo "$(date) fw$g missing, aborting"; exit 1; }
done
echo "$(date) new data ready"

# 2) new settings: 200B tokens (end step 381,000) + new data mix
python - << 'PY'
import glob, os
d = "/fsx/ajay-stu/ar-data/bin/"
eng = [d + "fineweb.bin"] + sorted(glob.glob(d + "fw*/fineweb.bin"))
code = [f"{d}stack_{l}.bin" for l in
        ["python", "javascript", "java", "cpp", "typescript", "go", "rust", "shell", "sql"]]
math = [d + "finemath.bin", d + "finemath3.bin"]
def split(files, share):
    s = [os.path.getsize(f) for f in files]
    return [round(share * x / sum(s), 4) for x in s]
files = eng + code + math
w = split(eng, 0.65) + split(code, 0.20) + [0.09, 0.06]
out = []
for l in open("run_args.sh").read().splitlines():
    if l.startswith("BIN="): l = 'BIN="' + " ".join(files) + '"'
    elif l.startswith("WEIGHTS="): l = 'WEIGHTS="' + " ".join(map(str, w)) + '"'
    elif l.startswith("STOP_AT="): l = "STOP_AT=381000"
    elif l.startswith("LR_HORIZON="): l = "LR_HORIZON=381000"
    out.append(l)
open("run_args.sh", "w").write("\n".join(out) + "\n")
print("settings updated:", len(files), "files, stop at 381000")
PY
sed -i 's/ckpt_185999/ckpt_380999/' watchdog.sh     # watchdog's new "finished" marker

# 3) pause the watchdog so it can't launch a duplicate during the switch
pkill -f watchdog.sh
J=$(squeue -u $USER -h -n ar-real -o %i | head -1)
if [ -z "$J" ]; then
  echo "$(date) training not running; watchdog will start the new job"
  nohup ./watchdog.sh >> logs/watchdog.log 2>&1 &
  exit 0
fi
NODES=$(squeue -j $J -h -o %N)

# 4) after a checkpoint, switch only if no waiting job could take our nodes (max ~6 tries)
tries=0
while true; do
  N=$(grep -c "\[ckpt\] saved" logs/real_$J.log)
  while [ "$(grep -c "\[ckpt\] saved" logs/real_$J.log)" -eq "$N" ]; do
    [ -z "$(squeue -j $J -h 2>/dev/null)" ] && break 2    # old job died meanwhile
    sleep 30
  done
  tries=$((tries + 1))
  risky=$(squeue -p dev -t PD -h -o "%u %D" | awk -v me=$USER '$1 != me && $2 <= 4' | wc -l)
  if [ "$risky" -eq 0 ] || [ $tries -ge 6 ]; then break; fi
  echo "$(date) $risky waiting jobs could take our nodes; trying at the next checkpoint"
done
sleep 30
scancel $J 2>/dev/null
sleep 20
sbatch --nodelist=$NODES train_real.sbatch
echo "$(date) switched: old job $J stopped, new loop job submitted on $NODES"
nohup ./watchdog.sh >> logs/watchdog.log 2>&1 &