#!/bin/bash
# 用 cml custom_train submit 提交 SFT 训练，训练完成后自动起 cml workflow run 评测
#
# 用法：
#   ./submit_sft_via_cml.sh <SFT_RUNDIC> <EVAL_RUNDIC> [PREV_RUNDIC]
#
# 推荐调用方式（runDic 由 resolve_run_ids.sh 决定，不要手敲）：
#   eval "$(./scripts/resolve_run_ids.sh)"
#   ./scripts/submit_sft_via_cml.sh "$SFT_RUNDIC" "$EVAL_RUNDIC"
#
# 语义（与 program.md §746 + §5.2 对齐）：
#   SFT_RUNDIC  = R{n-1}.runDic（augment / yaml / 归档全用这个，不 +1）
#   EVAL_RUNDIC = R{n-1}.runDic + 1（仅本步起 eval 时使用，仅在这一处 +1）
#   PREV_RUNDIC = SFT_RUNDIC - 1（默认；只用作上一轮 sft_output 改名）

set -euo pipefail

SFT_RUNDIC=${1:?usage: $0 <SFT_RUNDIC> <EVAL_RUNDIC> [PREV_RUNDIC]}
EVAL_RUNDIC=${2:?usage: $0 <SFT_RUNDIC> <EVAL_RUNDIC> [PREV_RUNDIC]}
PREV_RUNDIC=${3:-$((SFT_RUNDIC-1))}

# 防御：EVAL_RUNDIC 必须严格大于 SFT_RUNDIC，否则一定是调用方算错
if [ "$EVAL_RUNDIC" -le "$SFT_RUNDIC" ]; then
  echo "[cml-sft] ❌ EVAL_RUNDIC ($EVAL_RUNDIC) 必须 > SFT_RUNDIC ($SFT_RUNDIC)" >&2
  exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT=${AUTORESEARCH_ROOT:-/mnt/wangsenhao/autoresearch-zk}
MODEL_OLD=${MODEL_OLD:-/mnt/wangsenhao/verl_zk/qwen4b_cispo_wokl_add_bvt_2/global_step_5/actor/huggingface}
TPL=${SFT_TRAIN_JOB_TEMPLATE:-$SCRIPT_DIR/sft_train_job.yaml.tpl}
YAML=/tmp/sft_train_job_r${SFT_RUNDIC}.yaml

if [ -f ~/.cloudml-cli/.profile ]; then
  source ~/.cloudml-cli/.profile
else
  echo "[cml-sft] ❌ 未找到 ~/.cloudml-cli/.profile，请先安装并初始化 cml" >&2
  exit 1
fi

# 1. 渲染 yaml 模板（用 SFT_RUNDIC，不是 EVAL_RUNDIC）
sed \
  -e "s|{RUNDIC}|${SFT_RUNDIC}|g" \
  -e "s|{PREV_RUNDIC}|${PREV_RUNDIC}|g" \
  -e "s|{AUTORESEARCH_ROOT}|${ROOT}|g" \
  "$TPL" > "$YAML"
echo "[cml-sft] yaml: $YAML  (SFT_RUNDIC=$SFT_RUNDIC EVAL_RUNDIC=$EVAL_RUNDIC PREV_RUNDIC=$PREV_RUNDIC)"

# 2. 提交训练任务
SUBMIT_OUT=$(cml custom_train submit --filename "$YAML" 2>&1)
echo "$SUBMIT_OUT"
JOB_ID=$(echo "$SUBMIT_OUT" | grep -oE 't-[0-9]+-[a-z0-9]+' | head -1)
if [ -z "$JOB_ID" ]; then
  echo "[cml-sft] ❌ 提交失败" >&2
  exit 1
fi
echo "[cml-sft] ✅ JobID: $JOB_ID"

# 3. 后台 watcher：等训练成功 → 自动起评测（用 EVAL_RUNDIC）
WATCHER_LOG=/tmp/r${SFT_RUNDIC}_logs/cml_watcher.log
mkdir -p "$(dirname "$WATCHER_LOG")"

nohup bash -c '
JOB_ID='"$JOB_ID"'
SFT_RUNDIC='"$SFT_RUNDIC"'
EVAL_RUNDIC='"$EVAL_RUNDIC"'
LOG='"$WATCHER_LOG"'
ROOT='"$ROOT"'
MODEL_NEW="$ROOT/sft_output"
MODEL_OLD='"$MODEL_OLD"'
EVAL_WORKFLOW_ID=f-20260408161444-wu3pz
EVAL_VERSION=v28

source ~/.cloudml-cli/.profile

echo "[$(date)] 等待 cml job $JOB_ID 完成..." >> "$LOG"
while true; do
  STATE=$(cml custom_train describe "$JOB_ID" 2>/dev/null | grep -o "\"state\": \"[a-z]*\"" | head -1 | sed "s/.*: \"//;s/\"//")
  case "$STATE" in
    succeed)
      echo "[$(date)] 训练成功" >> "$LOG"
      break ;;
    failed|killed)
      echo "[$(date)] ❌ 训练 $STATE" >> "$LOG"
      exit 1 ;;
    *)
      sleep 60 ;;
  esac
done

# 验证产出
[ ! -f "$MODEL_NEW/_SUCCESS" ] && [ ! -f "$MODEL_NEW/config.json" ] && {
  echo "[$(date)] ❌ 产出缺失" >> "$LOG"; exit 1
}

# 起评测（这里且仅这里用 EVAL_RUNDIC）
echo "[$(date)] 启动 CML 评测 EVAL_RUNDIC=$EVAL_RUNDIC" >> "$LOG"
cml workflow run \
  --workflow_id $EVAL_WORKFLOW_ID --version $EVAL_VERSION \
  --global_inputs runDic=$EVAL_RUNDIC \
  --global_inputs model_path_new=$MODEL_NEW \
  --global_inputs model_path_old=$MODEL_OLD >> "$LOG" 2>&1
echo "[$(date)] cml workflow run 提交完成 (workflow$EVAL_RUNDIC)" >> "$LOG"
' > /tmp/r${SFT_RUNDIC}_logs/watcher_runner.log 2>&1 &

WATCHER_PID=$!
echo "[cml-sft] watcher PID: $WATCHER_PID（log: $WATCHER_LOG）"

# 4. 立即输出可查看命令
cat <<EOF

[cml-sft] 任务提交完成，关键命令：
  查看任务状态:  cml custom_train describe $JOB_ID
  查看实时日志:  cml custom_train logs $JOB_ID --follow
  停止任务:      cml custom_train kill $JOB_ID
  watcher log:   tail -f $WATCHER_LOG

训练完成后，evaluation workflow 会自动启动；评测产物在
  /mnt/xiaoai-zk-model-train-tj5/workflow5/workflow${EVAL_RUNDIC}/

JobID:        $JOB_ID
SFT_RUNDIC:   $SFT_RUNDIC   (yaml / sft_output 命名)
EVAL_RUNDIC:  $EVAL_RUNDIC  (workflow 评测目录)
EOF