#!/usr/bin/env bash
# HF Trainer only has one save_steps, so save every 500 and prune afterwards:
# keep every 500 up to 5000, then only multiples of 5000. Optimizer state lives in the
# newest checkpoint alone -- 8.5 GB each and only the latest can resume anyway.
# Never touches the highest-numbered checkpoint (it may still be mid-write).
set -u
R=/data/minha/pose6daug/train_robocasa/runs/pnpcountertocab_actaug_256_newemb_from60k_gbs64_30000steps_20260922_140904
LOG=/tmp/claude-1001/-home-nvidia/68c27ffb-706b-4cb3-8bf2-4618a200d378/scratchpad/keep_actaug.log
while true; do
    latest=$(ls -d "$R"/checkpoint-* 2>/dev/null | sed 's/.*checkpoint-//' | sort -n | tail -1)
    if [ -n "${latest:-}" ]; then
        for d in "$R"/checkpoint-*; do
            step=${d##*checkpoint-}
            [ "$step" = "$latest" ] && continue
            [ -f "$d/model.safetensors.index.json" ] || continue
            if [ "$step" -gt 5000 ] && [ $((step % 5000)) -ne 0 ]; then
                rm -rf "$d"; echo "$(date -u +%H:%M:%S) dropped checkpoint-$step (5k 간격 아님)" >> "$LOG"
                continue
            fi
            if [ -f "$d/optimizer.pt" ]; then
                rm -f "$d/optimizer.pt"; echo "$(date -u +%H:%M:%S) optimizer 제거 checkpoint-$step" >> "$LOG"
            fi
        done
    fi
    ps -eo args | grep -q "[g]r00t_finetune.py" || { sleep 120; ps -eo args | grep -q "[g]r00t_finetune.py" || break; }
    sleep 60
done
latest=$(ls -d "$R"/checkpoint-* 2>/dev/null | sed 's/.*checkpoint-//' | sort -n | tail -1)
for d in "$R"/checkpoint-*; do
    step=${d##*checkpoint-}; [ "$step" = "$latest" ] && continue
    [ "$step" -gt 5000 ] && [ $((step % 5000)) -ne 0 ] && rm -rf "$d" && continue
    [ -f "$d/optimizer.pt" ] && rm -f "$d/optimizer.pt"
done
echo "$(date -u +%H:%M:%S) pruner done; optimizer only in checkpoint-$latest" >> "$LOG"