caj 0.0.1__py3-none-any.whl

caj/__init__.py

@@ -0,0 +1,116 @@
+import functools
+import hashlib
+import os
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from typing import overload
+from warnings import warn
+
+if sys.version_info < (3, 15):
+    from typing_extensions import sentinel
+
+import jax
+from platformdirs import user_cache_path
+
+from ._cache import Cache
+from ._serialization import deserialize_pytree, serialize_jaxpr, serialize_pytree
+from ._writers import HashWriter, WriteLimitError
+
+DEFAULT_CACHE_DIR = user_cache_path(__package__)
+DEFAULT_CACHE_MAX_BYTES = 1_000_000_000
+
+_MISSING = sentinel("_MISSING")
+
+
+@overload
+def cache[**P, R](func: Callable[P, R], /) -> Callable[P, R]: ...
+
+
+@overload
+def cache[**P, R](
+    func: Callable[P, R],
+    /,
+    *,
+    dir: str | os.PathLike[str],
+    max_bytes: int | None = ...,
+) -> Callable[P, R]: ...
+
+
+@overload
+def cache[**P, R](
+    *, dir: str | os.PathLike[str], max_bytes: int | None = ...
+) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
+
+
+def cache[**P, R](
+    func: Callable[P, R] | _MISSING = _MISSING,
+    /,
+    *,
+    dir: str | os.PathLike[str] | _MISSING = _MISSING,
+    max_bytes: int | _MISSING = _MISSING,
+) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
+    if dir is not _MISSING:
+        dir = Path(dir).absolute()
+
+        if max_bytes is _MISSING:
+            max_bytes = DEFAULT_CACHE_MAX_BYTES
+        elif max_bytes is not None and max_bytes <= 0:
+            raise ValueError(f"max_bytes must be None or positive, got {max_bytes!r}")
+
+    elif max_bytes is not _MISSING:
+        raise TypeError("max_bytes is only allowed if dir is also specified")
+
+    else:
+        dir = DEFAULT_CACHE_DIR
+        max_bytes = DEFAULT_CACHE_MAX_BYTES
+
+    cache = Cache(dir, max_bytes=max_bytes)
+
+    def decorator(func: Callable[P, R], /) -> Callable[P, R]:
+        jaxpr_func = jax.make_jaxpr(func, return_shape=True)
+
+        @functools.wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            jaxpr, ret_shape = jaxpr_func(*args, **kwargs)
+
+            h = hashlib.blake2b(digest_size=16)
+
+            f = HashWriter(h)
+            serialize_jaxpr(f, jaxpr)
+            serialize_pytree(f, args)
+            serialize_pytree(f, kwargs)
+
+            key = h.hexdigest()
+
+            try:
+                with cache.read(key) as f:
+                    ret = deserialize_pytree(f, like=ret_shape)
+            except KeyError:
+                pass
+            except (OSError, RuntimeError) as e:
+                warn(
+                    f"{__package__}: failed to load cached entry: {e}",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+            else:
+                return ret
+
+            ret = func(*args, **kwargs)
+
+            try:
+                with cache.write(key) as f:
+                    serialize_pytree(f, ret, exc_workaround=WriteLimitError)
+            except (OSError, WriteLimitError) as e:
+                warn(
+                    f"{__package__}: failed to save cache to {cache.dir}: {e}",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+
+            return ret
+
+        return wrapper
+
+    return decorator(func) if func is not _MISSING else decorator

caj/_cache.py

@@ -0,0 +1,113 @@
+import contextlib
+import os
+import sys
+import time
+from collections.abc import Generator
+from pathlib import Path
+from warnings import warn
+
+if sys.version_info >= (3, 14):
+    from io import Writer
+else:
+    from typing_extensions import Writer
+
+from safewrite import atomic_write
+
+from ._typing import SupportsReadSeek
+from ._writers import BoundedWriter
+
+
+class Cache:
+    _SUFFIX = f".{__package__}"
+
+    def __init__(self, dir: Path, /, *, max_bytes: int | None = None) -> None:
+        self.dir = dir
+        self.max_bytes = max_bytes
+
+    @contextlib.contextmanager
+    def read(self, key: str, /) -> Generator[SupportsReadSeek[bytes], None, None]:
+        path = self._path(key)
+        try:
+            f = path.open("rb")
+        except FileNotFoundError as e:
+            raise KeyError(key) from e
+        try:
+            yield f
+        finally:
+            f.close()
+        self._hit(key)
+
+    @contextlib.contextmanager
+    def write(self, key: str, /) -> Generator[Writer[bytes], None, None]:
+        path = self._path(key)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with atomic_write(path, mode="wb") as f:
+            if self.max_bytes is not None:
+                f = BoundedWriter(f, max_bytes=self.max_bytes)
+            yield f
+        self._cull()
+
+    def _path(self, key: str, /) -> Path:
+        return self.dir / f"{key}{self._SUFFIX}"
+
+    def _hit(self, key: str, /) -> None:
+        path = self._path(key)
+        with contextlib.suppress(OSError):
+            os.utime(path, ns=(time.time_ns(), path.stat().st_mtime_ns))
+
+    def _cull(self) -> None:
+        if self.max_bytes is None:
+            return
+
+        cache_entries: list[tuple[Path, int, int]] = []
+
+        try:
+            with os.scandir(self.dir) as dir_entries:
+                for dir_entry in dir_entries:
+                    path = Path(dir_entry.path)
+
+                    if path.suffix != self._SUFFIX:
+                        continue
+
+                    try:
+                        if not dir_entry.is_file(follow_symlinks=False):
+                            continue
+                        st = dir_entry.stat(follow_symlinks=False)
+                    except OSError:
+                        continue
+
+                    cache_entries.append((path, st.st_atime_ns, st.st_size))
+        except OSError:
+            warn(
+                f"{__package__}: failed to cull cache in {self.dir} to {self.max_bytes} bytes; "
+                "unable to list files",
+                RuntimeWarning,
+                stacklevel=3,
+            )
+            return
+
+        total_size = sum(entry[2] for entry in cache_entries)
+
+        if total_size <= self.max_bytes:
+            return
+
+        cache_entries.sort(key=lambda entry: entry[1])
+
+        for cache_entry in cache_entries:
+            assert cache_entry[0].suffix == self._SUFFIX
+            try:
+                cache_entry[0].unlink(missing_ok=True)
+            except OSError:
+                continue
+
+            total_size -= cache_entry[2]
+
+            if total_size <= self.max_bytes:
+                return
+
+        warn(
+            f"{__package__}: failed to cull cache in {self.dir} to {self.max_bytes} bytes; "
+            f"current size is at least {total_size} bytes",
+            RuntimeWarning,
+            stacklevel=3,
+        )

caj/_serialization.py

@@ -0,0 +1,71 @@
+import sys
+from typing import Any
+
+if sys.version_info >= (3, 14):
+    from io import Writer
+else:
+    from typing_extensions import Writer
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import jaxlib
+
+from ._typing import SupportsReadSeek
+
+
+def serialize_jaxpr(f: Writer[bytes], jaxpr: Any, /) -> None:
+    f.write(jax.__version__.encode())
+    f.write(jaxlib.__version__.encode())
+    f.write(str(jaxpr).encode())
+    serialize_pytree(f, jaxpr.consts)
+
+
+# Workaround for https://github.com/patrick-kidger/equinox/issues/1255
+def serialize_filter_spec(f: Writer[bytes], x: object, /) -> None:
+    if isinstance(x, jax.Array) and jnp.issubdtype(x.dtype, jax.dtypes.prng_key):
+        x = jax.random.key_data(x)
+    return eqx.default_serialise_filter_spec(f, x)
+
+
+def serialize_pytree(
+    f: Writer[bytes],
+    pytree: object,
+    /,
+    *,
+    exc_workaround: type[Exception] | None = None,
+) -> None:
+    if exc_workaround is not None:
+        # Workaround for https://github.com/patrick-kidger/equinox/issues/1255
+        assert issubclass(exc_workaround, Exception)
+        assert not issubclass(exc_workaround, RuntimeError)
+
+        try:
+            eqx.tree_serialise_leaves(f, pytree, filter_spec=serialize_filter_spec)
+        except RuntimeError as e:
+            cause = e.__cause__
+            while isinstance(cause, RuntimeError):
+                cause = cause.__cause__
+            if isinstance(cause, exc_workaround):
+                raise cause from e  # ty: ignore[invalid-raise]
+            raise
+    else:
+        eqx.tree_serialise_leaves(f, pytree, filter_spec=serialize_filter_spec)
+
+
+# Workaround for https://github.com/patrick-kidger/equinox/issues/1255
+def deserialize_filter_spec(f: SupportsReadSeek[bytes], x: object, /) -> object:
+    ret = eqx.default_deserialise_filter_spec(f, x)
+    if isinstance(x, (jax.Array, jax.ShapeDtypeStruct)) and jnp.issubdtype(
+        x.dtype, jax.dtypes.prng_key
+    ):
+        return jax.random.wrap_key_data(
+            ret, impl=jax.random.key_impl(jnp.zeros((), x.dtype))
+        )
+    return ret
+
+
+def deserialize_pytree[T](f: SupportsReadSeek[bytes], /, *, like: T) -> T:
+    return eqx.tree_deserialise_leaves(
+        f, like=like, filter_spec=deserialize_filter_spec
+    )

caj/_typing.py

@@ -0,0 +1,16 @@
+import sys
+from collections.abc import Buffer
+from typing import Protocol
+
+if sys.version_info >= (3, 14):
+    from io import Reader
+else:
+    from typing_extensions import Reader
+
+
+class SupportsReadSeek[T](Reader[T], Protocol):
+    def seek(self, offset: int, whence: int, /) -> object: ...
+
+
+class SupportsUpdate(Protocol):
+    def update(self, data: Buffer, /) -> None: ...

caj/_writers.py

@@ -0,0 +1,47 @@
+import sys
+from collections.abc import Buffer
+from typing import override
+
+if sys.version_info >= (3, 14):
+    from io import Writer
+else:
+    from typing_extensions import Writer
+
+
+from ._typing import SupportsUpdate
+
+
+class HashWriter(Writer[Buffer]):
+    def __init__(self, hasher: SupportsUpdate, /) -> None:
+        self._hasher = hasher
+
+    @override
+    def write(self, data: Buffer, /) -> int:
+        data = memoryview(data)
+        self._hasher.update(data)
+        return data.nbytes
+
+
+class WriteLimitError(Exception):
+    pass
+
+
+class BoundedWriter(Writer[Buffer]):
+    def __init__(self, writer: Writer[Buffer], /, *, max_bytes: int) -> None:
+        self._writer = writer
+        self._max_bytes = max_bytes
+        self._bytes_written = 0
+
+    @override
+    def write(self, data: Buffer, /) -> int:
+        data = memoryview(data)
+        if self._bytes_written + data.nbytes > self._max_bytes:
+            raise WriteLimitError(
+                f"write would exceed limit of {self._max_bytes} bytes"
+            )
+
+        ret = self._writer.write(data)
+        assert 0 <= ret <= data.nbytes
+
+        self._bytes_written += ret
+        return ret

caj-0.0.1.dist-info/METADATA

@@ -0,0 +1,111 @@
+Metadata-Version: 2.4
+Name: caj
+Version: 0.0.1
+Summary: Automatic persistent caching for JAX-based computations
+Author: Gabriel S. Gerlero
+Author-email: Gabriel S. Gerlero <ggerlero@cimec.unl.edu.ar>
+License-Expression: Apache-2.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Typing :: Typed
+Requires-Dist: equinox>=0.13,<0.14
+Requires-Dist: jax>=0.5,<0.12
+Requires-Dist: platformdirs>=2.1,<5
+Requires-Dist: safewrite>=0.1.3,<0.2
+Requires-Dist: typing-extensions>=4.16,<5
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+<div align="center">
+  <a href="https://github.com/gerlero/caj"><img src="https://raw.githubusercontent.com/gerlero/caj/main/logo.png" alt="caj" width="200"/></a>
+
+  **Automatic persistent caching for JAX function invocations**
+
+  [![CI](https://github.com/gerlero/caj/actions/workflows/ci.yml/badge.svg)](https://github.com/gerlero/caj/actions/workflows/ci.yml)
+  [![Codecov](https://codecov.io/gh/gerlero/caj/branch/main/graph/badge.svg)](https://codecov.io/gh/gerlero/caj)
+  [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+  [![ty](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ty/main/assets/badge/v0.json)](https://github.com/astral-sh/ty)
+  [![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+  [![Publish](https://github.com/gerlero/caj/actions/workflows/pypi-publish.yml/badge.svg)](https://github.com/gerlero/caj/actions/workflows/pypi-publish.yml)
+  [![PyPI](https://img.shields.io/pypi/v/caj)](https://pypi.org/project/caj/)
+  [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/caj)](https://pypi.org/project/caj/)
+</div>
+
+**caj** is a simple persistent cache for JAX-based code.
+
+Decorate a function with `@cache`, and **caj** will store the return on disk. Later calls with the same inputs will load the result directly from the cache instead of performing the computation again.
+
+```python
+import jax.numpy as jnp
+from caj import cache
+
+
+@cache
+@jax.jit
+def compute(x):
+    return jnp.linalg.eigvalsh(x)
+
+
+x = jnp.eye(1000)
+
+y = compute(x)
+```
+
+Running the above code twice will only compute the eigenvalues once: the second time, **caj** will just load the result from disk.
+
+**caj** takes advantage of the computation tracing provided by JAX so it can detect changes in the decorated function or any other functions it calls and not return cached results that were computed with different code.
+
+## Installation
+
+Install **caj** from PyPI:
+
+```console
+pip install caj
+```
+
+## Usage
+
+For most uses, the default cache is enough:
+
+```python
+from caj import cache
+
+
+@cache
+def f(x): ...
+```
+
+By default, **caj** stores entries in a per-user cache directory and limits the cache to **1 GB**.
+
+A custom cache directory can be specified with `dir`:
+
+```python
+@cache(dir=".cache")
+def f(x): ...
+```
+
+The maximum size of a custom cache can also be configured via `max_bytes`:
+
+```python
+@cache(dir=".cache", max_bytes=2_000_000_000)
+def f(x): ...
+```
+
+Or, set `max_bytes=None` for no size limit:
+
+```python
+@cache(dir=".cache", max_bytes=None)
+def f(x): ...
+```
+
+When a size limit is enabled, **caj** will remove the least recently used entries as necessary to keep the cache within the target size.

caj-0.0.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+caj/__init__.py,sha256=eGwFjfCcbXANcflNUTlXYU1qjpt3ZuotGrp9HDVdCJQ,3241
+caj/_cache.py,sha256=2tW4htp4_qcDaEceZSLWltqt699_OSQrD4pnF_KZRgc,3336
+caj/_serialization.py,sha256=Bg2DuFQts3VZfCLqiDjrgZb9xbRoXi_4rGXI2WNu-_Y,2258
+caj/_typing.py,sha256=Izwkkk6vpLoP1Ay7f6-Tv-Ng5ZceGPPMF9rvm_WepwM,378
+caj/_writers.py,sha256=YuNzNxGwBhMjhjwqte8murAof_pPwk3AOawOl7X09Lk,1166
+caj/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+caj-0.0.1.dist-info/WHEEL,sha256=cmC5s21ojypbVslldL7IJq3hjZH-tINy4rziKePFsG0,81
+caj-0.0.1.dist-info/METADATA,sha256=6xcuacj9HXXYlCOQsA_GzHwS-drPahNoaegwmdaUg1g,3946
+caj-0.0.1.dist-info/RECORD,,

caj-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.23
+Root-Is-Purelib: true
+Tag: py3-none-any