pymc-cache 0.1.0__py3-none-any.whl

pymc_cache/__init__.py

@@ -0,0 +1,13 @@
+"""Model-aware result caching for standalone PyMC sampling."""
+from pymc_cache._version import __version__
+from pymc_cache.api import sample, sample_smc
+from pymc_cache.backends import CacheBackend, CacheEntry, DiskCache, MemoryCache, cleanup_cache
+from pymc_cache.identity import UncacheableError, cache_key, model_digest
+from pymc_cache.policy import CachePolicy, CacheWarning
+from pymc_cache.session import CacheLookup, CacheSession
+
+__all__ = [
+    "sample", "sample_smc", "CacheBackend", "MemoryCache", "DiskCache", "CacheSession",
+    "CachePolicy", "CacheEntry", "CacheLookup", "UncacheableError", "CacheWarning",
+    "cleanup_cache", "model_digest", "cache_key", "__version__",
+]

pymc_cache/_version.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

pymc_cache/api.py

@@ -0,0 +1,127 @@
+"""Explicit (non-monkeypatching) PyMC sampling with reusable results."""
+from __future__ import annotations
+
+import copy
+import importlib.util
+import os
+import warnings
+
+from pymc_cache.policy import CachePolicy, CacheWarning, report_from_idata
+from pymc_cache.session import CacheSession
+
+
+def _model(model):
+    import pymc as pm
+
+    return model if model is not None else pm.modelcontext(None)
+
+
+def _native_kwargs(model, kwargs: dict, *, operation: str) -> tuple[dict, dict, str]:
+    import pymc as pm
+
+    call = dict(kwargs)
+    key = dict(kwargs)
+    key.pop("progressbar", None)
+    if operation == "smc":
+        call.setdefault("draws", 2000)
+        call.setdefault("chains", max(2, min(4, os.cpu_count() or 1)))
+        key.update({"draws": call["draws"], "chains": call["chains"]})
+        kernel = call.get("kernel")
+        if kernel is not None and getattr(kernel, "__module__", "").startswith("pymc.smc"):
+            key["kernel"] = f"{kernel.__module__}.{kernel.__qualname__}"
+        return call, key, "smc"
+    call.setdefault("draws", 1000)
+    cores = call.get("cores") or min(4, os.cpu_count() or 1)
+    call.setdefault("cores", cores)
+    call.setdefault("chains", max(2, min(4, cores)))
+    sampler = call.get("nuts_sampler")
+    if sampler is None:
+        if int(pm.__version__.split(".")[0]) >= 6 and importlib.util.find_spec("nutpie"):
+            discrete = any(str(rv.dtype).startswith(("int", "uint", "bool"))
+                           for rv in model.free_RVs)
+            sampler = "pymc" if discrete else "nutpie"
+        else:
+            sampler = "pymc"
+        call["nuts_sampler"] = sampler
+    call["tune"] = call.get("tune") if call.get("tune") is not None else (
+        400 if int(pm.__version__.split(".")[0]) >= 6 and sampler == "nutpie" else 1000
+    )
+    key.update({k: call[k] for k in ("draws", "tune", "chains", "cores", "nuts_sampler")})
+    return call, key, sampler
+
+
+def _runner(model, call_kwargs, sample_fn, *, operation):
+    import pymc as pm
+
+    func = sample_fn if sample_fn is not None else (
+        pm.sample_smc if operation == "smc" else pm.sample
+    )
+    with model:
+        return func(model=model, **call_kwargs)
+
+
+def _call(*, model=None, cache=True, overwrite=False, policy=None, namespace="",
+          model_key=None, sampler_key=None, sample_fn=None, operation="sample", **sample_kwargs):
+    model = _model(model)
+    if policy is None:
+        policy = CachePolicy()
+    if not isinstance(policy, CachePolicy):
+        raise TypeError("policy must be CachePolicy")
+    call = dict(sample_kwargs)
+    key_kwargs = dict(call)
+    sampler_identity = operation
+    versions = None
+    hook = getattr(sample_fn, "__pymc_cache_normalize__", None)
+    if sample_fn is not None and hook is not None:
+        normalized = hook(model, copy.deepcopy(call))
+        if not isinstance(normalized, dict) or set(normalized) != {
+            "call_kwargs", "key_kwargs", "sampler_identity", "versions"
+        }:
+            raise TypeError("sample_fn normalization must return call_kwargs, key_kwargs, "
+                            "sampler_identity, versions")
+        call = dict(normalized["call_kwargs"])
+        key_kwargs = dict(normalized["key_kwargs"])
+        sampler_identity = normalized["sampler_identity"]
+        versions = dict(normalized["versions"])
+    elif sample_fn is None:
+        call, key_kwargs, sampler_identity = _native_kwargs(model, call, operation=operation)
+    elif cache is not False and sampler_key is None:
+        warnings.warn("sampling result is not cacheable: sample_fn requires sampler_key or "
+                      "__pymc_cache_normalize__", CacheWarning, stacklevel=2)
+        cache = False
+    if cache is not False and call.get("return_inferencedata") is False:
+        warnings.warn("sampling result is not cacheable: return_inferencedata=False",
+                      CacheWarning, stacklevel=2)
+        cache = False
+    session = CacheSession(cache, model=model, policy=policy, namespace=namespace,
+                           model_key=model_key, sampler_key=sampler_key)
+    lookup = session.lookup(key_kwargs, operation=operation, overwrite=overwrite,
+                            versions=versions, sampler_identity=sampler_identity)
+    if lookup.hit:
+        return lookup.result
+    result = _runner(model, call, sample_fn, operation=operation)
+    report = report_from_idata(result)
+    status = report.get("status", "complete") if report else "complete"
+    session.store(lookup, result, status=status, effective_sampler=sampler_identity)
+    return result
+
+
+def sample(*, model=None, cache=True, overwrite=False, policy=None, namespace="",
+           model_key=None, sampler_key=None, sample_fn=None, **sample_kwargs):
+    """Run ``pm.sample`` (or a trusted runner), reusing a matching result.
+
+    No global PyMC function is replaced. ``cache='disk'`` persists under
+    ``.pymc-cache``; ``cache=False`` always samples anew. A user-supplied
+    sample_fn must opt into safe identity or supply a trusted sampler_key.
+    """
+    return _call(model=model, cache=cache, overwrite=overwrite, policy=policy,
+                 namespace=namespace, model_key=model_key, sampler_key=sampler_key,
+                 sample_fn=sample_fn, operation="sample", **sample_kwargs)
+
+
+def sample_smc(*, model=None, cache=True, overwrite=False, policy=None, namespace="",
+               model_key=None, sampler_key=None, **sample_kwargs):
+    """Run and cache ``pm.sample_smc`` without conflating it with NUTS."""
+    return _call(model=model, cache=cache, overwrite=overwrite, policy=policy,
+                 namespace=namespace, model_key=model_key, sampler_key=sampler_key,
+                 operation="smc", **sample_kwargs)

