`LazyReferenceMapper` cache leaking memory

[honzaflash]

Preface: There could be reasons things are the way they currently are. I am not very familiar with the library.

That being said, I have observed a memory leak when periodically opening a parquet dataset via xarray+kerchunk
and I have traced it to the lru_cache usage on instance methods of LazyReferenceMapper.
Specifically, LazyReferenceMapper.listdir and LazyReferenceMapper._key_to_record.
Turns out lru_cache on methods is a common pitfall and it tends to result in memory leaks.

LazyReferenceMapper.open_refs is also cached BUT it is defined as a closure per-instance (in .setup) so no memory leakage there.

Assuming this is a bug, I would suggest giving listdir and _key_to_record a similar treatment as open_refs
or even better: applying the decorator in the constructor manually as per this stack overflow answer.

I am happy to work on a PR if desired.

(In case the reader is not aware of the peril of lru_cache on methods:

If a method is cached, the self instance argument is included in the cache

- https://docs.python.org/3/library/functools.html#functools.lru_cache

The two principal tools for caching methods are functools.cached_property() and functools.lru_cache(). The former stores results at the instance level and the latter at the class level.

- https://docs.python.org/3/faq/programming.html#faq-cache-method-calls

Therefore, the instance reference (self) is stored at the class level)

I can also provide an example to reproduce but it seems like this is a case of an issue that is well enough documented.

[TLCFEM]

A bit more discussion can be seen here: https://discuss.python.org/t/memoizing-methods-considered-harmful/24691

The memory usage indeed increases in long-running processes.
IMHO a manual cache would be better.

[martindurant]

Do I understand that you are opening the same dataset multiple times, or using ReferenceFS for different datasets?

[martindurant]

Either way, your point is valid: the thing that we actually care about caching is the URL->reference arrays, so that should be neither a method not closure, but a normal function.

The fsspec filesystem instances should not themselves be cached, because of the cachable=False instance attribute, but if a reference is preventing cleanup, that would indeed be a problem.

[honzaflash]

Do I understand that you are opening the same dataset multiple times

Yes.

I am not actually familiar with the lower-level kerchunk magic. We periodically check a dataset in our web service, ironically, to cache the result*. If this is of interest, the call looks something like this:

with xr.open_dataset(PTH, engine="kerchunk", chunks={}, decode_timedelta=True, cache=False) as ds:
    end_time = ds.cf["time"].max()

Adding cache clearing after resolved the memory usage growth:

fsspec.implementations.reference.LazyReferenceMapper.listdir.cache_clear()
fsspec.implementations.reference.LazyReferenceMapper._key_to_record.cache_clear()

the thing that we actually care about caching is the URL

I'm not sure I understand - are you talking about open_refs there? That function/method/closure does not actually take self so it doesn't "leak" like the other two I've mentioned. Sorry if I'm being confusing.

This is the fix that I had in mind based on the linked SO answer in my fork: honzaflash@5a3c61b

• pretty simple
• should fix the memory leak
• otherwise pretty transparent - the methods should work just like they did - since function decorator @ is just syntactic sugar for what the change does (except it does it inside the constructor and the self is already bound to the method)
• makes constructing a LazyReferenceMapper take longer, don't know if it's enough for you to care about

I haven't tested it at all, just wanted to share that idea. Maybe you are thinking of a better, smarter way of caching whatever actually needs caching in there.

EDIT: *yes, this is because the dataset file is updated every hour. However, for thoroughness I've tried it with a different, static dataset as well and the issue is still there as expected.