#!/bin/bash
# Utility script to run a training command and properly save logs.
# 导入函数方法：source run_cmd.sh

run_cmd()
{
  local cmd="$1"
  local log_dir="$2"
  local log_file="$3"
  local node_rank="$4"
  local node_num="$5"
  local cmd_exit_code=0

  if [[ ${node_rank} -eq 0 || ${node_num} -eq 1 ]]; then
    # mkdir
    if [[ ! -d ${log_dir} ]]; then
      echo "Creating log dir: ${log_dir}"
      mkdir -p "${log_dir}"
    fi
    # train
    echo "${cmd}" 2>&1 | tee "${log_file}"
    echo "[$(date "+%Y-%m-%d %H:%M:%S")] - Begin of Training Script." | tee -a "${log_file}"
    eval "${cmd}" 2>&1 | tee -a "${log_file}"
    cmd_exit_code=${PIPESTATUS[0]}
    echo "[$(date "+%Y-%m-%d %H:%M:%S")] - End of Training Script." | tee -a "${log_file}"
  else
    TIMESTAMP=$(date +%s)
    TMP_LOG_FILE=$(mktemp -p /tmp "${LOG_FILE_NAME}_${TIMESTAMP}_XXX.log.tmp")
    echo "${cmd}" 2>&1 | tee "${TMP_LOG_FILE}"
    echo "[$(date "+%Y-%m-%d %H:%M:%S")] - Begin of Training Script." | tee -a "${TMP_LOG_FILE}"
    eval "${cmd}" 2>&1 | tee -a "${TMP_LOG_FILE}"
    cmd_exit_code=${PIPESTATUS[0]}
    echo "[$(date "+%Y-%m-%d %H:%M:%S")] - End of Training Script." | tee -a "${TMP_LOG_FILE}"
    mv "${TMP_LOG_FILE}" "${log_file}"
    if [[ $? -ne 0 ]]; then
      echo "⚠️ 节点${node_rank}的日志保存于 ${TMP_LOG_FILE}。由于$(dirname ${log_file}) 目录不是共享存储，无法移动至此目录。"
      echo "⚠️ 请手动将日志文件移动到共享存储的日志目录中以便后续查看。"
    fi
  fi
  return ${cmd_exit_code}
}
