#!/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 throughput_test_qwen_7B_ar                                                 #job name
#SBATCH --gpus-per-node 8
#SBATCH -e workdir/throughput_error.log
#SBATCH -o workdir/throughput_out.log

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

# 模型路径 - 7B模型（非AR）
model_path=output_models/finetune_qwen2.5-7B-bd-size-32-lr-1e-5/checkpoint-2500

# 任务配置 - 使用gsm8k进行throughput测试
task=gsm8k

# Throughput测试结果保存路径
throughput_log_path=eval_results/qwen2_5_7B_throughput_test.log
throughput_results_dir=eval_results/throughput_results_7B

# 创建必要的目录
mkdir -p workdir
mkdir -p workdir/slurm_log/throughput_test_7B
mkdir -p eval_results
mkdir -p ${throughput_results_dir}

echo "开始执行7B模型Throughput测试..."
echo "任务: ${task}"
echo "模型路径: ${model_path}"
echo "结果保存路径: ${throughput_results_dir}"

# 定义要测试的batch sizes
batch_sizes=(1 2 4 8 16 32 64)

# 计数器
experiment_count=0
total_experiments=${#batch_sizes[@]}

echo "总共需要运行 ${total_experiments} 个throughput测试"
echo "开始时间: $(date)" | tee -a ${throughput_log_path}

# 遍历所有batch size进行throughput测试
for batch_size in "${batch_sizes[@]}"; do
    experiment_count=$((experiment_count + 1))
    
    echo ""
    echo "=========================================="
    echo "Throughput测试 ${experiment_count}/${total_experiments}"
    echo "任务: ${task}"
    echo "批次大小: ${batch_size}"
    echo "测试时间: $(date)"
    echo "=========================================="

    # 记录测试信息到日志
    echo "=========================================" >> ${throughput_log_path}
    echo "Throughput测试 ${experiment_count}/${total_experiments} - Batch Size: ${batch_size}" >> ${throughput_log_path}
    echo "开始时间: $(date)" >> ${throughput_log_path}
    echo "=========================================" >> ${throughput_log_path}

    # 构建throughput测试命令
    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=False,small_block_size=8,dynamic_threshold=False,factor=0.9,threshold=0.9 \
        --output_path ${throughput_results_dir}/batch_${batch_size} \
        --log_samples
    "

    echo "执行throughput测试命令..."
    echo "命令: $eval_cmd"

    # 记录开始时间
    start_time=$(date +%s)
    
    # 使用srun执行throughput测试命令
    srun -o workdir/slurm_log/throughput_test_7B/slurm_%A_%j_${task}_batch_${batch_size}.log \
         -e workdir/slurm_log/throughput_test_7B/error_%A_%j_${task}_batch_${batch_size}.log \
        bash -c "${eval_cmd}" >> ${throughput_log_path} 2>&1

    # 记录结束时间
    end_time=$(date +%s)
    duration=$((end_time - start_time))

    echo "Throughput测试 ${experiment_count} 完成！"
    echo "批次大小: ${batch_size}, 耗时: ${duration}秒" | tee -a ${throughput_log_path}
    echo "结束时间: $(date)" >> ${throughput_log_path}
    echo "" >> ${throughput_log_path}
done

echo ""
echo "=========================================="
echo "所有Throughput测试完成！"
echo "总共运行了 ${total_experiments} 个测试"
echo "结束时间: $(date)"
echo "结果保存在: ${throughput_results_dir}"
echo "日志文件: ${throughput_log_path}"
echo "=========================================="

# 在日志中记录完成信息
echo "=========================================" >> ${throughput_log_path}
echo "所有Throughput测试完成！总共运行了 ${total_experiments} 个测试" >> ${throughput_log_path}
echo "完成时间: $(date)" >> ${throughput_log_path}
echo "=========================================" >> ${throughput_log_path}

# 生成throughput总结报告
echo "正在生成Throughput总结报告..."
summary_file=${throughput_results_dir}/throughput_summary.txt

echo "7B模型Throughput测试总结报告" > ${summary_file}
echo "测试时间: $(date)" >> ${summary_file}
echo "模型路径: ${model_path}" >> ${summary_file}
echo "任务: ${task}" >> ${summary_file}
echo "=========================================" >> ${summary_file}

for batch_size in "${batch_sizes[@]}"; do
    echo "Batch Size: ${batch_size}" >> ${summary_file}
    # 这里可以添加更多的性能指标提取逻辑
    echo "详细结果请查看: ${throughput_results_dir}/batch_${batch_size}" >> ${summary_file}
    echo "---" >> ${summary_file}
done

echo "Throughput总结报告已生成: ${summary_file}"
