jpivarski
commented
9 months ago This is true, and the memory use in uproot.dask is over and above uproot.open (and getting the TTree metadata).
import gc
import psutil
import uproot
this_process = psutil.Process()
def memory_diff(task):
gc.disable()
gc.collect()
start_memory = this_process.memory_full_info().uss
task()
gc.collect()
stop_memory = this_process.memory_full_info().uss
gc.enable()
return stop_memory - start_memory
def task():
with uproot.open(
{"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"}
) as tree:
pass
for _ in range(200):
print(f"{memory_diff(task) * 1e-6:.3f} MB")reports
28.156 MB
1.483 MB
0.000 MB
0.008 MB
0.012 MB
0.004 MB
3.932 MB
0.262 MB
0.000 MB
0.000 MB
0.000 MB
-1.040 MB
0.000 MB
0.803 MB
0.246 MB
0.000 MB
0.000 MB
0.000 MB
0.000 MB
-1.040 MB
0.807 MB
0.242 MB
0.004 MB
0.000 MB
0.000 MB
0.000 MB
...Change the task to
def task():
tree = uproot.open(
{"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"}
)so that it leaks file handles, and it's
26.059 MB
1.499 MB
0.004 MB
0.012 MB
0.033 MB
0.004 MB
0.000 MB
-0.012 MB
6.046 MB
0.258 MB
0.004 MB
0.000 MB
-1.049 MB
0.008 MB
-1.049 MB
0.000 MB
0.000 MB
...We'd eventually run out of file handles this way, but apparently not memory (on the MB scale).
Now
def task():
lazy = uproot.dask(
{"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"}
)87.364 MB
27.988 MB
25.252 MB
28.987 MB
25.227 MB
26.903 MB
25.219 MB
30.024 MB
23.155 MB
27.898 MB
26.325 MB
28.975 MB
24.158 MB
28.991 MB
...This is a problem. (Also, it's noticeably slower, though there might be good reasons for that.)
Using Pympler,
>>> import gc
>>> import pympler.tracker
>>> import uproot
>>>
>>> summary_tracker = pympler.tracker.SummaryTracker()
>>>
>>> # run it once to get past the necessary first-time things (filling uproot.classes, etc.)
>>> lazy = uproot.dask(
... {"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"}
... )
>>> del lazy
>>> gc.collect()
0
>>> # run print_diff enough times to get to the quiescent state
>>> summary_tracker.print_diff()
...
>>> summary_tracker.print_diff()
types | # objects | total size
======= | =========== | ============
>>>
>>> # what does an Uproot Dask array bring in?
>>> lazy = uproot.dask(
... {"https://github.com/CoffeaTeam/coffea/raw/master/tests/samples/nano_dy.root": "Events"}
... )
>>> summary_tracker.print_diff()
types | # objects | total size
================================================= | =========== | ============
dict | 72059 | 11.54 MB
bytes | 3 | 5.66 MB
list | 24007 | 1.73 MB
numpy.int64 | 37491 | 1.14 MB
uproot.source.cursor.Cursor | 21000 | 984.38 KB
numpy.ndarray | 4501 | 492.30 KB
str | 4886 | 436.10 KB
uproot.models.TObject.Model_TObject | 5999 | 281.20 KB
tuple | 3385 | 147.00 KB
uproot.models.TObjArray.Model_TObjArray | 3000 | 140.62 KB
uproot.models.TNamed.Model_TNamed | 2999 | 140.58 KB
frozenset | 1 | 128.21 KB
int | 3492 | 95.51 KB
awkward._nplikes.typetracer.TypeTracerArray | 1878 | 88.03 KB
uproot.models.TTree.Model_ROOT_3a3a_TIOFeatures | 1500 | 70.31 KB
>>>
>>> # what goes away when we delete it?
>>> del lazy
>>> gc.collect()
14
>>> gc.collect()
0
>>> summary_tracker.print_diff()
types | # objects | total size
============================================== | =========== | ============
code | 0 | 37 B
aiohttp.helpers.TimerContext | -1 | -48 B
awkward.contents.recordarray.RecordArray | -1 | -48 B
dask.highlevelgraph.HighLevelGraph | -1 | -48 B
dask_awkward.utils.LazyInputsDict | -1 | -48 B
dask.blockwise.BlockwiseDepDict | -1 | -48 B
awkward.highlevel.Array | -1 | -48 B
dask_awkward.layers.layers.AwkwardInputLayer | -1 | -48 B
dask_awkward.lib.core.Array | -1 | -48 B
asyncio.trsock.TransportSocket | -2 | -80 B
ssl.SSLObject | -2 | -96 B
aiohttp.streams.StreamReader | -2 | -96 B
asyncio.sslproto.SSLProtocol | -2 | -96 B
asyncio.sslproto._SSLPipe | -2 | -96 B
bytearray | -2 | -112 BHardly anything goes away when lazy is deleted! That's not good!
This TTree has 1499 TBranches. So having approximately that many TIOFeatures, TypeTracerArray, twice that many Model_TNamed, Model_TObjArray (TBranch and TLeaf), and four times as many Model_TObject make sense.
There are only 3 bytes objects, but they comprise 5.66 MB. I don't know, offhand, what they could be, but I think they're more likely Uproot than Dask. There are a lot of big dicts, which is not too surprising, and I can't say offhand whether I expect more in Uproot or more in Dask.
The one, major problem is that del lazy followed by gc.collect() does not get rid of as many objects as were brought in. It may be reasonable for the dask-awkward array of a large TTree to be 30 MB, but it's not reasonable for it to still be around after deleting.
Who gets a reference to it and doesn't let go? It might be possible to find out with gc.get_referrers, but it might not be at the level of lazy (the ak.Array and dak.Array are listed as objects that go away). Let me think about that...
agoose77
lgray
reproducer:
This particular instance leaks ~30MB per open. This adds up very quickly if you need to extract the form of hundreds of files in a remote process as evident from https://github.com/CoffeaTeam/coffea/issues/1007 where this bug manifested pretty nastily.