User-pickling with dynamic sub-classing
Opened by BramVanroy (closed)

This is a continuation of the now closed PR in https://github.com/huggingface/datasets/pull/3206. The discussion there has shaped a new approach to do this.

In this PR, behavior of `pklregister` and `Pickler` is extended. Earlier, users were already able to register custom pickle functions. That is useful if they have objects that are not easily picklable with default methods. When one registers a custom function to a type, an object of that type will be pickled with the given function by `Pickler` which looks up the type in its `dispatch` table. The downside of this method, and of `pickle` in general, is that it is limited to direct type-matching and does not allow sub-classes. In many, default, cases that is not an issue. But when you are using external libraries where classes (e.g. parsers, models) are sub-classed this is not ideal. 

```python
from datasets.fingerprint import Hasher
from datasets.utils.py_utils import pklregister

class BaseParser:
    pass

class EnglishParser(BaseParser):
    pass

@pklregister(BaseParser)
def custom_pkl_func(pickler, obj):
    print(f"Called the custom pickle function for type {type(obj)}!")
    # do something with the obj and ultimately save with the pickler

base = BaseParser()
en = EnglishParser()

# Hasher.hash uses the Pickler behind the scenes
# `custom_pkl_func` called for base
Hasher.hash(base)
# `custom_pkl_func` not called for en :-(
Hasher.hash(en)
```

In the example above we'd want to sub-class `EnglishParser` to be handled in the same way as its super-class `BaseParser`. This PR solves that by allowing for a keyword-argument `allow_subclasses` in `pklregister` (default: `False`). 

```python
@pklregister(BaseParser, allow_subclasses=True)
```

When this option is enabled, we not only save the function in `Pickler.dispatch` but also save it in a custom table `Pickler.subclass_dispatch` **which allows us to dynamically add sub-classes of that class to the real dispatch table**. Then, if we want to pickle an  object `obj` with `Pickler.dump()` (which ultimately will call `Pickler.save()`) we _first_  check whether any of the object's super-classes exist in `Pickler.sublcass_dispatch` and get the related custom pickle function. If we find one, we add the type of `obj` alongside the function to `Pickler.dispatch`. All of this happens at the start of the call to `Pickler.save()`. _Only then_ dill.Pickler's `save` will be called, which in turn will call `pickle._Pickler.save` which handles everything. Here, the `Pickler.dispatch` table will be used to look up custom pickler functions - and it now also includes the function for `obj`, which was copied from its super-class, which we added at the very start of our custom `Pickler.save()`.

For edge cases and, especially, for testing, a contextmanager class `TempPickleRegistry` is included that resets the pickle registry on exit to its previous state.

```python
with TempPickleRegistry():
    @pklregister(MyObjClass)
    def pickle_registry_test_false(pickler, obj):
        pickler.save(obj.fancy_method())

    some_obj = MyObjClass()
    dumps(some_obj)
    # `MyObjClass` is in Pickler.dispatch

# ... `MyObjClass` is _not_ in Pickler.dispatch anymore
```

closes https://github.com/huggingface/datasets/issues/3178

To Do
====
- [x] Write tests
- [ ] Write documentation/examples?

Comments:
@lhoestq Feel free to have a look. The implementation is slightly different from what you suggested. I have opted to overwrite `save` instead of meddling with `save_global`. `save_global` is called very late down in dill/pickle so it is hard to control for what is happening there. I might be wrong. Pickling is more complex than I thought! 

The linked issue (`map` with spaCy) also works now!

```python
import pickle
import spacy
from spacy import Language
from datasets import load_dataset
from datasets.utils.py_utils import dumps, pklregister

@pklregister(Language, allow_subclasses=True)
def hash_spacy_language(pickler, nlp: Language):
    pickler.save(nlp.to_bytes())

def main():
    fin = r"large/file.txt"
    nlp = spacy.load("en_core_web_sm")

    def tokenize(l):
        return {"tok": [t.text for t in nlp(l["text"])]}

    ds =  load_dataset("text", data_files=fin)
    ds = ds["train"].map(tokenize)

    # Sanity check: load NLP from pickle created with our own `dumps`
    config = nlp.config
    lang_cls = spacy.util.get_lang_class(config["nlp"]["lang"])
    nlp2 = lang_cls.from_config(config)
    nlp2.from_bytes(pickle.loads(dumps(nlp)))

    assert isinstance(nlp2, type(nlp))
    assert dumps(nlp) == dumps(nlp2)

if __name__ == '__main__':
    main()
```

