upload_agillm4_checkpoints.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from huggingface_hub import HfApi


def iso_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def stamp_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def load_json(path: Path, default: Any) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return default


def save_json(path: Path, data: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2, sort_keys=True), encoding="utf-8")


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def upload_file(api: HfApi, repo_id: str, local_path: Path, remote_path: str, message: str) -> None:
    api.upload_file(
        repo_id=repo_id,
        path_or_fileobj=str(local_path),
        path_in_repo=remote_path,
        commit_message=message,
    )


def delete_remote_not_kept(api: HfApi, repo_id: str, remote_dir: str, keep_basenames: set[str]) -> list[str]:
    deleted: list[str] = []
    try:
        files = api.list_repo_files(repo_id=repo_id)
    except Exception as exc:
        print(f"[upload] WARN list_repo_files failed: {exc}", flush=True)
        return deleted
    prefix = remote_dir.rstrip("/") + "/"
    victims = []
    for file_path in files:
        if not file_path.startswith(prefix):
            continue
        name = Path(file_path).name
        base = name[:-7] if name.endswith(".sha256") else name
        if base not in keep_basenames:
            victims.append(file_path)
    if victims:
        try:
            api.delete_files(repo_id=repo_id, paths=victims, commit_message=f"Prune AGILLM4 uploads under {remote_dir}")
            deleted.extend(victims)
        except Exception as exc:
            print(f"[upload] WARN delete_files failed for {len(victims)} files: {exc}", flush=True)
    return deleted


def latest_file(glob_root: Path, pattern: str) -> Path | None:
    files = [p for p in glob_root.glob(pattern) if p.is_file()]
    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)],
        capture_output=True,
        text=True,
        timeout=60,
        check=False,
    )
    if result.returncode != 0:
        return {"checked_at": iso_now(), "error": result.stderr.strip() or result.stdout.strip()}
    try:
        return json.loads(result.stdout)
    except Exception:
        return {"checked_at": iso_now(), "error": "failed to parse status json", "raw": result.stdout[-4000:]}


def write_tail(src: Path, dst: Path, lines: int) -> None:
    dst.parent.mkdir(parents=True, exist_ok=True)
    if not src.exists():
        dst.write_text("", encoding="utf-8")
        return
    result = subprocess.run(["tail", "-n", str(lines), str(src)], capture_output=True, text=True, check=False)
    dst.write_text(result.stdout, encoding="utf-8", errors="replace")


def maybe_upload_large(
    api: HfApi,
    repo_id: str,
    state: dict[str, Any],
    kind: str,
    path: Path | None,
    remote_dir: str,
    interval_sec: int,
    keep: int,
) -> bool:
    if path is None or not path.exists():
        print(f"[upload] no {kind} checkpoint yet", flush=True)
        return False
    now = time.time()
    last_t = float(state.get(f"last_{kind}_upload_time", 0) or 0)
    identity = f"{path.name}:{path.stat().st_size}:{int(path.stat().st_mtime)}"
    if state.get(f"last_{kind}_identity") == identity:
        print(f"[upload] {kind} unchanged: {path.name}", flush=True)
        return False
    if last_t and now - last_t < interval_sec:
        remaining = int(interval_sec - (now - last_t))
        print(f"[upload] {kind} interval not due for {remaining}s: {path.name}", flush=True)
        return False

    digest = sha256_file(path)
    sha_path = path.with_suffix(path.suffix + ".upload.sha256")
    sha_path.write_text(f"{digest}  {path.name}\n", encoding="utf-8")
    remote_name = f"{stamp_now()}_{path.name}"
    remote_path = f"{remote_dir.rstrip('/')}/{remote_name}"
    print(f"[upload] uploading {kind}: {path} -> {repo_id}/{remote_path}", flush=True)
    upload_file(api, repo_id, path, remote_path, f"Upload AGILLM4 {kind} checkpoint {path.name}")
    upload_file(api, repo_id, sha_path, remote_path + ".sha256", f"Upload AGILLM4 {kind} checksum {path.name}")

    history = list(state.get(f"{kind}_uploads", []))
    history.append({"name": remote_name, "remote_path": remote_path, "sha256": digest, "uploaded_at": iso_now(), "size": path.stat().st_size})
    history = history[-max(1, keep):]
    state[f"{kind}_uploads"] = history
    state[f"last_{kind}_upload_time"] = now
    state[f"last_{kind}_identity"] = identity
    keep_names = {item["name"] for item in history}
    deleted = delete_remote_not_kept(api, repo_id, remote_dir, keep_names)
    if deleted:
        print(f"[upload] pruned {len(deleted)} remote {kind} files", flush=True)
    return True


def main() -> int:
    parser = argparse.ArgumentParser(description="Bounded AGILLM4 checkpoint uploader")
    parser.add_argument("--repo", default=os.environ.get("AGILLM4_UPLOAD_REPO", "OpenTransformer/AGILLM-4"))
    parser.add_argument("--prefix", default=os.environ.get("AGILLM4_UPLOAD_PREFIX", "training/agillm4_floor_v47_v4pro"))
    parser.add_argument("--save-dir", type=Path, default=Path(os.environ.get("AGILLM4_UPLOAD_SAVE_DIR", "/workspace/agillm4_4090_ckpts")))
    parser.add_argument("--log", type=Path, default=Path(os.environ.get("AGILLM4_UPLOAD_LOG", "/workspace/agillm4_floor_train.log")))
    parser.add_argument("--script", type=Path, default=Path(os.environ.get("AGILLM4_UPLOAD_SCRIPT", "/workspace/agillm-4/nB300_agillm4.py")))
    parser.add_argument("--state", type=Path, default=Path(os.environ.get("AGILLM4_UPLOAD_STATE", "/workspace/agillm4_upload_state_v4pro.json")))
    parser.add_argument("--stage", type=Path, default=Path(os.environ.get("AGILLM4_UPLOAD_STAGE", "/workspace/agillm4_upload_stage_v4pro")))
    parser.add_argument("--full-interval-sec", type=int, default=int(os.environ.get("AGILLM4_UPLOAD_FULL_INTERVAL_SEC", str(7 * 24 * 3600))))
    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()

    api = HfApi()
    prefix = args.prefix.strip("/")
    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")

    tail_path = args.stage / "train_tail.log"
    write_tail(args.log, tail_path, args.tail_lines)
    upload_file(api, args.repo, tail_path, f"{prefix}/logs/train_tail.log", "Update AGILLM4 training log tail")

    latest_json = args.save_dir / "latest.json"
    if latest_json.exists():
        shutil.copy2(latest_json, args.stage / "latest.json")
        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_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)

    state["last_status_upload_at"] = iso_now()
    save_json(args.state, state)
    manifest_path = args.stage / "upload_state.json"
    save_json(manifest_path, state)
    upload_file(api, args.repo, manifest_path, f"{prefix}/status/upload_state.json", "Update AGILLM4 upload state")
    print(f"[upload] done {iso_now()} repo={args.repo} prefix={prefix}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())