Private dispatch table not checked

[wuisawesome]

I'm attempting to register a custom reducer with cloudpickle. In my use case, I need multiple CloudPickler's to serialize an object with different behavior, so I can't override the class-level global dispatch table.

To demonstrate, assume I have the following class:

class CustomClass:
    def __init__(self, *args):
        self.x = len(args)

c = CustomClass()

I can define a custom reducer and register it with a private dispatch table:

def custom_reducer(obj):
    return CustomClass, (0,)

io_buffer = io.BytesIO()
custom_pickler = pickle.Pickler(io_buffer)
custom_pickler.dispatch_table = {CustomClass: custom_reducer}
custom_pickler.dump(c)

io_buffer.seek(0)
c1 = pickle.load(io_buffer)
assert c1.x == 1

I can also prove that this doesn't affect the global dispatch table:

c2 = pickle.loads(pickle.dumps(c))
assert c2.x == 0

Unfortunately, when I try the same behavior with cloudpickle, it doesn't respect the private dispatch table:

io_buffer = io.BytesIO()
custom_pickler = cloudpickle.CloudPickler(io_buffer)
custom_pickler.dispatch_table = {CustomClass: custom_reducer}
custom_pickler.dump(c)

io_buffer.seek(0)
c3 = cloudpickle.load(io_buffer)
assert c3.x == 1 # assertion fails becase c3.x == 0

(reference to how this was implemented in pickle)

[wuisawesome]

Btw one mitigation for this bug is to extend the CloudPickler instead, though if anyone knows how to fix this bug, that would be ideal.

class CustomCloudPickler(cloudpickle.CloudPickler):
    custom_reducer_table = {CustomClass: custom_reducer}
    dispatch_table = ChainMap(
        custom_reducer_table,
        cloudpickle.CloudPickler.dispatch_table)
    dispatch = dispatch_table

io_buffer = io.BytesIO()
custom_pickler = CustomCloudPickler(io_buffer)
custom_pickler.dispatch_table = {CustomClass: custom_reducer}
custom_pickler.dump(c)

io_buffer.seek(0)
c4 = cloudpickle.load(io_buffer)
assert c4.x == 1

[pierreglaser]

Hi @wuisawesome, thanks for the report,

The problem is known, and a PR has been submitted to solve the problem(#395), I'll try to review it soon.

In the meantime, a possible workaround is explained in joblib/loky#260 and related issues. I don't have time to re-explain it here now, but it should look like:

import cloudpickle
import io
import pickle

class CustomClass:
    def __init__(self, *args):
        self.x = len(args)

c = CustomClass()

def custom_reducer(obj):
    return CustomClass, (0,)


io_buffer = io.BytesIO()
custom_pickler = cloudpickle.CloudPickler(io_buffer)
dt = pickle.Pickler.dispatch_table.__get__(custom_pickler)
new_dt = dt.new_child({CustomClass: custom_reducer})
pickle.Pickler.dispatch_table.__set__(custom_pickler, new_dt)
custom_pickler.dump(c)

io_buffer.seek(0)
c3 = cloudpickle.load(io_buffer)
assert c3.x == 1 # assertion fails becase c3.x == 0

[pierreglaser]

Classifying this as an enhancement since IIRC we never claimed full compatibility with the Pickler API.