(Potential) memory leak with uproot.iterate

[pfackeldey]

(this is tested in uproot v5.6.5, likely present in other versions as well).

The issue & reproducer

When running the following snippet and benchmarking this with memray:

import uproot

def loop_iterate(rootfile):
    with uproot.open(rootfile, array_cache=None, object_cache=None) as f:
        tree = f["tree"]
        for batch in tree.iterate(step_size="200 MB", library="ak"):
            print(repr(batch))


if __name__ == "__main__":
    loop_iterate("~/Downloads/zlib9-jagged0.root")

I'm getting a memory consumption of the physical RAM usage (RSS) of up to 1.6GB (even though the step size is 200 MB):

This is already surprising, however another indication that there's something up is that an explicit gc.collect at the end of each iteration improves the RSS situation by ~2x, i.e.:

import uproot
+import gc

def loop_iterate(rootfile):
    with uproot.open(rootfile, array_cache=None, object_cache=None) as f:
        tree = f["tree"]
        for batch in tree.iterate(step_size="200 MB", library="ak"):
            print(repr(batch))
+           del batch
+           gc.collect()


if __name__ == "__main__":
    loop_iterate("~/Downloads/zlib9-jagged0.root")

which gives up to 800MB RSS consumption:

Why is this bad?

RSS is the physical RAM usage by this process, which dask monitors to decide if a worker should be killed due to OOM or not.

What I've found so far...

The memory usage grows by the following function: https://github.com/scikit-hep/uproot5/blob/main/src/uproot/behaviors/TBranch.py#L1440-L1452 and to be more explicit by this part of it: https://github.com/scikit-hep/uproot5/blob/main/src/uproot/behaviors/TBranch.py#L3421-L3428

What does work correctly is that the file arrays dictionary by the above function is ~200 MB, that's good! However, this _ranges_or_baskets_to_arrays still uses ~800 MB to fill the ~200 MB arrays dict and does not free it again.

Also, the "popper-trick" that @jpivarski introduced in #1305 does enable the manual gc.collect to help here (without it even that won't help).

So, my understanding right now is that uproot.iterate does yield correctly sized arrays, but it uses way to much memory while doing so and also doesn't free it properly.

Other implications

_ranges_or_baskets_to_arrays is used also in other loading functions, and some quick tests showed that:

