cache.py

"""Revision-pinned, checksum-validated ephemeral Dataset cache."""
from __future__ import annotations
import hashlib,json,os,re,shutil,threading,urllib.parse,urllib.request
from pathlib import Path

REVISION=re.compile(r"^[0-9a-f]{40}$")
class RevisionCache:
    def __init__(self,root:Path,manifest_url:str,dataset_id:str):
        self.root=Path(root);self.manifest_url=manifest_url;self.dataset_id=dataset_id;self._lock=threading.Lock();self._status={"status":"not_loaded","ready":False,"dataset_revision":None,"artifact_count":0,"cached_bytes":0,"error":None}
    def status(self): return dict(self._status)
    def _bytes(self,url):
        with urllib.request.urlopen(url,timeout=60) as response:return response.read()
    def _download(self,url,target):
        digest=hashlib.sha256();size=0
        with urllib.request.urlopen(url,timeout=120) as response,target.open("wb") as output:
            while True:
                block=response.read(1024*1024)
                if not block:break
                output.write(block);digest.update(block);size+=len(block)
        return size,digest.hexdigest()
    def _selected(self,manifest):
        result=[]
        for item in manifest.get("artifacts",[]):
            if item.get("storage")!="dataset":continue
            if item.get("content_type") not in {"application/vnd.apache.parquet","application/vnd.sqlite3"}:continue
            path=item.get("path","");parts=path.split("/")
            if not parts or any(not part or part in {".",".."} for part in parts):raise ValueError("unsafe artifact path")
            result.append(item)
        if not result:raise ValueError("manifest has no cache artifacts")
        return result
    def _valid(self,directory,manifest,items):
        ready=directory/"ready.json"
        if not ready.is_file():return False
        try:
            marker=json.loads(ready.read_text())
            if marker.get("manifest_sha256")!=hashlib.sha256(json.dumps(manifest,sort_keys=True,separators=(",", ":")).encode()).hexdigest():return False
            for item in items:
                path=directory/item["path"]
                if not path.is_file() or path.stat().st_size!=item["size_bytes"]:return False
                if hashlib.sha256(path.read_bytes()).hexdigest()!=item["checksum"]["value"]:return False
            return True
        except Exception:return False
    def bootstrap(self,verify_active:bool=True):
      with self._lock:
        previous=self.status();had_ready=previous.get("ready") is True
        if not had_ready:self._status.update({"status":"loading","ready":False,"error":None})
        try:
            manifest_raw=self._bytes(self.manifest_url);manifest=json.loads(manifest_raw);revision=manifest["dataset_revision"]
            if not REVISION.fullmatch(revision):raise ValueError("invalid Dataset revision")
            items=self._selected(manifest);self.root.mkdir(parents=True,exist_ok=True);final=self.root/revision
            canonical_manifest_sha=hashlib.sha256(json.dumps(manifest,sort_keys=True,separators=(",", ":")).encode()).hexdigest()
            if not verify_active and had_ready and previous.get("dataset_revision")==revision and Path(previous.get("cache_dir",""))==final and final.is_dir():
                return self.status()
            if final.exists() and not self._valid(final,manifest,items):shutil.rmtree(final)
            if not final.exists():
                stage=self.root/(".%s.partial-%d"%(revision,os.getpid()))
                if stage.exists():shutil.rmtree(stage)
                stage.mkdir()
                try:
                    for item in items:
                        target=stage/item["path"];target.parent.mkdir(parents=True,exist_ok=True)
                        url="https://huggingface.co/datasets/%s/resolve/%s/%s"%(self.dataset_id,revision,urllib.parse.quote(item["path"],safe="/=-._"))
                        size,digest=self._download(url,target)
                        if size!=item["size_bytes"] or digest!=item["checksum"]["value"]:raise ValueError("artifact checksum mismatch: "+item["path"])
                    (stage/"manifest.json").write_bytes(manifest_raw)
                    (stage/"ready.json").write_text(json.dumps({"dataset_revision":revision,"manifest_sha256":canonical_manifest_sha},sort_keys=True)+"\n")
                    os.replace(stage,final)
                except Exception:
                    if stage.exists():shutil.rmtree(stage)
                    raise
            self._status={"status":"ready","ready":True,"dataset_revision":revision,"artifact_count":len(items),"cached_bytes":sum(x["size_bytes"] for x in items),"cache_dir":str(final),"error":None}
            retained={revision,previous.get("dataset_revision")}
            for directory in self.root.iterdir():
                if directory.is_dir() and not directory.name.startswith(".") and directory.name not in retained:
                    try:shutil.rmtree(directory)
                    except OSError:pass
        except Exception as error:
            if had_ready:self._status=previous
            else:self._status.update({"status":"error","ready":False,"error":str(error)[:500]})
            raise
        return self.status()