scripts/snapshot_operator.py

#!/usr/bin/env python3
"""Copy a stable operator snapshot without stopping or locking the evaluator."""
from __future__ import annotations

import argparse
import json
from pathlib import Path
import shutil
import sys
import tempfile
import time

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))
from tsfm_bench.remote_state import HANDOFF_ROOTS, utc_now


def inventory(root):
    result = {}
    for prefix in HANDOFF_ROOTS:
        for path in (root / prefix).rglob("*"):
            if path.is_symlink():
                raise ValueError(f"Refusing symlink: {path.relative_to(root)}")
            if path.is_file():
                stat = path.stat()
                result[path.relative_to(root)] = (stat.st_size, stat.st_mtime_ns, stat.st_ino)
    return result


def snapshot(root, destination):
    if destination.exists():
        raise ValueError("Snapshot destination must not exist")
    before = inventory(root)
    with tempfile.TemporaryDirectory(prefix="livehouse-copy-", dir=destination.parent) as temp:
        stage = Path(temp) / "snapshot"
        stage.mkdir()
        for rel in sorted(before):
            target = stage / rel
            target.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(root / rel, target)
        if inventory(root) != before:
            raise RuntimeError("Operator changed during copying; retry between evaluation cycles")
        record = {"created_at":utc_now(), "files":len(before),
                  "source_bytes":sum(stat[0] for stat in before.values())}
        (stage / ".snapshot.json").write_text(json.dumps(record, indent=2) + "\n")
        stage.rename(destination)
    return record


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--destination", type=Path, required=True)
    args = parser.parse_args()
    destination = args.destination.resolve()
    destination.parent.mkdir(parents=True, exist_ok=True)
    for attempt in range(3):
        try:
            print(json.dumps(snapshot(ROOT, destination)), flush=True)
            return
        except (FileNotFoundError, RuntimeError) as exc:
            if attempt == 2:
                raise
            print(f"{exc}; retrying in 20 seconds", flush=True)
            time.sleep(20)


if __name__ == "__main__":
    main()