Impossible to operate with enums #1178
Description
There are actually multiple problems that prevent using enums with dask.distributed.
The first one is that cloudpickle doesn't support serializing enum.Enum subclasses:
import enum
from distributed.protocol.serialize import serialize
class MyEnum(enum.Enum):
SPAM = 'SPAM'
serialize(MyEnum.SPAM)The result would be: PicklingError: Could not pickle object as excessively deep recursion required..
pickle supports enums serialization. But if try to register it as a custom serializer and use it, the second problem arises: for some reason it's impossible to register a custom serializers for enum.Enum subclasses:
import enum
import dask
import pickle
import dask.distributed
from distributed.protocol.serialize import (
serialize,
deserialize,
register_serialization
)
class Spam(enum.Enum):
EGGS = 'Eggs'
def register():
register_serialization(
Spam,
lambda o: [{}, [pickle.dumps(o)]],
lambda _, f: pickle.loads(f[0])
)
register()
def dask_setup(_):
register()
assert (
deserialize(*serialize(Spam.EGGS)) == Spam.EGGS
)
client = dask.distributed.Client('tcp://127.0.0.1:1234')
# Doesn't work
client.compute(dask.delayed(Spam.EGGS), True)The last statement would result in KeyError: '__main__.Spam' error thrown on the worker. The same thing happens if use LocalCluster().