#!/bin/bash
# lab_train.sh — π0.5 の続きからの学習を Runpod で回し, 終わったらモデルと記録を持ち帰って Pod を消す (依頼 49)
#
#   ./lab_train.sh --log Log17             学習 → Hugging Face → KIOXIA → Pod を消す → Mac に通知
#   ./lab_train.sh --log Log17 --dry-run   Pod を借りずに流れだけ確かめる (API・SSH・hf・通知を差し替え,
#                                          書くのは一時フォルダだけ. 差し替えは lab_train_fakes.sh)
#   ./lab_train.sh --log Log17 --stop      この Log の学習の Pod のログを持ち帰って消す (Mac が途中で止まったときなど)
#
#   --gpu          最初に試す GPU (既定 A40). 借りられなければ RTX A6000 → L40S (どれも 48GB) を順に試す
#   --no-fallback  ほかの GPU を試さない (借りられなければ通知して終わる)
#   --limit-hours  上限の時間 (既定 5. Pod を作った時刻から数える). 超えたらログを持ち帰って消す
#   --model        作るモデル (依頼 57. 名前は lab_start.sh の --model と同じ. 既定 corr):
#                    corr  = 3999 から trim2_corr_v3 で続きを学習 → PixelstoMotors/pi05-mecharm-trim-corr-lora-v1
#                    corr2 = corr の 1499 から trim2_corr_v4b で続きを学習 → PixelstoMotors/pi05-mecharm-trim-corr-lora-v2
#                            (設定 pi05_mecharm_trim_corr_v4_lora. 3c で取るのは 3999 ではなく corr の 1499)
#                    corr3 = corr2 の 1499 から trim2_corr_v5 で続きを学習 → PixelstoMotors/pi05-mecharm-trim-corr-lora-v3
#                            (依頼 97. 設定 pi05_mecharm_trim_corr_v5_lora. 3c で取るのは corr2 の 1499)
#   --steps / --init / --exp-name   学習 (既定 1500 / 続きの元の params / trim_corr_v3・trim_corr_v4・trim_corr_v5)
#                    続きの元: corr は /workspace/ckpt/3999/params, corr2 は /workspace/ckpt/corr_v1_1499/params,
#                              corr3 は /workspace/ckpt/corr_v2_1499/params
#   --poll-min     進みを見る間 (既定 5 分)
#   --stall-min    学習のログが何分書かれなければ「止まった」とするか (既定 30 分)
#
# 1. Pod を作る (名前 p2m-train-<Log>. lab_start.sh の p2m-<Log> とは別). 記録は lab/<Log>/train/lab_pod.json
# 2. SSH が通るまで待ち, Pod 側に上限の見張りを置く (上限 + 60 分で, Runpod が Pod に入れる Pod 用の鍵で自分を消す)
# 3. 準備 (lab_start.sh 3a〜3f と同じ. サーバーは起動しない)
# 4. tar を /workspace/data/ へ送り, pod_prep_trim_corr.sh <steps> <init>
# 5. 学習を setsid nohup で始める (Mac とのつながりが切れても止まらない). ログ /workspace/logs/train_<exp>.log,
#    GPU の様子は 30 秒おきに /workspace/logs/gpu_train.csv
# 6. 数分おきにステップ・loss を出す
# 7. 終わったら: 最後のチェックポイントを Hugging Face (非公開) へ上げて一覧を確かめる → ログと一覧を持ち帰る
#    → Pod を消す → KIOXIA の checkpoints/ へ Hugging Face から取り, Pod 側の一覧と名前・バイト数を比べる.
#    上げられなければ, Pod から KIOXIA へ scp で取ってから消す
# 8. 落ちた・止まった・上限を超えた・Ctrl-C: ログを持ち帰ってから消す
# 9. Mac に通知を出す. まとめは lab/<Log>/train/result_<時刻>.txt, ターミナルの出力は lab_train_<時刻>.txt
#    由来 (どのデータ・コード・設定で学習したか. provenance.py) を lab/<Log>/train/train_provenance.json に置き,
#    KIOXIA に新しく置けたチェックポイントと Hugging Face にも同じものを置く (依頼 54. 落ちた・止めたときも lab/ 側には置く)
#
# 鍵 (Runpod・Hugging Face) は KIOXIA の keys/ から読み, 画面・記録には出さない (lab_lib.sh).
# Pod 側の見張りには Mac の鍵を送らない.

set -u
set -o pipefail

# 中身は全部 main の中に置き, 最後の1行で呼ぶ (依頼 47 と同じ. 動いている途中でこのファイルが書き換わっても壊れない)
main() {
HERE="$(cd "$(dirname "$0")" && pwd)"
ORIG_ARGS=("$@")            # 由来に残す (依頼 54)
# shellcheck source=lab_lib.sh
. "$HERE/lab_lib.sh"

GPU=A40; FALLBACK=1; LIMIT_HOURS=5; STEPS=1500; INIT=""; EXP=""; MODEL=corr    # INIT・EXP の既定は --model から
POLL_MIN=5; STALL_MIN=30; LOG=""; DRY_RUN=0; MODE=train
while [ $# -gt 0 ]; do
    case "$1" in
        --log) lab_need_val Log18 "$@"; LOG="$2"; shift 2 ;;
        --gpu) lab_need_val A40 "$@"; GPU="$2"; shift 2 ;;
        --no-fallback) FALLBACK=0; shift ;;
        --limit-hours) lab_need_val 5 "$@"; LIMIT_HOURS="$2"; shift 2 ;;
        --model) lab_need_val corr2 "$@"; MODEL="$2"; shift 2 ;;
        --steps) lab_need_val 1500 "$@"; STEPS="$2"; shift 2 ;;
        --init) lab_need_val /workspace/ckpt/3999/params "$@"; INIT="$2"; shift 2 ;;
        --exp-name) lab_need_val trim_corr_v3 "$@"; EXP="$2"; shift 2 ;;
        --poll-min) lab_need_val 5 "$@"; POLL_MIN="$2"; shift 2 ;;
        --stall-min) lab_need_val 30 "$@"; STALL_MIN="$2"; shift 2 ;;
        --stop) MODE=stop; shift ;;
        --dry-run) DRY_RUN=1; shift ;;
        -h|--help) sed -n '2,39p' "$0"; exit 0 ;;
        *) echo "[ERROR] 知らない引数: $1  (./lab_train.sh --help)" >&2; exit 2 ;;
    esac
done
[ -n "$LOG" ] || { echo "[ERROR] --log Log17 のように Log 番号を付けてください" >&2; exit 2; }
lab_train_def "$MODEL" || { echo "[ERROR] --model は ${LAB_TRAIN_MODELS// / か }: $MODEL" >&2; exit 2; }
[ -n "$INIT" ] || INIT=$TRAIN_INIT
[ -n "$EXP" ] || EXP=$TRAIN_EXP
num() { awk -v x="$1" 'BEGIN { exit !(x ~ /^[0-9]+(\.[0-9]+)?$/ && x > 0) }'; }
for v in LIMIT_HOURS POLL_MIN STALL_MIN; do num "${!v}" || { echo "[ERROR] $v は 0 より大きい数: ${!v}" >&2; exit 2; }; done
[[ "$STEPS" =~ ^[1-9][0-9]*$ ]] || { echo "[ERROR] --steps は 1 以上の整数: $STEPS" >&2; exit 2; }
[[ "$EXP" =~ ^[A-Za-z0-9_.-]+$ ]] || { echo "[ERROR] --exp-name は英数字と _ . - だけ: $EXP" >&2; exit 2; }
[[ "$INIT" =~ ^[A-Za-z0-9_./-]+$ ]] || { echo "[ERROR] --init に使えない文字がある: $INIT" >&2; exit 2; }
case "$GPU" in *" "*) GPU_ID="$GPU" ;; *) GPU_ID="NVIDIA $GPU" ;; esac
secs() { awk -v x="$1" -v m="$2" 'BEGIN { printf "%d", x * m }'; }
LIMIT_SEC=$(secs "$LIMIT_HOURS" 3600); POLL_SEC=$(secs "$POLL_MIN" 60); STALL_SEC=$(secs "$STALL_MIN" 60)
WD_MARGIN_SEC=3600          # Pod 側の見張りは Mac の上限のこれだけあと (Mac が先に片付けるのがふつう)
WRAP_EXTEND_SEC=10800       # 片付け (上げる・持ち帰る) を始めたら, Pod 側の見張りを今から 3 時間あとへ延ばす
CREATE_WAIT=30; SSH_WAIT=10; SSH_TRIES=90; LOGIN_WAIT=5; LOGIN_TRIES=60; GONE_WAIT=5; GONE_TRIES=30
FETCH_TRIES=3; FETCH_WAIT=10
FALLBACK_GPUS=("NVIDIA RTX A6000" "NVIDIA L40S")    # A40 と同じ 48GB. この順に試す

