Spaces:
Sleeping
Sleeping
File size: 10,783 Bytes
e317359 | 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 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | """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
|