def loop_manual(rootfile):
    with uproot.open(rootfile, array_cache=None, object_cache=None) as f:
        tree = f["tree"]

        starts = list(range(0, tree.num_entries, tree.num_entries // 20))
        stops = starts[1:] + [tree.num_entries]
        ranges = list(zip(starts, stops))
        for start, stop in ranges:
            print(f"entry {start} to {stop}")
            entry = tree.arrays(entry_start=start, entry_stop=stop, library="ak")
            print(repr(entry))


def loop_same_chunks(rootfile):
    with uproot.open(rootfile, array_cache=None, object_cache=None) as f:
        tree = f["tree"]

        chunk_starts = 0
        chunk_stops = 53687091
        for _ in range(10):
            print(f"entry {chunk_starts} to {chunk_stops}")
            entry = tree.arrays(entry_start=chunk_starts, entry_stop=chunk_stops, library="ak")
            print(repr(entry))

have a similar memory behavior, see e.g. the profile for loop_manual (the numerical values of the y axis is different of course because I can't exactly mirror "200 MB" steps by hand):

and for loop_same_chunks:

What I want to see / was expecting

The orange and blue line overlap and roughly follow a saw tooth shape with 200 MB jumps per iteration (and not much additional overhead in RAM).

---

This has been originally been found by @oshadura in the scope of the integration challenge, here I just attach a local reproducer with some first findings.

cc @oshadura @alexander-held

[ianna]

@pfackeldey - thanks for the detailed report! Indeed, it is a serious issue. I will have a look asap. Thanks!

[pfackeldey]

Thanks @ianna, I'll try debugging as well meanwhile.
I suspect this is another ref-cycle (or similar) issue like @jpivarski and me solved in #1305 which prevents this memory from being garbage collected. Those tend to be hard to find from my previous debugging experience 😅

[ikrommyd]

Good find @pfackeldey. I'm wondering, what is your interpretation for the fact that resident size flattens out after some iterations? It's like it "stops leaking" after a bit.

[pfackeldey]

@ikrommyd I do not know why it looks like this

[pfackeldey]

I checked 4 'versions' of {TTree,RNTuple}.iterate in addition with an open data NanoAOD file (~1GB decompressed in total, stepsize="20 MB").

1. TTree
2. TTree + explicit gc.collect() at the end of every loop iteration
3. RNTuple
4. RNTuple + explicit gc.collect() at the end of every loop iteration

The results are as follows:

TTree

TTree + explicit gc.collect() at the end of every loop iteration

RNTuple

RNTuple + explicit gc.collect() at the end of every loop iteration

Summary

In the TTree case an explicit gc.collect() does help apparently. RSS is piling up in every iteration and it looks to me like the yielded 20 MB array is cleaned up properly (the little peaks) but somehow RAM is still getting full.

In the RNTuple case I suspect that the whole array is read first and then chunks of that are yielded. Although that would only explain it growing to 1GB. Probably there are some copies?

In general I think it is expected that the memory consumption is somewhat larger than 20 MB because uproot will have the compressed and decompressed chunk in memory at the same time. So in the worst case it should be 2x the chunk size, so 40 MB in this case... but not 1.4GB in TTree, or 2.5GB in RNTuple.

@ariostas any ideas why RNTuple looks like this?

I'm honestly pretty stuck here, and after days of debugging I'm not sure where else to look into. Any ideas? (maybe also @nsmith- ?)

(fyi: changing the stepsize to an int for the number of events doesn't change anything)

---

TTree code

with uproot.open(TTREE_FILE, array_cache=None, object_cache=None) as f:
  tree = f["Events"]
  for i, batch in enumerate(tree.iterate(step_size="20 MB", library="ak")):
    print(repr(batch))
    # uncomment to enable explicit gc.collect()
    # del batch
    # gc.collect()

RNTuple code

with uproot.open(RNTUPLE_FILE, array_cache=None, object_cache=None) as f:
  rntuple = f["ntuple"]
  for i, batch in enumerate(rntuple.iterate(step_size="20 MB", library="ak")):
    print(repr(batch))
    # uncomment to enable explicit gc.collect()
    # del batch
    # gc.collect()

[pfackeldey]

Ok, so I found apparently the main reason (and some solutions) to this issue.

Apparently, we're suffering in uproot from memory fragmentation of arenas created by malloc. When I switch the allocator to jemalloc (which has a much tighter fragmentation behavior) on MacOS it looks way better, see:

The solutions to this are basically:

1. glibc only: adjust malloc options through env vars (e.g. MALLOC_ARENA_MAX or MALLOC_TRIM_THRESHOLD)
2. glibc only: call ctypes.CDLL("libc.so.6").malloc_trim(0) after yielding batches
3. use a different allocator with tighter fragmentation behavior (e.g. jemalloc)
4. avoid allocating as much as possible in uproot, i.e. can we reuse buffers e.g. for decompressing?
5. (write uproot in a low level language to have more control over allocations?)

I didn't test 1-2 because I don't have a linux machine, but that should help.

(Also this here should help for dask: https://distributed.dask.org/en/stable/worker-memory.html#memory-not-released-back-to-the-os)

[nsmith-]

Memory fragmentation has been a problem for us since ~forever, unfortunately. When we have many small long-lived allocations interspersed with very large array allocations, the arenas end up mostly empty. See this old coffea issue e.g. scikit-hep/coffea#249

[pfackeldey]

Memory fragmentation has been a problem for us since ~forever, unfortunately. When we have many small long-lived allocations interspersed with very large array allocations, the arenas end up mostly empty. See this old coffea issue e.g. scikit-hep/coffea#249

yes, that makes sense. We did quite some work in awkward to reduce small long-lived allocations, so that should hopefully be noticeable already. Do you think we could reuse some buffers explicitly in uproot @nsmith- (or do you have any other idea how that could be mitigated apart from "uproot-rs" :D)?

This issue is not a programmatic memory leak then, should we turn this into a discussion instead for future reference?