[WIP] Allow user-defined hash functions via a registry
Opened by BramVanroy (closed)

Inspired by the discussion on hashing in https://github.com/huggingface/datasets/issues/3178#issuecomment-959016329, @lhoestq suggested that it would be neat to allow users more control over the hashing process. Specifically, it would be great if users can specify specific hashing functions depending on the **class** of the object.

As an example, we found in the linked topic that loaded spaCy models (`Language` objects) have different hashes when `dump`'d, but their byte representation with `Language.to_bytes()` _is_ deterministic. It would therefore be great if we could specify that for `Language` objects, the hasher should hash the objects `to_bytes()` return value instead of the object itself.

This PR adds a new, but tiny, dependency to manage the registry, namely [`catalogue`](https://github.com/explosion/catalogue). 

Two files have been changed (apart from the added dependency in `setup.py`) and one file has been added.

**utils.registry** (added)

This file defines our custom Registry and builds a registry called "hashers". A Registry is basically dictionary from names (str) to functions. A function can be added to the registry by a decorator, e.g. 

```python
@hashers.register(spacy.Language)
def hash_spacy_language(nlp):
    return Hasher.hash(nlp.to_bytes())
```

You'll notice that `spacy.Language` is not a string, even though the registry holds a str->func mapping. To accomplish this with classes in a dynamic way, catalogue.Registry needed to be subclassed and modified as `DatasetsRegistry`. All methods that use a name as an input are now modified so that classes are deterministically converted in strings in such a way that we can later retrieve the actual class from the string (below).

**utils.py_utils** (modified)

Added two functions to deal with classes and their qualified names, that is, their full descriptive name including the module. On the one hand it allows us to retrieve a string from a given class, e.g. given `Module` class, return `torch.nn.Module` str. Conversly, a function is added to convert such a full qualified name into a class. For instance, given the string `torch.nn.Module`, return the `Module` class. These straightforward methods allow us to interchangeably use classes and strings without any needed user interaction - they can just register a class, and behind the scenes `DatasetsRegistry` converts these to deterministic strings.

**fingerprint** (modified)

Updated Hasher.hash so that if the object to hash is an instance of a class in the registry, the registered function is used to hash the object instead of the default behavior. To do so we iterate over the registry `hashers` and convert its keys (strings) into classes, and then we can use `isinstance`.

```python
# Check if the current object is an instance that is
# applicable to the user-defined hashers. If so, hash
# with the user-defined function
for full_module_name, func in hashers.get_all().items():
    registered_cls = get_cls_from_qualname(full_module_name)
    if isinstance(value, registered_cls):
        return func(value)
```

**Putting it all together**

To test this, you can try the following example with spaCy. First install spaCy from source and checkout a specific commit.

```shell
git clone https://github.com/explosion/spaCy.git
cd spaCy/
git checkout cab9209c3dfcd1b75dfe5657f10e52c4d847a3cf
cd ..

git clone https://github.com/BramVanroy/datasets.git
cd datasets
git checkout registry
pip install -e .
pip install ../spaCy
spacy download en_core_web_sm
```

Now you can run the following script. By default it will use the custom hasher function for the Language object. You can enable the default behavior by commenting out `@hashers.register...`.

```python
import spacy

from datasets.fingerprint import Hasher
from datasets.utils.registry import hashers

# Register a function so that when the Hasher encounters a spacy.Language object
# it uses this custom function to hash instead of the default
@hashers.register(spacy.Language)
def hash_spacy_language(nlp):
    return Hasher.hash(nlp.to_bytes())


def main():
    print(hashers.get_all())
    nlp = spacy.load("en_core_web_sm")
    dump1 = Hasher.hash(nlp)
    nlp = spacy.load("en_core_web_sm")
    dump2 = Hasher.hash(nlp)
    print(dump1)
    # succeeds when using the registered custom function
    # fails if using the default
    assert dump1 == dump2


if __name__ == '__main__':
    main()
```

To do
====
- The above is just a proof-of-concept. I am open to changes/suggestions
- Tests still need to be written
- We should consider whether we can make `DatasetsRegistry` very restrictive and ONLY allowing classes. That would make testing easier - otherwise we also need to test for other sorts of objects.
- Maybe the `hashers` definition is better suited in `fingerprint`?
- Documentation/examples need to be updated
- Not sure why the logger is not working in `hash()`
- `get_cls_from_qualname` might need a fail-safe: is it possible for a full_qualname to not have a module, and if so how do we deal with that?


Comments:
Hi @BramVanroy, thanks for your PR.

There was a bug in TensorFlow/Keras. We have made a temporary fix in master branch. Please, merge master into your PR branch, so that the CI tests pass.

```
git checkout registry
git fetch upstream master
git merge upstream/master
```

@albertvillanova Done. Although new tests will need to be added. I am looking for some feedback on my initial proposal in this PR. Reviews and ideas welcome!

Hi ! Thanks for diving into this :)

