uproot4.iterate with XRootD - std::bad_alloc?

[ryuwd]

Currently I am using uproot4.iterate with 223 ROOT files (~400G total) over XRootD to filter through entries and only save certain columns in a parquet file.

I am using the latest development version of uproot as of the time of posting at commit 7e3353c.
I was using version 4.0.2, but ran into virtual memory problems (i.e. filling to max available + running into MemoryError), and figured that maybe using #281 would help.

Here's the code I'm using:

    arrays = uproot.iterate(
        data.get_list(tree), # a list of xrootd paths formatted like this: "{path}:{ROOT_tree_location}"
        filter_name=branches, # a set of 35 column / branch names
        library="pd", # my data is not jagged, so I'm using pandas
    )
    total = 0
    total_saved = 0
    with IncrementalPqWriter(output) as selection: # lets me write row groups to a parquet file from dataframes
        for df in arrays:
            mask = xicp_presel_mask(df, use_rectangular_cuts=use_rect_cuts) # returns boolean mask
            df["foldNumber"] = df["eventNumber"] % 2
            selection.write_from_df(df[mask])
            
            total += df.shape[0]
            total_saved += np.sum(mask)
            log.info(
                f" {np.sum(mask)} / {df.shape[0]} events saved. "
                f"Total: {total_saved} saved / {total} processed"
            )

  xicp_preselection  INFO     | Loading tree Xicp_ToPpKmPip/DecayTree from 223 ROOT file(s). 35 columns will be saved.
  xicp_preselection  INFO     |  10172 / 47818 events saved. Total: 10172 saved / 47818 processed
  xicp_preselection  INFO     |  7412 / 35161 events saved. Total: 17584 saved / 82979 processed
  xicp_preselection  INFO     |  6339 / 30022 events saved. Total: 23923 saved / 113001 processed
  xicp_preselection  INFO     |  6169 / 28699 events saved. Total: 30092 saved / 141700 processed
  xicp_preselection  INFO     |  8061 / 37881 events saved. Total: 38153 saved / 179581 processed
  xicp_preselection  INFO     |  5939 / 28098 events saved. Total: 44092 saved / 207679 processed
  xicp_preselection  INFO     |  10060 / 47716 events saved. Total: 54152 saved / 255395 processed
  xicp_preselection  INFO     |  4562 / 21916 events saved. Total: 58714 saved / 277311 processed
  xicp_preselection  INFO     |  9492 / 44276 events saved. Total: 68206 saved / 321587 processed
  xicp_preselection  INFO     |  6724 / 31274 events saved. Total: 74930 saved / 352861 processed
  xicp_preselection  INFO     |  8489 / 39955 events saved. Total: 83419 saved / 392816 processed
  xicp_preselection  INFO     |  8082 / 38127 events saved. Total: 91501 saved / 430943 processed
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
/bin/bash: line 1: 34478 Aborted                 GIT_REVISION="$(git rev-parse HEAD)" PYTHONPATH="/.../" python /.../.snakemake/scripts/tmpreua4su2.wrapper.py

I think this error is maybe related to #281. In 4.0.2 I was just getting MemoryError - which seems to be gone now.

[jpivarski]

MemoryError and std::bad_alloc are pretty similar: the first is Python's exception, the second is C++'s exception. The allocation that actually fails to allocate memory, whether it's in Python or C++, is not where the real problem lies. It's a matter of other things using up the memory.

If it's running out of memory "legitimately," which is to say, you just don't have enough memory to do what you're trying to do, then it is difficult to diagnose on a different computer with a different amount of memory available. If it's a memory leak (which #281 was not), then it happens on any computer if you wait long enough.

What you're doing in your code doesn't involve Pandas, just some masking, so you can remove a big performance and memory hog by replacing library="pd" with library="np" and do everything with plain NumPy arrays. (What you've named "df" would come out as a dict of NumPy arrays, which can be selected with square brackets and strings, just as you have it.)

The Python garbage collector is supposed to get invoked when you start running low on memory, though it could only do that if the limit is approached while the Python code is active (i.e. a MemoryError, not std::bad_alloc). In #281, however, it wasn't getting invoked before a Pandas DataFrame construction, but that could be because Pandas has its own C++ code to hide allocations from the garbage collector. So removing Pandas might help that, or you could import gc and gc.collect() at the end of your loop.

Actually, thinking about this more deeply, the arrays you make in one step of iteration (mask, df, etc.) are still in scope when the next iteration starts, so either put all the logic in the body of the loop or explicitly del mask, del df, etc. when done, maybe in addition to gc.collect(). (The garbage collector can only remove things that are out of scope.) As it is, there has to be two iterations' worth of arrays in memory at the time that the next iteration makes arrays—the old and the new ones. So that can also help.

Beyond that, the resource utilization of XRootD are a bit of a mystery to me. @nsmith- recommends

uproot.open.defaults["xrootd_handler"] = uproot.source.xrootd.MultithreadedXRootDSource

and this should probably become the default. The current default does a vector-read, which has caused a number of problems.

Beyond that, maybe break your work up into smaller processes. (That's what I'd do if I run out of all other options.)

I'm going to convert this into a Discussion because it's not really a bug, that I know of.

[ryuwd]

Hi @jpivarski,

Thanks for the detailed explanation and your recommendations, it's really appreciated.

MemoryError and std::bad_alloc are pretty similar: the first is Python's exception, the second is C++'s exception.

Ah, I naively thought that if I was still dealing with an OOM problem that python would throw another MemoryError. It didn't occur to me that it could also present as std::bad_alloc. Thanks!

For some context: I was running on Scientific Linux 7 in an interactive job on a cluster, the virtual memory constraint in this case was 16GB, with 8 cores available. I guess I didn't notice this problem when I was running this on my local PC, as there were a lot of allocations but no complaints from the OS causing an allocation error, since there was plenty of swap (in contrast to the hard 16GB Vmem limit on the cluster).

Actually, thinking about this more deeply, the arrays you make in one step of iteration (mask, df, etc.) are still in scope when the next iteration starts, so either put all the logic in the body of the loop or explicitly del mask, del df, etc. when done, maybe in addition to gc.collect().

Thanks for pointing this out! I moved the operations into a function. I also implemented all the other suggestions (library="np", etc) and now the virtual memory doesn't exceed the Vmem upper-limit.

I realised that a previous call earlier on in this code to uproot.lazy was filling up close to my 16GB available Vmem, although I specified 500 MB as the array cache size. I was running ak.sum over 223 doubles in the tree across the 223 files (i.e. 1 entry per file, but these are the same files holding the tree with 35+ columns), so I guess that something is going on with XRootD there.
When I changed this code to use uproot.iterate(..., library='np') with gc.collect the usage went right down. Now my Vmem usage doesn't exceed 1.5 GB over the whole process.

Anyway - thanks! The suggestions helped a lot.