#!/bin/bash

run_type="$1"
model="$2"
gpu_cnt="$3"
max_steps="$4"

current_datetime=$(date +%Y%m%d%H%M%S)

if [ "${run_type}"="lora_sft" ]; then
    run_name="${run_type}_${model}_${gpu_cnt}_gpu_${max_steps}_step_${current_datetime}"
    
else
    run_name="${run_type}_${model}_${gpu_cnt}_gpu_${current_datetime}"
fi

output_dir="./results/${run_name}"

if [ ! -d "$output_dir" ]; then
    mkdir -p "$output_dir"
    echo "output_dir created: $output_dir"
else
    echo "output_dir exists: $output_dir"
fi

# echo "${run_type} ${model} ${gpu_cnt} ${max_steps} ${run_name} ${output_dir}"
python prepare_yaml_file.py ${run_type} ${model} ${max_steps} ${run_name} ${output_dir}

export USE_MODELSCOPE_HUB=1

# echo "Start recording gpu status "
# # 0 means not printing gpu status
# python gpu_status.py ${output_dir} 1 10 &
# gpu_status_pid=$!
# echo "${gpu_status_pid}"

if [ "${gpu_cnt}"="1" ]; then
   ASCEND_RT_VISIBLE_DEVICES=0 llamafactory-cli train ${output_dir}/${run_name}.yml | tee "${output_dir}/log.txt" &
   train_pid=$!
   echo "Start train"
else
   FORCE_TORCHRUN=1 llamafactory-cli train ${output_dir}/${run_name}.yml | tee "${output_dir}/log.txt" &
   train_pid=$!
   echo "Start train"
fi

wait $train_pid
echo "Train ended"

# sleep 90
# kill $gpu_status_pid
# echo "Gpu status ended"