#!/bin/bash
#SBATCH --job-name=ar-real
#SBATCH --partition=dev
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=96
#SBATCH --exclusive
#SBATCH --exclude=ip-10-0-136-14,ip-10-0-136-94,ip-10-0-136-253,ip-10-0-138-178,ip-10-0-139-113,ip-10-0-140-247
#SBATCH --time=30-00:00:00
#SBATCH --signal=B:USR1@4200
#SBATCH --output=logs/real_%j.log
#SBATCH --error=logs/real_%j.err

# Training runs in a loop inside this one job, so the nodes are kept:
#   logs/RELOAD -> stop after the next checkpoint, re-read run_args.sh, continue
#   logs/STOP   -> stop after the next checkpoint and end the job cleanly
#   ~70 min before a time limit Slurm sends USR1 -> we request STOP (nothing lost)
#   crash -> retry up to 3 times here before giving the nodes back
set -uo pipefail
cd /fsx/ajay-stu/AR
source .venv/bin/activate
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n1)
export FI_PROVIDER=efa
export NCCL_PROTO=simple
export NCCL_SOCKET_IFNAME=enp73s0
export NCCL_IB_DISABLE=0
export FI_EFA_USE_DEVICE_RDMA=1
export LD_LIBRARY_PATH=/fsx/ajay-stu/AR/.venv/lib/python3.10/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH:-}
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TOKENIZERS_PARALLELISM=false
# keep ~524K tokens per step for any node count
case $SLURM_NNODES in
  1) BS=4; GA=4 ;;
  2) BS=4; GA=2 ;;
  3|4) BS=4; GA=1 ;;
  5|6|7|8) BS=2; GA=1 ;;     # 8 nodes: 64 GPUs x 2 seqs x 4096 = 524K tokens
  *) BS=1; GA=1 ;;
esac
#       # 4 nodes -> 1 micro-batch, 2 nodes -> 2 (same step size)

# handover from a smaller job: wait (max 30 min) until our other training job has ended
for i in $(seq 90); do
  squeue -u $USER -h -n ar-real -t RUNNING -o %i | grep -vxq "$SLURM_JOB_ID" || break
  sleep 20
done
if squeue -u $USER -h -n ar-real -t RUNNING -o %i | grep -vxq "$SLURM_JOB_ID"; then
  echo "=== other training job still running, giving nodes back ==="; exit 4
fi

rm -f logs/STOP logs/RELOAD    # clear leftovers from an earlier job
trap 'echo "=== time limit near: stopping after the next checkpoint ==="; touch logs/STOP logs/RELOAD' USR1

round=0; fails=0
while true; do
  round=$((round + 1))
  source run_args.sh                                   # latest settings every round
  ./preflight.sh || { echo "PREFLIGHT FAILED: busy GPUs"; exit 3; }
  export MASTER_PORT=$((29500 + (SLURM_JOB_ID + round) % 1000))
  echo "=== round $round: nodes=$SLURM_NNODES stop_at=$STOP_AT lr_horizon=$LR_HORIZON ==="
  echo "=== free space: $(df -h /fsx | tail -1) ==="

  # run in the background so the time-limit warning can be handled mid-training
  srun bash -c "torchrun --nnodes=\$SLURM_NNODES --nproc_per_node=8 --node_rank=\$SLURM_NODEID \
    --rdzv_id=${SLURM_JOB_ID}_$round --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    -m ar.train.train_fsdp --bin $BIN --weights $WEIGHTS --model_size full \
    --batch_size $BS --seq_len 4096 --max_steps $STOP_AT --lr_horizon $LR_HORIZON \
    --warmup_steps 2000 --peak_lr $PEAK_LR --grad_accum $GA \
    --ckpt_every 2000 --ckpt_dir /fsx/ajay-stu/ar-data/checkpoints" &
  pid=$!
  while true; do
    wait $pid; rc=$?
    kill -0 $pid 2>/dev/null || break                  # training ended; rc is its exit code
  done

  if [ -f logs/STOP ]; then
    rm -f logs/STOP logs/RELOAD
    echo "=== stopped cleanly after a checkpoint ==="; exit 0
  fi
  if [ -f logs/RELOAD ]; then
    rm -f logs/RELOAD; fails=0
    echo "=== reload requested: restarting with new settings, same nodes ==="
    continue
  fi
  [ $rc -eq 0 ] && { echo "=== training finished ==="; break; }
  fails=$((fails + 1))
  echo "=== training crashed (rc=$rc), retry $fails/3 on the same nodes ==="
  [ $fails -ge 3 ] && exit $rc
  sleep 60
done