diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 625480cba..5eb63fe57 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -53,7 +53,7 @@ jobs: - name: Set python test settings run: | - echo "INSTALL_EXTRAS='[dev,parsl,dask]'" >> $GITHUB_ENV + echo "INSTALL_EXTRAS='[dev,caches,parsl,dask]'" >> $GITHUB_ENV - name: Install dependencies (Linux) if: startsWith( matrix.os, 'ubuntu' ) @@ -66,7 +66,7 @@ jobs: uv pip install torch --index-url https://download.pytorch.org/whl/cpu uv pip install xgboost # install checked out coffea - uv pip install -q '.[dev,parsl,dask,triton]' --upgrade + uv pip install -q '.[dev,caches,parsl,dask,triton]' --upgrade uv pip list - name: Install dependencies (MacOS) if: matrix.os == 'macOS-latest' @@ -79,7 +79,7 @@ jobs: uv pip install torch uv pip install xgboost # install checked out coffea - uv pip install -q '.[dev,dask]' --upgrade + uv pip install -q '.[dev,caches,dask]' --upgrade uv pip list - name: Install dependencies (Windows) if: matrix.os == 'windows-latest' @@ -90,7 +90,7 @@ jobs: uv pip install torch uv pip install xgboost # install checked out coffea - uv pip install -q '.[dev,dask]' --upgrade + uv pip install -q '.[dev,caches,dask]' --upgrade uv pip list - name: Start triton server with example model diff --git a/pyproject.toml b/pyproject.toml index 1e15ef001..1b3a00ea7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,10 @@ xrootd = [ triton = [ "tritonclient[grpc,http]>=2.56.0", ] +caches = [ + "numcodecs>=0.13.1", + "zict>=3.0.0", +] dev = [ "pre-commit", "flake8", diff --git a/src/coffea/nanoevents/__init__.py b/src/coffea/nanoevents/__init__.py index d80766de2..4ff4b5ad5 100644 --- a/src/coffea/nanoevents/__init__.py +++ b/src/coffea/nanoevents/__init__.py @@ -1,6 +1,7 @@ """NanoEvents and helpers""" from coffea.nanoevents.factory import NanoEventsFactory +from coffea.nanoevents.mapping import BufferCache, NoCompressionCodec from coffea.nanoevents.schemas import ( FCC, BaseSchema, @@ -19,7 +20,9 @@ __all__ = [ "NanoEventsFactory", "BaseSchema", + "BufferCache", "NanoAODSchema", + "NoCompressionCodec", "PFNanoAODSchema", "TreeMakerSchema", "PHYSLITESchema", diff --git a/src/coffea/nanoevents/mapping/__init__.py b/src/coffea/nanoevents/mapping/__init__.py index 4aa20b46f..670697a9a 100644 --- a/src/coffea/nanoevents/mapping/__init__.py +++ b/src/coffea/nanoevents/mapping/__init__.py @@ -1,3 +1,7 @@ +from .buffer_cache import ( + BufferCache, + NoCompressionCodec, +) from .parquet import ParquetSourceMapping, TrivialParquetOpener from .preloaded import ( PreloadedOpener, @@ -7,6 +11,8 @@ from .uproot import TrivialUprootOpener, UprootSourceMapping __all__ = [ + "BufferCache", + "NoCompressionCodec", "TrivialUprootOpener", "UprootSourceMapping", "TrivialParquetOpener", diff --git a/src/coffea/nanoevents/mapping/buffer_cache.py b/src/coffea/nanoevents/mapping/buffer_cache.py new file mode 100644 index 000000000..cfa4ec97a --- /dev/null +++ b/src/coffea/nanoevents/mapping/buffer_cache.py @@ -0,0 +1,216 @@ +import dataclasses +import typing as tp +from collections.abc import MutableMapping + +import numpy as np + + +@dataclasses.dataclass(slots=True, frozen=True) +class ShapeDTypeStruct: + dtype: np.dtype + shape: tuple[int, ...] + strides: tuple[int, ...] + + +ByteBuffer: tp.TypeAlias = bytes + + +@tp.runtime_checkable +class Codec(tp.Protocol): + def encode(self, arr: np.ndarray) -> tuple[ByteBuffer, ShapeDTypeStruct]: ... + def decode(self, buffer: ByteBuffer, struct: ShapeDTypeStruct) -> np.ndarray: ... + + +class NoCompressionCodec(Codec): + def encode(self, arr: np.ndarray) -> tuple[ByteBuffer, ShapeDTypeStruct]: + struct = ShapeDTypeStruct(dtype=arr.dtype, shape=arr.shape, strides=arr.strides) + return arr.tobytes(), struct + + def decode(self, buffer: ByteBuffer, struct: ShapeDTypeStruct) -> np.ndarray: + arr = np.frombuffer(buffer, struct.dtype) + return np.lib.stride_tricks.as_strided(arr, struct.shape, struct.strides) + + +class NumCodecsWrapper: + __slots__ = ("_codec",) + + def __init__(self, codec: tp.Any) -> None: + self._codec = codec + + def encode(self, arr: np.ndarray) -> tuple[ByteBuffer, ShapeDTypeStruct]: + struct = ShapeDTypeStruct(dtype=arr.dtype, shape=arr.shape, strides=arr.strides) + encoded = self._codec.encode(arr.tobytes()) + return encoded, struct + + def decode(self, buffer: ByteBuffer, struct: ShapeDTypeStruct) -> np.ndarray: + decoded = self._codec.decode(buffer) + arr = np.frombuffer(decoded, struct.dtype) + return np.lib.stride_tricks.as_strided(arr, struct.shape, struct.strides) + + +ByteBufferCache: tp.TypeAlias = MutableMapping[tp.Hashable, ByteBuffer] +ShapeDTypeStructCache: tp.TypeAlias = MutableMapping[tp.Hashable, ShapeDTypeStruct] + + +class CodecAwareCache(MutableMapping): + __slots__ = ("_cache", "_meta", "_codec") + + def __init__(self, cache: ByteBufferCache, codec: Codec): + self._cache: ByteBufferCache = cache + self._meta: ShapeDTypeStructCache = {} + + if not isinstance(codec, Codec): + raise TypeError(f"codec must be an instance of Codec, got {type(codec)}") + self._codec = codec + + @property + def cache(self) -> ByteBufferCache: + return self._cache + + @property + def meta(self) -> ShapeDTypeStructCache: + return self._meta + + @property + def codec(self) -> Codec: + return self._codec + + def __setitem__(self, key: tp.Hashable, value: np.ndarray) -> None: + buf, struct = self.codec.encode(value) + self.cache[key] = buf + self.meta[key] = struct + + def __getitem__(self, key: tp.Hashable) -> np.ndarray: + return self.codec.decode(buffer=self.cache[key], struct=self.meta[key]) + + def __delitem__(self, key: tp.Hashable): + del self.cache[key] + del self.meta[key] + + def __iter__(self) -> tp.Iterator[tp.Hashable]: + return iter(self.cache) + + def __len__(self) -> int: + return len(self.cache) + + +# can't type hint without importing it, so we do this instead +NumCodecsCodec: tp.TypeAlias = tp.Any + + +def BufferCache( + cache: ByteBufferCache | None, + codec: Codec | NumCodecsCodec | None, # noqa: F821 +) -> MutableMapping: + """ + A compressed buffer cache. Supports all numcodecs.abc.Codec types. + + ## In-memory buffer cache + + Buffer caches give you more fine-grained control over internal + memory management of an awkward Array (here: NanoEvents). One powerful + feature is for example to compress the buffers in-memory to reduce + the total memory footprint. Buffers are decompressed upon use (`__getitem__`) + and compressed upon `__setitem__`. In a scenario where you have many buffers in + an awkward Array this can be highly beneficial because most arrays are then + compressed in RAM, while only a few at a time will be decompressed for a specific + operation. + + Example (in-memory no compression) + ------- + >>> buffer_cache=BufferCache(cache=None, codec=None) # or `NoCompressionCodec()` + >>> NanoEventsFactory.from_root(..., buffer_cache=buffer_cache) + + + Example (in-memory compressed) + ------- + >>> from numcodecs import Blosc + >>> codec = Blosc("zstd", clevel=1, shuffle=Blosc.BITSHUFFLE) + >>> buffer_cache=BufferCache(cache=None, codec=codec) + >>> NanoEventsFactory.from_root(..., buffer_cache=buffer_cache) + + + Example (LRU-backed compressed in-memory) + ------- + >>> from numcodecs import Blosc + >>> import zict + >>> codec = Blosc("zstd", clevel=1, shuffle=Blosc.BITSHUFFLE) + >>> capacity = 500_000_000 # 500 MB + >>> # len gives the number of bytes in the bytebuffer + >>> cache = zict.LRU(n=capacity, d={}, weight=lambda k,v: len(v)) + >>> buffer_cache=BufferCache(cache=cache, codec=codec) + >>> NanoEventsFactory.from_root(..., buffer_cache=buffer_cache) + + + ## On-disk buffer cache + + The on-disk buffer cache is the most aggressive way to offload buffers from RAM. + A simple on-disk buffer cache example is as follows: + + Example (on-disk compressed) + ------- + >>> from numcodecs import Blosc + >>> import zict + >>> codec = Blosc("zstd", clevel=1, shuffle=Blosc.BITSHUFFLE) + >>> buffer_cache=BufferCache(cache=zict.File("my_cache"), codec=codec) + >>> NanoEventsFactory.from_root(..., buffer_cache=buffer_cache) + + .. caution:: + + The comes with some caveats though: + + 1. The directory for the on-disk cache should be chosen to be as close as possible + to the CPU. That means that NFS backed paths (e.g. `/afs/` or `/eos/` at CERN) are + highly disouraged for this cache. A better choice would be `/tmp/...` on the worker. + + 2. It's probably good to cleanup this cache once it isn't needed anymore. For dask usage + with the coffea Executors one can use the `cachestrategy` argument of the Executor class + to make sure the on-disk cache is created in the local temp directory of the dask worker itself. + (see: https://distributed.dask.org/en/stable/worker.html#api-documentation) + + + ## Other examples + + Example (hierarchical) + ------- + >>> import zict + >>> cache = zict.Buffer( + >>> fast={}, + >>> slow=zict.File("mycache"), + >>> n=100, + >>> weight=lambda k,v: len(v), # len gives the number of bytes in the bytebuffer + >>> ) + >>> buffer_cache=BufferCache(cache=cache, codec=None) + >>> NanoEventsFactory.from_root(..., buffer_cache=buffer_cache) + """ + if cache is None: + cache = {} + + if not isinstance(cache, MutableMapping): + raise TypeError( + f"cache must be an instance of MutableMapping, got {type(cache)}" + ) + + if codec is not None: + try: + import numcodecs + except ModuleNotFoundError as err: + raise ModuleNotFoundError("""to use BufferCache, you must install numcodecs: + +pip install numcodecs + +or + +conda install -c conda-forge numcodecs""") from err + + # auto-wrap for numcodecs.abc.Codec + if isinstance(codec, numcodecs.abc.Codec): + codec = NumCodecsWrapper(codec=codec) + + # at this point we expect a proper Codec instance + if not isinstance(codec, Codec): + raise TypeError(f"codec must be an instance of Codec, got {type(codec)}") + + return CodecAwareCache(cache=cache, codec=codec) + + return cache diff --git a/tests/test_buffer_cache.py b/tests/test_buffer_cache.py new file mode 100644 index 000000000..5cc08f870 --- /dev/null +++ b/tests/test_buffer_cache.py @@ -0,0 +1,134 @@ +from collections.abc import MutableMapping + +import awkward as ak +import pytest + +from coffea.nanoevents import NanoAODSchema, NanoEventsFactory +from coffea.nanoevents.mapping import BufferCache +from coffea.nanoevents.util import unquote + + +def _make_events_with_cache(path: str, cache: MutableMapping) -> ak.Array: + factory = NanoEventsFactory.from_root( + {path: "Events"}, + schemaclass=NanoAODSchema, + mode="virtual", + buffer_cache=cache, + ) + return factory.events() + + +def _check_cache(events) -> None: + cache = events.attrs["@events_factory"].buffer_cache + assert len(cache) == 0 + + # materialize something and check that this is now properly populated in the cache + ak.materialize(events.Jet.pt) + + cache = events.attrs["@events_factory"].buffer_cache + assert len(cache) == 2 + + keys = [*map(unquote, events.attrs["@events_factory"].buffer_cache.keys())] + assert frozenset(keys) == frozenset( + [ + # nJet + "a9490124-3648-11ea-89e9-f5b55c90beef//Events;1/0-40/offsets/nJet,!load,!counts2offsets,!skip,!offsets", + # Jet_pt + "a9490124-3648-11ea-89e9-f5b55c90beef//Events;1/0-40/data/Jet_pt,!load,!content", + ] + ) + + +def test_buffer_cache(tests_directory): + pytest.importorskip("zict") + + events = _make_events_with_cache( + path=f"{tests_directory}/samples/nano_dy.root", + cache=BufferCache(cache=None, codec=None), + ) + + _check_cache(events) + + +def test_compressed_buffer_cache_in_memory(tests_directory): + pytest.importorskip("numcodecs") + pytest.importorskip("zict") + + from numcodecs import Blosc + + codec = Blosc("zstd", clevel=1, shuffle=Blosc.BITSHUFFLE) + events = _make_events_with_cache( + path=f"{tests_directory}/samples/nano_dy.root", + cache=BufferCache(cache=None, codec=codec), + ) + + _check_cache(events) + + +def test_compressed_buffer_cache_on_disk(tests_directory): + pytest.importorskip("numcodecs") + pytest.importorskip("zict") + + import zict + from numcodecs import Blosc + + codec = Blosc("zstd", clevel=1, shuffle=Blosc.BITSHUFFLE) + ondisk = zict.File(f"{tests_directory}/mycache") + events = _make_events_with_cache( + path=f"{tests_directory}/samples/nano_dy.root", + cache=BufferCache(cache=ondisk, codec=codec), + ) + + _check_cache(events) + + # clean up + import os + + ondisk.clear() # rm's all files in mycache + os.rmdir(f"{tests_directory}/mycache") + + +def test_buffer_cache_lru(tests_directory): + zict = pytest.importorskip("zict") + + # large enough to succeed + cache = zict.LRU(n=100_000_000, d={}, weight=lambda k, v: len(v)) + events = _make_events_with_cache( + path=f"{tests_directory}/samples/nano_dy.root", + cache=BufferCache(cache=cache, codec=None), + ) + + _check_cache(events) + + # small enough to fail (lru cache too small -> keys got evicted -> _check_cache fails) + cache = zict.LRU(n=100, d={}, weight=lambda k, v: len(v)) + events = _make_events_with_cache( + path=f"{tests_directory}/samples/nano_dy.root", + cache=BufferCache(cache=cache, codec=None), + ) + + with pytest.raises(AssertionError): + _check_cache(events) + + +def test_buffer_cache_hierarchical(tests_directory): + zict = pytest.importorskip("zict") + + hierarchical_cache = zict.Buffer( + fast={}, + slow=zict.File(f"{tests_directory}/mycache"), + n=100, + weight=lambda k, v: len(v), + ) + events = _make_events_with_cache( + path=f"{tests_directory}/samples/nano_dy.root", + cache=BufferCache(cache=hierarchical_cache, codec=None), + ) + + _check_cache(events) + + # clean up + import os + + hierarchical_cache.clear() # rm's all files in mycache + os.rmdir(f"{tests_directory}/mycache")