With this approach you get the right hash when doing `Hasher.hash(nlp)` but if you try to hash an object that has `nlp` as one of its attributes for example you will get different hashes every time.

This is because `Hasher.hash` is not recursive itself. Indeed what happens when you try to hash an object is that:
1. it is dumped with our custom `dill` pickler (which is recursive)
2. the bytes of the dump are hashed

To fix this we must integrate the custom hashing as a custom pickler dumping instead.

Note that we're only using the `pickler.dumps` method and not `pickler.loads` since we only use it to get hashes, so it doesn't matter if `loads` doesn't reconstruct the object exactly. What's important it only to capture all the necessary information that defines how the object transforms the data (here `nlp.to_bytes()` determines how the spacy pipeline transforms the text).

Our pickler already has a registry and you can register new dump functions with:
```python
import dill
import spacy
from datasets.utils.py_utils import pklregister

@pklregister(spacy.Language)
def _save_spacy_language(pickler, nlp):
    pickler.save_reduce(...)  # I think we can use nlp.to_bytes() here
    dill._dill.log.info(...)
```

You can find some examples of custom dump functions in `py_utils.py`

Ah, darn it. Completely missed that register. Time wasted, unfortunately. 

To better understand what you mean, I figured I'd try the basis of your snippet and I've noticed quite an annoying side-effect of how the pickle dispatch table seems to work. It explicitly uses an object's [`type()`](https://github.com/python/cpython/blob/87032cfa3dc975d7442fd57dea2c6a56d31c911a/Lib/pickle.py#L557-L558), which makes sense for pickling some (primitive) types it is not ideal for more complex ones, I think. `Hasher.hash` has the same issue as far as I can tell.

https://github.com/huggingface/datasets/blob/d21ce54f2c2782f854f975eb1dc2be6f923b4314/src/datasets/fingerprint.py#L187-L191

This is very restrictive, and won't work for subclasses. In the case of spaCy, for instance, we register `Language`, but `nlp` is an instance of `English`, which is a _subclass_ of `Language`. These are different types, and so they will not match in the dispatch table. Maybe this is more general approach to cover such cases? Something like this is a start but too broad, but ideally a hierarchy is constructed and traversed of all classes in the table and the lowest class is selected to ensure that the most specific class function is dispatched.

```python
    def hash(cls, value: Any) -> str:
        # Try to match the exact type
        if type(value) in cls.dispatch:
            return cls.dispatch[type(value)](cls, value)

        # Try to match instance (superclass)
        for type_cls, func in cls.dispatch.items():
            if isinstance(value, type_cls):
                return cls.dispatch[type_cls](cls, value)

        return cls.hash_default(value)
```

