Spaces:
Running
Running
File size: 4,498 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 | #!/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()
|