DATA_ROOT=/Volumes/KIOXIA/r2r2r-data
NOTIFY_TITLE="lab_train $LOG"
if [ $DRY_RUN = 1 ]; then
    # shellcheck source=lab_train_fakes.sh
    . "$HERE/lab_train_fakes.sh"      # DATA_ROOT・LAB_ROOT・鍵を一時フォルダへ, API・SSH・hf・通知を偽物に, 待ちを短く
fi
CKPT_ROOT="$DATA_ROOT/checkpoints"
ADD_TRIM="$DATA_ROOT/lab/Log11/pod/add_trim_config.py"
use_model() {               # use_model corr|corr2|corr3 → 作るモデルで決まる値 (lab_lib.sh の lab_train_def. 依頼 57 / 97)
    lab_train_def "$1" || return 1
    DATA_NAME=$TRAIN_DATA                                   # corr: ..._trim2_corr_v3 / corr2: ..._trim2_corr_v4b (依頼 58) / corr3: ..._trim2_corr_v5
    TAR="$DATA_ROOT/datasets/$DATA_NAME/${DATA_NAME}_raw_episodes.tar"
    CONFIG=$TRAIN_CONFIG                                    # corr: pi05_mecharm_trim_corr_lora / corr2: ..._v4_lora / corr3: ..._v5_lora
    HF_REPO_IN=$TRAIN_REPO_IN                               # 続きの元 (corr: 3999 / corr2: corr の 1499 / corr3: corr2 の 1499. lab_start.sh と同じ)
    HF_REPO_OUT=$TRAIN_REPO_OUT                             # 上げる先 (非公開). 名前は checkpoints/ のフォルダ名に合わせる
    CKPT_NAME=$TRAIN_CKPT_NAME                              # KIOXIA の checkpoints/<これ>/<ステップ>/
    CKPT_IN=$TRAIN_CKPT_IN                                  # 3c で続きの元を置く場所
    CKPT_IN_KIOXIA="$CKPT_ROOT/$TRAIN_FROM_KIOXIA"          # HF_REPO_IN と同じもの (由来の指紋を取る. 依頼 54)
    PREP_VARIANT=$TRAIN_VARIANT                             # pod_prep_trim_corr.sh の 3 つ目 (v3 / v4 / v5)
}
use_model "$MODEL"
POD_EXP_DIR=/workspace/openpi/checkpoints/$CONFIG/$EXP
TRAIN_LOG=/workspace/logs/train_$EXP.log
GPU_LOG=/workspace/logs/gpu_train.csv
RUN_LOG=/workspace/logs/train_run.log
EXIT_FILE=/workspace/logs/train_exit_code
PREP_LOG=/workspace/logs/prep_trim_corr.log
LAST_STEP=$((STEPS - 1))      # openpi は 0 から数え, 最後のステップ (num_train_steps - 1) を必ず保存する

# ---- 書く前に KIOXIA のマウントを確かめる ----
if [ ! -d "$DATA_ROOT" ]; then
    echo "[ERROR] SSD 未マウント: $DATA_ROOT" >&2; exit 1
fi
LOGDIR="$LAB_ROOT/$LOG/train"
STATE="$LOGDIR/lab_pod.json"
mkdir -p "$LOGDIR" || { echo "[ERROR] $LOGDIR を作れない" >&2; exit 1; }
LOGFILE="$LOGDIR/lab_train_$(date +%Y%m%d_%H%M%S)$([ $MODE = stop ] && echo _stop).txt"
trap '' PIPE                # ターミナルが閉じて tee が終わっても, 書けないだけで続ける (Pod を消すところまで行く)
exec > >(tee -a "$LOGFILE") 2>&1
echo "[LOG   ] ターミナルの出力を残します → $LOGFILE"

lab_need_key "$RUNPOD_KEY_FILE" "Runpod の API キー"
lab_need_key "$HF_TOKEN_FILE" "Hugging Face の鍵"
[ -s "$POD_SSH_KEY" ] || { echo "[ERROR] Pod に入る SSH の鍵がありません: $POD_SSH_KEY" >&2; exit 1; }
[ -s "$POD_SSH_KEY.pub" ] || { echo "[ERROR] SSH の公開鍵がありません: $POD_SSH_KEY.pub" >&2; exit 1; }
lab_ssh_opts
pssh() { pod_ssh "$@" < /dev/null; }      # 標準入力を Pod へ渡さない ssh (時間の上限は pod_ssh. 長いものは LAB_TO=秒 を前に付ける)

# ---- 段階ごとの秒数 ----
T_NAMES=(); T_SECS=()
CUR_STAGE="(始める前)"; STAGE_T0=$SECONDS
POD_ID=""; POD_T0=""; COST_PER_HR=0; GPU_USED=""; WRAPPING=0; TRAIN_T0=""
stage() {                   # stage 名前 [nolimit]  (nolimit: 上限の確かめをしない. 見守りは様子を見てから上限を見る)
    CUR_STAGE="$1"; STAGE_T0=$SECONDS; echo; echo "=== $1  ($(date '+%H:%M:%S'))"
    if [ "${2:-}" != nolimit ] && [ $WRAPPING = 0 ] && [ -n "$POD_T0" ] && [ $(( $(date +%s) - POD_T0 )) -ge "$LIMIT_SEC" ]; then
        wrap_fail timeout "上限 $LIMIT_HOURS 時間を超えた (「$1」の前)"
    fi
}
stage_done() { T_NAMES+=("$CUR_STAGE"); T_SECS+=($((SECONDS - STAGE_T0))); echo "    ($((SECONDS - STAGE_T0)) 秒)"; }
summary() {
    local i
    echo; echo "=== かかった時間"
    for i in "${!T_NAMES[@]}"; do printf '  %-36s %6d 秒\n' "${T_NAMES[$i]}" "${T_SECS[$i]}"; done
}
hm() { local x=$1; [ "$x" -lt 0 ] && x=0; printf '%d:%02d' $((x / 3600)) $((x % 3600 / 60)); }

