uproot's memory usage

[pfackeldey]

This post is an update on uproot's memory usage

The recent PR #1305 improved uproot's memory usage significantly. Before, memory was accumulating where it shouldn't, a comparison is shown in #1305 (comment).

The memory usage can be measured with the following setup:

Requirements

• uproot (including PR perf: removing output arrays from reference cycles so they don't have to wait for GC #1305)
• memray
• absl-py (for CLI args)

Script

import uproot

from absl import app
from absl import flags

FLAGS = flags.FLAGS
flags.DEFINE_string("rootfile", "~/Downloads/zlib9-jagged0.root", "path to root file.")
flags.DEFINE_string("loop", None, "choose loop to run.")
flags.DEFINE_boolean("gc", True, "turns on pythons GC.")


def loop_tbranch_array(rootfile):
    tree = uproot.open(rootfile, array_cache=None, object_cache=None)["tree"]

    how_often = 8
    # repeat the first one
    branch = list(tree.keys())[0]
    for _ in range(how_often):
        batch = tree[branch].array(library="np")
        print(repr(batch))


def loop_ttree_arrays(rootfile):
    tree = uproot.open(rootfile, array_cache=None, object_cache=None)["tree"]

    how_often = 8
    for _ in range(how_often):
        batch = tree.arrays(library="np")
        print(repr(batch))


def loop_iterate(rootfile):
    how_often = 8
    filenames = [{rootfile: "tree"}] * how_often
    # choosing (very large) step_size to be 10 GB just to make sure everything 
    # is read in one go (same as in the other loops)
    for batch in uproot.iterate(filenames, step_size="10 GB", library="np"):
        print(repr(batch))


def main(argv):
    del argv

    if not FLAGS.gc:
        import gc

        gc.disable()

    loop = {l.__name__: l for l in [loop_tbranch_array, loop_ttree_arrays, loop_iterate]}[FLAGS.loop]
    # run
    loop(FLAGS.rootfile)


if __name__ == "__main__":
    flags.mark_flag_as_required("loop")
    app.run(main)

How to run (with GC disabled)

for l in 'loop_tbranch_array' 'loop_ttree_arrays' 'loop_iterate'; do\
  echo memray run -o $l.nogc.bin issue.py --loop=$l --nogc &&\
  memray run -o $l.nogc.bin issue.py --loop=$l --nogc &&\
  memray flamegraph $l.nogc.bin;\
  done

The resulting flame graph (with GC disabled) is as follows for each loop kind:

Uproot's memory consumption is stable (not growing) and follows an expected sawtooth structure 🎉