#!/bin/bash
# 多卡跑一批算子的六轮 A/B：每卡一个 ab_repeat.sh，卡内串行。
#
# 它只做编排，不做测量。ab_repeat.sh 一行不动——那是已经跑出结论的、验证过的串行
# 路径，不该为了并行去改它。分片交给 partition.py（按实测耗时贪心，比 hash 分片在
# 这批活上快 50 分钟）。
#
# 并行必须隔离 tuner cache：8 个进程写同一个 sqlite 是要避开的事。代价是 pin 下来
# 的 config 与共享缓存下跑的批次不同，两批的绝对延迟不可互引——batch_summary.py
# 的口径守卫会自己拒绝这种对比，不用记着。
#
# 用法：
#   SWITCH=FLAGTREE_MUSA_RLC_ENHANCE DEVICES=0,1,2,3,4,5,6,7 \
#     SCREEN=runs/ab_rlc_XXX/screen.tsv \
#     ZL_BACKEND=musa FLAGGEMS_DIR=/path/to/FlagGems bash campaign.sh
#
#   DRY_RUN=1      只打印分片方案，不起跑
#   EXTRA_OPS=".." 追加筛选表里没有的算子（如 benchmark 刚修好的），塞给最空的卡
#   ROUNDS=6  CELL_TIMEOUT=2400  VERDICT=IR_CHANGED
set -uo pipefail

cd "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

: "${SWITCH:?SWITCH is required, e.g. FLAGTREE_MUSA_RLC_ENHANCE}"
: "${SCREEN:?SCREEN is required, 指向 ab_rlc_test.sh 的 screen.tsv}"
DEVICES="${DEVICES:-0}"
ROUNDS="${ROUNDS:-6}"
CELL_TIMEOUT="${CELL_TIMEOUT:-2400}"
VERDICT="${VERDICT:-IR_CHANGED}"
EXTRA_OPS="${EXTRA_OPS:-}"
DRY_RUN="${DRY_RUN:-0}"
TAB=$'\t'

[[ -f "$SCREEN" ]] || { echo "!!! 筛选表不存在: $SCREEN" >&2; exit 2; }