# ---- 由来 (依頼 54. provenance.py. 書けなくても学習の片付けは止めない) ----
PROV_PROCESS=""; PROV_ERROR=""; PROV_CKPT_DIR=""; PROV_HF=0; PROV_DONE=0
OPENPI_JSON='[]'            # 3f のあとに Pod の openpi のコミットと変更の有無を入れる
iso_of() { [ -n "${1:-}" ] && [ "$1" != 0 ] && date -r "$1" '+%Y-%m-%dT%H:%M:%S%z' | sed -E 's/([0-9]{2})$/:\1/'; }
prov_inputs() {             # 入力 (学習データ・続きの元 (corr は 3999・corr2 は corr の 1499・corr3 は corr2 の 1499)・設定を足すスクリプト) を spec の inputs の形で出す
    local init_ck=""
    case "$INIT" in "$CKPT_IN"/*) init_ck="$CKPT_IN_KIOXIA" ;; esac
    jq -n --arg data "$DATA_NAME" --arg data_dir "$DATA_ROOT/datasets/$DATA_NAME" \
        --arg init "$INIT" --arg init_ck "$init_ck" --arg init_id "${CKPT_IN_KIOXIA#"$CKPT_ROOT"/}" --arg repo_in "$HF_REPO_IN" \
        --arg add_trim "$ADD_TRIM" \
        '{inputs: ([{kind: "dataset", id: $data, path: $data_dir}]
                  + [{kind: "checkpoint", id: (if $init_ck == "" then (if $init == "" then null else $init end) else $init_id end),
                      path: (if $init_ck == "" then null else $init_ck end),
                      hf_repo: (if $init_ck == "" then null else $repo_in end), pod_path: (if $init == "" then null else $init end)}]
                  + [{kind: "config", id: "add_trim_config.py", path: $add_trim}])}'
}
PROV_CODE_ARGS=(--repo-dir "${LAB_PROV_REPO:-$HERE}" --code-path lab_train.sh --code-path lab_lib.sh --code-path provenance.py --code-path pi05)
prov_prepare() {            # 始める前にコードと入力の指紋を測って PROV_START に持つ (終わったときに測ると, 途中で変えたものになる)
    local out
    out=$(prov_inputs | lab_timeout 120 python3 "$HERE/provenance.py" prepare --data-root "$DATA_ROOT" "${PROV_CODE_ARGS[@]}" 2>/dev/null) \
        && jq -e '.code and .inputs' <<<"$out" >/dev/null 2>&1 && PROV_START=$(jq -c . <<<"$out") && return 0
    echo "[WARN  ] 始める時点のコードと入力の指紋を測れなかった (由来は終わったときに測る)"
    PROV_START=""
    return 1
}

prov_spec() {               # 由来の spec (lab_train.sh が知っていること) を JSON で出す
    local t_end rented start
    t_end=$(jq -r '.terminated_epoch // empty' "$STATE" 2>/dev/null)
    rented=$(jq -r '.rented_sec // empty' "$STATE" 2>/dev/null)
    start=${PROV_START:-}
    if [ -z "$start" ] && [ $MODE = stop ]; then
        # --stop で始めたときの記録がない (前の形の記録など): 止めた今のものを測らず「分からない」とする (codex R2)
        start=$(prov_inputs | jq -c '{code: {repo: "r2r2r-scripts", commit: null, dirty: null, error: "始めたときの記録がない (--stop)"},
                 inputs: [.inputs[] | . + {fingerprint: null, fingerprint_error: "始めたときの記録がない (--stop)", measured_at: null}]}') || return 1
    elif [ -z "$start" ]; then
        start=$(prov_inputs) || return 1        # 始める前に測れなかった: ここで測る (measured_at が終わったときの時刻になる)
    fi
    jq -n --argjson start "$start" --arg script lab_train.sh --argjson argv "$(printf '%s\0' ${ORIG_ARGS[@]+"${ORIG_ARGS[@]}"} | jq -Rs 'split("\u0000")[:-1]')" \
        --arg proc "$PROV_PROCESS" --arg err "$PROV_ERROR" --arg log "$LOG" --arg dry "$DRY_RUN" \
        --arg data "$DATA_NAME" --arg init "$INIT" --argjson openpi "$OPENPI_JSON" \
        --arg config "$CONFIG" --arg exp "$EXP" --arg steps "$STEPS" --arg limit "$LIMIT_HOURS" \
        --arg repo_out "$HF_REPO_OUT" --arg ckpt_name "$CKPT_NAME" --arg step "${CKPT_STEP:-}" --arg ck_dir "$PROV_CKPT_DIR" \
        --arg hf "${HF_MSG:-}" --arg kx "${KIOXIA_MSG:-}" --arg prov_hf "$PROV_HF" \
        --arg pod "$POD_ID" --arg gpu "$GPU_USED" --arg cph "$COST_PER_HR" --arg yen "$YEN_PER_USD" \
        --arg t0 "$(iso_of "$POD_T0")" --arg t1 "$(iso_of "$t_end")" --arg tt "$(iso_of "$TRAIN_T0")" --arg rented "$rented" \
        --arg cost_line "${COST_LINE:-}" \
        'def num: if . == "" then null else tonumber end; def str: if . == "" then null else . end;
        {producer: {script: $script, argv: $argv, dry_run: ($dry == "1")},
         external_code: $openpi,
         inputs: $start.inputs,
         code: ($start.code // null),
         outputs: (if $ck_dir == "" then [] else [{kind: "checkpoint", id: ($ckpt_name + "/" + $step), path: $ck_dir,
                   hf_repo: (if $prov_hf == "1" then $repo_out else null end)}] end),     # 指紋 files-v2 は由来のファイルを数えない (依頼 56)
         config: {name: $config, exp_name: $exp, num_train_steps: ($steps | num), init_params: ($init | str), data_name: $data,
                  limit_hours: ($limit | num), checkpoint_step: ($step | str)},
         compute: {provider: "runpod", pod_id: ($pod | str), gpu: ($gpu | str), started_at: ($t0 | str), train_started_at: ($tt | str),
                   ended_at: ($t1 | str), duration_sec: ($rented | num), cost_per_hr_usd: ($cph | num),
                   cost_usd: (if $rented == "" then null else (($rented | tonumber) / 3600 * ($cph | tonumber) * 100 | round / 100) end),
                   yen_per_usd: ($yen | num), cost_line: ($cost_line | str)},
         numbering: {series: "フィジカルAI研究室", series_log: $log, article_no: ""},
         status: {process: $proc, task: "not_applicable"},
         error: ($err | str),
         notes: {hugging_face: ($hf | str), kioxia: ($kx | str)}}'
}
PROV_TO=${LAB_PROV_TO:-300} # 由来を書く全体 (spec・書く・Hugging Face) の時間の上限 (通知を遅らせすぎない. codex R1・R2)
kill_group() {              # kill_group pgid → そのプロセスグループ (由来を書く子と, その子孫全部) に TERM, 2 秒たって残っていれば KILL
    # グループごと送るので, たどったあとに生まれた子孫も取りこぼさない. 番号の使い回しは, グループの全員が終わったあとにしか起きない
    # (KILL の前にまだ誰か残っているかを見る) (codex R3)
    kill -TERM -- "-$1" 2>/dev/null
    local i
    for i in 1 2 3 4; do pgrep -g "$1" >/dev/null 2>&1 || return 0; sleep 0.5; done
    kill -KILL -- "-$1" 2>/dev/null
    return 0
}
prov_write() {              # lab/<Log>/train/ (と, 今回新しく置いたチェックポイント・Hugging Face) に由来を置く. RESULT_LINES に1行足す
    local f pid t0 line
    [ $PROV_DONE = 1 ] && return 0
    PROV_DONE=1
    echo; echo "=== 由来 (train_provenance.json) を書く"
    [ -n "$PROV_PROCESS" ] || PROV_PROCESS=failed
    f=$(mktemp)
    # 自分のプロセスグループで流す (set -m). 時間切れのときにグループごと止める. 中の python と hf にはそれぞれ 120 秒の上限がある
    set -m
    ( trap - INT TERM HUP; prov_write_body ) > "$f" 2>&1 &
    pid=$!
    set +m
    t0=$SECONDS
    while kill -0 "$pid" 2>/dev/null && [ $((SECONDS - t0)) -lt "$PROV_TO" ]; do sleep 1; done
    if kill -0 "$pid" 2>/dev/null; then
        kill_group "$pid"; wait "$pid" 2>/dev/null
        echo "PROVLINE 由来: $PROV_TO 秒で書き終わらなかったので打ち切った (書けたものだけ残る)" >> "$f"
    else
        wait "$pid" 2>/dev/null
    fi
    grep -v '^PROVLINE ' "$f"
    line=$(sed -n 's/^PROVLINE //p' "$f" | tail -1); rm -f "$f"
    [ -n "$line" ] || line="由来: 書く途中で終わった"
    echo "[PROV  ] $line"
    RESULT_LINES=$(printf '%s\n%s' "${RESULT_LINES:-}" "$line")
}
prov_write_body() {         # 子のシェルで流す. 画面に出すものと, 最後に「PROVLINE <まとめの1行>」を出す
    local spec out lab_f ck_f="" outs line
    if [ ! -d "$DATA_ROOT" ] || [ ! -d "$LOGDIR" ]; then
        echo "PROVLINE 由来: KIOXIA がマウントされていない ($DATA_ROOT). 書かなかった"; return 1
    fi
    lab_f=$(dest_of train_provenance.json)
    outs=(--out "$lab_f")
    if [ -n "$PROV_CKPT_DIR" ] && [ -d "$PROV_CKPT_DIR" ]; then
        ck_f="$PROV_CKPT_DIR/train_provenance.json"
        outs+=(--out "$ck_f")
    fi
    spec=$(prov_spec) || spec=""
    if [ -z "$spec" ]; then
        echo "PROVLINE 由来: spec を作れなかった (jq). 書かなかった"; return 1
    fi
    out=$(printf '%s' "$spec" | lab_timeout 120 python3 "$HERE/provenance.py" write --data-root "$DATA_ROOT" "${outs[@]}" "${PROV_CODE_ARGS[@]}" \
              --secret-file "$RUNPOD_KEY_FILE" --secret-file "$HF_TOKEN_FILE" 2>&1)
    echo "$out" | sed 's/^/[PROV  ] /'
    line="由来: $([ -f "$lab_f" ] && echo "$lab_f" || echo "lab 側に書けなかった")"
    if [ -n "$ck_f" ]; then
        if grep -qxF -- "$ck_f" <<<"$out"; then line="$line ・ $ck_f"; else line="$line ・ チェックポイントには置けなかった"; fi
    fi
    if [ "$PROV_HF" = 1 ] && [ -f "$lab_f" ]; then
        # Hugging Face (今回 hf_upload_ckpt.py が上げたリポジトリ) にも同じものを足す. だめでも止めない
        if lab_hf_upload_file_mac "$HF_REPO_OUT" "$lab_f" train_provenance.json "lab_train.sh $LOG: train_provenance.json" | sed 's/^/[PROV  ] hf: /'; then
            line="$line ・ Hugging Face ($HF_REPO_OUT) にも置いた"
        else
            line="$line ・ Hugging Face には置けなかった"
        fi
    fi
    echo "PROVLINE $line"
}

# ---- 通知とまとめ ----
RESULT="$LOGDIR/result_$(date +%Y%m%d_%H%M%S).txt"
finish() {                  # finish 終了コード 通知の文 (まとめを書き, 通知を出して終わる)
    local rc="$1" msg="$2"
    FINISHED=1
    prov_write
    summary
    {
        echo "lab_train.sh $LOG  $(date '+%Y-%m-%d %H:%M:%S')"
        echo "結果: $msg"
        echo "${RESULT_LINES:-}"
    [ "${WD_OK:-1}" = 1 ] || echo "※ Pod 側の見張りは働いていなかった (${wd:-答えなし}). ./lab_stop.sh --check で Pod が残っていないか確かめる"
        echo "ターミナルの記録: $LOGFILE"
    } > "$RESULT"
    echo; echo "[RESULT] $msg"; echo "[RESULT] まとめ → $RESULT"
    lab_notify "$NOTIFY_TITLE" "$msg" "$rc"
    exit "$rc"
}
early_fail() {              # Pod を作る前に止まる (借りていないので消すものはない)
    echo; echo "[STOP  ] 「${CUR_STAGE}」で止まりました: $*"
    PROV_PROCESS=failed; PROV_ERROR="「${CUR_STAGE}」: $*"
    RESULT_LINES="Pod は借りていません"
    finish 1 "失敗: $* (Pod は借りていません)"
}

# ---- 同じ Log の lab_train.sh を2つ同時に動かさない (lab_start.sh と同じ形の印. 別のフォルダなのでぶつからない) ----
LOCK="$LOGDIR/.lab_train.lock"
LOCKED=0
FINISHED=0
cleanup() {
    # 思わぬところで終わった (set -u の止まりなど) ときも, 借りた Pod は残さない
    if [ "$FINISHED" = 0 ] && [ -n "$POD_ID" ]; then
        FINISHED=1
        echo "[ERROR] lab_train.sh が思わぬところで終わりました (「${CUR_STAGE}」). ログを持ち帰って Pod を消します"
        ( wrap_fail unexpected "lab_train.sh が思わぬところで終わった (「${CUR_STAGE}」)" )
    fi
    if [ $LOCKED = 1 ] && [ "$(cat "$LOCK/pid" 2>/dev/null)" = "$$" ]; then rm -rf "$LOCK"; fi
    [ -n "${CAFF_PID:-}" ] && kill "$CAFF_PID" 2>/dev/null
}
trap cleanup EXIT
if ! mkdir "$LOCK" 2>/dev/null; then
    other=$(cat "$LOCK/pid" 2>/dev/null)
    if [ -n "$other" ] && kill -0 "$other" 2>/dev/null; then
        echo "[ERROR] $LOG の lab_train.sh がもう動いています (pid $other). 止めるならそのターミナルで Ctrl-C (ログを持ち帰って Pod を消します)"
    else
        echo "[ERROR] $LOG の lab_train.sh の印が残っています (pid ${other:-?}・動いていないように見える)"
        echo "        ほかに動いていないことを確かめてから, 印を消して打ち直してください: rm -r $LOCK"
    fi
    exit 1
fi
LOCKED=1
echo $$ > "$LOCK/pid" || { echo "[ERROR] $LOCK/pid を書けない"; exit 1; }

save_state() {              # save_state jq の式 [--arg/--argjson ...]
    local f="$1"; shift
    jq "$@" "$f" "$STATE" > "$STATE.tmp" && mv "$STATE.tmp" "$STATE"
}

# ---- 片付け (持ち帰る・消す) ----
dest_of() {                 # dest_of 名前 → <train>/名前 (あれば 名前_<時刻>)
    local d="$LOGDIR/$1"
    if [ -e "$d" ]; then
        case "$1" in *.*) d="$LOGDIR/${1%.*}_$(date +%Y%m%d_%H%M%S).${1##*.}" ;; *) d="${d}_$(date +%Y%m%d_%H%M%S)" ;; esac
    fi
    echo "$d"
}

fetch_logs() {              # 学習のログ・GPU のログなどを持ち帰る. Pod に入れなければ FETCH_TRIES 回まで待ってやり直す
    local i f dest got=0 listing
    for i in $(seq 1 "$FETCH_TRIES"); do
        if lab_pod_addr && LAB_TO=$LAB_TO_SHORT pssh true >/dev/null; then
            for f in "$TRAIN_LOG" "$GPU_LOG" "$RUN_LOG" "$PREP_LOG" /workspace/logs/watchdog.log; do
                dest=$(dest_of "$(basename "$f")")
                if pod_scp "root@$POD_IP:$f" "$dest" >/dev/null; then
                    echo "[LOG   ] $(basename "$f") を持ち帰った → $dest ($(wc -l < "$dest" | tr -d ' ') 行)"; got=$((got + 1))
                else
                    rm -f "$dest"; echo "[LOG   ] $(basename "$f") は Pod にない (持ち帰れなかった)"
                fi
            done
            # Pod 側のチェックポイントの一覧 (全部のステップ. train_state も含む. 名前とバイト数)
            dest=$(dest_of pod_checkpoints.tsv)
            listing=$(LAB_TO=$LAB_TO_SHORT pssh ": lt=ls_all; cd $POD_EXP_DIR 2>/dev/null && find . -type f -printf '%P\t%s\n' | LC_ALL=C sort | sed 's/^/F\t/'")
            grep $'^F\t' <<<"$listing" | cut -f2- > "$dest"
            if [ -s "$dest" ]; then echo "[LOG   ] Pod 側のチェックポイントの一覧 → $dest ($(wc -l < "$dest" | tr -d ' ') ファイル)"
            else rm -f "$dest"; echo "[LOG   ] Pod 側にチェックポイントはない ($POD_EXP_DIR)"; fi
            FETCHED=$got
            return 0
        fi
        echo "[LOG   ] ($i/$FETCH_TRIES) Pod に入れない (状態 ${POD_STATUS:-?})"
        [ "$i" -lt "$FETCH_TRIES" ] && { echo "[LOG   ] $FETCH_WAIT 秒待ってやり直します"; sleep "$FETCH_WAIT"; }
    done
    FETCHED=0
    return 1
}

delete_pod() {              # Pod を消し (停止ではなく削除), 一覧からも消えたのを API で確かめ, 費用を出す. 消せなければ 1
    local tries pods gone=0 t_end s sent=1
    echo
    echo "=== Pod を消す (id $POD_ID)  ($(date '+%H:%M:%S'))"
    lab_pod_kind "$POD_ID"; local known=$?
    if [ $known = 0 ] && [ "$LAB_POD_KIND" = gone ]; then
        echo "[STOP  ] Pod (id $POD_ID) はもう API にない"
        t_end=$(jq -r '.terminated_epoch // empty' "$STATE" 2>/dev/null); [ -n "$t_end" ] || t_end=$(date +%s)
    else
        sent=0
        for tries in 1 2 3; do
            if runpod_api DELETE "/pods/$POD_ID" >/dev/null; then sent=1; break; fi
            echo "[STOP  ] ($tries/3) Pod を消せない (HTTP $RUNPOD_HTTP)"
            [ $tries -lt 3 ] && sleep "$GONE_WAIT"
        done
        t_end=$(date +%s)
        [ $sent = 1 ] && echo "[STOP  ] 削除 (Terminate) を送った ($(date -r "$t_end" '+%H:%M:%S'))"
    fi
    for _ in $(seq 1 "$([ "${sent:-1}" = 1 ] && echo "$GONE_TRIES" || echo 1)"); do
        if pods=$(runpod_api GET /pods) && lab_valid_pods "$pods" \
           && ! jq -e --arg id "$POD_ID" 'any(.[]; .id == $id and .desiredStatus != "TERMINATED")' <<<"$pods" >/dev/null; then
            gone=1; break
        fi
        sleep "$GONE_WAIT"
    done
    s=$((t_end - POD_T0))
    COST_LINE=$(lab_cost_line "$s" "$COST_PER_HR")
    save_state '. + {terminated_epoch:$t, terminated_at:($t|strflocaltime("%Y-%m-%d %H:%M:%S")), rented_sec:$s, gone_confirmed:($g == 1)}' \
        --argjson t "$t_end" --argjson s "$s" --argjson g "$gone"
    if [ $gone = 1 ]; then
        echo "[STOP  ] Pod (id $POD_ID) が消えたことを API で確かめた"
    else
        echo "[ERROR] Pod (id $POD_ID) が API の一覧に残っています$([ $sent = 0 ] && echo " (削除を受け付けなかった)"). Runpod の画面で確かめてください"
    fi
    echo "[COST  ] $COST_LINE  GPU=$GPU_USED"
    POD_GONE=$gone
    [ $gone = 1 ]
}

pod_left_msg() { echo "Pod (id $POD_ID) が残っているかもしれません. ./lab_train.sh --log $LOG --stop か Runpod の画面で消してください"; }

wrap_fail() {               # wrap_fail 種類 わけ → ログを持ち帰ってから Pod を消し, 通知して終わる (終了コード 1)
    local kind="$1" why="$2" logs
    WRAPPING=1; trap '' INT TERM HUP     # ここからは Ctrl-C でも止めない (Pod を消すところまで行く. 二重には入らない)
    echo; echo "[STOP  ] 「${CUR_STAGE}」で止めます: $why"
    case "$kind" in signal|stop|timeout) PROV_PROCESS=interrupted ;; *) PROV_PROCESS=failed ;; esac      # 由来 (依頼 54)
    [ "${TRAIN_DONE:-0}" = 1 ] && PROV_PROCESS=completed      # 学習は最後まで終わっていて, そのあとの片付けで止まった
    PROV_ERROR="$kind: 「${CUR_STAGE}」: $why"
    if [ -z "$POD_ID" ]; then RESULT_LINES="Pod は借りていません"; finish 1 "失敗: $why (Pod は借りていません)"; fi
    # 学習を止めてからログを持ち帰る (消す前に書き終わったログを取る. 止められなくても続ける)
    if [ "$kind" != crash ] && [ -n "$TRAIN_T0" ]; then
        LAB_TO=$LAB_TO_SHORT pssh "$(lab_train_kill_cmd)" >/dev/null
    fi
    if fetch_logs; then logs="ログを持ち帰った ($FETCHED 個)"; else logs="ログは持ち帰れなかった (Pod に入れない)"; fi
    if delete_pod; then
        RESULT_LINES=$(printf '%s\n' "わけ: $why" "$logs → $LOGDIR" "Pod (id $POD_ID) は消した (API で確かめた)" "$COST_LINE")
        finish 1 "失敗: $why. $logs. Pod は消しました. ${COST_LINE#借りていた時間 }"
    else
        RESULT_LINES=$(printf '%s\n' "わけ: $why" "$logs → $LOGDIR" "$(pod_left_msg)" "${COST_LINE:-}")
        finish 1 "失敗: $why. $logs. $(pod_left_msg)"
    fi
}

on_signal() { wrap_fail signal "$1 で止めた"; }
trap 'on_signal Ctrl-C' INT
trap 'on_signal "TERM (kill)"' TERM
trap 'on_signal "ターミナルが閉じた (HUP)"' HUP

# ---- KIOXIA の checkpoints/ に置く ----
POD_LIST=""                  # Pod 側の最後のチェックポイントの一覧 (パス<TAB>バイト数. train_state なし)
compare_with_pod() {        # compare_with_pod フォルダ → Pod 側の一覧と名前・バイト数が全部同じなら 0
    local mine d rc
    mine=$(mktemp)
    # 由来 (train_provenance.json) は Mac があとで足すので数えない (依頼 54)
    if ! { lab_list_files "$1" | { grep -v $'^train_provenance\.json\t' || true; }; } > "$mine" || [ ! -s "$POD_LIST" ]; then
        echo "[KIOXIA] 一覧を作れない (KIOXIA 側 $1 か Pod 側 $POD_LIST). 合ったとはしない"; rm -f "$mine"; return 1
    fi
    d=$(diff "$mine" "$POD_LIST"); rc=$?; rm -f "$mine"
    [ $rc = 0 ] && return 0          # diff の終了コードが 0 のときだけ「合った」(2 = 読めない も合わないとする)
    echo "[KIOXIA] Pod 側の一覧と合わない (diff の終了コード $rc. < KIOXIA / > Pod. 先頭 20 行):"
    head -20 <<<"$d" | sed 's/^/         /'
    return 1
}
place_ckpt() {              # place_ckpt hf|scp → KIOXIA の checkpoints/<CKPT_NAME>/<ステップ>/ に置いて確かめる. KIOXIA_MSG に結果
    local how="$1" parent="$CKPT_ROOT/$CKPT_NAME" dest tmp n bytes
    dest="$parent/$CKPT_STEP"
    if [ ! -d "$DATA_ROOT" ]; then KIOXIA_MSG="KIOXIA がマウントされていない (置けなかった)"; return 1; fi
    if [ -e "$dest" ] || [ -L "$dest" ]; then         # 前からあるものは上書き・移動・削除しない. 中身が同じかだけを見る
        if compare_with_pod "$dest"; then KIOXIA_MSG="$dest はもとからあり, Pod 側と全部同じ (書き換えていない)"; return 0; fi
        KIOXIA_MSG="$dest がもとからあり, Pod 側と違う. 上書きしなかった"; return 1
    fi
    mkdir -p "$parent" || { KIOXIA_MSG="$parent を作れない"; return 1; }
    tmp="$parent/.incoming_${CKPT_STEP}_$(date +%Y%m%d_%H%M%S)_$$"
    mkdir "$tmp" || { KIOXIA_MSG="$tmp を作れない"; return 1; }
    if [ "$how" = hf ]; then
        echo "[KIOXIA] Hugging Face ($HF_REPO_OUT) から取ります → $tmp"
        lab_hf_download_mac "$HF_REPO_OUT" "$tmp" || { KIOXIA_MSG="Hugging Face から取れなかった (途中のものは $tmp)"; return 1; }
    else
        echo "[KIOXIA] Pod から scp で取ります → $tmp"
        LAB_TO=$LAB_TO_LONG pod_scp -r "root@$POD_IP:$POD_CKPT/params" "root@$POD_IP:$POD_CKPT/assets" \
                   "root@$POD_IP:$POD_CKPT/_CHECKPOINT_METADATA" "$tmp/" \
            || { KIOXIA_MSG="Pod から scp で取れなかった (途中のものは $tmp)"; return 1; }
    fi
    rm -rf "${tmp:?}/.cache" "${tmp:?}/.gitattributes"     # hf download の印 (自分で作った一時フォルダの中だけ)
    rm -f "${tmp:?}/train_provenance.json"                 # Hugging Face にあった前の回の由来は持ってこない (今回のものをあとで置く. codex R1)
    if ! compare_with_pod "$tmp"; then KIOXIA_MSG="取ったものが Pod 側の一覧と合わない (そのまま $tmp に置いた)"; return 1; fi
    # mkdir は「なければ作る」を1回で行う (あれば失敗). 作れたら中身を移す (確かめと移すの間に別の実行が作っても上書きしない)
    if ! mkdir "$dest" 2>/dev/null; then KIOXIA_MSG="$dest が途中でできた. 上書きしない (取ったものは $tmp)"; return 1; fi
    # .DS_Store なども移す (codex R3). ._* (exFAT が付ける) は移さない (依頼 52): exFAT では ._<名前> は本体と一緒に
    # 動くので, 本体のあとに ._<名前> を mv すると「No such file」で止まる (依頼 49 の初回). 残った ._* は rm_incoming で片付ける
    if ! ( shopt -s dotglob nullglob; set --; for f in "$tmp"/*; do case "${f##*/}" in ._*) ;; *) set -- "$@" "$f" ;; esac; done
           [ $# = 0 ] || mv "$@" "$dest"/ ); then KIOXIA_MSG="$tmp から $dest へ移せない (途中まで移したかもしれない)"; return 1; fi
    rm_incoming "$tmp" || echo "[KIOXIA] $tmp に ._ 以外のものが残っている, または消せない. 消さずに残します"
    compare_with_pod "$dest" || { KIOXIA_MSG="$dest に移したあとの一覧が Pod 側と合わない"; return 1; }
    n=$(wc -l < "$POD_LIST" | tr -d ' '); bytes=$(awk -F'\t' '{s += $2} END {printf "%d", s}' "$POD_LIST")
    KIOXIA_MSG="$dest に置いた ($n ファイル $bytes バイト. Pod 側と名前・バイト数が全部合った)"
    PROV_CKPT_DIR="$dest"       # 今回新しく置いたものにだけ由来を置く (前からあるものには書き足さない. 依頼 54)
}

