Add AGILLM4 live local prune guard

agillm4_local_prune_guard.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from dataclasses import asdict, dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Iterable
+
+
+def iso_now() -> str:
+    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+
+
+@dataclass(frozen=True)
+class TrainerIdentity:
+    pid: int
+    start_ticks: int
+    cmdline: str
+
+
+def read_cmdline(proc_dir: Path) -> list[str]:
+    try:
+        return [
+            part.decode("utf-8", errors="replace")
+            for part in (proc_dir / "cmdline").read_bytes().split(b"\0")
+            if part
+        ]
+    except Exception:
+        return []
+
+
+def read_start_ticks(proc_dir: Path) -> int | None:
+    try:
+        stat = (proc_dir / "stat").read_text(encoding="utf-8", errors="replace")
+        rest = stat[stat.rfind(")") + 2 :].split()
+        return int(rest[19])
+    except Exception:
+        return None
+
+
+def find_trainers(script_name: str, required_arg: str) -> list[TrainerIdentity]:
+    matches: list[TrainerIdentity] = []
+    for proc_dir in Path("/proc").iterdir():
+        if not proc_dir.name.isdigit():
+            continue
+        args = read_cmdline(proc_dir)
+        if not args or required_arg not in args:
+            continue
+        if not any(Path(arg).name == script_name for arg in args):
+            continue
+        start_ticks = read_start_ticks(proc_dir)
+        if start_ticks is None:
+            continue
+        matches.append(TrainerIdentity(int(proc_dir.name), start_ticks, " ".join(args)))
+    return sorted(matches, key=lambda item: item.pid)
+
+
+def path_artifacts(path: Path) -> list[Path]:
+    return [
+        path,
+        path.with_suffix(".sha256"),
+        path.with_suffix(path.suffix + ".upload.sha256"),
+        path.with_suffix(path.suffix + ".dtmp"),
+    ]
+
+
+def delete_artifacts(path: Path) -> list[str]:
+    deleted: list[str] = []
+    for artifact in path_artifacts(path):
+        try:
+            if artifact.exists():
+                artifact.unlink()
+                deleted.append(str(artifact))
+        except Exception as exc:
+            deleted.append(f"FAILED {artifact}: {type(exc).__name__}: {exc}")
+    return deleted
+
+
+def matching_checkpoints(save_dir: Path, pattern: str, excluded_name_parts: Iterable[str]) -> list[Path]:
+    excluded = tuple(excluded_name_parts)
+    return sorted(
+        [
+            path
+            for path in save_dir.glob(pattern)
+            if path.is_file()
+            and path.stat().st_size > 0
+            and not any(part in path.name for part in excluded)
+        ],
+        key=lambda path: path.stat().st_mtime,
+    )
+
+
+def prune_to_keep(
+    save_dir: Path,
+    pattern: str,
+    keep: int,
+    excluded_name_parts: Iterable[str] = (),
+) -> list[str]:
+    if keep < 0:
+        return []
+    checkpoints = matching_checkpoints(save_dir, pattern, excluded_name_parts)
+    victims = checkpoints[: max(0, len(checkpoints) - keep)]
+    deleted: list[str] = []
+    for victim in victims:
+        deleted.extend(delete_artifacts(victim))
+    return deleted
+
+
+def prune_stale_tmp(save_dir: Path, older_than_sec: int) -> list[str]:
+    if older_than_sec <= 0:
+        return []
+    cutoff = time.time() - older_than_sec
+    deleted: list[str] = []
+    for tmp in save_dir.glob("*.dtmp"):
+        try:
+            if tmp.is_file() and tmp.stat().st_mtime < cutoff:
+                tmp.unlink()
+                deleted.append(str(tmp))
+        except Exception as exc:
+            deleted.append(f"FAILED {tmp}: {type(exc).__name__}: {exc}")
+    return deleted
+
+
+def prune_orphan_sidecars(save_dir: Path) -> list[str]:
+    deleted: list[str] = []
+    sidecars = sorted(set(save_dir.glob("*.sha256")) | set(save_dir.glob("*.pt.upload.sha256")))
+    for sidecar in sidecars:
+        name = sidecar.name
+        if name.endswith(".pt.upload.sha256"):
+            base = save_dir / name[: -len(".upload.sha256")]
+        elif name.endswith(".sha256"):
+            base = save_dir / name[: -len(".sha256")]
+        else:
+            continue
+        try:
+            if not base.exists():
+                sidecar.unlink()
+                deleted.append(str(sidecar))
+        except Exception as exc:
+            deleted.append(f"FAILED {sidecar}: {type(exc).__name__}: {exc}")
+    return deleted
+
+
+def free_gb(path: Path) -> float:
+    stat = os.statvfs(path)
+    return (stat.f_bavail * stat.f_frsize) / (1024**3)
+
+
+def prune_uploaded_delta_for_reserve(save_dir: Path, reserve_free_gb: float) -> list[str]:
+    if reserve_free_gb <= 0 or free_gb(save_dir) >= reserve_free_gb:
+        return []
+    deltas = matching_checkpoints(save_dir, "*_delta_step*.pt", ())
+    deleted: list[str] = []
+    for delta in deltas:
+        upload_marker = delta.with_suffix(delta.suffix + ".upload.sha256")
+        if not upload_marker.exists():
+            continue
+        deleted.extend(delete_artifacts(delta))
+        if free_gb(save_dir) >= reserve_free_gb:
+            break
+    return deleted
+
+
+def write_state(path: Path, payload: dict) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
+
+
+def same_trainer(current: list[TrainerIdentity], initial: TrainerIdentity) -> bool:
+    return any(item.pid == initial.pid and item.start_ticks == initial.start_ticks for item in current)
+
+
+def run_once(args: argparse.Namespace, initial: TrainerIdentity) -> dict:
+    current = find_trainers(args.script_name, args.required_arg)
+    if not same_trainer(current, initial):
+        return {
+            "checked_at": iso_now(),
+            "action": "exit_training_restarted",
+            "initial_trainer": asdict(initial),
+            "current_trainers": [asdict(item) for item in current],
+        }
+
+    deleted = {
+        "delta": prune_to_keep(args.save_dir, "*_delta_step*.pt", args.keep_delta),
+        "full": prune_to_keep(args.save_dir, "*_step*.pt", args.keep_full, ("_delta_step",)),
+        "uploaded_delta_for_reserve": prune_uploaded_delta_for_reserve(args.save_dir, args.reserve_free_gb),
+        "orphan_sidecars": prune_orphan_sidecars(args.save_dir),
+        "stale_tmp": prune_stale_tmp(args.save_dir, args.stale_tmp_sec),
+    }
+    state = {
+        "checked_at": iso_now(),
+        "action": "pruned",
+        "save_dir": str(args.save_dir),
+        "free_gb": round(free_gb(args.save_dir), 3),
+        "initial_trainer": asdict(initial),
+        "deleted": deleted,
+        "policy": {
+            "keep_delta": args.keep_delta,
+            "keep_full": args.keep_full,
+            "interval_sec": args.interval_sec,
+            "reserve_free_gb": args.reserve_free_gb,
+            "stale_tmp_sec": args.stale_tmp_sec,
+            "stop_on_training_restart": True,
+        },
+    }
+    return state
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Local AGILLM4 checkpoint pruning guard")
+    parser.add_argument("--save-dir", type=Path, default=Path(os.environ.get("AGILLM4_SAVE_DIR", "/workspace/agillm4_4090_ckpts")))
+    parser.add_argument("--state", type=Path, default=Path(os.environ.get("AGILLM4_PRUNE_STATE", "/workspace/agillm4_local_prune_guard_state.json")))
+    parser.add_argument("--script-name", default=os.environ.get("AGILLM4_TRAIN_SCRIPT_NAME", "nB300_agillm4.py"))
+    parser.add_argument("--required-arg", default=os.environ.get("AGILLM4_TRAIN_REQUIRED_ARG", "train"))
+    parser.add_argument("--keep-delta", type=int, default=int(os.environ.get("AGILLM4_LOCAL_KEEP_DELTA", "1")))
+    parser.add_argument("--keep-full", type=int, default=int(os.environ.get("AGILLM4_LOCAL_KEEP_FULL", "1")))
+    parser.add_argument("--interval-sec", type=int, default=int(os.environ.get("AGILLM4_PRUNE_INTERVAL_SEC", "300")))
+    parser.add_argument("--reserve-free-gb", type=float, default=float(os.environ.get("AGILLM4_PRUNE_RESERVE_GB", "6")))
+    parser.add_argument("--stale-tmp-sec", type=int, default=int(os.environ.get("AGILLM4_PRUNE_STALE_TMP_SEC", "3600")))
+    parser.add_argument("--once", action="store_true")
+    args = parser.parse_args()
+
+    trainers = find_trainers(args.script_name, args.required_arg)
+    if len(trainers) != 1:
+        state = {
+            "checked_at": iso_now(),
+            "action": "exit_bad_trainer_count",
+            "trainer_count": len(trainers),
+            "trainers": [asdict(item) for item in trainers],
+        }
+        write_state(args.state, state)
+        print(json.dumps(state, indent=2, sort_keys=True), flush=True)
+        return 2
+
+    initial = trainers[0]
+    print(f"[prune-guard] guarding trainer pid={initial.pid} start_ticks={initial.start_ticks}", flush=True)
+    while True:
+        state = run_once(args, initial)
+        write_state(args.state, state)
+        print(json.dumps(state, indent=2, sort_keys=True), flush=True)
+        if state["action"].startswith("exit_"):
+            return 0
+        if args.once:
+            return 0
+        time.sleep(max(5, args.interval_sec))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())