Add AGILLM4 uploader local pruning

upload_agillm4_checkpoints.py

@@ -83,6 +83,56 @@ def latest_file(glob_root: Path, pattern: str) -> Path | None:
     return max(files, key=lambda p: p.stat().st_mtime) if files else None
 
 
+def latest_file_excluding(glob_root: Path, pattern: str, excluded_name_parts: tuple[str, ...]) -> Path | None:
+    files = [
+        p for p in glob_root.glob(pattern)
+        if p.is_file() and not any(part in p.name for part in excluded_name_parts)
+    ]
+    return max(files, key=lambda p: p.stat().st_mtime) if files else None
+
+
+def checkpoint_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 prune_local_checkpoints(
+    save_dir: Path,
+    pattern: str,
+    keep: int,
+    label: str,
+    excluded_name_parts: tuple[str, ...] = (),
+) -> list[str]:
+    if keep < 0:
+        return []
+    files = sorted(
+        [
+            p for p in save_dir.glob(pattern)
+            if p.is_file()
+            and p.stat().st_size > 0
+            and not any(part in p.name for part in excluded_name_parts)
+        ],
+        key=lambda p: p.stat().st_mtime,
+    )
+    victims = files[:max(0, len(files) - keep)]
+    deleted: list[str] = []
+    for path in victims:
+        for artifact in checkpoint_artifacts(path):
+            try:
+                if artifact.exists():
+                    artifact.unlink()
+                    deleted.append(str(artifact))
+            except Exception as exc:
+                print(f"[upload] WARN local {label} prune failed for {artifact}: {exc}", flush=True)
+    if deleted:
+        print(f"[upload] pruned {len(deleted)} local {label} artifacts", flush=True)
+    return deleted
+
+
 def status_json(script: Path, log: Path, save_dir: Path) -> dict[str, Any]:
     result = subprocess.run(
         [sys.executable, "-u", str(script), "status", "--json", "--log", str(log), "--save_dir", str(save_dir)],
@@ -167,6 +217,8 @@ def main() -> int:
     parser.add_argument("--delta-interval-sec", type=int, default=int(os.environ.get("AGILLM4_UPLOAD_DELTA_INTERVAL_SEC", str(24 * 3600))))
     parser.add_argument("--keep-full", type=int, default=int(os.environ.get("AGILLM4_UPLOAD_KEEP_FULL", "2")))
     parser.add_argument("--keep-delta", type=int, default=int(os.environ.get("AGILLM4_UPLOAD_KEEP_DELTA", "2")))
+    parser.add_argument("--local-keep-full", type=int, default=int(os.environ.get("AGILLM4_LOCAL_KEEP_FULL", "1")))
+    parser.add_argument("--local-keep-delta", type=int, default=int(os.environ.get("AGILLM4_LOCAL_KEEP_DELTA", "1")))
     parser.add_argument("--tail-lines", type=int, default=int(os.environ.get("AGILLM4_UPLOAD_TAIL_LINES", "5000")))
     args = parser.parse_args()
 
@@ -175,14 +227,22 @@ def main() -> int:
     args.stage.mkdir(parents=True, exist_ok=True)
     state = load_json(args.state, {})
 
+    local_pruned = {
+        "delta": prune_local_checkpoints(args.save_dir, "*_delta_step*.pt", args.local_keep_delta, "delta"),
+        "full": prune_local_checkpoints(args.save_dir, "*_step*.pt", args.local_keep_full, "full", ("_delta_step",)),
+    }
+
     status = status_json(args.script, args.log, args.save_dir)
     status["upload_policy"] = {
         "full_interval_sec": args.full_interval_sec,
         "delta_interval_sec": args.delta_interval_sec,
         "keep_full_current_files": args.keep_full,
         "keep_delta_current_files": args.keep_delta,
+        "local_keep_full_files": args.local_keep_full,
+        "local_keep_delta_files": args.local_keep_delta,
         "note": "Small status/log tail uploads are frequent; multi-GB deltas/full checkpoints are rate-limited for HF public storage.",
     }
+    status["local_pruned"] = local_pruned
     status_path = args.stage / "status.json"
     save_json(status_path, status)
     upload_file(api, args.repo, status_path, f"{prefix}/status/status.json", "Update AGILLM4 training status")
@@ -197,7 +257,7 @@ def main() -> int:
         upload_file(api, args.repo, args.stage / "latest.json", f"{prefix}/status/latest.json", "Update AGILLM4 latest checkpoint metadata")
 
     newest_delta = latest_file(args.save_dir, "*_delta_step*.pt")
-    newest_full = latest_file(args.save_dir, "*_step*.pt")
+    newest_full = latest_file_excluding(args.save_dir, "*_step*.pt", ("_delta_step",))
     maybe_upload_large(api, args.repo, state, "delta", newest_delta, f"{prefix}/checkpoints/deltas", args.delta_interval_sec, args.keep_delta)
     maybe_upload_large(api, args.repo, state, "full", newest_full, f"{prefix}/checkpoints/full", args.full_interval_sec, args.keep_full)