social-data-attribution-w2/scripts/unlearning/prune_intermediates.py

#!/usr/bin/env python3
"""Keep only the latest checkpoint-N per faithful multiseed dir; delete older
intermediates. Insurance against the scratch-quota OSError-122 job-kill during a
large overnight grid. Safe for live jobs: the trainer never reads old checkpoints,
and the highest-numbered checkpoint (current/just-written) is always kept.
"""

from pathlib import Path
import shutil

MS = Path.home() / "scratch/n16_selectivity/multiseed"
pruned = 0
for d in MS.glob("faithful_*"):
    if not d.is_dir():
        continue
    cks = [c for c in d.glob("checkpoint-*") if c.name.split("-")[-1].isdigit()]
    if len(cks) <= 1:
        continue
    latest = max(cks, key=lambda c: int(c.name.split("-")[-1]))
    for c in cks:
        if c != latest:
            shutil.rmtree(c, ignore_errors=True)
            pruned += 1
print(f"pruned {pruned} intermediate checkpoints")