Implement dynamic class provenance tracking to fix isinstance semantics and add support for dynamically defined enums

[ogrisel]

This is a fix for #244 (and #101) to add support for dynamically defined Enum subclasses.

Properly adding support for dynamic Enums required to more broadly fix the isinstance semantics as initially requested in #195.

The proposed solution involves tracking the provenance of pickled dynamic class definitions with a pair of weakref.WeakKeyDictionary / weakref.WeakValueDictionary protected by a threading.Lock.

[codecov-io]

Codecov Report

Merging #246 into master will increase coverage by 0.93%.
The diff coverage is 96.42%.

@@            Coverage Diff             @@
##           master     #246      +/-   ##
==========================================
+ Coverage   88.22%   89.15%   +0.93%     
==========================================
  Files           1        1              
  Lines         552      627      +75     
  Branches      112      129      +17     
==========================================
+ Hits          487      559      +72     
- Misses         42       44       +2     
- Partials       23       24       +1

Impacted Files	Coverage Δ
cloudpickle/cloudpickle.py	89.15% <96.42%> (+0.93%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6ab5eed...0cbd071. Read the comment docs.

[ogrisel]

@pitrou @pierreglaser if you have any suggestions / comments on this solution, please let me know.

[ogrisel]

I had forgotten about #105 which implemented similar support for Enum classes in __main__ as first requested in #101.

However I also think this is a dangerous feature in its current implementation. I think we need to better think trough when an how to reconcile dynamic types (including enums for also other dynamically defined types) on a cluster with some explicit scoping mechanism.

[ogrisel]

Here is another valid use case that cannot be support by this PR:

from enum import Enum
from dask.distributed import Client


class Color(Enum):
    RED = 1
    GREEN = 2


client = Client()
data_futures = client.scatter([Color.RED, Color.GREEN])


def is_red(x):
    return x is Color.RED


for f in client.map(is_red, data_futures):
     print(f.result())

data_futures and is_red are pickled each by their own CloudPickler instances with their own memoizer and global_ref namespaces. The reconstructed Color classes on the worker are therefore independent and is_red can never return True under those conditions.

[ogrisel]

@mrocklin you might be interested in this PR. Basically it makes it possible to properly support enums in dask and fix the isinstance semantics for objects passed as futures.

@pitrou @mrocklin I would also love to get your feedback on the solution I used to tackle this issue that involves a WeakKeyDictionary / WeakValueDictionary pair protected by a lock. Maybe there is a better way to achieve a similar result.

@pcmoritz @HyukjinKwon I would love it if you could test this PR with Ray and Spark to check that it does not introduce any regression.

[ogrisel]

@mrocklin when launching the dask-distributed tests against this PR, test_prometheus fails:

https://travis-ci.org/cloudpipe/cloudpickle/jobs/494254377

It does not seem to be related to the change of this PR but I am not 100% sure and it seems deterministic.

[mrocklin]

Resolved here: dask/distributed#2533

[ogrisel]

Alright: all downstream CI tests pass (dask, loky, joblib). I removed the hard dependency to the enum module so that it should fix a broken test on Py 2.7 on Ray. Waiting for more feedback.

[ogrisel]

@llllllllll @ssanderson I would also love to get your input on this PR.

[pcmoritz]

From the Ray side, this PR passes the test suite, so looking forward to seeing it merged!

Thanks for all the hard work :)

[llllllllll]

Sorry it took so long for me to look at this. I will try to be more responsive here in the future.

I really like the type class tracking work, this seems super useful and fixes hard to explain isinstance bugs. I think there are a few cases that we still don't support with the enum serialization. It's pretty late right now so I am not sure yet what the correct fix is. I posted some tests in the comments below that show the behavior I am talking about. Tomorrow I can PR those tests against this branch and then look into making them pass.

[llllllllll]

how can this be None if we are always passing in _ensure_tracking(obj) on line 590. The enum case also appears to pass this as a non-None value.

[ogrisel]

It used to be the case in a former version of this PR. However I wonder if it's not a good idea to make it possible to disable tracking by passing None for forward compatibility.

I am not sure. Would this be a YAGNI?

[llllllllll]

Hmm, this seems fine given that it is already implemented.

[llllllllll]

There is currently a case that we don't support, which we may want to consider. Currently the enum module has code that adds a failing reduce to the enum subclass if the mixed in type isn't (default) picklable.

        # If a custom type is mixed into the Enum, and it does not know how
        # to pickle itself, pickle.dumps will succeed but pickle.loads will
        # fail.  Rather than have the error show up later and possibly far
        # from the source, sabotage the pickle protocol for this class so
        # that pickle.dumps also fails.
        #
        # However, if the new class implements its own __reduce_ex__, do not
        # sabotage -- it's on them to make sure it works correctly.  We use
        # __reduce_ex__ instead of any of the others as it is preferred by
        # pickle over __reduce__, and it handles all pickle protocols.
        if '__reduce_ex__' not in classdict:
            if member_type is not object:
                methods = ('__getnewargs_ex__', '__getnewargs__',
                        '__reduce_ex__', '__reduce__')
                if not any(m in member_type.__dict__ for m in methods):
                    _make_class_unpicklable(enum_class)

This means I get the following error:

In [18]: class A:
    ...:     def a(self): return 'a'
    ...:     
    ...:     

In [19]: class AEnum(A, enum.Enum):
    ...:     value = 1
    ...:     
    ...:     

In [20]: cp.loads(cp.dumps(AEnum.value))
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-20-24ee7c65848f> in <module>()
----> 1 cp.loads(cp.dumps(AEnum.value))

