Assigning virtual arrays as a field to an array materializes the virtual arrays

[jrueb]

Version of Awkward Array

1.4.0 (and git main)

Description and code to reproduce

When inserting a virtual array into an array created with ak.from_buffers with lazy=True, the generator function of the virtual array gets called right away. This however happens only for the first virtual array being inserted. I would expect the generator function not being called at all, as the virtual arrays content is not required at this point.
Code example:

import numpy as np
import awkward as ak


def some_function1(length):
    print("called1")
    return np.ones(length)


def some_function2(length):
    print("called2")
    return np.ones(length)


array = ak.from_buffers(*ak.to_buffers(ak.Array({"something": [1,2,3]})), lazy=True)
array["virtual"] = ak.virtual(some_function1, (len(array),), length=len(array))
array["virtual"] = ak.virtual(some_function2, (len(array),), length=len(array))

This will print called1, even though it shouldn't. If lazy=False in the ak.from_buffers call, it doesn't print anything, like it should.

In the same context i noticed a continuous increase in memory usage, in case the array is also used as part of the args parameter. To me this looks a lot like something isn't being freed. Code example:

import os
import psutil
import gc
import numpy as np
import awkward as ak


def some_function(length, something):
    print("called")
    return np.ones(length)


while True:
    gc.collect()
    big = ak.Array({"something": np.random.rand(100000000)})
    array = ak.from_buffers(*ak.to_buffers(big), lazy=True)
    array["virtual"] = ak.virtual(some_function, (len(array), array), length=len(array))
    array["virtual"] = ak.virtual(some_function, (len(array), array), length=len(array))
    print(psutil.Process(os.getpid()).memory_info().rss)

The resident set size, which is printed, will continue to increase with every loop. This does not happen (there is almost no change to the number) if array["virtual"] is only set once. It also does not happen if big instead of array is used in the args parameter.

[agoose77]

RE the lazy aspect to this question, I assume what's happening here is that ak.with_field is materialising the virtual array when broadcasting.

[jpivarski]

RE the lazy aspect to this question, I assume what's happening here is that ak.with_field is materialising the virtual array when broadcasting.

That's right.

Assigning an array x as a new field of another array y necessarily evaluates any generators in x and y because it needs to broadcast them to see how to restructure x to fit y. That's so that something like

array_of_3vectors["mass"] = 125.0

can set a "mass" field with the same value in all records in array_of_3vectors, regardless of whether array_of_3vectors is an array of records, an array of lists of records, or something deeper. Similarly,

array_of_lists_of_lists["new_field"] = array_of_lists

will broadcast each element from the shallower array_of_lists into the corresponding array_of_lists_of_lists.

This is a fundamental problem with our design of lazy arrays: the idea is that it will be materialized when needed, but it can be hard for users to guess when it's needed. (Maybe you didn't know that field-assignment broadcasts, for instance.) There's a long history of bug-fixes for lazy arrays being materialized "too early," but without a clear definition of when is "too early." In some cases, it was obvious: some code was just checking something that could be determined from type alone, such as the value of the "__record__" parameter. Materializing a VirtualArray node for that would be wrong, because the "__record__" parameter is expressed in the VirtualArray's Form, and we can search down that tree instead. Other cases have been more borderline: should a slice materialize a VirtualArray? It's a calculation on an array, but most people didn't expect it to materialize an array, especially when the slice was just picking out one field (e.g. my_array.field). Now a slice of a VirtualArray creates a new VirtualArray that would cascade through evaluating the slice, then evaluating the original array.

