#!/bin/bash
# Run all evals on THIS node, on one GPU nobody is using (for an interactive session on a compute node).
# Usage: ./eval_here.sh [model.pt]    (default: export the newest complete checkpoint first)
set -uo pipefail
cd /fsx/ajay-stu/AR
source .venv/bin/activate
export PYTHONPATH=$PWD HF_DATASETS_OFFLINE=1
H=$(hostname)

# 1) never on our own training nodes (their GPUs are fully used by training)
TRAIN=$(scontrol show hostnames $(squeue -u $USER -n ar-real -h -o %N | head -1) 2>/dev/null)
echo "$TRAIN" | grep -qx "$H" && { echo "STOP: $H is one of our training nodes"; exit 1; }

# 2) one free GPU (inside a Slurm session, use the GPU Slurm gave us)
if [ -z "${CUDA_VISIBLE_DEVICES:-}" ]; then
  G=$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | awk -F', ' '$2 < 2000 {print $1; exit}')
  [ -z "$G" ] && { echo "no free GPU on $H"; exit 1; }
  export CUDA_VISIBLE_DEVICES=$G
fi
used=$(nvidia-smi -i ${CUDA_VISIBLE_DEVICES%%,*} --query-gpu=memory.used --format=csv,noheader,nounits)
[ "$used" -ge 2000 ] && { echo "GPU $CUDA_VISIBLE_DEVICES already in use ($used MiB)"; exit 1; }
echo "=== on $H, GPU $CUDA_VISIBLE_DEVICES ==="

# 3) model: the given file, or export the newest complete checkpoint once (~5 min, CPU)
C=/fsx/ajay-stu/ar-data/checkpoints
if [ -n "${1:-}" ]; then
  M=$1
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
  [ -f $M ] || python -m ar.eval.export_ckpt --ckpt $CK --out $M || { echo "export failed"; exit 1; }
fi
echo "=== evaluating $(basename $M)  $(date -u) ==="

# 4) all evals; keep the score lines
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
echo "=== done  $(date -u) ==="