#!/bin/bash
#SBATCH -A nvr_elm_llm                                                                #account
#SBATCH -p batch_block1,batch_block3,batch_block4                                                         #partition
#SBATCH -t 04:00:00                                                                   #wall time limit, hr:min:sec
#SBATCH -N 1                                                                           #number of nodes
#SBATCH -J batch_eval_4_configs                                                       #job name
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/error.log
#SBATCH -o workdir/out.log
#SBATCH --exclusive

# Set the environment variables first before running the command.
export HF_ALLOW_CODE_EVAL=1
export HF_DATASETS_TRUST_REMOTE_CODE=true

# 任务配置
task=gsm8k
# task=mbpp
# task=humaneval
# task=minerva_math

# 模型路径
model_path=models/Fast_dLLM_1.5B

# 评估参数
batch_size=1
block_cache_refresh_interval=16
small_block_size=16
dynamic_threshold=False
factor=1

# 创建必要的目录
mkdir -p workdir
mkdir -p workdir/slurm_log/eval_throughput

echo "开始执行4种配置的batch evaluation..."
echo "任务: ${task}"
echo "模型路径: ${model_path}"
echo "批次大小: ${batch_size}"

# 定义4种配置
configs=(
    "True 0.9"   # 使用block cache, 阈值0.9
    "True 1.0"   # 使用block cache, 阈值1.0
    "False 0.9"  # 不使用block cache, 阈值0.9
    "False 1.0"  # 不使用block cache, 阈值1.0
)

# 循环执行4种配置
for i in "${!configs[@]}"; do
    IFS=' ' read -r use_block_cache threshold <<< "${configs[$i]}"
    
    echo ""
    echo "=========================================="
    echo "执行配置 $((i+1))/4:"
    echo "使用块缓存: ${use_block_cache}"
    echo "阈值: ${threshold}"
    echo "=========================================="
    
    # 构建评估命令
    eval_cmd='
    export HF_ALLOW_CODE_EVAL=1;
    export HF_DATASETS_TRUST_REMOTE_CODE=true;
    accelerate launch chengyue_script/eval_mdm_qwen_batch.py \
        --tasks '"${task}"' \
        --batch_size '"${batch_size}"' \
        --confirm_run_unsafe_code \
        --model dfm_qwen \
        --fewshot_as_multiturn \
        --apply_chat_template \
        --model_args model_path='"${model_path}"',show_speed=True,use_block_cache='"${use_block_cache}"',block_cache_refresh_interval='"${block_cache_refresh_interval}"',small_block_size='"${small_block_size}"',dynamic_threshold='"${dynamic_threshold}"',factor='"${factor}"',threshold='"${threshold}"'
    '
    
    echo "执行评估命令..."
    echo $eval_cmd
    
    # 使用srun执行评估命令，日志路径包含关键参数信息
    srun -o workdir/slurm_log/eval_throughput/slurm_%A_%j_${task}_batch${batch_size}_config$((i+1))_block_cache_${use_block_cache}_thresh${threshold}.log -e workdir/slurm_log/eval_throughput/error_%A_%j_${task}_batch${batch_size}_config$((i+1))_block_cache_${use_block_cache}_thresh${threshold}.log \
        bash -c "${eval_cmd}"
    
    echo "配置 $((i+1)) 评估完成！"
    
    # 在配置之间稍作等待
    sleep 10
done

echo ""
echo "=========================================="
echo "所有4种配置的评估完成！"
echo "=========================================="
echo "配置总结："
echo "1. 使用block cache, 阈值0.9"
echo "2. 使用block cache, 阈值1.0"
echo "3. 不使用block cache, 阈值0.9"
echo "4. 不使用block cache, 阈值1.0"
echo "=========================================="