But for broadcasting, we can't delay that. To delay a calculation, we have to be able to say what its Form is going to be without actually performing the calculation. There are even some slices in which that's not possible: the Form of the output depends on the specific values of the array. (That is, if some lists are empty or well-aligned or something, the output would have a different Form than if they weren't, and in order to make a new VirtualArray, we have to predict the Form it would have upon evaluation.)

Dask solves the lazy array problem differently, and that's where we're directing new effort on this:

One difference is that a Dask collection has a .compute() method, so you get to say when it's the right time to start a calculation. As a consequence, the Dask DAG can be a different thing from an ordinary array (i.e. .compute() returns a different type, mapping DAG → in-memory array), so it doesn't need to support all the same operations that an ordinary array supports. Perhaps there would be no way to ask a DAG what the array type will be without computing it, for instance. But this would relax constraints on what can go into the DAG so that any operation can be lazy.

Development of an Awkward-Dask collection is starting next month.

---

As for solving your problem: are you trying to make a record array with virtual fields? Instead of assigning new fields into an existing object, you could create the array all at once by calling the ak.Array constructor:

array = ak.Array({"field1": ak.virtual(...), "field2": ak.virtual(...)})

or ak.zip, limited to the non-virtual, list-like dimensions:

array = ak.zip({"field1": ak.virtual(...), "field2": ak.virtual(...)}, depth_limit=1)

If the fields were lists of virtual data (i.e. the list offsets are not virtual, but the contents are), then you could use depth_limit=2 to zip those lists together without materializing the virtual data. depth_limit=1 is equivalent to the ak.Array constructor.

[jrueb]

I forgot to mention this: The array in my case is a NanoEvents generated by Coffea. I don't think the ak.Array constructor is usable here. I would have to loop over all the fields, possibly recursively, most likely breaking something about the structure.

Something I still don't understand and which made me really sure this a bug in laziness is that the second insertion does not cause the function to be called. If the first call is on purpose, I would assume the second insertion to have the same behavior. Why doesn't the second array need to be broadcasted?
If this is really how it is supposed to be, I can workaround this by just inserting a virtual array that doesn't need any heavy computation first, followed by my virtual arrays that actually need the heavy computation. Only the first one will be called right away anyway. But this seemed too arbitrary to me.

Should I maybe create a separate issue for the memory leak? I thought these two are probably connected because I only could reproduce the leak in this exact same context.

[jpivarski]

Why doesn't the second array need to be broadcasted?

Broadcasting depends on all arrays involved in the broadcast. In each line like

array["virtual"] = ak.virtual(some_function1, (len(array),), length=len(array))

there are two arrays: the array on the left-hand side and the ak.virtual(...) on the right-hand side. The left-hand side is changed by this operation—it contains virtual arrays and they're materialized by the broadcast, after which they hold data in a cache. It happens to be the case in this instance that

array["virtual2"] = ak.virtual(some_function2, (len(array),), length=len(array))

(or assigning to the same field name) doesn't follow a code path that evaluates the right-hand side because the left-hand side is different now, but that's a fragile situation that I wouldn't rely on.

You can always collect many arrays and pass them all to an ak.Array constructor or ak.zip, even if you had to go through a loop or recursive function to find them: this is a common functional pattern, and it works well with Awkward Arrays because—apart from this field-assignment and the internal caching of VirtualArrays—they're immutable. In fact, nearly every function in the Awkward codebase is a recursive walk through an unknown tree, returning a new tree with this functional pattern. I would especially recommend it if you have a large number of fields. The broadcasting in a field-assignment is a performance cost (even potential broadcasting, as it has to check to see if broadcasting is necessary).

About the memory leak, let me look at that again. I think I know what it is (and just edited away a lot of speculation). I'll post in the next reply.

[jpivarski]

The key thing about the memory leak is that it happens when a lazy array is used as a function argument but not when a non-lazy array is used as a function argument. I think what's happening is that Python can't see reference cycles like A → B → A when A is a Python object and B is a C++ object. All of the mallocs and frees are appropriately paired up (because the C++ code uses std::shared_ptr throughout), but that doesn't work for cycles.

@nsmith- found a similar issue when a VirtualArray's cache contained the VirtualArray itself. The cache is a Python mapping (to let you use cachetools, a plain dict, etc.) but the VirtualArray and its components (ArrayGenerator and ArrayCache) are C++ objects that were designed to be independent of Python. When they wrap a Python object, their std::shared_ptr calls PyDECREF, so mallocs and frees are paired up, but Python's garbage collector can't traverse the object graph through a C++ object to know that a std::shared_ptr points to something that contains a Python object. Not only has Python's garbage collector not been instrumented to do this, but also the type information about what needs to be traversed and how to interpret it does not exist after compile-time. You need a language that can reflect all of its types to be able to do that.

Our solution with the VirtualArray cache problem was to make the ArrayCache hold only a weak reference to the Python cache object, and then connect the strong reference through ak.Array (a hidden attribute named _caches). It works: the cycle only goes through Python objects and the garbage collector is able to detect cycles and delete them, but at a considerable cost in complexity. The bug-fix to make this reference weak (#427) was followed by many other bug-fixes to deal with the fact that it is weak and the reference can get lost (#428, #479, #560).

What you've found is the same story, but with ArrayGenerator's hold on Python objects as arguments to the function, rather than ArrayCache's hold on a Python object as a MutableMapping. I think @nsmith- predicted yet another way of forming cycles: by the ArrayGenerator's function holding a closure to the array itself, but holding arguments is effectively the same. That sounded unlikely to come up in practice: why would the function to fill in an array need to reference the holder where that function would go? Furthermore, in a purely immutable world, it wouldn't be possible to set up that situation: you wouldn't be able to reference a VirtualArray in one of the arguments that constructs the VirtualArray because it hasn't been made yet. @nsmith-'s example of a closure referencing itself is possible because Python's global namespaces are mutable: that's how you can write a recursive function—Python doesn't attempt to evaluate a reference to f inside the definition of f until f is called, which is after it has been fully defined. (See this if you're interested.) Your example of the arguments referencing itself is only possible because of field-assignment, which is another exception from the general rule against mutability. If you weren't assigning

array["virtual"] = something_involving(array)

but building it all at once with the ak.Array constructor or ak.zip, then it wouldn't be possible to make this cyclic reference.

Just as in the vagueness about when a lazy array materializes, this memory leak is a consequence of a bad design decision: if we implemented the array nodes in Python, rather than C++, then none of this would have ever come up. As the attempt to patch VirtualArray caches shows, trying to patch this with weak references will bring in its own problems. Just as the vagueness of lazy array materialization will be solved by relying instead on Dask, the problems introduced by C++ are being addressed by refactoring Awkward Array's middle layer from C++ to Python. There's a talk on that; it's an ongoing project, but the gist is below:

We replace the C++ box with a Python box:

The motivation described in that talk was for JAX (hence the "differentiated kernels" that becomes possible), but many other issues boiled down to the same thing. JAX and Dask can both be supported better if those libraries' tracer objects can "see" down to the level of kernels, but the long-standing problem of reference cycles would also be fixed by letting Python's garbage collector "see" down to this level. Kernel functions, by the way, do zero memory management (purely borrowed references).

So the memory leak you found is real, but our fix for it is going to be this refactoring.

[jpivarski]

Someday, I should give a talk on "Python/C++ integration: just because you can doesn't mean you should." The tooling is great (e.g. pybind11), but the problems are architectural, six steps ahead of the actual coding and hard to predict. I certainly didn't see this far ahead.

Sometimes, of course, you have to include C++ in a Python project for performance or interoperability with another C++ codebase. But then, the lessons would be like the lessons about mutability from functional programming: if we can't be purists, at least we can learn to be wary and only introduce mutability—or C++ integration—when necessary, knowing which problems to look out for when doing so. With mutability, we have to worry about the distinction between copying and viewing. With C++ integration, we have to worry about integrating memory management. Just matching shared pointers to PyINCREF and PyDECREF (which pybind11 does automatically, by the way) is not enough: reference cycles are important, too.

[nsmith-]

Yet another example of creating reference cycles is by creating a virtual array that closes over a highlevel array using the same cache: scikit-hep/coffea#487
There are so many ways to make cycles 😢

[jrueb]

Thanks for the insightful discussion. i decided to give the ak.Array constructor a chance as @jpivarski suggested and it turned out to actually solve not only the first part but both parts of the issue, even with NanoEvents. Everything put into a simple class looks like this:

class VirtualArrayCopier:
    def __init__(self, array):
        self.data = {f: array[f] for f in ak.fields(array)}
        self.behavior = array.behavior

    def __setitem__(self, key, value):
        self.data[key] = value

    def get(self):
        array = ak.Array(self.data)
        array.behavior = self.behavior
        return array

    def wrap_with_copy(self, func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(self.get(), *args, **kwargs)
        return wrapper

By using the __setitem__ method followed by the get method, one can add fields of virtual arrays without accidentally triggering the materialization. At the same time using the data attribute (or simpler wrap_with_copy) inside the args parameter, one avoids the cyclic reference. Probably there still are some things to improve here.