src/tsfm_bench/remote_state.py

"""Durable operator checkpoints in an organization-owned private HF Dataset."""
from __future__ import annotations

import gzip
import hashlib
import json
import os
import shutil
import tarfile
import tempfile
from datetime import datetime, timezone
from pathlib import Path

DATA = Path("space/vendor/ts_bench/data/sample_week")
RESULTS = Path("space/results")
STATE_MANIFEST = "state/manifest.json"
STATE_ARCHIVE = "state/runtime.tar.gz"
HANDOFF_ROOTS = (RESULTS, Path("space/vendor/ts_bench/data"), Path("outputs"),
                 Path("paper"), Path("data"), Path("logs"))


def utc_now():
    return datetime.now(timezone.utc).isoformat()


def digest(path: Path) -> str:
    with path.open("rb") as handle:
        return hashlib.file_digest(handle, "sha256").hexdigest() if hasattr(hashlib, "file_digest") else _digest_stream(handle)


def _digest_stream(handle):
    result = hashlib.sha256()
    for block in iter(lambda: handle.read(1024 * 1024), b""):
        result.update(block)
    return result.hexdigest()


def runtime_files(root: Path):
    """Keep authoritative state and reachable inputs, excluding regenerable caches."""
    files = []
    for directory in (RESULTS, DATA / "normalized"):
        for path in (root / directory).rglob("*"):
            rel = path.relative_to(root)
            if (path.is_file() and path.suffix in {".json", ".jsonl", ".csv", ".tex"}
                    and "canonical_all_history" not in rel.parts):
                files.append(path)
    inputs = root / DATA / "model_inputs_latest"
    manifest = inputs / "manifest.json"
    if manifest.exists():
        files.append(manifest)
        if (inputs / "manifest.csv").exists():
            files.append(inputs / "manifest.csv")
        for entry in json.loads(manifest.read_text()):
            task_id = entry["task_id"]
            if not task_id or task_id in {".", ".."} or Path(task_id).name != task_id:
                raise ValueError("Invalid input task path")
            files.extend(p for p in (inputs / task_id).glob("*") if p.is_file())
    summary = root / DATA / "summary.json"
    if summary.exists():
        files.append(summary)
    return sorted(set(files))


def pack_files(root: Path, files: list[Path], destination: Path) -> dict:
    """Deterministic archives with a per-file integrity manifest and no symlinks."""
    secrets = [value.encode() for key, value in os.environ.items()
               if value and len(value) >= 12 and any(word in key for word in ("TOKEN", "API_KEY", "PASSWORD", "SECRET"))]
    entries = []
    with destination.open("wb") as raw, gzip.GzipFile(filename="", fileobj=raw, mode="wb", mtime=0) as compressed:
        with tarfile.open(fileobj=compressed, mode="w") as archive:
            for path in files:
                if path.is_symlink() or not path.resolve().is_relative_to(root.resolve()):
                    raise ValueError(f"Refusing symlink: {path}")
                name = path.relative_to(root).as_posix()
                if any(part.startswith(".env") for part in Path(name).parts):
                    raise ValueError("Environment files belong in a secrets manager")
                content = path.read_bytes()
                if any(secret in content for secret in secrets):
                    raise ValueError(f"Configured credential detected in {name}")
                info = tarfile.TarInfo(name)
                info.size = len(content)
                info.mode = 0o600
                import io
                archive.addfile(info, io.BytesIO(content))
                entries.append({"path": name, "bytes": len(content), "sha256": hashlib.sha256(content).hexdigest()})
    return {"sha256": digest(destination), "bytes": destination.stat().st_size, "files": entries}


def restore_archive(root: Path, archive: Path, manifest: dict, *, allowed_roots=(RESULTS, DATA)):
    if digest(archive) != manifest["sha256"]:
        raise ValueError("Remote checkpoint checksum mismatch")
    expected = {entry["path"]: entry for entry in manifest["files"]}
    root = root.resolve()
    with tarfile.open(archive, "r:gz") as bundle:
        members = bundle.getmembers()
        if len(members) != len(expected) or {m.name for m in members} != set(expected):
            raise ValueError("Checkpoint inventory mismatch")
        for member in members:
            path = Path(member.name)
            if (not member.isfile() or path.is_absolute() or ".." in path.parts
                    or not (root / path).resolve().is_relative_to(root)
                    or not any(path.is_relative_to(prefix) for prefix in allowed_roots)):
                raise ValueError(f"Unsafe checkpoint member: {member.name}")
            with bundle.extractfile(member) as source:
                content = source.read()
            entry = expected[member.name]
            if len(content) != entry["bytes"] or hashlib.sha256(content).hexdigest() != entry["sha256"]:
                raise ValueError(f"Checkpoint file checksum mismatch: {member.name}")
        # Verify the entire archive before replacing any local state.
        for member in members:
            target = root / member.name
            target.parent.mkdir(parents=True, exist_ok=True)
            with bundle.extractfile(member) as source, target.open("wb") as destination:
                shutil.copyfileobj(source, destination)