# 设备可见性变量名从 backends 层拿，不写死 MUSA_VISIBLE_DEVICES
VIS_KEY=$(python -c "import sys; sys.path.insert(0, '.')
from backends import get_backend
print(get_backend().visibility_env_keys[0])" 2>/dev/null) || VIS_KEY=""
[[ -n "$VIS_KEY" ]] || { echo "!!! 取不到后端的可见性变量名，检查 ZL_BACKEND" >&2; exit 2; }

echo ">>> switch=$SWITCH  rounds=$ROUNDS  devices=$DEVICES  可见性变量=$VIS_KEY"
python partition.py --screen "$SCREEN" --devices "$DEVICES" \
  --verdict "$VERDICT" --rounds "$ROUNDS" || exit 2

# 每卡一行 "<device><TAB><OPS><TAB><权重秒>"
mapfile -t PLAN < <(python partition.py --screen "$SCREEN" --devices "$DEVICES" \
  --verdict "$VERDICT" --format shell) || exit 2

if [[ -n "$EXTRA_OPS" ]]; then
  # 追加到当前负载最小的卡。按算子个数选是错的——一个重算子能顶三个轻的，
  # partition.py 的第三列给的就是权重秒数
  lightest=0; least=""
  for i in "${!PLAN[@]}"; do
    IFS="$TAB" read -r _d _o _w <<< "${PLAN[$i]}"
    if [[ -z "$least" ]] || (( ${_w:-0} < least )); then least=${_w:-0}; lightest=$i; fi
  done
  IFS="$TAB" read -r d o w <<< "${PLAN[$lightest]}"
  PLAN[$lightest]="${d}${TAB}${o:+$o }${EXTRA_OPS}${TAB}${w}"
  echo ">>> EXTRA_OPS 追加到卡${d}（当前最空，$(( ${w:-0} * ROUNDS / 60 )) 分钟）: $EXTRA_OPS"
fi

if [[ "$DRY_RUN" == "1" ]]; then
  echo ">>> DRY_RUN=1，不起跑。逐卡将执行："
  for row in "${PLAN[@]}"; do
    IFS="$TAB" read -r d o w <<< "$row"
    printf '    卡%s  %4s 分钟  OPS="%s"\n' "$d" "$(( ${w:-0} * ROUNDS / 60 ))" "$o"
  done
  exit 0
fi

ROOT="runs/campaign_${SWITCH}_$(date +%Y%m%d-%H%M%S)"
mkdir -p "$ROOT"
printf '%s\n' "${PLAN[@]}" > "$ROOT/plan.tsv"
echo ">>> 产物: $(pwd)/$ROOT"

declare -a PIDS=() CARDS=()
for row in "${PLAN[@]}"; do
  IFS="$TAB" read -r dev ops _w <<< "$row"
  [[ -n "$ops" ]] || { echo ">>> 卡$dev 没有分到活，跳过"; continue; }
  out="$ROOT/card${dev}"
  mkdir -p "$out"
  # AB_ROOT 让每卡有独立产物根：八个 ab_repeat.sh 同一秒启动，默认的时间戳目录名会撞
  env "$VIS_KEY=$dev" ZL_DEVICE="$dev" ZL_ISOLATE_TUNER_CACHE=1 \
      AB_ROOT="$out" SWITCH="$SWITCH" ROUNDS="$ROUNDS" \
      CELL_TIMEOUT="$CELL_TIMEOUT" OPS="$ops" NO_COLOR=1 \
      bash ab_repeat.sh > "$ROOT/card${dev}.log" 2>&1 &
  PIDS+=($!); CARDS+=("$dev")
  echo ">>> 卡$dev pid=$! ops=$ops"
done

(( ${#PIDS[@]} )) || { echo "!!! 没有一张卡分到活" >&2; exit 2; }
echo ">>> ${#PIDS[@]} 张卡已启动，等待全部结束……"

fail=0
for i in "${!PIDS[@]}"; do
  if wait "${PIDS[$i]}"; then
    echo ">>> 卡${CARDS[$i]} 结束 rc=0"
  else
    rc=$?; fail=$((fail+1))
    echo "!!! 卡${CARDS[$i]} 结束 rc=$rc —— 看 $ROOT/card${CARDS[$i]}.log"
  fi
done

# 合并逐格记录：卡号加在最前面，便于事后按卡追
{
  first=1
  for d in "${CARDS[@]}"; do
    f="$ROOT/card${d}/status.tsv"
    [[ -f "$f" ]] || continue
    if (( first )); then printf 'device%s%s\n' "$TAB" "$(head -1 "$f")"; first=0; fi
    tail -n +2 "$f" | sed "s/^/${d}${TAB}/"
  done
} > "$ROOT/status.tsv"

echo
echo "########## 汇总 ##########"
echo ">>> 失败的卡: $fail / ${#PIDS[@]}"
if [[ -s "$ROOT/status.tsv" ]]; then
  echo ">>> 逐格记录: $(pwd)/$ROOT/status.tsv  ($(( $(wc -l < "$ROOT/status.tsv") - 1 )) 格)"
  awk -F"$TAB" 'NR>1 && ($6+0)!=0 {n++} END{printf ">>> 非零退出的格: %d\n", n+0}' "$ROOT/status.tsv"
  awk -F"$TAB" 'NR>1 {fb+=$10; ng+=$11; po+=$12}
       END{printf ">>> 三个标记合计: replay_fallback=%d no_cudagraph=%d capture_poisoned=%d\n",
           fb+0, ng+0, po+0}' "$ROOT/status.tsv"
fi
echo ">>> 逐算子报告在各卡的 summary.log 里；单独重出："
echo "    python batch_summary.py --ab-rounds $ROOT/card<N>/<op> --label <op>"
echo ">>> 注意：本批 ZL_ISOLATE_TUNER_CACHE=1，绝对延迟不可与共享缓存的批次互引"