feat: add buffer caches

[pfackeldey]

This PR adds buffer caches that are used at runtime during the processing of a single chunk. This PR is enabled through #1507 (and the latest awkward release).

The purpose of these caches is to reduce memory usage. By default events will keep all buffers in memory until a Processor(...).process (or the wrapped versions of it from the Runner interface) finishes. That's suboptimal, e.g. some columns may only be used once and never again, or we could make use of in-memory compression, ... Essentially ways to reduce our memory footprint to allow 'cranking up the chunksize to 11'.
The point here is that we typically have hundreds if not thousands of buffers in events, but per awkward operation we only need to have 1-2 (maybe a handful) typically decompressed in memory as arrays at the same time.

Benchmark Overview

I'll add some runtime measurements and benchmarks in the following, but first some information:
This PR adds 4 types of buffer caches:

1. a dict-like cache, BufferCache: that's essentially no different than not using a cache at all, but it allows to manually delete buffers which isn't possible right now
2. a in-memory compressed cache, CompressedBufferCache: that one auto-supports codecs from numcodecs, but following the Codec-protocol one can implement their own custom codec.
3. a HDF5-based cache, HDF5BufferCache: that one stores all buffers in a h5py.Group as datasets. They may live on disk, or in-memory, and may be compressed or not (one can provide compression opts using e.g. hdf5plugin.
4. a hierarchical cache, hiearchical_cache: that one gets the above caches as inputs with byte-limits and implements an LRU logic with several cache layers (similar to: https://github.com/scikit-hep/coffea/blob/master/src/coffea/processor/dask/__init__.py#L19-L37).

The following benchmarks have been made with:

• my local Mac (M3 Max, 36GB RAM, Apple SSD [write speed=2670408352 bytes/sec; read speed=10731944948 bytes/sec)
• using mimalloc as the memory allocator to avoid memory fragmentation overheads as discussed in High memory overhead in uproot.iterate uproot5#1535
• a NanoAOD MC file with ~920k events (locally on the SSD, no streaming involved)

The benchmark task is to read/materialize all Jet columns in order of the full file (chunksize~=920k), but re-materialize the kinematic fields "pt", "eta", "phi", "mass" in between every column read everytime. The idea is to simulate a workflow that reads many different columns but reads certain columns multiple times (the kinematic ones), the logic is roughly like this:

def analysis(events):
    kinematics = ["pt", "eta", "phi", "mass"]
    for f in events.Jet.fields:
        ak.materialize(events.Jet[f])
        for k in kinematics:
             ak.materialize(events.Jet[k])

BufferCache

Peak RSS: 1.51 GB
Runtime: 13.476896047592163 s

CompressedBufferCache

Peak RSS: 1.01 GB
Runtime: 14.100797891616821 s

HDF5BufferCache (no compression, on disk)

Peak RSS: 0.73 GB
Runtime: 13.34046983718872 s

HierarchicalCache (100 MB: BufferCache -> 500 MB: CompressedBufferCache -> HDF5BufferCache)

Peak RSS: 1.07 GB
Runtime: 12.970969200134277 s

Summary:

• In-memory compression reduces peak RSS by ~500 MB with barely any overhead compared to the pure IO time
• On-disk HDF5 cache without compression reduces peak RSS by ~800 with no measurable runtime overhead (note: this is using a fast SSD)
• a mixture of in-memory compression and on-disk recovers some of the compression/decompression overhead while still reducing the memory footprint by ~500 MB

These techniques reduce the RSS usage significantly with no/barely any runtime overhead. They allow users to increase their chunk sizes significantly as awkward typically only uses 1-2 buffers at a time for an operation which means we only need those to be uncompressed in-memory at the same time. Increasing chunk sizes is of major importance: it reduces effectively python overhead, makes fully use of compiled kernels and SIMD, decreases the number of tasks/jobs.

This could also be interesting for infrastructure maintainers: having a SSD backed HDF5BufferCache shows no runtime overhead while we can essentially extend our RAM per worker by SSDs (which should be much less expensive $ than buying RAM).

In the future we can also use these caches to hold intermediates. Often intermediate arrays may only be needed once (e.g. for the eval of an ML model), until they're needed they could be in such compressed/ondisk caches to reduce the peak RSS usage.

A big portion of this work has been driven and enabled by @ikrommyd, thanks! (also looking forward to your review :P)

[pfackeldey]

More benchmark results for running locally (the above mentioned Mac, or on lxplus with eos) can be found here for a ttbar hadronic open data nanoaod file (1M events): https://github.com/pfackeldey/coffea_buffercache_benchmarks?tab=readme-ov-file#results

Even on lxplus using a HDF5 on disk buffer cache doesn't seem to introduce significant overhead (although those are HDDs), while reducing the memory footprint (peak RSS) from 2.81 GB to 0.93 GB (~3x memory reduction) for materializing all Jet & Muon columns (interleaved with repeated materialization of kinematic columns).

[lgray]

This is incredibly interesting data! I would like to see higher scale testing, but we should inform people like @PerilousApricot about these kinds of insights coming down the pipeline for directly attached local storage considerations.

It would be very useful to get these kinds of metrics on a full analysis soon, to understand buffer access multiplicity, etc. in the realistic scenarios. This could probably be checked quickly with HiggsDNA, @ikrommyd.

Some future thoughts: is blosc working on GPU implementations? could we fit something blosc-like into what's going on with uproot + nVidiaGDS @fstrug. Similarly, since we now know that awkward is using at most 2-3 buffers at once most of the time should we be thinking about more aggressive GPU memory memory management in general. If we can structure PCIe transfers correctly we can hide all the latency and process enormous arrays in GPU memory @kmohrman.

[nsmith-]

How do these compare to a simple LRU cache for the buffers? LRU cache misses imply extra loading/decompressing, but perhaps the extra compute and network is not much for an appropriately large LRU while still reducing memory.

[pfackeldey]

How do these compare to a simple LRU cache for the buffers? LRU cache misses imply extra loading/decompressing, but perhaps the extra compute and network is not much for an appropriately large LRU while still reducing memory.

A simple LRU cache also limits memory, the problem with that is that streaming again the data (or even reading it again with uproot) is expensive. Dumping the raw bytes (no compression) into a hdf5 dataset and just reading them again is much faster (this is local on my laptop):

# reading with uproot (first time)
In []: %timeit -n1 -r1 ak.materialize(events.Jet.pt)
410 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

# reading from hdf5 cache (second time)
In []: %timeit -n1 -r1 ak.materialize(events.Jet.pt)
10.9 ms ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

The problem with a plain LRU is that usually people stream their data or read it from some slow filesystem, so evicting a column and later reading it again can be very costly (and error prone given all the streaming issues we've experienced in the past).

Initializing an LRU cache with this PR is as simple as wrapping them into a zict.LRU, e.g.:

bc = BufferCache()

# 500 MB capacity
lru = zict.LRU(n=500_000_000, d=bc, weight=bc.get_nbytes)

NanoEventsFactory.from_root(..., buffer_cache=lru)

or since 8292c31:

from coffea.nanoevents.mapping.buffer_cache import lru_cache

# 500 MB capacity
lru = lru_cache(capacity=500_000_000)

NanoEventsFactory.from_root(..., buffer_cache=lru)

[pfackeldey]

Some future thoughts: is blosc working on GPU implementations? could we fit something blosc-like into what's going on with uproot + nVidiaGDS @fstrug. Similarly, since we now know that awkward is using at most 2-3 buffers at once most of the time should we be thinking about more aggressive GPU memory memory management in general. If we can structure PCIe transfers correctly we can hide all the latency and process enormous arrays in GPU memory @kmohrman.

I'm not aware of a Blosc codec GPU implementation, it's also not fully clear yet what the best codec is (it's different between columns of course...). I started some studies for NanoAOD to compare how different columns compare with different codecs (and their settings). I'm happy to share that if people are interested to see this.

Blosc generally seems to be highly optimized for CPUs, e.g., the chunking & block size determination happens dynamically depending on the cache sizes and levels of your specific CPU. Here, Blosc is beneficial because it can add bit-shuffling on top of zstd which turns out (in first tests) to be an effective filter for our data usually, but nothing else from Blosc is used here (i.e. we don't use any chunking here).

I think this is important for running on GPUs too, we can think about the memory limitation not being all needed columns of events but rather the largest (in bytes) combination of 2-3 buffers that awkward needs at the same time in memory. My hope is that we get into a state where kernels become the limiting performance factor for both GPU and CPU.

There can be also other future implementations, such as a user-agnostic buffer caches with a longer lifetime than the dask/taskvine cluster, similarly to the caching we've worked on in https://arxiv.org/abs/2207.08598 (but with more control on the buffer cache). That would let analysis groups work more efficiently together when they run over the same datasets repeatedly.

[ikrommyd]

A simple LRU cache also limits memory, the problem with that is that streaming again the data (or even reading it again with uproot) is expensive. Dumping the raw bytes (no compression) into a hdf5 dataset and just reading them again is much faster (this is local on my laptop):

I think it's good to have a simple LRU cache implemented by default too though. In the sense that it is the simplest cache you can have that can save you some memory. And let's say that if I have my data locally, I may not care about dumping bytes as that is a "copy" and I may be fine with reading from the original file.

[pfackeldey]

A simple LRU cache also limits memory, the problem with that is that streaming again the data (or even reading it again with uproot) is expensive. Dumping the raw bytes (no compression) into a hdf5 dataset and just reading them again is much faster (this is local on my laptop):

I think it's good to have a simple LRU cache implemented by default too though. In the sense that it is the simplest cache you can have that can save you some memory. And let's say that if I have my data locally, I may not care about dumping bytes as that is a "copy" and I may be fine with reading from the original file.

done in 8292c31

[ikrommyd]

I assume here we will introduce optional dependencies like pip install coffea[caches] or something similar right?

[pfackeldey]

I assume here we will introduce optional dependencies like pip install coffea[caches] or something similar right?

As you like to have it. For now, I put all imports into the specific cache implementation, such that it only fails upon use with missing dependencies.

[pfackeldey]

Btw, it's not too complicated to 'back' any awkward array (also intermediates) with such a cache (at the cost of re-wrapping with to/from buffers):

def with_backend_cache(arr: ak.Array, *, cache: tp.MutableMapping) -> ak.Array:
  form, length, buffers = ak.to_buffers(arr)

  # put into cache
  buffer_keys = list(buffers)
  for bk in buffer_keys:
    cache[bk] = buffers.pop(bk)
  
  def get(bk):
    print("Get", bk, "from", cache)
    return cache[bk]

  # lazify buffer container
  from functools import partial
  container = {bk: partial(get, bk) for bk in cache}

  return ak.from_buffers(
    form,
    length,
    container,
    enable_virtualarray_caching=lambda form_key, attribute: attribute != "data",
    attrs=arr.attrs,
    behavior=arr.behavior,
  )

arr = ak.zip({"foo": np.ones(100), "bar": np.ones(100)})

from coffea.nanoevents.mapping.buffer_cache import BufferCache
cache = BufferCache()
arr_with_cache = with_backend_cache(arr, cache=cache)
print(arr_with_cache.foo + 1)
>> "Get node1-data from BufferCache(nbuffers=2, nbytes=1600)"
>> [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, ..., 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]

That can be handy for certain intermediate awkward arrays. If you agree, I can make another PR and add such a helper function to coffea.

In general, I think it's probably better to think about a MutableMapping-like intermediate store that can handle any type (not just awkward arrays), so one would do something like:

DB = IntermediateStore(cache=...)

# ak.Arrays
DB["some_quantity"] = events.Jet.pt + 1
# np.ndarrays
DB["some_other_quantity"] = ak.to_numpy(events.genWeight) + 1

instead of storing everything on events...

But that's rather something for another PR. Not sure what you think about this? (This is dangerous though to use with an LRU cache as evicted buffers can't be recovered)

[ikrommyd]

Definitely for another PR as you said. I think this can be a nice-to-have addition but a difficult concept to teach to the users. It can be worth it though for low RAM workers because users tend to create a lot of intermediate "columns" that they either histogram or export to output NTuples. However, I don't think it's gonna be used that often if I'm being honest because it feels to me like a power-user thing and is also something that you do not find in typical array-oriented-programming guides.
We can have it though, I just wouldn't prioritize it.

[nsmith-]

As discussed in the meeting, having numcodecs be a core dependency is probably fine, and it might even be preferred over cramjam in uproot as well since it supports all the same compression algorithms and has a few more features such as pre-compression filters that may be useful for reading/writing RNTuples, while only depending on numpy (FYI @ariostas). For zict and h5py we can have coffea[caches].

[nsmith-]

From the meeting: we think this is ready for review. We need to add tests but input is appreciated on what kinds of tests would be most useful.

[ikrommyd]

I have not looked at code line by line again but will do. Regarding the dependencies, I agree with what Nick said above. Regarding testing, I think you can probably wrap the caches so that they can log whenever they caching something and when something is being extracted from the cache? Then you can assert that the right things happen at the right time?

[ikrommyd]

Oh and I think we'd need appropriate error messages when things try to get imported inside the caches that say for example "if you want to use this buffer cache, install h5py " or something. We do this in awkward a lot.

[ikrommyd]

Only left minor comments. I would really love a test that they just work at least. I know they do because I've been using your profiling repo @pfackeldey but if we can at least have a run through a chunk of events so show that no errors are being raises it would be good. Imagine some library changes something in their API. We woulldn't know if a cache does not work anymore.

Secondly @pfackeldey @nsmith- should we keep the executors out of the scope of this one and expose the caches through the executors get_cache in another PR?

[pfackeldey]

Only left minor comments. I would really love a test that they just work at least. I know they do because I've been using your profiling repo @pfackeldey but if we can at least have a run through a chunk of events so show that no errors are being raises it would be good. Imagine some library changes something in their API. We woulldn't know if a cache does not work anymore.

Secondly @pfackeldey @nsmith- should we keep the executors out of the scope of this one and expose the caches through the executors get_cache in another PR?

Ok, I'll come up with some tests tomorrow. Not sure what you mean with keeping executors out of scope? I think you can already pass those caches to the executor interface using the cachestrategy kwarg (https://github.com/scikit-hep/coffea/blob/master/src/coffea/processor/executor.py#L1096), no?

[ikrommyd]

Ok, I'll come up with some tests tomorrow. Not sure what you mean with keeping executors out of scope? I think you can already pass those caches to the executor interface using the cachestrategy kwarg (https://github.com/scikit-hep/coffea/blob/master/src/coffea/processor/executor.py#L1096), no?

Looks like it will work out of the box yes. That can be in the test too :)

[ariostas]

Just a quick heads up, it seems like the lz4 implementation in numcodecs might have issues. Maybe it doesn't matter here since you would be compressing and decompressing with numcodecs, but just so you know. scikit-hep/uproot5#1574 (comment)

[ikrommyd]

I honestly don't see anything worth commenting on here. We've tested the implementation and done benchmarks. It's in very very good condition for an introduction to buffer caches. @lgray or @nsmith- can you give it a once over in case you see anything?

[pfackeldey]

@ikrommyd and @nsmith- I've refactored the code to make better use of zict. I think it's much better now. Can you have a look again?

[lgray]

@ikrommyd @nsmith- ping ^_^

[nsmith-]

Just documentation tweaks. You can preview at https://coffea-hep--1508.org.readthedocs.build/en/1508/
edit: this currently isn't in the toctree anywhere

[ikrommyd]

Is this good to go?