#!/bin/bash
#SBATCH --account=reformo
#SBATCH --partition=booster
#SBATCH --nodes=1
#SBATCH --gres=gpu:4
#SBATCH --ntasks-per-node=4
#SBATCH --cpus-per-task=16
#SBATCH --time=12:00:00
#SBATCH --job-name=m2-cont-val-sweep
#SBATCH --output=/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/logs/cont_val_sweep_%j.out
#SBATCH --error=/e/scratch/reformo/schuhmann1_moss/out/m2_600m_ladder/logs/cont_val_sweep_%j.err
set -euo pipefail
SC=/e/scratch/reformo/schuhmann1_moss
CODE=/e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/continuous_large_talker_v1
SWEEP=$(realpath "$1")
source "$SC/code/fastgen.sh"
export HF_HOME="$SC/hfcache" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1
export PYTHONPATH="$CODE:/e/home/jusers/schuhmann1/jupiter/m2_600m_tts_training/cascade_v3_timed:$SC/code/small_tts:$SC/code/vc_env_pkgs:$PYTHONPATH"
export OMP_NUM_THREADS=4 PYTHONUNBUFFERED=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
PREV_STAGE=""
while IFS=$'\t' read -r STAGE CHECKPOINT; do
  [[ -n "$CHECKPOINT" ]] || continue
  if [[ -n "$PREV_STAGE" && "$STAGE" != "$PREV_STAGE" ]]; then
    python3 -u "$CODE/select_validation_checkpoint.py" --eval-root \
      "$SC/out/m2_600m_ladder/caption_curriculum_v1/validation_loss" --world 4
  fi
  RESULT="$SC/out/m2_600m_ladder/caption_curriculum_v1/validation_loss/$(basename "$CHECKPOINT")/full/SUMMARY.json"
  if [[ -f "$RESULT" ]] && python3 - "$RESULT" <<'PY'
import json, sys
raise SystemExit(0 if json.load(open(sys.argv[1])).get('status') == 'PASS' else 1)
PY
  then
    echo "VALIDATION_SWEEP_ALREADY_PASS $STAGE $CHECKPOINT"
    PREV_STAGE="$STAGE"
    continue
  fi
  echo "VALIDATION_SWEEP_BEGIN $CHECKPOINT"
  srun --ntasks=4 --ntasks-per-node=4 --gpus-per-task=1 --kill-on-bad-exit=1 --cpu-bind=none \
    bash -c 'cd "$1"; python3 -u eval_validation_loss.py --checkpoint "$2"' _ "$CODE" "$CHECKPOINT" </dev/null
  PREV_STAGE="$STAGE"
done < <(python3 - "$SWEEP" <<'PY'
import json, sys
for item in json.load(open(sys.argv[1]))['candidates']:
    print(item['stage'] + '\t' + item['checkpoint'])
PY
)
python3 -u "$CODE/select_validation_checkpoint.py" --eval-root \
  "$SC/out/m2_600m_ladder/caption_curriculum_v1/validation_loss" --world 4