Spaces:
Running on Zero
Running on Zero
File size: 5,480 Bytes
68c1777 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """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
|