#!/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 を作った時刻から数える). 超えたらログを持ち帰って消す
#   --steps / --init / --exp-name   学習 (既定 1500 / /workspace/ckpt/3999/params / trim_corr_v3)
#   --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
#
# 鍵 (Runpod・Hugging Face) は KIOXIA の keys/ から読み, 画面・記録には出さない (lab_lib.sh).
# Pod 側の見張りには Mac の鍵を送らない.

set -u
set -o pipefail

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

GPU=A40; FALLBACK=1; LIMIT_HOURS=5; STEPS=1500; INIT=/workspace/ckpt/3999/params; EXP=trim_corr_v3
POLL_MIN=5; STALL_MIN=30; LOG=""; DRY_RUN=0; MODE=train
while [ $# -gt 0 ]; do
    case "$1" in
        --log) LOG="${2:-}"; shift 2 ;;
        --gpu) GPU="${2:-}"; shift 2 ;;
        --no-fallback) FALLBACK=0; shift ;;
        --limit-hours) LIMIT_HOURS="${2:-}"; shift 2 ;;
        --steps) STEPS="${2:-}"; shift 2 ;;
        --init) INIT="${2:-}"; shift 2 ;;
        --exp-name) EXP="${2:-}"; shift 2 ;;
        --poll-min) POLL_MIN="${2:-}"; shift 2 ;;
        --stall-min) STALL_MIN="${2:-}"; shift 2 ;;
        --stop) MODE=stop; shift ;;
        --dry-run) DRY_RUN=1; shift ;;
        -h|--help) sed -n '2,29p' "$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; }
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"
DATA_NAME=mecharm_redblock_box_teach_v1_trim2_corr_v3
TAR="$DATA_ROOT/datasets/$DATA_NAME/${DATA_NAME}_raw_episodes.tar"
ADD_TRIM="$DATA_ROOT/lab/Log11/pod/add_trim_config.py"
CONFIG=pi05_mecharm_trim_corr_lora
HF_REPO_IN=PixelstoMotors/pi05-mecharm-trim-lora-v1           # 続きから学習する 3999 (lab_start.sh と同じ)
HF_REPO_OUT=PixelstoMotors/pi05-mecharm-trim-corr-lora-v1     # 上げる先 (非公開). 名前は checkpoints/ のフォルダ名に合わせる
CKPT_NAME=pi05_mecharm_trim_corr_lora_v1                      # KIOXIA の checkpoints/<これ>/<ステップ>/
CKPT_IN=/workspace/ckpt/3999
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

# ---- 段階ごとの秒数 ----
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)); }

# ---- 通知とまとめ ----
RESULT="$LOGDIR/result_$(date +%Y%m%d_%H%M%S).txt"
finish() {                  # finish 終了コード 通知の文 (まとめを書き, 通知を出して終わる)
    local rc="$1" msg="$2"
    FINISHED=1
    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}」で止まりました: $*"
    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 && 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=$(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"
    if [ -z "$POD_ID" ]; then RESULT_LINES="Pod は借りていません"; finish 1 "失敗: $why (Pod は借りていません)"; fi
    # 学習を止めてからログを持ち帰る (消す前に書き終わったログを取る. 止められなくても続ける)
    if [ "$kind" != crash ] && [ -n "$TRAIN_T0" ]; then
        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)
    if ! lab_list_files "$1" > "$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"
        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 の印 (自分で作った一時フォルダの中だけ)
    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
    # ._* (exFAT が付ける)・.DS_Store なども移す (残ると rmdir が失敗する. codex R3)
    if ! ( shopt -s dotglob nullglob; mv "$tmp"/* "$dest"/ ) || ! rmdir "$tmp"; then KIOXIA_MSG="$tmp から $dest へ移せない (途中まで移したかもしれない)"; return 1; fi
    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 側と名前・バイト数が全部合った)"
}

extend_watchdog() {         # Pod 側の見張りの時刻を, 今から WRAP_EXTEND_SEC あとへ延ばす (片付けの途中で消えないように)
    local t=$(( $(date +%s) + WRAP_EXTEND_SEC ))
    if 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
    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=$(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=$(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
    fi
    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")
    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 "         $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)"
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" \
      '{id:$id, kind:"train", 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"))}' > "$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 pssh true >/dev/null && break; sleep "$LOGIN_WAIT"; done
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))
pssh "mkdir -p /workspace/r2r2r /workspace/logs /workspace/ckpt /workspace/data" || wrap_fail prep "Pod にフォルダを作れない"
lab_watchdog_script | pod_ssh "cat > /workspace/r2r2r/lab_watchdog.sh" || wrap_fail prep "見張りのスクリプトを置けない"
out=$(pssh "$(lab_watchdog_start_cmd "$WD_AT")")
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. スクリプトを送る"
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 を入れ, 鍵を渡す"
pssh "$(lab_cmd_hf_install)" || wrap_fail prep "huggingface_hub を入れられない"
lab_send_hf_token || wrap_fail prep "Hugging Face の鍵を渡せない"
pssh 'hf auth whoami' || wrap_fail prep "Hugging Face に入れない (鍵が違う: $HF_TOKEN_FILE を確かめる)"
stage_done
stage "3c. 続きから学習するモデルを取る ($HF_REPO_IN → $CKPT_IN)"
pssh "$(lab_cmd_ckpt "$HF_REPO_IN" "$CKPT_IN")" || wrap_fail prep "hf download に失敗"
stage_done
stage "3d. openpi を取る ($LAB_OPENPI_COMMIT)"
pssh "$(lab_cmd_openpi "$LAB_OPENPI_COMMIT")" || wrap_fail prep "openpi を取れない"
stage_done
stage "3e. uv sync"
pssh "$(lab_cmd_uv_sync)" || wrap_fail prep "uv sync に失敗"
stage_done
stage "3f. addon・設定・assets"
pssh "$(lab_cmd_addon)" || wrap_fail prep "addon・設定・assets に失敗"
stage_done

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

# ---- 5. 学習を始める ----
stage "5. 学習を始める"
out=$(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=$(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 $?