~/projects/python/cloudpickle/cloudpickle/cloudpickle.py in dumps(obj, protocol)
   1015     try:
   1016         cp = CloudPickler(file, protocol=protocol)
-> 1017         cp.dump(obj)
   1018         return file.getvalue()
   1019     finally:

~/projects/python/cloudpickle/cloudpickle/cloudpickle.py in dump(self, obj)
    297         self.inject_addons()
    298         try:
--> 299             return Pickler.dump(self, obj)
    300         except RuntimeError as e:
    301             if 'recursion' in e.args[0]:

/usr/lib64/python3.6/pickle.py in dump(self, obj)
    407         if self.proto >= 4:
    408             self.framer.start_framing()
--> 409         self.save(obj)
    410         self.write(STOP)
    411         self.framer.end_framing()

/usr/lib64/python3.6/pickle.py in save(self, obj, save_persistent_id)
    494             reduce = getattr(obj, "__reduce_ex__", None)
    495             if reduce is not None:
--> 496                 rv = reduce(self.proto)
    497             else:
    498                 reduce = getattr(obj, "__reduce__", None)

~/.virtualenvs/cloudpickle/lib/python3.6/enum.py in _break_on_call_reduce(self, proto)
     44     """Make the given class un-picklable."""
     45     def _break_on_call_reduce(self, proto):
---> 46         raise TypeError('%r cannot be pickled' % self)
     47     cls.__reduce_ex__ = _break_on_call_reduce
     48     cls.__module__ = '<unknown>'

TypeError: <AEnum.value: 1> cannot be pickled

However, I am able to pickle A objects just fine with cloudpickle. I am not sure if we can fix this because it looks a lot like an explicit __reduce_ex__

[ogrisel]

Maybe that's what the type argument is for. I forgot to address @pierreglaser's comments above.

[ogrisel]

I had a deeper look with @pierreglaser and it seems that this is a fundamental limitation of enums that inherit from another Python class: this problem is already present for enums defined statically in importable modules. I therefore believe this cannot be fixed in cloudpickle but should be fixed upstream in Python's pickle or enum modules themselves.

[ogrisel]

To be more specific:

In [6]: cat mymodule.py
from enum import Enum


class A:
    def method_a(self):
        return "a"


class E(A, Enum):
    x = 1
    y = 2



In [7]: from mymodule import E

In [8]: import pickle

In [9]: pickle.dumps(E.x)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-9-5c82fc8d827e> in <module>()
----> 1 pickle.dumps(E.x)

/opt/python3.7/lib/python3.7/enum.py in _break_on_call_reduce(self, proto)
     42     """Make the given class un-picklable."""
     43     def _break_on_call_reduce(self, proto):
---> 44         raise TypeError('%r cannot be pickled' % self)
     45     cls.__reduce_ex__ = _break_on_call_reduce
     46     cls.__module__ = '<unknown>'

TypeError: <E.x: 1> cannot be pickled

The class A can be imported from mymodule so maybe it's actually a bug in Python itself.

[pierreglaser]

I think somehow this is not the case anymore in recent python versions (3.7.1 and later):

Python 3.7.1 (tags/v3.7.1:260ec2c36a, Jan 24 2019, 18:18:14)                                                                                                                                                       
Type 'copyright', 'credits' or 'license' for more information                                                                                                                                                      
IPython 7.5.0 -- An enhanced Interactive Python. Type '?' for help.                                                                                                                                                
                                                                                                                                                                                                                   
In [1]: class A:                                                                                                                                                                                                   
   ...:     def a(self): return "a"                                                                                                                                                                                
   ...: import enum                                                                                                                                                                                                
   ...: class AEnum(A, enum.Enum):                                                                                                                                                                                 
   ...:     value = 1                                                                                                                                                                                              
   ...: import pickle                                                                                                                                                                                              
   ...: pickle.dumps(AEnum.value)                                                                                                                                                                                  
Out[1]: b'\x80\x03c__main__\nAEnum\nq\x00K\x01\x85q\x01Rq\x02.'     

See: Enum: mixin classes don't mix well with already mixed Enums

[ogrisel]

@llllllllll thanks for the review. I think I addressed your comments.

[ogrisel]

It seems that my latest commits have broken some dask tests although I am not 100% sure because there is no output for 10min and travis kills the job without letting pytest report the failures. Will have to do that manually but I won't have time soon enough. If someone else wants to have a look don't hesitate.

[ogrisel]

@llllllllll I have rebased this PR to check that the downstream tests of dask distributed, ray, loky and joblib pass.

I now pass all the bases as you suggested so and also added support for methods defined directly on the enum class itself.

[philippotto]

Sorry for nagging here, but is there anything which blocks this PR? The fix would be pretty useful for my team.

[ogrisel]

I think we might want to release cloudpickle 1.0 without this change and 1.1 with this PR merged so that we have a pretty uncontroversial 1.0 release for cloudpickle as @mrocklin recently asked.

[ogrisel]

@llllllllll I would like to merge this but I would really appreciate it if you could confirm that I took your comments into account as intended.

[llllllllll]

Sorry about that again. I just got back from C++now and did another pass. All of my comments have been addressed.

Thanks for working on this feature, this is really amazing work!

[llllllllll]

Hmm, this seems fine given that it is already implemented.

[ogrisel]

Thanks you @llllllllll for your review. I am merging.

[ogrisel]

@philippotto cloudpickle 1.1.0 is out with this fix included.

[philippotto]

Awesome, thank you!