pymc_cache/backends.py

@@ -0,0 +1,262 @@
+"""Cache backends: bounded in-memory LRU and atomic, inspectable NetCDF disk storage."""
+from __future__ import annotations
+
+import copy
+import json
+import os
+import re
+import threading
+import uuid
+import warnings
+from collections import OrderedDict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Protocol
+
+from pymc_cache.io import get_group, load_inference_data, save_inference_data
+from pymc_cache.naming import model_slug, params_label
+from pymc_cache.policy import CacheWarning
+
+_ENTRY_SCHEMA = 3
+_ENTRY_META = "_pymc_cache_entry_v3"
+
+
+class CacheBackend(Protocol):
+    """Minimal backend interface; extra management operations are optional."""
+
+    def load(self, key: str, **kwargs): ...
+
+    def save(self, key: str, idata, **kwargs) -> None: ...
+
+
+class MemoryCache:
+    """Thread-safe, copy-isolated LRU of posterior results."""
+
+    def __init__(self, max_entries: int = 8):
+        if isinstance(max_entries, bool) or not isinstance(max_entries, int) or max_entries < 1:
+            raise ValueError("max_entries must be a positive integer")
+        self._max_entries = max_entries
+        self._store: OrderedDict[str, object] = OrderedDict()
+        self._lock = threading.RLock()
+
+    def load(self, key: str, **kwargs):
+        with self._lock:
+            if key not in self._store:
+                return None
+            self._store.move_to_end(key)
+            return copy.deepcopy(self._store[key])
+
+    def save(self, key: str, idata, **kwargs) -> None:
+        snapshot = copy.deepcopy(idata)
+        with self._lock:
+            self._store[key] = snapshot
+            self._store.move_to_end(key)
+            while len(self._store) > self._max_entries:
+                self._store.popitem(last=False)
+
+    def clear(self) -> int:
+        """Discard stored entries and return the number removed."""
+        with self._lock:
+            removed = len(self._store)
+            self._store.clear()
+            return removed
+
+    def __len__(self):
+        with self._lock:
+            return len(self._store)
+
+
+@dataclass(frozen=True)
+class CacheEntry:
+    """Metadata for an owned, valid disk entry (without loading posterior arrays)."""
+
+    key: str
+    model_slug: str
+    operation: str
+    created_at: float
+    size_bytes: int
+    status: str
+    path: Path
+    software_versions: dict[str, str]
+
+
+class DiskCache:
+    """One-file, full-key-verified atomic NetCDF result cache.
+
+    Entries live at ``<root>/<model>/<parameters>-<full-sha256>.nc``. Contents
+    have an independently checked full key and schema; legacy files are never
+    reused. The directory is a trusted local cache, not an untrusted file reader.
+    """
+
+    def __init__(self, base_dir: str | Path | None = None, model=None, *,
+                 max_entries_per_model: int = 20):
+        if (isinstance(max_entries_per_model, bool) or
+                not isinstance(max_entries_per_model, int) or max_entries_per_model < 1):
+            raise ValueError("max_entries_per_model must be a positive integer")
+        self._base = Path(base_dir) if base_dir is not None else Path(".pymc-cache")
+        self._model_slug = model_slug(model)
+        self._keep = max_entries_per_model
+
+    def _path(self, key: str, sample_kwargs: dict | None = None) -> Path:
+        if not isinstance(key, str) or not re.fullmatch(r"[a-f0-9]{64}", key):
+            raise ValueError("cache key must be a full 64-character SHA-256 hexadecimal digest")
+        path = self._base / self._model_slug / f"{params_label(sample_kwargs)}-{key}.nc"
+        if not path.resolve().is_relative_to(self._base.resolve()):
+            raise ValueError("cache path escapes the cache root")
+        return path
+
+    @staticmethod
+    def _read_metadata(path: Path) -> dict:
+        import h5py
+
+        with h5py.File(path, "r") as f:
+            # Both ArviZ 0 and ArviZ 1 write group attrs to the HDF5 group.
+            raw = f["posterior"].attrs[_ENTRY_META]
+        if isinstance(raw, bytes):
+            raw = raw.decode("utf-8")
+        data = json.loads(raw)
+        if data.get("schema") != _ENTRY_SCHEMA:
+            raise ValueError("unknown cache schema")
+        return data
+
+    def load(self, key: str, sample_kwargs: dict | None = None):
+        path = self._path(key, sample_kwargs)
+        if self._base.is_symlink() or path.parent.is_symlink() or path.is_symlink() or not path.exists():
+            return None
+        try:
+            metadata = self._read_metadata(path)
+            if metadata.get("key") != key or metadata.get("status") not in {"complete", "met"}:
+                raise ValueError("cache entry identity or status does not match")
+            idata = load_inference_data(path)
+            post = get_group(idata, "posterior")
+            post.attrs.pop(_ENTRY_META, None)
+            return idata
+        except (OSError, ValueError, KeyError, TypeError, RuntimeError) as exc:
+            warnings.warn(f"pymc-cache: ignoring unreadable cache file {path} ({exc})",
+                          CacheWarning, stacklevel=2)
+            return None
+
+    def save(self, key: str, idata, sample_kwargs: dict | None = None, *,
+             operation="sample", status="complete", software_versions=None) -> None:
+        path = self._path(key, sample_kwargs)
+        if status not in ("complete", "met"):
+            raise ValueError("only complete or met results may be stored")
+        if self._base.is_symlink() or path.parent.is_symlink() or path.is_symlink():
+            raise ValueError("refusing to write through a cache symlink")
+        path.parent.mkdir(parents=True, exist_ok=True)
+        if self._base.is_symlink() or path.parent.is_symlink():
+            raise ValueError("refusing to write through a cache symlink")
+        if path.exists():
+            import h5py
+
+            try:
+                with h5py.File(path, "r") as existing:
+                    legacy = "posterior" in existing and any(
+                        f"_pymc_cache_entry_v{version}" in existing["posterior"].attrs
+                        for version in (1, 2)
+                    )
+            except OSError:
+                legacy = False
+            if legacy:
+                raise ValueError("refusing to overwrite a schema-1/2 cache entry")
+        import time
+
+        metadata = {"key": key, "schema": _ENTRY_SCHEMA, "status": status,
+                    "operation": operation, "created_at": time.time(),
+                    "software_versions": software_versions or {}}
+        tmp = path.with_name(f"{path.name}.{os.getpid()}.{uuid.uuid4().hex}.tmp")
+        try:
+            save_inference_data(idata, tmp, attrs={_ENTRY_META: json.dumps(metadata)})
+            os.replace(tmp, path)
+        finally:
+            tmp.unlink(missing_ok=True)
+        self.prune(self._keep)
+
+    def entries(self) -> list[CacheEntry]:
+        """List this format's entries without materializing posterior data."""
+        results = []
+        if not self._base.exists() or self._base.is_symlink():
+            return results
+        for directory in self._base.iterdir():
+            if not directory.is_dir() or directory.is_symlink():
+                continue
+            for path in directory.glob("*.nc"):
+                if path.is_symlink() or not path.is_file():
+                    continue
+                try:
+                    metadata = self._read_metadata(path)
+                    key = metadata["key"]
+                    if not re.fullmatch(r"[0-9a-f]{64}", key) or not path.name.endswith(
+                        f"-{key}.nc"
+                    ):
+                        continue
+                    results.append(CacheEntry(key, directory.name, metadata["operation"],
+                                              float(metadata["created_at"]), path.stat().st_size,
+                                              metadata["status"], path,
+                                              metadata.get("software_versions", {})))
+                except (OSError, ValueError, KeyError, TypeError):
+                    continue
+        return sorted(results, key=lambda entry: entry.created_at, reverse=True)
+
+    def delete(self, key: str) -> bool:
+        """Remove matching owned entries, never arbitrary files or symlinks."""
+        matches = [entry for entry in self.entries() if entry.key == key]
+        for entry in matches:
+            entry.path.unlink()
+        return bool(matches)
+
+    def prune(self, keep: int | None = None) -> int:
+        """Keep the N newest entries per model; return count removed."""
+        if keep is None:
+            keep = self._keep
+        if isinstance(keep, bool) or not isinstance(keep, int) or keep < 0:
+            raise ValueError("keep must be a nonnegative integer")
+        counts: dict[str, int] = {}
+        removed = 0
+        for entry in self.entries():
+            current = counts.get(entry.model_slug, 0)
+            counts[entry.model_slug] = current + 1
+            if current >= keep:
+                entry.path.unlink()
+                removed += 1
+        return removed
+
+    def clear(self) -> int:
+        """Delete only format-owned entries, preserving other files."""
+        entries = self.entries()
+        for entry in entries:
+            entry.path.unlink()
+        return len(entries)
+
+
+_default_memory_cache: MemoryCache | None = None
+_default_lock = threading.Lock()
+
+
+def get_default_cache() -> MemoryCache:
+    """Lazily construct the process-wide memory cache."""
+    global _default_memory_cache
+    with _default_lock:
+        if _default_memory_cache is None:
+            _default_memory_cache = MemoryCache()
+        return _default_memory_cache
+
+
+def resolve_cache(cache_arg, model=None, *, default_directory=".pymc-cache") -> CacheBackend | None:
+    """Resolve bool, memory/disk alias, directory, or custom backend."""
+    if cache_arg is False:
+        return None
+    if cache_arg is True or cache_arg == "memory":
+        return get_default_cache()
+    if cache_arg == "disk":
+        return DiskCache(default_directory, model=model)
+    if isinstance(cache_arg, (str, Path)):
+        return DiskCache(cache_arg, model=model)
+    if callable(getattr(cache_arg, "load", None)) and callable(getattr(cache_arg, "save", None)):
+        return cache_arg
+    raise TypeError(f"cache must be bool | 'disk' | str | Path | CacheBackend; got {cache_arg!r}")
+
+
+def cleanup_cache(base_dir: str | Path | None = None) -> int:
+    """Remove only new-format entries in a directory; do not delete the root."""
+    return DiskCache(base_dir).clear()