If this all looks good to you, I'll start writing on some documentation and examples.


One more thing. This is a reduction function for SpaCy Language that should work with the new API:
```python
@pklregister(Language, allow_subclasses=True)
def hash_spacy_language(pickler, obj):
    def create_language(config, bytes_data):
        lang_cls = spacy.util.get_lang_class(config["nlp"]["lang"])
        nlp = lang_cls.from_config(config)
        return nlp.from_bytes(bytes_data)

    args = (obj.config, obj.to_bytes())
    pickler.save_reduce(create_language, args, obj=obj)
```
so IMO we are missing a test with `pickler.save_reduce`. 

> One more thing. This is a reduction function for SpaCy Language that should work with the new API:
> 
> ```python
> @pklregister(Language, allow_subclasses=True)
> def hash_spacy_language(pickler, obj):
>     def create_language(config, bytes_data):
>         lang_cls = spacy.util.get_lang_class(config["nlp"]["lang"])
>         nlp = lang_cls.from_config(config)
>         return nlp.from_bytes(bytes_data)
> 
>     args = (obj.config, obj.to_bytes())
>     pickler.save_reduce(create_language, args, obj=obj)
> ```
> 
> so IMO we are missing a test with `pickler.save_reduce`.

Sure that seems a good idea, but I do not quite understand what `save_reduce` does. Could you give some more info about what reduce functions do and how they differ from regular `save` and `save_global`? I've read about it but the docs nor the built-in `pickle` code seem really helpful.

I'm no pickle expect, but here is my understanding. I believe the pickler uses the reduce function when you do `loads` to reconstructs the original object from the parameters/arguments that were saved with `dumps`.

For example your sanity check could be simplified from
```python
    config = nlp.config
    lang_cls = spacy.util.get_lang_class(config["nlp"]["lang"])
    nlp2 = lang_cls.from_config(config)
    nlp2 = nlp2.from_bytes(pickle.loads(dumps(nlp)))
```
to
EDIT: 
pickle.loads(pickle.dumps(nlp))

```python
    nlp2 = loads(dumps(nlp))  # using our custom pickler
```

Though note that while it can be convenient for tests, we actually don't care about the reconstruction of the object since we're only using the pickler for `dumps` to compute hashes.

> I'm no pickle expect, but here is my understanding. I believe the pickler uses the reduce function when you do `loads` to reconstructs the original object from the parameters/arguments that were saved with `dumps`.
> 
> For example your sanity check could be simplified from
> 
> ```python
>     config = nlp.config
>     lang_cls = spacy.util.get_lang_class(config["nlp"]["lang"])
>     nlp2 = lang_cls.from_config(config)
>     nlp2 = nlp2.from_bytes(pickle.loads(dumps(nlp)))
> ```
> 
> to
> 
> ```python
>     nlp2 = pickle.loads(pickle.dumps(nlp))
> ```
> 
> Though note that while it can be convenient for tests, we actually don't care about the reconstruction of the object since we're only using the pickler for `dumps` to compute hashes.

Yes, the sanity check can be simplified like that _if_ we use `pickle.dumps` - but that would not test our own `dumps` functionality and would do a naive dump instead of using `to_bytes`. It won't work if we use our own `dumps`, exactly because of the reason that we want custom pickling and being able to call `to_bytes`. To reconstruct the object from the pickled bytes from `to_bytes` we need `from_bytes`. The result of pickle/dill loads will therefore always be a `bytes` object and not a `Language` object.

