WIP: Support pickling lru_cache on CPython

[ssanderson]

I was poking through the issue tracker over the weekend and noticed #178. This is an initial pass on what I think is feasible to support without aggressive hackery:

TL;DR:

• This PR enables support for pickling functions decorated with lru_cache on CPython.
• This PR does not pickle an lru_cached function's associated cache.
• This PR preserves the value of cached_func.maxsize, but loses the value of typed that was passed to the cached function. I don't think typed is used very often, but it seems unfortunate that this PR silently discards that info. I don't have a good idea for how to fix that without upstream changes in CPython.
• This PR doesn't change the current behavior on PyPy which is to just deep-serialize the entire function closure, including the cache. It's unclear to me whether we should be willing to pay an extra runtime cost (though, admittedly, a pretty trivial one) to make this behavior consistent.

---

Following @ogrisel's comment in #178 (comment), we can support simple usage of lru_cache in CPython straightforwardly, because lru-cached functions are instances of an extension type that provides enough of an API (via .__wrapped__ and .cache_info()) to allow us to extract the cache's maxsize.

Unfortunately, we do not have a way to recover the typed parameter that was added in python 3.3, which causes the decorated function to use separate caches for different input types. This might be enough of a concern that we don't want to merge this PR (or, at least, wouldn't want to enable it without some kind of explicit opt-in from the user). As-is, this PR just silently discards the typed param, which is not great.

On PyPy, things are a bit trickier. There, lru_cache uses the pure-python implementation from the stdlib, which just returns a closure that closes over the cache state. Since we can already serialize closures, this "works", but it introduces a new issue, which is that PyPy also serializes the cache, which might not be desirable (at least, it seems unfortunate that the behavior would be different between CPython and PyPY). On the other hand, it's not clear how we could prevent this: pickle's dispatching is all type based, and to any external viewer an lru-cached function is just a regular python function on pypy, so there isn't really a way for us to treat lru-cached functions specially unless we want to add a special case to check on all function serializations (e.g. we could branch on if hasattr(func, 'cache_clear') or similar). Whether that cost is worth the effort to ensure consistency between CPython and PyPy is an interesting question. I'm seeing several unrelated test failures on pypy on master anyway, so maybe we don't care?

[csadorf]

I just remembered that I looked at this while back, but never completed it. But might be a good point of reference nonetheless: master...csadorf:fix/issue-178

[pierreglaser]

Thanks for the patch. You also need to implement this behavior for cloudpickle_fast.py (this is why travis fails on 3.8). This should just be some copy-pasting here and there.
This reducer duplication is clunky and we should probably do some refactoring work to avoid it (see #284)

[ssanderson]

@pierreglaser I'm happy to port this to cloudpickle_fast if we decide this is a reasonable path forward, but I think the first order question to answer is whether the behavior added here (in particular, losing the value of typed) is acceptable. If not, I think that essentially means that lru_cache is unsupportable without upstream changes unless we want to resort to memory-unsafe ctypes hacks.

cc @ogrisel and @csadorf, any opinions on ⬆️?

[pierreglaser]

Ah, right. Is there any good reason for CPython not to expose the typed parameter?
As for the PyPy, that means that even right now, cached functions are serializable by cloudpickle? In this case the behavior right now is not consistent either — adding support to serialize CPython’s cached functions would not really worsen the situation.

[ssanderson]

Is there any good reason for CPython not to expose the typed parameter?

Not that I can think of besides a possible general aversion to exposing implementation details. But given that lru_cache-decorated functions already expose other cache statistics, it seems reasonable to expose typed as well.

As for the PyPy, that means that even right now, cached functions are serializable by cloudpickle?

That seems to be the case yes, although I'm a bit confused how that works b/c the pure python implementation in the stdlib uses RLocks to protect concurrent calls to cached functions from separate threads, which means the closure contains an RLock, which I didn't think we supported. Did that change recently?

[pierreglaser]

Not that I can think of besides a possible general aversion to exposing implementation details.

Then my opinion is that it is reasonable to open an issue in upstream to start a discussion. It could be exposed as a read-only attribute, it should not hurt much.

Did that change recently?

No, we do not support RLock serialization explicitly, but it may be that PyPy RLocks know how to pickle themselves (although I haven’t checked)

[csadorf]

@pierreglaser I'm happy to port this to cloudpickle_fast if we decide this is a reasonable path forward, but I think the first order question to answer is whether the behavior added here (in particular, losing the value of typed) is acceptable. If not, I think that essentially means that lru_cache is unsupportable without upstream changes unless we want to resort to memory-unsafe ctypes hacks.

cc @ogrisel and @csadorf, any opinions on ?

I think that we cannot move forward without also supporting the typed value. That is because from an application point of view this is not just a pure implementation detail, but could potentially lead to very different behavior that would be extremely difficult to debug.

[ssanderson]

I think that we cannot move forward without also supporting the typed value.

I think this is my inclination as well. If so, then this PR can probably be closed until upstream provides a way to expose typed to python.

[ssanderson]

I posted an issue to bugs.python.org here: https://bugs.python.org/issue38565. Will post a PR to expose the attribute in a bit. Any advice on who to prod for getting those reviewed/triaged would be helpful.

[ogrisel]

The typed value is now exposed in upstream Python 3.9 (dev): python/cpython@051ff52

Feel free to update this PR accordingly.

[ogrisel]

Thinking about this a bit more, lru_cache could typically be used to decorate an interactively function called on a dedicated process worker in loky, dask (and probably ray too). For instance let's consider the following piece of code using the loky / dask API:

@lru_cache(maxsize=1000)
def expensive_computation(a):
    return a ** 2

for result in executor.map(expensive_computation, [0, 1, 42] * 10000):
    print(result)

the lru_cache will never kick-in, even if the call ends up running on the same worker process consecutively, because each time the definition of the square interactively defined function is shipped over and over again and gets garbage collected in between calls.

This behavior might be surprising to the users but I am not sure exactly how we could fix this. We could annotate the LRUCache instance with a UUID at pickling time a bit like what we do for class definition and enums (see _lookup_class_or_track and _get_or_create_tracker_id) and use it to redecorate interactively defined function at unpickling time and then forcibly keep those LRUCache instances around when the function defintions itself gets garbage collected in case a future remote call will need to reuse it again. But then I am not sure when we should collect the cache in case the original function and its cache are collected in the caller process? Never? Maybe we could have a limit on the number of LRUCache instances we keep around and collect the least recently used (some kind LRUCache or LRUCache instances?).

[NeilGirdhar]

There's no way to plug into extract_func_data to teach it how to pickle custom method types, is there?