hsbatakurki/HuggingMes

1#!/usr/bin/env python32"""HuggingMes Hermes state backup via Hugging Face Datasets."""3 4import hashlib5import json6import logging7import os8import shutil9import signal10import random11import socket12import sys13import tempfile14import threading15import time16from pathlib import Path17 18os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")19os.environ.setdefault("HF_HUB_VERBOSITY", "error")20os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "300")   # 5 min — default 10s causes timeout on large state21os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")      # faster multipart upload/download (replaces deprecated HF_HUB_ENABLE_HF_TRANSFER)22 23from huggingface_hub import HfApi, snapshot_download, upload_folder24from huggingface_hub.errors import HfHubHTTPError, RepositoryNotFoundError25 26logging.getLogger("huggingface_hub").setLevel(logging.ERROR)27 28HERMES_HOME = Path(os.environ.get("HERMES_HOME", "/opt/data"))29STATUS_FILE = Path("/tmp/huggingmes-sync-status.json")30STATE_FILE = HERMES_HOME / ".huggingmes-sync-state.json"31INTERVAL = int(os.environ.get("SYNC_INTERVAL", "600"))32INITIAL_DELAY = int(os.environ.get("SYNC_START_DELAY", "10"))33HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()34HF_USERNAME = os.environ.get("HF_USERNAME", "").strip()35SPACE_AUTHOR_NAME = os.environ.get("SPACE_AUTHOR_NAME", "").strip()36BACKUP_DATASET_NAME = os.environ.get("BACKUP_DATASET_NAME", "huggingmes-backup").strip()37INCLUDE_ENV = os.environ.get("SYNC_INCLUDE_ENV", "").strip().lower() in {"1", "true", "yes"}38MAX_FILE_SIZE_BYTES = int(os.environ.get("SYNC_MAX_FILE_BYTES", str(50 * 1024 * 1024)))39 40EXCLUDED_DIRS = {41    ".cache",42    ".git",43    ".npm",44    ".venv",45    "__pycache__",46    "node_modules",47    "venv",48}49EXCLUDED_TOP_LEVEL = {"logs", STATE_FILE.name}50if not INCLUDE_ENV:51    EXCLUDED_TOP_LEVEL.add(".env")52 53HF_API = HfApi(token=HF_TOKEN) if HF_TOKEN else None54STOP_EVENT = threading.Event()55_REPO_ID_CACHE: str | None = None56 57 58def write_status(status: str, message: str, fingerprint: str | None = None, marker: tuple[int, int, int] | None = None) -> None:59    timestamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())60    payload = {61        "status": status,62        "message": message,63        "timestamp": timestamp,64    }65    66    # Update temporary status file for health checks/dashboard67    tmp_path = STATUS_FILE.with_suffix(".tmp")68    try:69        tmp_path.write_text(json.dumps(payload), encoding="utf-8")70        tmp_path.replace(STATUS_FILE)71    except OSError:72        pass73 74    # Update persistent state file in HERMES_HOME75    if fingerprint or marker:76        state = {}77        if STATE_FILE.exists():78            try:79                state = json.loads(STATE_FILE.read_text(encoding="utf-8"))80            except Exception:81                pass82        83        if fingerprint:84            state["last_fingerprint"] = fingerprint85        if marker:86            state["last_marker"] = list(marker)87        state["last_sync"] = timestamp88        89        try:90            STATE_FILE.write_text(json.dumps(state), encoding="utf-8")91        except OSError:92            pass93 94 95def resolve_backup_repo() -> str:96    global _REPO_ID_CACHE97    if _REPO_ID_CACHE:98        return _REPO_ID_CACHE99 100    namespace = HF_USERNAME or SPACE_AUTHOR_NAME101    if not namespace and HF_API is not None:102        whoami = HF_API.whoami()103        namespace = whoami.get("name") or whoami.get("user") or ""104 105    namespace = str(namespace).strip()106    if not namespace:107        raise RuntimeError("Could not determine HF username. Set HF_USERNAME or use an account HF_TOKEN.")108 109    _REPO_ID_CACHE = f"{namespace}/{BACKUP_DATASET_NAME}"110    return _REPO_ID_CACHE111 112 113def ensure_repo_exists() -> str:114    repo_id = resolve_backup_repo()115    try:116        HF_API.repo_info(repo_id=repo_id, repo_type="dataset")117    except RepositoryNotFoundError:118        HF_API.create_repo(repo_id=repo_id, repo_type="dataset", private=True)119    return repo_id120 121 122def should_exclude(rel_posix: str, path: Path) -> bool:123    parts = Path(rel_posix).parts124    if not parts:125        return False126    if parts[0] in EXCLUDED_TOP_LEVEL:127        return True128    if any(part in EXCLUDED_DIRS for part in parts):129        return True130    # Exclude SQLite temporary files (shm, wal, journal)131    if path.is_file():132        name_lower = path.name.lower()133        if name_lower.endswith((".db-shm", ".db-wal", ".db-journal")):134            return True135        try:136            return path.stat().st_size > MAX_FILE_SIZE_BYTES137        except OSError:138            return True139    return False140 141 142def metadata_marker(root: Path) -> tuple[int, int, int]:143    if not root.exists():144        return (0, 0, 0)145    file_count = 0146    total_size = 0147    newest_mtime = 0148    for path in root.rglob("*"):149        if not path.is_file():150            continue151        rel = path.relative_to(root).as_posix()152        if should_exclude(rel, path):153            continue154        try:155            stat = path.stat()156        except OSError:157            continue158        file_count += 1159        total_size += int(stat.st_size)160        newest_mtime = max(newest_mtime, int(stat.st_mtime_ns))161    return (file_count, total_size, newest_mtime)162 163 164def fingerprint_dir(root: Path) -> str:165    hasher = hashlib.sha256()166    if not root.exists():167        return hasher.hexdigest()168    for path in sorted(p for p in root.rglob("*") if p.is_file()):169        rel = path.relative_to(root).as_posix()170        if should_exclude(rel, path):171            continue172        hasher.update(rel.encode("utf-8"))173        with path.open("rb") as handle:174            for chunk in iter(lambda: handle.read(1024 * 1024), b""):175                hasher.update(chunk)176    return hasher.hexdigest()177 178 179def create_snapshot_dir(source_root: Path) -> Path:180    staging_root = Path(tempfile.mkdtemp(prefix="huggingmes-sync-"))181    for path in sorted(source_root.rglob("*")):182        rel = path.relative_to(source_root)183        rel_posix = rel.as_posix()184        if should_exclude(rel_posix, path):185            continue186        target = staging_root / rel187        if path.is_dir():188            target.mkdir(parents=True, exist_ok=True)189            continue190        target.parent.mkdir(parents=True, exist_ok=True)191        try:192            shutil.copy2(path, target)193        except OSError:194            # File may have been deleted by the application since enumeration195            continue196    return staging_root197 198 199def restore() -> bool:200    if not HF_TOKEN:201        write_status("disabled", "HF_TOKEN is not configured.")202        return False203 204    repo_id = resolve_backup_repo()205    write_status("restoring", f"Restoring Hermes state from {repo_id}")206    try:207        with tempfile.TemporaryDirectory() as tmpdir:208            snapshot_download(repo_id=repo_id, repo_type="dataset", token=HF_TOKEN, local_dir=tmpdir)209            tmp_path = Path(tmpdir)210            if not any(tmp_path.iterdir()):211                write_status("fresh", "Backup dataset is empty. Starting fresh.")212                return True213 214            HERMES_HOME.mkdir(parents=True, exist_ok=True)215            for child in tmp_path.iterdir():216                if should_exclude(child.name, child):217                    continue218                target = HERMES_HOME / child.name219                if target.is_dir():220                    shutil.rmtree(target, ignore_errors=True)221                elif target.exists():222                    target.unlink()223                if child.is_dir():224                    shutil.copytree(child, target)225                else:226                    shutil.copy2(child, target)227 228        write_status("restored", f"Restored Hermes state from {repo_id}")229        return True230    except RepositoryNotFoundError:231        write_status("fresh", f"Backup dataset {repo_id} does not exist yet.")232        return True233    except HfHubHTTPError as exc:234        if exc.response is not None and exc.response.status_code == 404:235            write_status("fresh", f"Backup dataset {repo_id} does not exist yet.")236            return True237        write_status("error", f"Restore failed: {exc}")238        print(f"Restore failed: {exc}", file=sys.stderr)239        return False240    except Exception as exc:241        write_status("error", f"Restore failed: {exc}")242        print(f"Restore failed: {exc}", file=sys.stderr)243        return False244 245 246def sync_once(last_fingerprint: str | None = None, last_marker: tuple[int, int, int] | None = None):247    # If no state provided, try to load from persistent state file248    if last_fingerprint is None and last_marker is None:249        if STATE_FILE.exists():250            try:251                state = json.loads(STATE_FILE.read_text(encoding="utf-8"))252                last_fingerprint = state.get("last_fingerprint")253                m = state.get("last_marker")254                if m and len(m) == 3:255                    last_marker = tuple(m)256            except Exception:257                pass258 259    repo_id = ensure_repo_exists()260    current_marker = metadata_marker(HERMES_HOME)261    if last_marker is not None and current_marker == last_marker:262        write_status("synced", "No Hermes state changes detected (marker match).")263        return (last_fingerprint or "", current_marker)264 265    current_fingerprint = fingerprint_dir(HERMES_HOME)266    if last_fingerprint is not None and current_fingerprint == last_fingerprint:267        write_status("synced", "No Hermes state changes detected (fingerprint match).")268        return (last_fingerprint, current_marker)269 270    hostname = socket.gethostname()271    write_status("syncing", f"Uploading Hermes state to {repo_id} from {hostname}")272    snapshot_dir = create_snapshot_dir(HERMES_HOME)273    try:274        upload_folder(275            folder_path=str(snapshot_dir),276            repo_id=repo_id,277            repo_type="dataset",278            token=HF_TOKEN,279            commit_message=f"HuggingMes sync [{hostname}] {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}",280            ignore_patterns=[".git/*", ".git"],281        )282    finally:283        shutil.rmtree(snapshot_dir, ignore_errors=True)284 285    write_status("success", f"Uploaded Hermes state to {repo_id}", fingerprint=current_fingerprint, marker=current_marker)286    return (current_fingerprint, current_marker)287 288 289def handle_signal(_sig, _frame) -> None:290    STOP_EVENT.set()291 292 293def loop() -> int:294    signal.signal(signal.SIGTERM, handle_signal)295    signal.signal(signal.SIGINT, handle_signal)296    try:297        repo_id = resolve_backup_repo()298        write_status("configured", f"Backup loop active for {repo_id} with {INTERVAL}s interval.")299    except Exception as exc:300        write_status("error", str(exc))301        print(f"Hermes sync error: {exc}")302        return 1303 304    last_fingerprint = fingerprint_dir(HERMES_HOME)305    last_marker = metadata_marker(HERMES_HOME)306    time.sleep(INITIAL_DELAY)307    print(f"Hermes state sync started: every {INTERVAL}s -> {repo_id}")308 309    while not STOP_EVENT.is_set():310        try:311            last_fingerprint, last_marker = sync_once(last_fingerprint, last_marker)312        except Exception as exc:313            write_status("error", f"Sync failed: {exc}")314            print(f"Hermes sync failed: {exc}")315        316        # Add 10% jitter to interval to avoid synchronized commits from multiple containers317        jitter = random.uniform(0.9, 1.1)318        if STOP_EVENT.wait(INTERVAL * jitter):319            break320    return 0321 322 323def main() -> int:324    HERMES_HOME.mkdir(parents=True, exist_ok=True)325    if len(sys.argv) < 2:326        return loop()327    command = sys.argv[1]328    if command == "restore":329        return 0 if restore() else 1330    if command == "sync-once":331        try:332            sync_once()333            return 0334        except Exception as exc:335            write_status("error", f"Shutdown sync failed: {exc}")336            print(f"Hermes sync: shutdown sync failed: {exc}")337            return 1338    if command == "loop":339        return loop()340    print(f"Unknown command: {command}", file=sys.stderr)341    return 1342 343 344if __name__ == "__main__":345    raise SystemExit(main())346