Memory consumption during uproot.iterate

[tomeichlersmith]

I am interested in processing a O(2 GB) file using uproot.iterate rather than loading the full array into memory. I want to do this because I would like to process many files in parallel, but do not have the required memory to load the entire contents of each file into memory all at once. This led me down the rabbit hole of investigating the memory consumption of a test script where I noticed some odd behavior. While the memory used when using uproot.iterate is significantly lower than the memory used when loading the entire file (2.5GB compared to >6GB), it does not have the shape I would expect - it continuously increases until the entire file is closed. I know from reading #926 that using a step size of a certain memory unit is complicated to interpret, but I'm wondering

• Is this monotonically increasing behavior expected/understood? What is the explanation?
• If so, is there a way I can estimate the actual memory usage? (or the inverse problem, how to best select a step_size to keep the process analyzing a file under a certain soft memory limit).

Below, I have uploaded the memory-usage vs time plots as generated by mprof. I am using awkward 2.4.5 and uproot 5.0.12. I copied my test script for reference - the argument to the test script refers to the function that is run (and thus profiled).

Data Shape Complexities

Not sure if these are pertinent, but I think all of the details are helpful when talking about memory usage.

• I am only selecting certain branches of the TTree I am loading (as can be seen in the test script).
• One of these branches (HcalSimHits) is doubly-ragged (it is a variable length array whose elements have variable length arrays as members), one of these branches ('EcalRecHits`) is singly-ragged (it is a variable length array), and the others are not ragged (one value per entry in the tree).
• The objects in these branches originated as C++ classes that were serialized in a C++/ROOT framework.

Runs for Reference

First, (for context), I'd like to show some plots where I don't use uproot.iterate so we can have them as reference when seeing how uproot.iterate behaves.

Just have numpy generate large chunks

def np_generate_chunks():
    for i in range(10):
        yield np.full(25_000_000, 1, dtype=np.int32)


@profile
def numpy():
    for events in np_generate_chunks():
        process(events)

Have uproot load entire file

def uproot_whole_file():
    with uproot.open('test.root') as f:
        events = f['LDMX_Events'].arrays(**selection_kw)
        process(events)

Tests

Naive solution from docs

def uproot_file_iterate():
    for events in uproot.iterate('test.root:LDMX_Events', **selection_kw):
        process(events)

Different attempt from docs

def uproot_tree_iterate():
    with uproot.open('test.root') as f:
        for events in f['LDMX_Events'].iterate(**selection_kw):
            process(events)

Try answer from #648

def uproot_tree_iterate_no_object_cache():
    with uproot.open('test.root', object_cache=None) as f:
        for events in f['LDMX_Events'].iterate(**selection_kw):
            process(events)

Force a specific step size that (I think) should keep the chunk-in-memory smaller than the default.

def uproot_iterate_small_chunk():
    for events in uproot.iterate('test.root:LDMX_Events', step_size=10000, **selection_kw):
        process(events)

Full test.py Script

import numpy as np
import uproot
import time

selection_kw = dict(
    filter_name = [
        'PEFF**',
        'HcalSimHits**',
        'EcalRecHits**',
        'EventHeader**'
    ]
)


@profile
def process(events):
    time.sleep(1)


@profile
def uproot_whole_file():
    """as a bench mark, lets just load the whole file into memory"""
    with uproot.open('test.root') as f:
        events = f['LDMX_Events'].arrays(**selection_kw)
        process(events)


@profile
def uproot_file_iterate():
    """naive solution, first thing I tried after looking at docs"""
    for events in uproot.iterate('test.root:LDMX_Events', **selection_kw):
        process(events)


@profile
def uproot_tree_iterate():
    """noticed docs have an iterate method on the tree itself, try that"""
    with uproot.open('test.root') as f:
        for events in f['LDMX_Events'].iterate(**selection_kw):
            process(events)


@profile
def uproot_tree_iterate_no_object_cache():
    """saw this option in the discussion
    https://github.com/scikit-hep/uproot5/discussions/648
    """
    with uproot.open('test.root', object_cache=None) as f:
        for events in f['LDMX_Events'].iterate(**selection_kw):
            process(events)


@profile
def uproot_iterate_small_chunk():
    """force a really small chunk to see an extreme"""
    for events in uproot.iterate('test.root:LDMX_Events', step_size=10000, **selection_kw):
        process(events)


def np_generate_chunks():
    for i in range(10):
        yield np.full(25_000_000, 1, dtype=np.int32)


@profile
def numpy():
    """another benchmark, just have numpy generate ~100MB chunks for me"""
    for events in np_generate_chunks():
        process(events)


def main():
    import sys
    globals()[sys.argv[1]]()

if __name__ == '__main__':
    main()

[agoose77]

The default uproot source for reading files from disk uses memory mapping (mmap). Memory mapping is useful to avoid an extra copy of the data when reading, to share (re-use) memory between processes, and to improve performance of random access (c.f. https://stackoverflow.com/a/6383253). However, it has some drawbacks; namely that the OS is responsible for managing the memory (and also, the cache is populated via page faulting). In this case, I think the behaviour that you're seeing is a direct consequence of the kernel caching memory pages, and choosing not to free them (c.f. https://stackoverflow.com/a/1972889).

If you want to avoid this caching, you can open the file yourself and pass the file-handle to uproot. This bypasses np.memmap, and instead performs single-threaded Python reads. This may be slower (no caching, more copies!), but the OS memory manager should be more aggressive in freeing memory pages.

[tomeichlersmith]

Using the file-handle directly to bypass mmap seems to have a similar effect as the fsspec source mentioned by @lobis although it is back to the time-performance from before. This reminds me that I'm reading this file over NFS and so trusting the times I'm showing here is not a good idea since I haven't logged whether NFS mount is warm or cold.

def pyhandle():
    with open('test.root','rb') as f:
        with uproot.open(f) as rf:
            for events in rf['LDMX_Events'].iterate(**selection_kw):
                process(events)

[lobis]

An alternative to modifying your code to pass the file-handle would be to use the v5.2.0rc1 pre-release. In the next release fsspec will become a required dependency and it will also become the default uproot source for most cases, currently this also includes local filesystem access. This can also be tested in the latest release (5.1.2) by passing the handler=uproot.source.fsspec.FSSpecSource to the iterate/open method. (The default source for local files in that version would be uproot.source.file.MemmapSource).

I would be really interested in seeing the difference in performance between the two sources!

[tomeichlersmith]

Interestingly, it does look to behave better on the memory side, but @agoose77 is correct, I do observe the reading take measurably longer than before (by ~30% - from ~33s reading to ~43s reading).

def uproot_fsspec_source():
    with uproot.open('test.root', handler=uproot.source.fsspec.FSSpecSource) as f:
        for events in f['LDMX_Events'].iterate(**selection_kw):
            process(events)

[lobis]

Thanks for the plot!

I'm wondering if we can use fsspec to behave as uproot.source.file.MemmapSource.

In your case you are using fsspec with the loca file protocol (same as setting file://test.root).

You could also load the whole file into memory

with open("test.root", "rb") as f:
        contents = f.read()
        with fsspec.open("memory://test.root", "wb") as ff:
            ff.write(contents)

and then read it with uproot using fsspec (passing memory://test.root instead of test.root to uproot).

I'm guessing this would have the fastest read performance but obviously you would need to have the whole file in memory and also you need to read it whole once to load it.

I guess it should be possible to use fsspec to behave as a memory mapped file. I tried protocol chaining (memory::file, file::memory) but couldn't make it work. Perhaps @martindurant can comment on this.

[martindurant]

There is no "memory cache" filesystem - they all save to local disk. The memoryFS doesn't expect to pull its data from any other FS for those kinds of URL chains.

The bytes caching mechanism (i.e., in-file) has one mode which allows to load a whole file prospectively into memory with open(cache_type="all"), but this only applies to subclasses of AbstractBufferedFile, not ones that have native implementations such as local files.