This does not solve the problem for pickling, though. That is quite unfortunate IMO because that implies that users always have to specify the most specific class, which is not always obvious. (For instance, `spacy.load`'s signature returns `Language`, but as said before a subclass might be returned.)

Second, I am trying to understand `save_reduce` but I can find very little documentation about it, only the source code which is quite cryptic. Can you explain it a bit? The required arguments are not very clear to me and there is no docstring.

```python
    def save_reduce(self, func, args, state=None, listitems=None, dictitems=None, obj=None):
```

Here is an example illustrating the problem with sub-classes.

```python
import spacy

from spacy import Language
from spacy.lang.en import English

from datasets.utils.py_utils import Pickler, pklregister

# Only useful in the registry (matching with `nlp`)
# if you swap it out for very specific `English`
@pklregister(English)
def hash_spacy_language(pickler, nlp):
    pass


def main():
    print(Pickler.dispatch)
    nlp = spacy.load("en_core_web_sm")
    print(f"NLP type {type(nlp)} in dispatch table? ", type(nlp) in Pickler.dispatch)


if __name__ == '__main__':
    main()
```

Indeed that's not ideal.
Maybe we could integrate all the subclasses directly in `datasets`. That's simple to do but the catch is that if users have new subclasses of `Language` it won't work.

Otherwise we can see how to make the API simpler for users by allowing subclasses
```python
# if you swap it out for very specific `English`
@pklregister(Language, allow_subclasses=True)
def hash_spacy_language(pickler, nlp):
    pass
```

Here is an idea how to make this work, let me know what you think:

When `Pickler.dumps` is called, it uses `Pickler.save_global` which is a method that is going to be called recursively on all the objects. We can customize this part, and make it work as we want when it encounters a subclass of `Language`.

For example when it encounters a subclass of `Language`, we can dynamically register the hashing function for the subclass (`English` for example) in `Pickler.save_global`, right before calling the actual `dill.Pickler.save_global(self, obj, name=name)`:
```python
pklregister(type(obj))(hash_function_registered_for_parent_class)
dill.Pickler.save_global(self, obj, name=name)
```

In practice that means we can have an additional dispatch dictionary (similar to `Pickler.dispatch`) to store the hashing functions when `allow_subclasses=True`, and use this dictionary in `Pickler.save_global` to check if we need to use a hashing function registered with `allow_subclasses=True` and get `hash_function_registered_for_parent_class`.

If I understood you correctly, I do not think that that is enough because you are only doing this for a type and its direct parent class. You could do this for all superclasses (so traverse all ancestors and find the registered function for the first that is encountered). I can work on that, if you agree. The one thing that I am not sure about is how you want to create the secondary dispatch table. An empty dict as class variable in Pickler? (It doesn't have to be a true dispatcher, I think.)

I do not think that dynamic registration is the ideal situation (it feels a bit hacky). An alternative would be to subclass Pickle and Dill to make sure that instead of just type() checking in the dispatch table also superclasses are considered. But that is probably overkill.

> You could do this for all superclasses (so traverse all ancestors and find the registered function for the first that is encountered)

That makes sense indeed !

> The one thing that I am not sure about is how you want to create the secondary dispatch table. An empty dict as class variable in Pickler? (It doesn't have to be a true dispatcher, I think.)

Sure, let's try to not use too complicated stuff

> I do not think that dynamic registration is the ideal situation (it feels a bit hacky). An alternative would be to subclass Pickle and Dill to make sure that instead of just type() checking in the dispatch table also superclasses are considered. But that is probably overkill.

Indeed that would feel less hacky, but maybe it's too complex just for this. I feel like this part of the library is already hard to understand when you're not familiar with pickle. IMO having only a few changes that are simpler to understand is better than having a rewrite of `dill`'s core code.

Thanks a lot for your insights, it looks like we're going to have something that works well and that unlocks some nice flexibility for users :) Feel free to ping me anytime if I can help on this

Sure, thanks for brainstorming! I'll try to work on it this weekend. Will also revert the current changes in this PR and rename it. 

It seems like this is going in the right direction :). 

@BramVanroy Just one small suggestion for future contributions: instead of using `WIP` in the PR title, you can create a draft PR if you're still working on it.

Maybe I should just create a new (draft) PR then, seeing that I'll have to rename and revert the changes anyway? I'll link to this PR so that the discussion is at least referenced.

I can convert this PR to a draft PR. Let me know what would you prefer.

I think reverting my previous commits would make for a dirty (or confusing) commit history, so I'll just create a new one. Thanks.