Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """Initialize, restore, or archive the organization's durable operator state.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| from pathlib import Path | |
| import sys | |
| import tempfile | |
| ROOT = Path(__file__).resolve().parents[1] | |
| sys.path[:0] = [str(ROOT), str(ROOT / "src")] | |
| def main(): | |
| from dotenv import load_dotenv | |
| from huggingface_hub import CommitOperationAdd, hf_hub_download | |
| from tsfm_bench.remote_state import HANDOFF_ROOTS, RemoteState, pack_files, restore_archive, utc_now | |
| from scripts.daemon_service import owned_process | |
| load_dotenv(ROOT / ".env") | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("action", choices=["initialize", "restore", "archive-local", "restore-handoff"]) | |
| parser.add_argument("--repo", default=os.getenv("HF_STATE_REPO", "ThinkcatLab/LiveHouse-TS-state")) | |
| parser.add_argument("--root", type=Path, default=ROOT) | |
| args = parser.parse_args() | |
| root = args.root.resolve() | |
| pid_path = root / ".daemon.pid" | |
| if pid_path.exists() and owned_process(int(pid_path.read_text()), root): | |
| raise SystemExit("Stop this checkout's daemon before snapshotting or restoring state.") | |
| token = os.getenv("HF_TOKEN") | |
| if not token: | |
| raise SystemExit("HF_TOKEN is required") | |
| state = RemoteState(root, args.repo, token) | |
| if args.action == "initialize": | |
| manifest = state.checkpoint(initialize=True, archive_raw=True) | |
| print(json.dumps({"revision":state.revision, "files":len(manifest["files"]), | |
| "sha256":manifest["sha256"], "checkpointed_at":manifest["checkpointed_at"]})) | |
| elif args.action == "restore": | |
| manifest = state.restore() | |
| print(f"Restored {len(manifest['files'])} verified files from {state.revision}") | |
| elif args.action == "restore-handoff": | |
| if any((root / path).exists() for path in HANDOFF_ROOTS): | |
| raise SystemExit("Use an empty --root for historical handoff restore; existing data must not be replaced.") | |
| info = state.api.dataset_info(args.repo) | |
| if not info.private: | |
| raise SystemExit("Handoff archives must be private") | |
| def download(name): | |
| return Path(hf_hub_download(args.repo, name, repo_type="dataset", revision=info.sha, token=token)) | |
| pointer = json.loads(download("handoff/latest.json").read_text()) | |
| manifest = json.loads(download(pointer["manifest"]).read_text()) | |
| restore_archive(root, download(pointer["archive"]), manifest, allowed_roots=HANDOFF_ROOTS) | |
| print(f"Restored {len(manifest['files'])} verified historical files from {info.sha}") | |
| else: | |
| info = state.api.dataset_info(args.repo) | |
| if not info.private: | |
| raise SystemExit("Handoff archives must be private") | |
| # Include all available runtime archives, old inputs, papers and outputs. | |
| # Virtual environments/caches are reproducible; credentials use secrets. | |
| files = [] | |
| for name in HANDOFF_ROOTS: | |
| files.extend(p for p in (root / name).rglob("*") if p.is_file()) | |
| with tempfile.TemporaryDirectory(prefix="livehouse-handoff-") as temp: | |
| archive = Path(temp) / "operator.tar.gz" | |
| manifest = pack_files(root, sorted(files), archive) | |
| manifest.update(created_at=utc_now(), source_revision=os.getenv("LIVEHOUSE_SOURCE_REVISION", "local-migration")) | |
| prefix = "handoff/" + manifest["sha256"] | |
| commit = state.api.create_commit(args.repo, repo_type="dataset", parent_commit=info.sha, | |
| commit_message="Archive complete available operator data for organization handoff", | |
| operations=[ | |
| CommitOperationAdd(path_in_repo=prefix + ".tar.gz", path_or_fileobj=archive), | |
| CommitOperationAdd(path_in_repo=prefix + ".json", path_or_fileobj=json.dumps(manifest).encode()), | |
| CommitOperationAdd(path_in_repo="handoff/latest.json", path_or_fileobj=json.dumps({ | |
| "archive":prefix + ".tar.gz", "manifest":prefix + ".json", | |
| "created_at":manifest["created_at"], "sha256":manifest["sha256"], | |
| }).encode()), | |
| ]) | |
| print(json.dumps({"revision":commit.oid, "archive":prefix + ".tar.gz", "files":len(files), | |
| "bytes":manifest["bytes"], "sha256":manifest["sha256"]})) | |
| if __name__ == "__main__": | |
| main() | |