129 lines
3.5 KiB
Python
129 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import pickle
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
CURRENT_CACHE_VERSIONS = {
|
|
"next_step": 3,
|
|
"all_future": 5,
|
|
}
|
|
|
|
|
|
def infer_cache_kind(path: Path) -> str | None:
|
|
name = path.name
|
|
for kind in CURRENT_CACHE_VERSIONS:
|
|
marker = f"_{kind}_cache_"
|
|
if marker in name:
|
|
return kind
|
|
return None
|
|
|
|
|
|
def read_cache_version(path: Path) -> int | None:
|
|
try:
|
|
with path.open("rb") as f:
|
|
payload: Any = pickle.load(f)
|
|
except Exception:
|
|
return None
|
|
if not isinstance(payload, dict):
|
|
return None
|
|
version = payload.get("_cache_version")
|
|
try:
|
|
return int(version)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def should_remove(path: Path, remove_all: bool) -> tuple[bool, str]:
|
|
kind = infer_cache_kind(path)
|
|
if kind is None:
|
|
return False, "not a DeepHealth dataset cache"
|
|
|
|
if remove_all:
|
|
return True, "remove all dataset caches"
|
|
|
|
version = read_cache_version(path)
|
|
expected = CURRENT_CACHE_VERSIONS[kind]
|
|
if version is None:
|
|
return True, f"{kind} cache is unreadable or missing _cache_version"
|
|
if version != expected:
|
|
return True, f"{kind} cache version {version} != current {expected}"
|
|
return False, f"{kind} cache version {version} is current"
|
|
|
|
|
|
def iter_cache_files(root: Path, recursive: bool):
|
|
pattern = "**/*_cache_*.pkl" if recursive else "*_cache_*.pkl"
|
|
yield from root.glob(pattern)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description=(
|
|
"Remove obsolete DeepHealth dataset cache files. Defaults to dry-run; "
|
|
"pass --apply to delete."
|
|
)
|
|
)
|
|
parser.add_argument(
|
|
"--data_dir",
|
|
type=Path,
|
|
default=Path("."),
|
|
help="Directory containing dataset cache files, usually the data_prefix directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--recursive",
|
|
action="store_true",
|
|
help="Search recursively under data_dir.",
|
|
)
|
|
parser.add_argument(
|
|
"--all",
|
|
action="store_true",
|
|
help="Remove all recognized DeepHealth dataset caches, including current ones.",
|
|
)
|
|
parser.add_argument(
|
|
"--apply",
|
|
action="store_true",
|
|
help="Actually delete files. Without this flag, only prints what would be removed.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
root = args.data_dir.expanduser().resolve()
|
|
if not root.exists() or not root.is_dir():
|
|
raise SystemExit(f"data_dir is not a directory: {root}")
|
|
|
|
files = sorted(p for p in iter_cache_files(root, args.recursive) if p.is_file())
|
|
if not files:
|
|
print(f"No *_cache_*.pkl files found under {root}")
|
|
return
|
|
|
|
kept = 0
|
|
removed = 0
|
|
candidates = 0
|
|
for path in files:
|
|
remove, reason = should_remove(path, remove_all=bool(args.all))
|
|
rel = os.path.relpath(path, root)
|
|
if remove:
|
|
candidates += 1
|
|
action = "DELETE" if args.apply else "WOULD_DELETE"
|
|
print(f"{action}\t{rel}\t{reason}")
|
|
if args.apply:
|
|
path.unlink()
|
|
removed += 1
|
|
else:
|
|
kept += 1
|
|
print(f"KEEP\t{rel}\t{reason}")
|
|
|
|
if args.apply:
|
|
print(f"Removed {removed} cache file(s); kept {kept}.")
|
|
else:
|
|
print(
|
|
f"Dry run: {candidates} cache file(s) would be removed; kept {kept}. "
|
|
"Re-run with --apply to delete."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|