src/services/persistence.py

"""Persistence plane (spec §11): local-first experiment store, mirrored to a
private HF dataset repo when a write token is available. The manifest is the
authoritative record; index.jsonl is a rebuildable dashboard cache. Restart-safe
(spec P3): on boot the store pulls any experiments present on the Hub.
"""

import json
import os
import pathlib

from src.schemas import ExperimentManifest, artifact_envelope

LOCAL_ROOT = pathlib.Path(os.environ.get("MLOL_DATA_DIR", "mlol_data"))
HUB_REPO = os.environ.get("MLOL_EXPERIMENTS_REPO", "finpy1789/mlol-experiments")


class ExperimentStore:
    def __init__(self):
        self.root = LOCAL_ROOT / "experiments"
        self.root.mkdir(parents=True, exist_ok=True)
        self._hub_ok = None  # lazily determined

    # ---------- hub mirroring (best-effort; local always wins for reads) ----------

    def _api(self):
        from huggingface_hub import HfApi
        return HfApi()

    def hub_available(self) -> bool:
        if self._hub_ok is None:
            try:
                api = self._api()
                api.whoami()
                api.create_repo(HUB_REPO, repo_type="dataset", private=True, exist_ok=True)
                self._hub_ok = True
            except Exception:  # noqa: BLE001
                self._hub_ok = False
        return self._hub_ok

    def _hub_push(self, run_id: str, fname: str, local_path: pathlib.Path):
        if not self.hub_available():
            return
        try:
            self._api().upload_file(
                path_or_fileobj=str(local_path),
                path_in_repo=f"experiments/{run_id}/{fname}",
                repo_id=HUB_REPO, repo_type="dataset",
            )
        except Exception as e:  # noqa: BLE001
            print(f"[persist] hub push failed ({fname}): {e}")

    def recover_from_hub(self):
        """Pull manifests that exist on the Hub but not locally (Space restart)."""
        if not self.hub_available():
            return 0
        try:
            from huggingface_hub import hf_hub_download
            files = self._api().list_repo_files(HUB_REPO, repo_type="dataset")
        except Exception:  # noqa: BLE001
            return 0
        n = 0
        for f in files:
            parts = f.split("/")
            if len(parts) == 3 and parts[0] == "experiments":
                run_id, fname = parts[1], parts[2]
                local = self.root / run_id / fname
                if not local.exists():
                    local.parent.mkdir(parents=True, exist_ok=True)
                    try:
                        got = hf_hub_download(HUB_REPO, f, repo_type="dataset")
                        local.write_bytes(pathlib.Path(got).read_bytes())
                        n += 1
                    except Exception as e:  # noqa: BLE001
                        print(f"[persist] recover failed ({f}): {e}")
        return n

    # ---------- manifests ----------

    def save_manifest(self, m: ExperimentManifest):
        d = self.root / m.run_id
        d.mkdir(parents=True, exist_ok=True)
        p = d / "manifest.json"
        p.write_text(m.model_dump_json(indent=2))
        self._hub_push(m.run_id, "manifest.json", p)
        self.rebuild_index()

    def load_manifest(self, run_id: str) -> ExperimentManifest | None:
        p = self.root / run_id / "manifest.json"
        if not p.exists():
            return None
        return ExperimentManifest.model_validate_json(p.read_text())

    def list_runs(self) -> list[ExperimentManifest]:
        out = []
        for d in sorted(self.root.iterdir(), reverse=True):
            if (d / "manifest.json").exists():
                try:
                    out.append(ExperimentManifest.model_validate_json((d / "manifest.json").read_text()))
                except Exception as e:  # noqa: BLE001
                    print(f"[persist] unreadable manifest {d.name}: {e}")
        return out

    def rebuild_index(self):
        """index.jsonl is derived, never authoritative (spec §11)."""
        lines = [
            json.dumps({"run_id": m.run_id, "title": m.title, "state": m.state,
                        "domain": m.domain, "model": m.model_repo, "updated_at": m.updated_at})
            for m in self.list_runs()
        ]
        (self.root / "index.jsonl").write_text("\n".join(lines) + ("\n" if lines else ""))

    # ---------- artifacts ----------

    def save_artifact(self, run_id: str, name: str, payload: dict):
        m = self.load_manifest(run_id)
        if m is None:
            raise ValueError(f"unknown run {run_id}")
        p = self.root / run_id / name
        p.write_text(json.dumps(artifact_envelope(run_id, m, payload), indent=2, default=str))
        self._hub_push(run_id, name, p)

    def load_artifact(self, run_id: str, name: str) -> dict | None:
        p = self.root / run_id / name
        if not p.exists():
            return None
        data = json.loads(p.read_text())
        return data.get("payload", data)

    def save_binary(self, run_id: str, name: str, data: bytes):
        p = self.root / run_id / name
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_bytes(data)
        self._hub_push(run_id, name, p)
        return p

    def artifact_path(self, run_id: str, name: str) -> pathlib.Path:
        return self.root / run_id / name


_STORE = None


def get_store() -> ExperimentStore:
    global _STORE
    if _STORE is None:
        _STORE = ExperimentStore()
    return _STORE