rm_incoming() {             # rm_incoming 一時フォルダ → ._* (exFAT が付ける) しか残っていなければ, それとフォルダを消す. ほかのものがあれば何もせず 1
    # 自分で作った .incoming_* にだけ使う (前からある .incoming_* には触らない)
    ( shopt -s dotglob nullglob
      for f in "$1"/*; do case "${f##*/}" in ._*) [ -f "$f" ] && [ ! -L "$f" ] || exit 1 ;; *) exit 1 ;; esac; done
      for f in "$1"/*; do rm -f "$f" || exit 1; done
      rmdir "$1" )
}

extend_watchdog() {         # Pod 側の見張りの時刻を, 今から WRAP_EXTEND_SEC あとへ延ばす (片付けの途中で消えないように)
    local t=$(( $(date +%s) + WRAP_EXTEND_SEC ))
    if LAB_TO=$LAB_TO_SHORT pssh ": lt=deadline; echo $t > /workspace/logs/.lab_deadline.tmp && mv /workspace/logs/.lab_deadline.tmp /workspace/logs/.lab_deadline" >/dev/null; then
        echo "[WATCH ] Pod 側の見張りを $(date -r "$t" '+%H:%M') まで延ばした (上げる・持ち帰る間に消えないように)"
    else
        echo "[WARN  ] Pod 側の見張りを延ばせなかった. 前の時刻 ($(date -r "$WD_AT" '+%H:%M')) を過ぎると片付けの途中でも Pod が消える"
    fi
}

wrap_success() {            # 学習が最後まで終わった: 上げる → 持ち帰る → 消す → KIOXIA へ
    local hf_ok=0 hf_line out k_ok=0 k_how="" logs msg rc
    WRAPPING=1; trap '' INT TERM HUP
    TRAIN_DONE=1                # このあとの片付けで止まっても, 由来は completed (codex R1)
    CKPT_STEP=$S_CKPT
    POD_CKPT="$POD_EXP_DIR/$CKPT_STEP"
    echo; echo "[TRAIN ] 学習が終わりました (最後のチェックポイント $CKPT_STEP, loss ${S_LOSS})"
    extend_watchdog

    stage "7a. Pod 側の一覧 ($CKPT_STEP)"
    POD_LIST=$(dest_of "pod_files_$CKPT_STEP.tsv")
    out=$(LAB_TO=$LAB_TO_SHORT pssh "$(lab_train_list_cmd "$POD_CKPT")"); local list_rc=$?
    grep $'^F\t' <<<"$out" | cut -f2- > "$POD_LIST"
    if [ $list_rc = 0 ] && [ -s "$POD_LIST" ] && grep -q $'^_CHECKPOINT_METADATA\t' "$POD_LIST"; then
        echo "[LIST  ] $(wc -l < "$POD_LIST" | tr -d ' ') ファイル $(awk -F'\t' '{s += $2} END {printf "%d", s}' "$POD_LIST") バイト (train_state なし) → $POD_LIST"
    else
        stage_done; wrap_fail list "Pod 側のチェックポイント ($POD_CKPT) の一覧を取れない"
    fi
    stage_done

    stage "7b. Hugging Face へ上げる (${HF_REPO_OUT}・非公開)"
    out=$(LAB_TO=$LAB_TO_LONG pssh "cd /workspace && HF_HUB_DISABLE_PROGRESS_BARS=1 python3 /workspace/r2r2r/pi05/openpi_addon/hf_upload_ckpt.py --dir $POD_CKPT --repo $HF_REPO_OUT --message 'lab_train.sh $LOG: $CONFIG $EXP step $CKPT_STEP'")
    echo "$out" | grep -v '^$'
    hf_line=$(grep '^HFUP ' <<<"$out" | tail -1)
    case "$hf_line" in "HFUP ok "*|"HFUP same "*) hf_ok=1 ;; esac
    HF_MSG=${hf_line#HFUP }; [ -n "$HF_MSG" ] || HF_MSG="答えがない (ssh が切れたなど)"
    stage_done

    if [ $hf_ok = 0 ]; then         # 上げられなければ, Pod を消す前に Pod から KIOXIA へ取る
        stage "7c. KIOXIA へ Pod から取る (Hugging Face に上げられなかったため)"
        extend_watchdog
        place_ckpt scp && k_ok=1; k_how=scp
        echo "[KIOXIA] $KIOXIA_MSG"
        stage_done
    fi

    stage "7d. ログを持ち帰る"
    if fetch_logs; then logs="ログを持ち帰った"; else logs="ログは持ち帰れなかった"; fi
    stage_done
    delete_pod; local gone=$?

    if [ $hf_ok = 1 ]; then         # Hugging Face にあるので, Pod を消してから取る (取れなくても Hugging Face に残る)
        stage "7e. KIOXIA へ Hugging Face から取る"
        place_ckpt hf && k_ok=1; k_how=hf
        echo "[KIOXIA] $KIOXIA_MSG"
        stage_done
    fi

    RESULT_LINES=$(printf '%s\n' "学習: $CONFIG --exp-name=$EXP  最後のチェックポイント $CKPT_STEP  最後の loss ${S_LOSS} (Step ${S_LOGSTEP})" \
        "Hugging Face ($HF_REPO_OUT): $HF_MSG" "KIOXIA: $KIOXIA_MSG" "$logs → $LOGDIR" \
        "$([ $gone = 0 ] && echo "Pod (id $POD_ID) は消した (API で確かめた)" || pod_left_msg)" "${COST_LINE:-}")
    if [ $hf_ok = 1 ] && [ $k_ok = 1 ] && [ $gone = 0 ]; then
        rc=0; msg="成功: 学習が終わりました (step $CKPT_STEP, loss ${S_LOSS}). Hugging Face と KIOXIA に置き, Pod を消しました. ${COST_LINE#借りていた時間 }"
    else
        rc=1; msg="学習は終わりました (step $CKPT_STEP) が,"
        [ $hf_ok = 0 ] && msg="$msg Hugging Face に上げられなかった ($HF_MSG)."
        if [ $k_ok = 1 ]; then msg="$msg KIOXIA には置いた ($k_how)."; else msg="$msg KIOXIA に置けなかった ($KIOXIA_MSG)."; fi
        if [ $gone = 0 ]; then msg="$msg Pod は消しました."; else msg="$msg $(pod_left_msg)"; fi
        PROV_ERROR="学習のあとの片付け: $msg"
    fi
    PROV_PROCESS=completed; [ $hf_ok = 1 ] && [ "${hf_line#HFUP ok }" != "$hf_line" ] && PROV_HF=1     # Hugging Face へは今回上げたときだけ足す
    finish "$rc" "$msg"
}

# ---- --stop: 残っている学習の Pod のログを持ち帰って消す ----
if [ $MODE = stop ]; then
    CUR_STAGE="--stop"
    [ -f "$STATE" ] || { echo "[ERROR] $STATE がありません (lab_train.sh で作った Pod の記録). ./lab_stop.sh --check で Pod の一覧を見る"; exit 1; }
    POD_ID=$(jq -r '.id // ""' "$STATE"); POD_T0=$(jq -r '.created_epoch // 0' "$STATE")
    COST_PER_HR=$(jq -r '.costPerHr // 0' "$STATE"); GPU_USED=$(jq -r '.gpu // ""' "$STATE")
    TRAIN_T0=$(jq -r '.train_started_epoch // ""' "$STATE")
    EXP=$(jq -r --arg e "$EXP" '.exp // $e' "$STATE")
    # 依頼 57: 作っていたモデル (記録がなければ corr. 依頼 57 より前は corr だけ) の値に合わせる. 設定の名前も記録から
    MODEL=$(jq -r '.model // "corr"' "$STATE")
    use_model "$MODEL" || { echo "[ERROR] $STATE のモデルが分からない: $MODEL"; exit 1; }
    CONFIG=$(jq -r --arg c "$CONFIG" '.config // $c' "$STATE")
    # 由来 (依頼 54. codex R1): 始めたときの設定・コード・入力を記録から戻す. 記録にないものは「分からない」(空) にする
    STEPS=$(jq -r '.steps // ""' "$STATE"); INIT=$(jq -r '.init // ""' "$STATE"); LIMIT_HOURS=$(jq -r '.limit_hours // ""' "$STATE")
    PROV_START=$(jq -c '.prov_start // empty' "$STATE")
    OPENPI_JSON=$(jq -c '.openpi // [{name: "openpi", commit: null, dirty: null, error: "記録にない (--stop)"}]' "$STATE") || OPENPI_JSON='[]'
    POD_EXP_DIR=/workspace/openpi/checkpoints/$CONFIG/$EXP; TRAIN_LOG=/workspace/logs/train_$EXP.log
    [ -n "$POD_ID" ] || { echo "[ERROR] $STATE に Pod の id がない"; exit 1; }
    echo "[STOP  ] $LOG の学習の Pod: id $POD_ID  GPU=$GPU_USED  作った時刻 $(jq -r '.created_at // "?"' "$STATE")"
    wrap_fail stop "--stop で止めた"
fi

echo "[LAB   ] $(date '+%Y-%m-%d %H:%M:%S')  学習  GPU=$GPU_ID$([ $FALLBACK = 1 ] && echo " (借りられなければ ${FALLBACK_GPUS[*]})")"
echo "         --model $MODEL ($HF_REPO_IN → $HF_REPO_OUT)  データ $DATA_NAME"
echo "         $CONFIG --exp-name=$EXP  steps=$STEPS  init=$INIT  上限 $LIMIT_HOURS 時間  記録 $LOGDIR$([ $DRY_RUN = 1 ] && echo '  ドライラン (Pod は借りない)')"

# ---- 0. 送るものと Mac の道具 ----
stage "0. 送るものと Mac の道具を確かめる"
[ -f "$TAR" ] || early_fail "tar がない: $TAR"
[ -f "$ADD_TRIM" ] || early_fail "add_trim_config.py がない: $ADD_TRIM"
[ -f "$HERE/pi05/openpi_addon/pod_prep_trim_corr.sh" ] && [ -f "$HERE/pi05/openpi_addon/hf_upload_ckpt.py" ] \
    || early_fail "pi05/openpi_addon/ の pod_prep_trim_corr.sh・hf_upload_ckpt.py がない"
lab_have_hf_cli || early_fail "Mac に hf (huggingface_hub) がない. KIOXIA へ取れないので借りずに止めます"
if [ -e "$CKPT_ROOT/$CKPT_NAME/$LAST_STEP" ]; then
    echo "[WARN  ] $CKPT_ROOT/$CKPT_NAME/$LAST_STEP がもうあります. 学習のあと上書きはせず, 中身が同じかだけを見ます"
fi
echo "[DATA  ] tar $(stat -f '%z' "$TAR" 2>/dev/null || wc -c < "$TAR") バイト  ($TAR)"
prov_prepare && echo "[PROV  ] 始める時点のコード ($(jq -r '.code.commit // "?" | .[0:7]' <<<"$PROV_START")$(jq -r 'if .code.dirty then "・変更あり" else "" end' <<<"$PROV_START")) と入力 $(jq '.inputs | length' <<<"$PROV_START") つの指紋を測った"
lab_keep_awake
stage_done

# ---- 1. Pod を作る ----
stage "1. Pod を作る"
if [ -f "$STATE" ]; then
    OLD_ID=$(jq -r '.id // ""' "$STATE" 2>/dev/null)
    if [ -n "$OLD_ID" ] && [ -z "$(jq -r '.gone_confirmed // empty | select(. == true)' "$STATE" 2>/dev/null)" ]; then
        lab_pod_kind "$OLD_ID" || early_fail "前の学習の Pod (id $OLD_ID) の様子を API で確かめられない (HTTP ${RUNPOD_HTTP:-?}). 二重に借りないよう止めます"
        case "$LAB_POD_KIND" in
            running|stopped)
                early_fail "前の学習の Pod (id ${OLD_ID}・$LAB_POD_STATUS) が残っている. ./lab_train.sh --log $LOG --stop で持ち帰って消してから打ち直す" ;;
        esac
    fi
    lab_archive_state "$STATE" || early_fail "古い $STATE を残せない"
    echo "[POD   ] 前の記録は $LAB_ARCHIVED に残しました"
fi
others=$(runpod_api GET /pods) || early_fail "Pod の一覧を API から取れない (HTTP $RUNPOD_HTTP)"
lab_valid_pods "$others" || early_fail "Pod の一覧の答えが読めない. 二重に借りないよう止めます"
same=$(jq -r --arg name "p2m-train-$LOG" '.[] | select(.name == $name and .desiredStatus != "TERMINATED") | .id' <<<"$others")
[ -z "$same" ] || early_fail "記録にない p2m-train-$LOG の Pod (id $(echo $same)) が残っている. 二重に借りないよう止めます (./lab_stop.sh --check)"
n_others=$(jq 'length' <<<"$others")
if [ "$n_others" -gt 0 ]; then
    echo "[WARN  ] ほかに Pod が $n_others 個あります (消し忘れなら ./lab_stop.sh --check で確かめる):"
    jq -r '.[] | "         \(.id)  \(.name)  \(.desiredStatus)  \(.costPerHr) ドル/時"' <<<"$others"
fi
GPUS=("$GPU_ID")
if [ $FALLBACK = 1 ]; then for g in "${FALLBACK_GPUS[@]}"; do [ "$g" = "$GPU_ID" ] || GPUS+=("$g"); done; fi
for gi in "${!GPUS[@]}"; do
    g=${GPUS[$gi]}
    tries=$([ "$gi" = 0 ] && echo 5 || echo 2)
    echo "[POD   ] $g を借ります (空きがなければ $CREATE_WAIT 秒おきに $tries 回まで)"
    body=$(jq -n --arg name "p2m-train-$LOG" --arg gpu "$g" --arg tpl "$LAB_TEMPLATE_ID" --arg img "$LAB_IMAGE" \
              --arg pk "$(cat "$POD_SSH_KEY.pub")" \
              '{name:$name, templateId:$tpl, imageName:$img, gpuTypeIds:[$gpu], gpuCount:1,
                cloudType:"SECURE", containerDiskInGb:200, volumeInGb:0, ports:["22/tcp"],
                env:{PUBLIC_KEY:$pk}}')
    busy=0
    for attempt in $(seq 1 "$tries"); do
        out=$(mktemp)
        runpod_api POST /pods "$body" > "$out"; ok=$?        # $(...) に入れない (RUNPOD_HTTP を残す)
        resp=$(cat "$out"); rm -f "$out"
        if [ $ok = 0 ]; then
            POD_ID=$(jq -r '.id // ""' <<<"$resp")
            [ -n "$POD_ID" ] || early_fail "API の答えに Pod の id がない"
            COST_PER_HR=$(jq -r '.costPerHr // 0' <<<"$resp"); GPU_USED=$g
            break 2
        fi
        msg=$(jq -r '.error // .message // .' <<<"$resp" 2>/dev/null || echo "$resp")
        echo "[POD   ] 作れなかった ($g $attempt/$tries, HTTP $RUNPOD_HTTP): $msg"
        case "$msg" in
            *resources*|*"no longer any instances"*|*"not available"*|*"no instances"*) busy=1 ;;
            *) early_fail "Pod を作れない: $msg" ;;
        esac
        [ "$attempt" -lt "$tries" ] && { echo "[POD   ] $CREATE_WAIT 秒待って作り直します"; sleep "$CREATE_WAIT"; }
    done
    [ $busy = 1 ] && echo "[POD   ] $g は空きがない"
done
[ -n "$POD_ID" ] || early_fail "どの GPU (${GPUS[*]}) も借りられなかった"
POD_T0=$(date +%s)
jq -n --arg id "$POD_ID" --arg gpu "$GPU_USED" --arg log "$LOG" --argjson cost "$COST_PER_HR" --argjson t0 "$POD_T0" \
      --arg exp "$EXP" --argjson steps "$STEPS" --arg config "$CONFIG" --argjson lim "$LIMIT_SEC" \
      --arg init "$INIT" --arg lh "$LIMIT_HOURS" --arg ps "${PROV_START:-}" --arg model "$MODEL" \
      '{id:$id, kind:"train", model:$model, gpu:$gpu, log:$log, costPerHr:$cost, created_epoch:$t0,
        created_at:($t0|strflocaltime("%Y-%m-%d %H:%M:%S")), config:$config, exp:$exp, steps:$steps,
        limit_sec:$lim, limit_at:(($t0 + $lim)|strflocaltime("%Y-%m-%d %H:%M:%S")),
        init:$init, limit_hours:$lh, prov_start:(if $ps == "" then null else ($ps | fromjson) end)}' > "$STATE"
echo "[POD   ] 作った: id $POD_ID  GPU=$GPU_USED  1時間 $COST_PER_HR ドル  (記録 $STATE)"
echo "[POD   ] 上限: $(date -r $((POD_T0 + LIMIT_SEC)) '+%m-%d %H:%M') ($LIMIT_HOURS 時間) を過ぎたら, ログを持ち帰って消します"
stage_done

# ---- 2. SSH が通るまで待ち, Pod 側に上限の見張りを置く ----
stage "2. Pod が立ち上がるまで待つ"
for _ in $(seq 1 "$SSH_TRIES"); do lab_pod_addr && break; sleep "$SSH_WAIT"; done
[ -n "${POD_IP:-}" ] && [ -n "${SSH_PORT:-}" ] || wrap_fail prep "待っても SSH の IP とポートが出ない"
echo "[POD   ] SSH の IP とポートを API から受け取った (<POD_IP>:<SSH_PORT>)"
for _ in $(seq 1 "$LOGIN_TRIES"); do LAB_TO=$LAB_TO_SHORT pssh true >/dev/null && break; sleep "$LOGIN_WAIT"; done
LAB_TO=$LAB_TO_SHORT pssh 'echo "[POD   ] 入れた: $(hostname)"; nvidia-smi --query-gpu=name,memory.total --format=csv,noheader' \
    || wrap_fail prep "Pod に SSH で入れない"
WD_AT=$((POD_T0 + LIMIT_SEC + WD_MARGIN_SEC))
LAB_TO=$LAB_TO_SHORT pssh "mkdir -p /workspace/r2r2r /workspace/logs /workspace/ckpt /workspace/data" || wrap_fail prep "Pod にフォルダを作れない"
lab_watchdog_script | LAB_TO=$LAB_TO_SHORT pod_ssh "cat > /workspace/r2r2r/lab_watchdog.sh" || wrap_fail prep "見張りのスクリプトを置けない"
out=$(LAB_TO=$LAB_TO_SHORT pssh "$(lab_watchdog_start_cmd "$WD_AT")")      # 依頼 52: 初回はここで ssh が返らなかった
wd=$(grep -E '^WATCHDOG ' <<<"$out" | tail -1)
WD_OK=0
case "$wd" in
    "WATCHDOG run=1 pod=1 key=1"*)
        WD_OK=1
        echo "[WATCH ] Pod 側の見張りを置いた: $(date -r "$WD_AT" '+%m-%d %H:%M') を過ぎたら Pod が自分を消す (Pod 用の鍵. Mac の鍵は送っていない)" ;;
    *)
        # 止めずに続ける (Pod 用の鍵が入らない Pod だと毎回学習できなくなるため). 画面・通知・まとめで目立たせる
        echo "[WARN  ] Pod 側の見張りが働かないかもしれない (${wd:-答えなし}). Mac 側の上限だけで消します. Mac を止めないでください" ;;
esac
save_state '. + {watchdog_epoch:$w, watchdog:$s}' --argjson w "$WD_AT" --arg s "${wd:-なし}"
stage_done

# ---- 3. 準備 (lab_start.sh 3a〜3f と同じ. サーバーは起動しない) ----
stage "3a. スクリプトを送る"
LAB_TO=$LAB_TO_LONG pod_scp -r "$HERE/pi05" "root@$POD_IP:/workspace/r2r2r/" || wrap_fail prep "pi05/ を送れない"
pod_scp "$ADD_TRIM" "root@$POD_IP:/workspace/r2r2r/" || wrap_fail prep "add_trim_config.py を送れない"
stage_done
stage "3b. huggingface_hub を入れ, 鍵を渡す"
LAB_TO=$LAB_TO_LONG pssh "$(lab_cmd_hf_install)" || wrap_fail prep "huggingface_hub を入れられない"
lab_send_hf_token || wrap_fail prep "Hugging Face の鍵を渡せない"
LAB_TO=$LAB_TO_SHORT pssh 'hf auth whoami' || wrap_fail prep "Hugging Face に入れない (鍵が違う: $HF_TOKEN_FILE を確かめる)"
stage_done
stage "3c. 続きから学習するモデルを取る ($HF_REPO_IN → $CKPT_IN)"
LAB_TO=$LAB_TO_LONG pssh "$(lab_cmd_ckpt "$HF_REPO_IN" "$CKPT_IN")" || wrap_fail prep "hf download に失敗"
stage_done
stage "3d. openpi を取る ($LAB_OPENPI_COMMIT)"
LAB_TO=$LAB_TO_LONG pssh "$(lab_cmd_openpi "$LAB_OPENPI_COMMIT")" || wrap_fail prep "openpi を取れない"
stage_done
stage "3e. uv sync"
LAB_TO=$LAB_TO_LONG pssh "$(lab_cmd_uv_sync)" || wrap_fail prep "uv sync に失敗"
stage_done
stage "3f. addon・設定・assets"
pssh "$(lab_cmd_addon_train "$MODEL")" || wrap_fail prep "addon・設定・assets に失敗"
# 由来 (依頼 54): 学習に使う Pod の openpi のコミットと変更の有無 (addon・設定を入れたあと). 取れなくても続ける
# git status が最後まで通ったときだけ OPENPI_END を出す. ssh の終了コードと両方を見る (途中で切れたものを「変更なし」にしない. codex R1)
out=$(LAB_TO=$LAB_TO_SHORT pssh ": lt=openpi_rev; cd /workspace/openpi && c=\$(git rev-parse HEAD) && s=\$(git status --porcelain=v1 --untracked-files=all) && echo \"OPENPI \$c\" && printf '%s\\n' \"\$s\" | sed -n 's/^\\(..\\) /OPF \\1 /p' && echo OPENPI_END"); orc=$?
if oc=$(sed -n 's/^OPENPI \([0-9a-f]\{40\}\)$/\1/p' <<<"$out" | tail -1) && [ -n "$oc" ] && [ $orc = 0 ] && grep -qx OPENPI_END <<<"$out"; then
    OPENPI_JSON=$(grep '^OPF ' <<<"$out" | head -200 | cut -c8- | jq -Rn --arg c "$oc" --arg want "$LAB_OPENPI_COMMIT" \
        '[inputs] as $f | [{name: "openpi", commit: $c, dirty: ($f | length > 0), dirty_files: $f, expected_commit: $want}]') \
        || OPENPI_JSON='[]'
    echo "[PROV  ] openpi $oc ($(jq -r '.[0].dirty_files | length' <<<"$OPENPI_JSON") ファイル変更・addon と設定を入れたあと)"
else
    OPENPI_JSON=$(jq -n --arg want "$LAB_OPENPI_COMMIT" '[{name: "openpi", commit: null, dirty: null, expected_commit: $want, error: "Pod から読めなかった"}]')
    echo "[WARN  ] Pod の openpi のコミットを読めなかった (由来には「読めなかった」と残す)"
fi
save_state '. + {openpi:$o}' --argjson o "$OPENPI_JSON"       # --stop で戻す (codex R2)
stage_done

# ---- 4. データを送り, 学習の準備 ----
stage "4a. tar を送る"
LAB_TO=$LAB_TO_LONG pod_scp "$TAR" "root@$POD_IP:/workspace/data/" || wrap_fail prep "tar を送れない"
LAB_TO=$LAB_TO_SHORT pssh "ls -l /workspace/data/$(basename "$TAR")" || wrap_fail prep "送った tar が見えない"
stage_done
stage "4b. pod_prep_trim_corr.sh $STEPS $INIT $PREP_VARIANT"
LAB_TO=$LAB_TO_LONG pssh "set -o pipefail; bash /workspace/r2r2r/pi05/openpi_addon/pod_prep_trim_corr.sh $STEPS $INIT $PREP_VARIANT 2>&1 | tee $PREP_LOG" \
    || wrap_fail prep "pod_prep_trim_corr.sh に失敗 (上の出力を見る)"
stage_done

# ---- 5. 学習を始める ----
stage "5. 学習を始める"
out=$(LAB_TO=$LAB_TO_SHORT pssh "$(lab_train_start_cmd "$CONFIG" "$EXP" "$TRAIN_LOG" "$GPU_LOG" "$RUN_LOG" "$EXIT_FILE")")
grep -q '^STARTED' <<<"$out" || wrap_fail prep "学習を始められない ($(tail -3 <<<"$out" | tr '\n' ' '))"
TRAIN_T0=$(date +%s)
save_state '. + {train_started_epoch:$t, train_started_at:($t|strflocaltime("%Y-%m-%d %H:%M:%S"))}' --argjson t "$TRAIN_T0"
echo "[TRAIN ] 学習を始めました (Pod の中で setsid nohup. Mac とのつながりが切れても続く)"
echo "         ログ $TRAIN_LOG  GPU $GPU_LOG (30 秒おき)"
stage_done

# ---- 6. 見守る ----
stage "6. 学習を見守る ($POLL_MIN 分おき)" nolimit       # 上限は下のループで, 終わっているかを見てから
told=0; ssh_ng=0
while :; do
    # 上限を超えていても, 先に様子を見る (確かめる間に学習が終わっていれば, 消さずに持ち帰る)
    out=$(LAB_TO=$LAB_TO_SHORT pssh "$(lab_train_status_cmd "$TRAIN_LOG" "$EXIT_FILE" "$POD_EXP_DIR" "$STEPS")")
    now=$(date +%s)
    if lab_train_parse_status "$out" && [ "$S_EXIT" = 0 ] && [ "$S_CKPT" = "$LAST_STEP" ] && [ "$S_CM" = 1 ]; then
        :       # 終わっている → 下で成功へ
    elif [ $((now - POD_T0)) -ge "$LIMIT_SEC" ]; then
        wrap_fail timeout "上限 $LIMIT_HOURS 時間を超えた (学習を始めてから $(hm $((now - TRAIN_T0))))"
    fi
    if lab_train_parse_status "$out"; then
        ssh_ng=0
        step=$S_BAR; [ "$step" = - ] && step=$S_LOGSTEP
        printf '[TRAIN ] %s  step %s/%s  loss %s%s  経過 %s  上限まで %s\n' "$(date '+%H:%M')" "$step" "$STEPS" "$S_LOSS" \
            "$([ "$S_LOGSTEP" != - ] && echo " (Step $S_LOGSTEP の行)")" "$(hm $((now - TRAIN_T0)))" "$(hm $((POD_T0 + LIMIT_SEC - now)))"
        if [ "$S_EXIT" != - ]; then
            if [ "$S_EXIT" = 0 ] && [ "$S_CKPT" = "$LAST_STEP" ] && [ "$S_CM" = 1 ]; then
                stage_done; wrap_success
            fi
            wrap_fail crash "学習が終わったが, 終了コード ${S_EXIT}・最後のチェックポイント $S_CKPT (待っていたのは $LAST_STEP)"
        fi
        [ "$S_RUN" = 0 ] && wrap_fail crash "学習のプロセスが終わりの印を書かずに消えた"
        if [ "$S_AGE" != - ] && [ "$S_AGE" -ge "$STALL_SEC" ]; then
            wrap_fail stall "学習のログが $((S_AGE / 60)) 分書かれていない (止まった. $STALL_MIN 分で止めると決めている)"
        fi
        if [ $told = 0 ]; then
            told=1
            echo
            echo "========================================================================"
            echo " 学習は Pod の中で回っています. ここまで来れば寝てかまいません."
            echo " このターミナルは閉じない・Mac のふたは閉じない (電源につないだまま)."
            echo " 終わると通知が出ます. まとめは $LOGDIR/result_*.txt"
            [ $WD_OK = 1 ] || echo " ※ Pod 側の見張りがない. Mac が止まると Pod が残る (朝に ./lab_stop.sh --check で確かめる)"
            echo "========================================================================"
            echo
        fi
    else
        ssh_ng=$((ssh_ng + 1))
        echo "[WARN  ] $(date '+%H:%M') Pod の様子を見られない ($ssh_ng 回続けて)"
        if lab_pod_kind "$POD_ID" && [ "$LAB_POD_KIND" != running ]; then
            wrap_fail gone "Pod が動いていない (API の状態 ${LAB_POD_STATUS:-なし})"
        fi
        lab_pod_addr || true      # IP とポートが変わっていれば取り直す
    fi
    sleep "$POLL_SEC"
done
}

main "$@"; exit $?
