#!/usr/bin/env bash
# Keep optimizer.pt only in the newest checkpoint of this run. 10 checkpoints x 8.55 GB of
# optimizer state is 77 GB that only matters for resuming, and only the latest can resume.
# Never touches the highest-numbered checkpoint (it may still be mid-write) or the weights.
set -u
R=/data/minha/pose6daug/train_robocasa/runs/pnpcountertocab_mimicgen_natural_256_short_newemb_from60k_gbs64_5000steps_20260922_020158
LOG=/tmp/claude-1001/-home-nvidia/68c27ffb-706b-4cb3-8bf2-4618a200d378/scratchpad/prune_optimizers.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                 # newest: leave it alone
            [ -f "$d/model.safetensors.index.json" ] || continue # incomplete: leave it alone
            if [ -f "$d/optimizer.pt" ]; then
                sz=$(du -h "$d/optimizer.pt" | cut -f1)
                rm -f "$d/optimizer.pt"
                echo "$(date -u +%H:%M:%S) removed checkpoint-$step/optimizer.pt ($sz)" >> "$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
# final pass once training has exited
latest=$(ls -d "$R"/checkpoint-* 2>/dev/null | sed 's/.*checkpoint-//' | sort -n | tail -1)
for d in "$R"/checkpoint-*; do
    [ "${d##*checkpoint-}" = "$latest" ] && continue
    [ -f "$d/optimizer.pt" ] && rm -f "$d/optimizer.pt" && echo "$(date -u +%H:%M:%S) final: removed ${d##*/}/optimizer.pt" >> "$LOG"
done
echo "$(date -u +%H:%M:%S) pruner done; optimizer kept only in checkpoint-$latest" >> "$LOG"