openpi/scripts/prune_stepcmp_checkpoints.py

#!/usr/bin/env python3

import argparse
import pathlib
import shutil
import time
from datetime import datetime, timezone


def utc_ts() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")


def prune_once(roots: list[pathlib.Path], keep_steps: set[str]) -> int:
    removed = 0
    for root in roots:
        if not root.is_dir():
            continue
        for child in root.iterdir():
            if not child.is_dir():
                continue
            if child.name.startswith("tmp_"):
                continue
            if not child.name.isdigit():
                continue
            if child.name in keep_steps:
                continue
            shutil.rmtree(child, ignore_errors=True)
            print(f"[{utc_ts()}] pruned {child}", flush=True)
            removed += 1
    return removed


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--interval-seconds", type=int, default=30)
    parser.add_argument("--keep-steps", nargs="+", default=["100", "500", "2000"])
    parser.add_argument("roots", nargs="+")
    args = parser.parse_args()

    roots = [pathlib.Path(root) for root in args.roots]
    keep_steps = set(args.keep_steps)
    print(
        f"[{utc_ts()}] retention pruner started interval_s={args.interval_seconds} keep_steps={sorted(keep_steps)}",
        flush=True,
    )
    while True:
        prune_once(roots, keep_steps)
        time.sleep(args.interval_seconds)


if __name__ == "__main__":
    main()