How to ensure arrays are loaded before AsdfFile is closed?

[lgarrison]

In a project I work on, I often want to read one value or array out of an ASDF file and then close the file. My first instinct is to write something like:

with asdf.open('myfile.asdf') as af:
    a = af['arr1']

But then I remember that the arrays are lazy-loaded (and mem-mapped), so I use [:] to force a non-lazy load and copy_arrays=True to ensure this isn't just creating a mem-mapped view that will become invalid when af is closed:

with asdf.open('myfile.asdf', copy_arrays=True) as af:
    a = af['arr1'][:]

I could simplify this by using .copy() instead of [:] and removing copy_arrays=True:

with asdf.open('myfile.asdf') as af:
    a = af['arr1'].copy()

Except I sometimes work with compressed data (which can't be mem-mapped), so I think this would first decompress the array, then create a copy of it (correct me if I'm wrong here!), and I'd rather not eat the performance/peak memory penalty.

Furthermore, when reading my collaborator's code, I'm not always sure what's an array and what's not. So if I see a pattern like:

with asdf.open('myfile.asdf') as af:
    v = af['AmbiguousName']

I'm not always sure if it's a latent bug, or valid code (if it's referring to, e.g., an int or other non-array data).

So my question is: is there a pattern I can use that will ensure array data is immediately loaded (and not memory-mapped) on first reference? I can't always do lazy_load=False, copy_arrays=True, because I don't always want all arrays to be loaded—just referenced ones. How do others approach this problem?

Thanks!

[braingram]

Thanks for starting this discussion. I believe you've captured all the options that currently exist. These range from "safe and slow" (lazy_load=False and copy_arrays=True) to "dangerous and fast" (the current default lazy_load=True and copy_arrays=False).

Thanks for bringing up the point about compression complicating these options. For your example:

with asdf.open('myfile.asdf', copy_arrays=True) as af:
    a = af['arr1'][:]

The array would be decompressed and copied to memory when any element of af['arr1'] is accessed and the [:] would result in a copy of that already decompressed data.

with asdf.open('myfile.asdf') as af:
    a = af['arr1'].copy()

has the same result.

with asdf.open('myfile.asdf', lazy_load=False, copy_arrays=True) as af:
    a = af['arr1']

avoids the extra copy. However as you noted this loads every array in the file into memory during asdf.open.

Keeping the file open can avoid the extra copy.

with asdf.open('myfile.asdf') as af:
    a = af['arr1']
    # do stuff with a, no extra copy is needed

However scope of af has to be considered as

def get_arr():
    with asdf.open('myfile.asdf') as af:
        return af['arr1']   # note that af falls out of scope

a = get_arr()
print(a[0])   # error!

the above example will fail with a OSError: Attempt to read block data from missing block if the return result is used as it was never accessed with af was open.

[lgarrison]

Keeping the file open can avoid the extra copy.

with asdf.open('myfile.asdf') as af:
    a = af['arr1']
    # do stuff with a, no extra copy is needed

I do lean on this pattern whenever possible, but sometimes the array is returned (as in your last example), or attached to some other object with a longer lifetime than the AsdfFile.

Relatedly, I've seen ASDF newcomers who are trying to use this pattern get tripped up in situations like:

with asdf.open('myfile.asdf') as af:
    a = af['arr1'] * 2  # no error
    b = af['arr2']  # latent bug!
    c = af['SomethingElse']  # maybe a bug?

Doing an array operation loads the array, but if the next array happens to not be multiplied by 2, then that's a hidden bug. And whether c has a hidden bug depends on whether af['SomethingElse'] is an array or not.

I do wonder if some intermediate solution would make sense, where every array that is explicitly named gets eagerly loaded; something like lazy_load='eager'? I realize this wouldn't fix every lifetime issue, because one could easily keep a reference to the tree and then find arrays that have exceeded the AsdfFile lifetime that way, but I think that's where you would want to use lazy_load=False anyway.

Or perhaps if a reliable way existed to load an array without triggering an extra copy, then one could just use that? For example:

with asdf.open('myfile.asdf') as af:
    a = af['arr1'].load()  # syntax 1
    a = af.load('arr1')  # syntax 2

The load() function here would have "load if not already loaded" semantics. Then the rule of thumb for developers (at least in my project) could be to always use load() when referencing arrays. Which still has the problem that one can't tell what's an array by reading the code (lazy_load='eager' would sidestep that problem), but at least it provides a clear mechanism by which to (1) avoid lifetime bugs, and (2) avoid extra copies.

[braingram]

Thanks for the write-up and the thought you've put into this.

For load, NDArrayType (the class that handles the lazy array loading) does have a (private) _make_array function that I think accomplishes what you've described. Perhaps we should expose this as a public method? One issue this won't address is how to trigger loading of entire sub-trees. Let's say arr1 and arr2 are in a data dictionary in the asdf file. It might be nice to call

with asdf.open('myfile.asdf') as af:
    a = af['data'].load()  # to trigger loading of af['data']['arr1'] and af['data']['arr2']

However this would mean that af['data'] could no longer be a standard python dictionary. This is related to another discussion about making more parts of the tree lazy: #1705 There is an open PR adding lazy_tree however it does not include the above load method. This might help to address some of the concerns about using builtin classes vs custom classes (as long as load always returns a "builtin" or deserialized type, in this case ideally a standard numpy array).

One "solution" I'm envisioning is an asdf file opened with lazy_tree=True and extending the lazy object API to include a load. Do you think that would work for your uses? I'm curious to hear your thoughts on if the "lazy_tree" would make this situation better or worse (as now even dict-like and list-like containers are linked to the lifetime of the AsdfFile object).

[lgarrison]

Thanks; I did wonder about loading entire sub-trees, too, but I wasn't sure if promoting the tree from a regular dict to an ASDF-aware type was within scope.

I think lazy_tree combined with load() would be okay. In my project, we don't have that much non-array data or complicated trees, so I would be inclined to keep using lazy_load since I don't mind if non-array data gets eagerly loaded.

However, I did a little more testing of this pattern:

with asdf.open('myfile.asdf', copy_arrays=True) as af:
    a = af['arr1'][:]

And it looks like for both compressed and non-compressed arrays, the max RSS is as expected (i.e. only a single copy of the [decompressed] array is in memory). I think [:] is just creating a view. So I may not even need _make_array()/load(), although the latter might be more readable.

I guess that just leaves me with the af['AmbiguousName'] problem. A generic load() that's safe to tack on to any part of the tree (unlike [:]) may be a good solution. Then the rule of thumb for lifetime safety in my project would be: any bare af['somekey'] that's not part of a larger expression should have a .load() attached. And always use copy_arrays=True (which I understand will become unnecessary in ASDF 4; #1648).

[ketozhang]

Do all options of explicit loading give NDArrayType? Is there any way to get an actual np.ndarray?

(Dealing with ketozhang/asdf-pydantic#13)

[braingram]

I'm not aware of one. You could use the [:] view technique or asanyarray to convert the NDArrayType to an array (which won't copy the data but will produce a ndarray). I opened an issue to track this: #1787