#!/usr/bin/env bash
# CP_CKPT_KEEP prunes .ckpt files but NOT their .zero_g1 sidecars. At
# CP_CKPT_FREQ=25 that leaked ~49 GB of orphans and filled a 193 GB disk, after
# which every save failed with only a [WARN] while training carried on. This
# deletes sidecars and .tmp files that no surviving .ckpt refers to.
set -uo pipefail
cd "$(dirname "$0")"
L=b3long_logs/janitor.log
while true; do
  freed=0
  for r in 0 1; do
    d=checkpoints_bluscriptcp/rank_$r
    [ -d "$d" ] || continue
    for f in "$d"/*.zero_g1; do
      [ -e "$f" ] || continue
      stem=$(basename "$f" .zero_g1)
      # keep if a matching .ckpt exists, or if a preserved copy references it
      [ -f "$d/$stem.ckpt" ] && continue
      ls "$d"/*"$stem".ckpt >/dev/null 2>&1 && continue
      sz=$(stat -c%s "$f" 2>/dev/null || echo 0)
      rm -f "$f" && freed=$(( freed + sz ))
    done
    # .tmp older than 10 min is a failed save, never an in-flight one
    find "$d" -name '*.tmp' -mmin +10 -delete 2>/dev/null
  done
  avail=$(df -BG --output=avail /root | tail -1 | tr -dc '0-9')
  [ "$freed" -gt 0 ] && echo "[$(date -u +%H:%M:%S)] freed $(( freed/1000000 )) MB, ${avail}G avail" >> "$L"
  [ "$avail" -lt 25 ] && echo "[$(date -u +%H:%M:%S)] WARNING only ${avail}G avail -- a save needs 17.3G" >> "$L"
  sleep 300
done