Memory leak in pyroot backend

[sjoerd-bouma]

I have a (by now ancient) script to extract trigger rates that has regularly been broken and unbroken by mattak updates, and although it runs again since the recent fixes to skip_incomplete, I noticed that the memory usage has skyrocketed. Changing the backend to uproot limits the memory usage to something reasonable, so it is pyroot-specific. Some quick debugging in Python hasn't led me anywhere yet, and I'm not sure that I would get anywhere trying to debug the root/c++ stuff.

Here's a MWE that demonstrates the memory usage ballooning

import gc
import mattak.Dataset
import time
import logging
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel(logging.INFO)

n_runs = 0
for run in range(0,2000):

    try:
        d = mattak.Dataset.Dataset(
            station=21, run=run,
            skip_incomplete=False,
            read_daq_status=False, read_run_info=False,
            backend='pyroot'
        )
        d.setEntries((0, d.N()))
        einfo = d.eventInfo()

        n_runs += 1
    except (IOError, ReferenceError, KeyError) as e:
        logger.warning('run error', exc_info=e)
        continue

[fschlueter]

@cozzyd maybe the new access to the sampling rate is the problem?

[sjoerd-bouma]

@cozzyd maybe the new access to the sampling rate is the problem?

for the record, the trigger rate script has been broken for a year or so, so I'm not sure at what time this would have been introduced. If I can think of a good way for python to self-report the problem I'm happy to run a git bisect : )

[sjoerd-bouma]

... so it turns out actually Python has a really neat memory tracking tool!

[ Top 10 ]
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/mattak/backends/pyroot/dataset.py:167: size=97.6 MiB, count=2132917, average=48 B
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/numpy/lib/function_base.py:959: size=23.2 MiB, count=329872, average=74 B
<string>:3: size=14.6 MiB, count=131946, average=116 B
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/mattak/backends/pyroot/dataset.py:182: size=7491 KiB, count=195917, average=39 B
<frozen importlib._bootstrap_external>:672: size=5615 KiB, count=39503, average=146 B
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/mattak/backends/pyroot/dataset.py:174: size=3093 KiB, count=65974, average=48 B
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/mattak/backends/pyroot/dataset.py:181: size=1997 KiB, count=65973, average=31 B
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/mattak/backends/pyroot/dataset.py:175: size=1797 KiB, count=65716, average=28 B
/home/sb/Python/.venv/nu-radio/lib/python3.10/site-packages/mattak/backends/pyroot/dataset.py:183: size=1666 KiB, count=60927, average=28 B
<frozen importlib._bootstrap>:241: size=1618 KiB, count=14650, average=113 B

so e.g. the main offender is l.167 of backends/pyroot/dataset.py, though it looks like a lot of other data is also persisted across multiple datasets. I am not sure why, though...

[CamphynR]

I also encounter a similar issue when using readRNOGDataMattak.py (from NuRadioReco/modules/io/RNO_G) which uses mattak datasets. In principle the readRNODataMattak reader should iterate over run directories, storing and clearing datasets while iterating. When trying this out over forced trigger data of season22, the university's cluster shuts down the job due to excessive RAM usage. The issue seems to be at least suppressed by using the uproot backend instead, which does allow for iterating over a full run season.

[cozzyd]

@CamphynR did this work before recently? (Trying to bisect)

[cozzyd]

I cannot reproduce this with Sjoerd's script. But I don't have it set up in a way to find calibration files. Is it finding calibration files in your case? That might be the problem.

[sjoerd-bouma]

I cannot reproduce this with Sjoerd's script. But I don't have it set up in a way to find calibration files. Is it finding calibration files in your case? That might be the problem.

Neither do I, and I do get 'could not find any calibration files' messages, though I thought they didn't really affect anything.

It would be really strange if the memory leak is present only on some platforms, right?

[cozzyd]

how much memory usage are you seeing?

[sjoerd-bouma]

how much memory usage are you seeing?

With the test script, I hit 1 GB after some 700 runs, it looks like (after starting with O(200 MB))

[cozzyd]

So after further testing, I think it's a memory fragmentation problem. I've verified there is no leak in mattak::Dataset and that the destructor is being called properly.

If I commented out the numpy.copy calls, the memory grows a lot slower, suggesting that it's the allocation of many small numpy arrays (that get put together in a list) that's the problem. Note that I also see memory increasing in uproot, though it's a bit slower. I added a modified version of your script to the tests directory with the ability to switch backends. Some preliminary memory profiling with memray also strongly suggests fragmentation based on how most of the time the memory usage is much less than the high water mark, if I'm interpreting correctly, though I should probably learn how to use memray properly.

Not sure why the benchmark test is failing now... it doesn't seem like it should be returning non-zero...

[fschlueter]

@cozzyd cant we do this is a branch?

[CamphynR]

In my case it is looking for and finding calibration files but not applying them (I am using the linear calibration). It has always been struggling with RAM usage but I only noticed now when running over a very big run list so I cannot say it worked properly at a previous point.

[cozzyd]

@cozzyd cant we do this is a branch?

I'm pretty sure the test failure is because this is the first time the benchmark has something running on the main branch :) I'm not really a fan of a ton of 1-2 commit branches for adding a test or a compatibility bug fix since then you end up with a gazillion branches.

@CamphynR I'll investigate voltage calibration loading, I wouldn't be surprised at all if there are memory leaks there.

For fixing the memory fragmentation issue, I think the thing to do is to allocate a single large numpy array for the start windows and for the readout delays for all events being read and return slices to it. This is closer to what uproot is essentially doing and it fragments more slowly. You will still fragment quite a bit when looping over many runs because the number of events varies a lot, but hopefully the memory usage from these particular things will be limited to about twice the maximum needed.

Obviously loading eventinfos in smaller batches should use much less memory (though may still fragment quite a bit, we should test). It's not like you can load 100,000 waveforms into memory anyway for the longest runs, that would require tens of GB of RAM before you did anything with them. Limiting the batch size to something reasonable like a few hundred is probably better (and it's what the iterate interface more or less does, I think...).