pymc_cache/cli.py

@@ -0,0 +1,50 @@
+"""Inspect and manage only entries owned by pymc-cache."""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+from pymc_cache.backends import DiskCache
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="pymc-cache", description=__doc__)
+    parser.add_argument("--directory", "--dir", dest="directory", default=".pymc-cache",
+                        help="cache directory (default: .pymc-cache)")
+    sub = parser.add_subparsers(dest="command", required=True)
+    listed = sub.add_parser("list", help="inspect owned cache entries")
+    listed.add_argument("--json", action="store_true", help="machine-readable metadata")
+    prune = sub.add_parser("prune", help="keep the newest entries per model")
+    prune.add_argument("--keep", type=int, default=20)
+    clear = sub.add_parser("clear", help="delete only entries owned by pymc-cache")
+    clear.add_argument("--yes", action="store_true", required=True,
+                       help="confirm permanent removal of owned entries")
+    for command in (listed, prune, clear):
+        command.add_argument("--directory", "--dir", dest="directory",
+                             default=argparse.SUPPRESS, help="cache directory")
+    args = parser.parse_args(argv)
+    cache = DiskCache(Path(args.directory))
+    if args.command == "list":
+        entries = cache.entries()
+        if args.json:
+            print(json.dumps([{"key": e.key, "model": e.model_slug,
+                               "operation": e.operation, "status": e.status,
+                               "path": str(e.path), "size_bytes": e.size_bytes,
+                               "created_at": e.created_at,
+                               "software_versions": e.software_versions}
+                              for e in entries], indent=2))
+        else:
+            for entry in entries:
+                print(f"{entry.key}  {entry.model_slug}  {entry.status}  "
+                      f"{entry.size_bytes} B  {entry.path}")
+            print(f"{len(entries)} entries")
+    elif args.command == "prune":
+        print(f"Removed {cache.prune(args.keep)} entries")
+    else:
+        print(f"Removed {cache.clear()} entries")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())