#!/bin/bash
#SBATCH -J leanRL
#SBATCH -p dpg
#SBATCH -A dpg
#SBATCH -N 1
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=4
#SBATCH --mem=100G
#SBATCH --time=15-00:00:00
#SBATCH --constraint="gpu_40g+"
#SBATCH --output=/beegfs/scratch/user/<anonymized>/experiments/lean/logs/%A_%a.log
#SBATCH --error=/beegfs/scratch/user/<anonymized>/experiments/lean/logs/%A_%a.err
#SBATCH --array=1

source ~/miniconda3/bin/activate
scontrol show job ${SLURM_JOB_ID}
nvidia-smi
nvidia-smi topo -m

source ~/.bashrc
conda activate verl


MODELS=("DeepSeek-Prover-1.5-7B-js-128-4" "DeepSeek-Prover-1.5-7B-grpo-128-5")
MODEL=${MODELS[$SLURM_ARRAY_TASK_ID]}
N=4
# Base checkpoint path
BASE_PATH=/beegfs/scratch/user/<anonymized>/fcdpg-verl/verl/checkpoints/LeanProver_FCDPG


echo "Processing model: $MODEL"

iteration=$(cat "${BASE_PATH}/${MODEL}/latest_checkpointed_iteration.txt")
echo "Latest iteration: $iteration"

LOCAL_DIR="${BASE_PATH}/${MODEL}/global_step_${iteration}/actor"
TARGET_DIR="${BASE_PATH}/${MODEL}/global_step_${iteration}/actor/huggingface_checkpoint_${iteration}"


if [ -d "$TARGET_DIR" ]; then
    echo "Target directory $TARGET_DIR already exists — skipping merge."
else
    echo "Merging into $TARGET_DIR..."
    python -m verl.model_merger merge \
        --backend fsdp \
        --local_dir ${LOCAL_DIR} \
        --target_dir ${TARGET_DIR}
fi

echo "Running evaluation..."
python -m eval.eval \
    --model_name ${TARGET_DIR} \
    --save_name ${MODEL} \
    --num_samples ${N}