File size: 1,952 Bytes
7034da5 | 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 | """Download the pinned accession index from bucket storage once per runtime."""
from contextlib import closing
import fcntl
import json
from pathlib import Path
import sqlite3
from huggingface_hub import HfApi, get_token
ROOT = Path(__file__).resolve().parent
def catalog_path():
pointer = ROOT / "data/catalog_snapshot.json"
if not pointer.exists():
return ROOT / "data/catalog.sqlite"
spec = json.loads(pointer.read_text())
cache = ROOT / ".cache/catalog"
cache.mkdir(parents=True, exist_ok=True)
target = cache / (spec["xet_hash"] + ".sqlite")
with (cache / "download.lock").open("a") as lock:
fcntl.flock(lock, fcntl.LOCK_EX)
if target.exists() and target.stat().st_size == spec["size"]:
return target
api = HfApi(token=get_token())
entries = list(api.get_bucket_paths_info(spec["bucket_id"], [spec["path"]]))
if len(entries) != 1 or entries[0].xet_hash != spec["xet_hash"] or entries[0].size != spec["size"]:
raise RuntimeError("The pinned accession snapshot is missing or has changed.")
temporary = target.with_suffix(".downloading")
print("Downloading accession index snapshot from bucket storage", flush=True)
api.download_bucket_files(spec["bucket_id"], [(entries[0], temporary)], raise_on_missing_files=True)
if temporary.stat().st_size != spec["size"]:
raise RuntimeError("Incomplete accession index download.")
with closing(sqlite3.connect(temporary.resolve().as_uri() + "?mode=ro", uri=True)) as connection:
manifest = json.loads(connection.execute("SELECT value FROM metadata WHERE key='manifest'").fetchone()[0])
if manifest.get("inventory_sha256") != spec["inventory_sha256"] or not manifest.get("full_snapshot"):
raise RuntimeError("Accession index does not match the published inventory.")
temporary.replace(target)
return target
|