hf-persist.py

#!/usr/bin/env python3
"""HuggingFace Dataset persistence for Mission Control .data/ directory.

On startup:  restore .data/ from HF Dataset
Periodically: upload .data/ to HF Dataset (every SYNC_INTERVAL seconds)
On shutdown:  final upload

Env vars:
  HF_TOKEN           - HuggingFace API token (required)
  HF_DATASET_REPO    - Dataset repo ID (default: auto from SPACE_ID)
  SYNC_INTERVAL      - Seconds between syncs (default: 120)
"""
import os, sys, signal, time, subprocess, threading, logging

logging.basicConfig(level=logging.INFO, format="[hf-persist] %(message)s")
log = logging.getLogger(__name__)

HF_TOKEN = os.environ.get("HF_TOKEN", "")
SPACE_ID = os.environ.get("SPACE_ID", "")
DATA_DIR = os.environ.get("MISSION_CONTROL_DATA_DIR", "/app/.data")
DATASET_PATH = "mission-control"
SYNC_INTERVAL = int(os.environ.get("SYNC_INTERVAL", "120"))

# Derive dataset repo from SPACE_ID (e.g. tao-shen/HuggingClaw-MissionControl -> tao-shen/HuggingClaw-MissionControl-data)
if os.environ.get("HF_DATASET_REPO"):
    HF_REPO_ID = os.environ["HF_DATASET_REPO"]
elif SPACE_ID:
    HF_REPO_ID = f"{SPACE_ID}-data"
else:
    HF_REPO_ID = ""

IGNORE_PATTERNS = ["*.lock", "*.tmp", "*.pid", "__pycache__", "*.wal", "*.shm"]


def ensure_dataset():
    """Create dataset repo if it doesn't exist."""
    from huggingface_hub import HfApi
    api = HfApi(token=HF_TOKEN)
    try:
        api.repo_info(repo_id=HF_REPO_ID, repo_type="dataset")
        log.info(f"Dataset {HF_REPO_ID} exists")
    except Exception:
        log.info(f"Creating dataset {HF_REPO_ID}")
        api.create_repo(repo_id=HF_REPO_ID, repo_type="dataset", private=True)
        log.info(f"Created dataset {HF_REPO_ID}")


def restore():
    """Download .data/ from HF Dataset on startup."""
    from huggingface_hub import snapshot_download
    import tempfile, shutil

    log.info(f"Restoring from {HF_REPO_ID}/{DATASET_PATH}/ ...")
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            snapshot_download(
                repo_id=HF_REPO_ID,
                repo_type="dataset",
                allow_patterns=f"{DATASET_PATH}/**",
                local_dir=tmpdir,
                token=HF_TOKEN,
            )
            src = os.path.join(tmpdir, DATASET_PATH)
            if os.path.isdir(src):
                count = 0
                for root, dirs, files in os.walk(src):
                    for f in files:
                        src_file = os.path.join(root, f)
                        rel = os.path.relpath(src_file, src)
                        dst_file = os.path.join(DATA_DIR, rel)
                        os.makedirs(os.path.dirname(dst_file), exist_ok=True)
                        shutil.copy2(src_file, dst_file)
                        count += 1
                log.info(f"Restored {count} files to {DATA_DIR}")
            else:
                log.info("No data found in dataset (fresh start)")
    except Exception as e:
        log.warning(f"Restore failed (starting fresh): {e}")


def save():
    """Upload .data/ to HF Dataset."""
    from huggingface_hub import HfApi
    from datetime import datetime

    if not os.path.isdir(DATA_DIR) or not os.listdir(DATA_DIR):
        log.info("Nothing to save (empty .data/)")
        return

    api = HfApi(token=HF_TOKEN)
    try:
        api.upload_folder(
            folder_path=DATA_DIR,
            path_in_repo=DATASET_PATH,
            repo_id=HF_REPO_ID,
            repo_type="dataset",
            commit_message=f"Sync mission-control data — {datetime.now().isoformat()}",
            ignore_patterns=IGNORE_PATTERNS,
        )
        log.info("Saved to dataset")
    except Exception as e:
        log.error(f"Save failed: {e}")


def register_agents(port):
    """Register HuggingClaw agents after server is ready."""
    import urllib.request, urllib.error, json

    agents = [
        {"name": "Adam", "role": "agent", "capabilities": ["parenting", "coding", "A2A"], "framework": "openclaw"},
        {"name": "Eve", "role": "agent", "capabilities": ["parenting", "coding", "A2A"], "framework": "openclaw"},
        {"name": "God", "role": "agent", "capabilities": ["supervision", "architecture", "A2A"], "framework": "openclaw"},
        {"name": "Cain", "role": "agent", "capabilities": ["learning", "coding"], "framework": "openclaw"},
    ]

    # Wait for server to be ready
    for i in range(30):
        try:
            req = urllib.request.Request(f"http://localhost:{port}/api/status?action=health")
            urllib.request.urlopen(req, timeout=2)
            break
        except Exception:
            time.sleep(2)
    else:
        log.warning("Server not ready after 60s, skipping agent registration")
        return

    api_key = os.environ.get("API_KEY", "")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"

    for agent in agents:
        try:
            data = json.dumps(agent).encode()
            req = urllib.request.Request(
                f"http://localhost:{port}/api/agents/register",
                data=data,
                headers=headers,
                method="POST",
            )
            resp = urllib.request.urlopen(req, timeout=10)
            log.info(f"Registered agent: {agent['name']} ({resp.status})")
        except urllib.error.HTTPError as e:
            body = e.read().decode()
            log.warning(f"Failed to register {agent['name']}: {e.code} {body}")
        except Exception as e:
            log.warning(f"Failed to register {agent['name']}: {e}")


def main():
    if not HF_TOKEN or not HF_REPO_ID:
        log.info("HF_TOKEN or HF_DATASET_REPO not set — persistence disabled, starting server directly")
        os.execvp("node", ["node", "server.js"])
        return

    port = os.environ.get("PORT", "7860")

    # Step 1: Ensure dataset exists
    ensure_dataset()

    # Step 2: Restore from dataset
    restore()

    # Step 3: Start server as subprocess
    proc = subprocess.Popen(["node", "server.js"], cwd="/app")
    log.info(f"Server started (PID {proc.pid})")

    # Step 4: Register agents in background
    reg_thread = threading.Thread(target=register_agents, args=(port,), daemon=True)
    reg_thread.start()

    # Step 5: Periodic sync
    stop_event = threading.Event()

    def sync_loop():
        while not stop_event.is_set():
            stop_event.wait(SYNC_INTERVAL)
            if not stop_event.is_set():
                save()

    sync_thread = threading.Thread(target=sync_loop, daemon=True)
    sync_thread.start()

    # Step 6: Handle shutdown
    def shutdown(signum, frame):
        log.info(f"Received signal {signum}, shutting down...")
        stop_event.set()
        proc.terminate()
        try:
            proc.wait(timeout=10)
        except subprocess.TimeoutExpired:
            proc.kill()
        save()  # Final save
        sys.exit(0)

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    # Wait for server process
    ret = proc.wait()
    stop_event.set()
    save()  # Final save
    sys.exit(ret)


if __name__ == "__main__":
    main()