class RemoteState:
    def __init__(self, root: Path, repo_id: str, token: str):
        from huggingface_hub import HfApi
        self.root = root.resolve()
        self.repo_id = repo_id
        self.api = HfApi(token=token)
        self.token = token
        self.revision = None

    def restore(self):
        from huggingface_hub import hf_hub_download
        info = self.api.dataset_info(self.repo_id)
        if not info.private:
            raise ValueError("Operator state repository must be private")
        revision = info.sha
        manifest = json.loads(Path(hf_hub_download(self.repo_id, STATE_MANIFEST, repo_type="dataset",
                                                  revision=revision, token=self.token)).read_text())
        archive = Path(hf_hub_download(self.repo_id, STATE_ARCHIVE, repo_type="dataset",
                                       revision=revision, token=self.token))
        # Extract into an empty staging root first; existing stale files must not
        # survive a restore and silently become future checkpoint contents.
        with tempfile.TemporaryDirectory(prefix="livehouse-restore-") as temp:
            stage = Path(temp)
            restore_archive(stage, archive, manifest)
            for rel in (RESULTS, DATA / "normalized", DATA / "model_inputs_latest"):
                if not (stage / rel).is_dir():
                    raise ValueError(f"Incomplete operator checkpoint: {rel}")
            for rel in (RESULTS, DATA / "normalized", DATA / "model_inputs_latest"):
                shutil.rmtree(self.root / rel, ignore_errors=True)
                (self.root / rel).parent.mkdir(parents=True, exist_ok=True)
                shutil.copytree(stage / rel, self.root / rel)
            if (stage / DATA / "summary.json").exists():
                shutil.copy2(stage / DATA / "summary.json", self.root / DATA / "summary.json")
        self.revision = revision
        return manifest

    def checkpoint(self, *, initialize=False, archive_raw=True):
        from huggingface_hub import CommitOperationAdd
        if self.revision is None:
            if not initialize:
                raise RuntimeError("Restore the remote checkpoint before evaluating")
            info = self.api.dataset_info(self.repo_id)
            if not info.private:
                raise ValueError("Operator state repository must be private")
            existing = self.api.list_repo_files(self.repo_id, repo_type="dataset", revision=info.sha)
            if STATE_MANIFEST in existing:
                raise ValueError("Remote operator state already exists; refusing initialization overwrite")
            self.revision = info.sha
        for required in (RESULTS / "evaluation_metrics_canonical.jsonl",
                         DATA / "normalized/observations.csv", DATA / "model_inputs_latest/manifest.json"):
            if not (self.root / required).is_file():
                raise ValueError(f"Incomplete operator state: {required}")
        with tempfile.TemporaryDirectory(prefix="livehouse-checkpoint-") as temp:
            temp = Path(temp)
            state_file = temp / "runtime.tar.gz"
            manifest = pack_files(self.root, runtime_files(self.root), state_file)
            manifest.update(schema_version=1, checkpointed_at=utc_now(),
                            source_revision=os.getenv("LIVEHOUSE_SOURCE_REVISION", "local-migration"))
            raw_files = sorted(p for p in (self.root / DATA / "raw").rglob("*") if p.is_file()) if archive_raw else []
            operations = [CommitOperationAdd(path_in_repo=".gitattributes", path_or_fileobj=b"*.tar.gz filter=lfs diff=lfs merge=lfs -text\nraw/*.json filter=lfs diff=lfs merge=lfs -text\nhandoff/*.json filter=lfs diff=lfs merge=lfs -text\n"),
                          CommitOperationAdd(path_in_repo=STATE_ARCHIVE, path_or_fileobj=state_file)]
            if raw_files:
                raw_file = temp / "raw.tar.gz"
                raw_manifest = pack_files(self.root, raw_files, raw_file)
                prefix = "raw/" + raw_manifest["sha256"]
                operations.extend([
                    CommitOperationAdd(path_in_repo=prefix + ".tar.gz", path_or_fileobj=raw_file),
                    CommitOperationAdd(path_in_repo=prefix + ".json", path_or_fileobj=json.dumps(raw_manifest).encode()),
                ])
                manifest["raw_batch"] = prefix
            operations.append(CommitOperationAdd(path_in_repo=STATE_MANIFEST, path_or_fileobj=json.dumps(manifest).encode()))
            # A stale worker must not reload a newer parent and overwrite it.
            commit = self.api.create_commit(self.repo_id, repo_type="dataset", operations=operations,
                                           parent_commit=self.revision, commit_message="Checkpoint live evaluation state")
            self.revision = commit.oid
            # Raw data are removed only inside the cloud runtime after an atomic
            # verified commit. The local migration backup is never deleted here.
            if os.getenv("LIVEHOUSE_CLOUD") == "1":
                for path in raw_files:
                    path.unlink()
            status_path = self.root / ".cloud-state/checkpoint.json"
            status_path.parent.mkdir(exist_ok=True)
            status_path.write_text(json.dumps({"checkpointed_at":manifest["checkpointed_at"], "revision":commit.oid}) + "\n")
            return manifest