#!/bin/bash
# Export the newest COMPLETE checkpoint, then run all evals on one GPU of a FULLY FREE node.
# (On this cluster "1 GPU anywhere" often lands on GPUs someone else is already using.)
# Run: nohup ./eval_latest.sh > logs/eval_$(date +%m%d_%H%M).log 2>&1 &
# Optional: MODEL=<existing export .pt> to skip the export step
set -uo pipefail
cd /fsx/ajay-stu/AR
source .venv/bin/activate
export PYTHONPATH=$PWD:${PYTHONPATH:-}
export HF_DATASETS_OFFLINE=1          # benchmarks are already saved; some nodes have no internet
C=/fsx/ajay-stu/ar-data/checkpoints
EXC="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"
SKIP=$( { echo $EXC | tr ',' '\n'; cat logs/bad_nodes.txt logs/bad_nodes_perm.txt 2>/dev/null; } | sort -u )

# 1) which model: the given export, or the newest complete checkpoint (exported once)
if [ -n "${MODEL:-}" ]; then
  M=$MODEL
else
  CK=$(for d in $C/ckpt_*; do if [ -f $d/.metadata ]; then echo $d; fi; done | sort -t_ -k2 -n | tail -1)
  M=/fsx/ajay-stu/ar-data/eval/model_$(basename $CK).pt
  mkdir -p /fsx/ajay-stu/ar-data/eval
  [ -f $M ] || srun -p dev -N1 -n1 --gpus=0 --cpus-per-task=16 --exclude=$EXC \
    python -m ar.eval.export_ckpt --ckpt $CK --out $M || { echo "export failed"; exit 1; }
fi
ls -lh $M || exit 1
echo "=== evaluating $(basename $M)  $(date -u) ==="

# 2) candidates: Slurm-idle nodes with all 8 GPUs empty, minus colleague and bad nodes
CANDS=$(./find_nodes.sh | grep -vxF -f <(echo "$SKIP") | head -10)
[ -z "$CANDS" ] && { echo "no fully free node right now, try again later"; exit 1; }
echo "candidate nodes: $(echo $CANDS)"

# 3) try them one by one until we get a GPU that is really free
rc=1
for NODE in $CANDS; do
  srun -p dev -N1 -n1 -w $NODE --gpus=1 --cpus-per-task=8 --immediate=120 bash -c '
    used=$(nvidia-smi -i ${CUDA_VISIBLE_DEVICES:-0} --query-gpu=memory.used --format=csv,noheader,nounits)
    if [ "$used" -ge 2000 ]; then echo "GPU already in use ($used MiB) on $(hostname), trying the next node"; exit 3; fi
    echo "=== on $(hostname), GPU ${CUDA_VISIBLE_DEVICES:-0} ==="
    python -m ar.eval.val_loss --model '"$M"' --data /fsx/ajay-stu/ar-data/heldout/stack_python.bin /fsx/ajay-stu/ar-data/heldout/fineweb.bin
    python -m ar.eval.humaneval --model '"$M"' | tail -1
    python -m ar.eval.gsm8k --model '"$M"' | tail -1
  '
  rc=$?
  [ $rc -eq 0 ] && break
done
echo "=== done (exit $rc)  $(date -u) ==="