#!/bin/bash
# B1 — pi0.5 LoRA SFT on turning_on_radio, PyTorch + peft, with live HF uploads.
#
# Validated by job 2135: LoRA applies (trainable 758.3M / 3.682B = 20.60%, the bulk
# being the fully-trainable action expert — RLinf's recipe wraps only the VLM), 200
# steps ran clean, loss trended down, TRAIN_PT_rc=0.
#
# save_interval is 2,000 steps (~2h at the measured 3.4 s/it), NOT the config's
# 10,000: training is moving to a B200 host in ~8h, and at 10,000 the first
# checkpoint would land after ~13h — i.e. nothing to carry over. 2,000 gives four
# checkpoints inside that window.
#
# Uploads run in a *separate* process watching the output dir, so a Hub outage or
# rate limit cannot kill the training job.
#SBATCH --job-name=b1_sft
#SBATCH --partition=baram
#SBATCH --gres=gpu:baram:2,gpu:debug:2
#SBATCH --cpus-per-task=32
#SBATCH --mem=110G
#SBATCH --time=4-00:00:00
#SBATCH --output=/data2/hyeongjinkim/behavior-challenge/logs/b1_sft_%j.log
#SBATCH --error=/data2/hyeongjinkim/behavior-challenge/logs/b1_sft_%j.log

set -uo pipefail
# Invalid HF_TOKEN in the login env shadows ~/.cache/huggingface/token.
unset HF_TOKEN

ROOT=/data2/hyeongjinkim/behavior-challenge
CFG=${CFG:-pi05_b1k_pytorch_local}
EXP_NAME=${EXP_NAME:-b1_sft_$(date +%Y%m%d)}
BATCH=${BATCH:-16}
STEPS=${STEPS:-50000}
SAVE_INTERVAL=${SAVE_INTERVAL:-2000}
LOG_INTERVAL=${LOG_INTERVAL:-100}
HF_REPO=${HF_REPO:-madokalif/pi05-b1k-turning-on-radio-lora}
NGPU=$(echo "$CUDA_VISIBLE_DEVICES" | tr ',' '\n' | wc -l)

export OPENPI_LORA=${OPENPI_LORA:-1}
export OPENPI_LORA_RANK=${OPENPI_LORA_RANK:-32}

CKPT_DIR=$ROOT/repos/openpi-sft/outputs/checkpoints/$CFG/$EXP_NAME

echo "=== job $SLURM_JOB_ID  $EXP_NAME ==="
echo "cfg=$CFG gpus=$NGPU batch=$BATCH steps=$STEPS save_every=$SAVE_INTERVAL"
echo "lora=$OPENPI_LORA rank=$OPENPI_LORA_RANK"
echo "hf_repo=$HF_REPO"
echo "ckpt_dir=$CKPT_DIR"
nvidia-smi --query-gpu=index,name,memory.total --format=csv

cd "$ROOT/repos/openpi-sft"

# Uploader first, so it is already watching when the first checkpoint lands.
.venv/bin/python "$ROOT/upload_ckpt_hf.py" \
    --watch-dir "$CKPT_DIR" \
    --repo-id "$HF_REPO" \
    --interval 300 > "$ROOT/logs/hf_upload_${SLURM_JOB_ID}.log" 2>&1 &
UPLOADER_PID=$!
trap 'echo "stopping uploader $UPLOADER_PID"; kill $UPLOADER_PID 2>/dev/null' EXIT
echo "uploader pid=$UPLOADER_PID -> logs/hf_upload_${SLURM_JOB_ID}.log"

.venv/bin/torchrun --standalone --nnodes=1 --nproc_per_node="$NGPU" \
    scripts/train_pytorch.py "$CFG" \
    --exp_name="$EXP_NAME" \
    --overwrite \
    --batch-size="$BATCH" \
    --num-train-steps="$STEPS" \
    --save-interval="$SAVE_INTERVAL" \
    --log-interval="$LOG_INTERVAL"
rc=$?
echo "TRAIN_B1_SFT_rc=$rc"

# Catch anything written after the uploader's last scan before the job exits.
echo "--- final upload sweep ---"
.venv/bin/python "$ROOT/upload_ckpt_hf.py" --watch-dir "$CKPT_DIR" --repo-id "$HF_REPO" --once 2>&1 | tail -5
exit $rc