But `save_reduce` is called when saving, right? Not when loading, AFAICT. I am just not sure what exactly it is saving. It is _potentially_ called [at the end of `save`](https://github.com/python/cpython/blob/24af9a40a8f85af813ea89998aa4e931fcc78cd9/Lib/pickle.py#L603) but only if we haven't returned by then. I just can't figure out what that base case is.

I don't think we expect users to write the reduce function that isn't going to be used anyway. So maybe let's stick with `save` ?

@BramVanroy 
As I understand `save_reduce` is very similar to `copyreg.pickle`, so I'd suggest you to check the following links:
* https://docs.python.org/3/library/copyreg.html#copyreg.pickle
* https://docs.python.org/3/library/pickle.html#object.__reduce__


@lhoestq 
> I don't think we expect users to write the reduce function that isn't going to be used anyway. So maybe let's stick with save ?

I agree. 

`save_reduce` is very similar to `copyreg.pickle` and `object.__reduce__`, which are part of public API (and `save` isn't), so I expect more advanced users to know how to write their own reduction functions. But, as you say, `pklregister` should also work with `save` (even though I think `save` is a bit lower-level, and harder to understand than `save_reduce`).

All our examples in `py_utils` that use `pklregister`  also use `save_reduce` in the last step, so my reduction for SpaCy is meant to be added there, and not to be written by users (because SpaCy is very popular, so the official support by us makes sense :)).

And in the tests, let's ignore the reconstruction part of pickle/dill, because it's not important for us, and focus on the generated dumps. What do you think?

@mariosasko What exactly do you mean with "isn't part of the public API"? It is [a public method](https://github.com/python/cpython/blob/24af9a40a8f85af813ea89998aa4e931fcc78cd9/Lib/pickle.py#L535) in base pickle, just like `dump` is but maybe you mean something else.

@BramVanroy Oh sorry, it's public (not prefixed with `"_"`)  but it's not documented in the docs. `save_reduce` is also not in the docs, but its signature/functionality is similar to `copyreg.pickle` and I see it more often being used in the projects on GH, so it's seems "more public" to me. 

Unfortunately I feel that pickle in general is under-documented. 😄 

For the documentation, I can add a brief example, maybe under "How-to Guides"? The only thing that isn't immediately obvious to me is how I can add that doc page to the TOC?

Yes great idea ! To add that doc page to the TOC, you just have to add it to the index.rst file in the "How-to guides" TOC section

@mariosasko @lhoestq Feel free to make any edits or suggestions in the text!

Hi @mariosasko. I wish you'd told me sooner, as I spent quite some time writing on this.

I'm also not sure whether it is too advanced to have in the documentation. The spaCy use-case seems potentially frequent. Or do you wish to add that case to the defaults, and whenever new issues come up that seem like frequent/obvious cases, add those internally as well?

Documenting the internal `pklregister` is overkill IMO (and it can be kept in docstrings), but we can document something higher level like `register_hash_func` once it's implemented.

So we keep the nice documentation you've written (thank you!), except we can rename it to "Advanced caching" and show an API that is similar to
```python
>>> @register_hash_func(Language, allow_subclasses=True)
>>> def hash_spacy_language(nlp: Language):
>>>     return (nlp.to_bytes(),)
```
This way we keep the documentation centered around the public API rather than the internals that may evolve/be too complicated to fit only one section.

> Or do you wish to add that case to the defaults, and whenever new issues come up that seem like frequent/obvious cases, add those internally as well?

Let's add it to the defaults since it's a frequent use-case. And also allow users to control the hashing using the API mentioned above if they face other non-trivially-hashable objects

Sure, I can have a go at implementing spaCy as a built-in. Should it be included in the tests? (Therefore adding spaCy to the tests requirements.)

Next, from your example, it seems that the return value of `register_hash_func` will be used in pickler.save automatically (calling pklregister a bit deeper). Any reason why it returns a tuple? I can work on this as well, if needed.

> Sure, I can have a go at implementing spaCy as a built-in. Should it be included in the tests? (Therefore adding spaCy to the tests requirements.)

That would be perfect !

> Next, from your example, it seems that the return value of register_hash_func will be used in pickler.save automatically (calling pklregister a bit deeper). 

Yes I think so. For example register_hash_func can call pklregister with the user's function, but wrapped to use pickler.save.

> Any reason why it returns a tuple? I can work on this as well, if needed.

It can either return an arbitrary object or a tuple. I like it a bit better if it's a tuple, so users understand more easily how to make the function take into account more than one item for the hash. It's also consistent with the streamlit caching functions, that also require a tuple. No strong opinion on this though

Let me know if I can help with anything

@lhoestq I do not have the time anymore to work on this. Can someone else pick this up?

Hi ! Sure someone else can continue this PR (either someone from HF, or other contributors can fork the PR).
I think I can work on this next week or the week after, but if anyone wants to work on this earlier feel free to comment here :)