Dask PicklingError function not found - dask

Question

I am trying do some parallelized data validation with pydantic & dask.
import dask.bag as db
from pydantic import BaseModel
class MyData(BaseModel):
id: int
name: str
def validate_data(data):
return MyData(**data)
data = [
{'id': 1, 'name': 'Foo'},
{'id': 2, 'name': 'Bar'}
]
bag = db.from_sequence(data)
bag.map(validate_data).compute()
This raises the following pickling error (full stack trace available here):
~/Library/Caches/pypoetry/virtualenvs/domi-IWOYYLRr-py3.7/lib/python3.7/site-packages/cloudpickle/cloudpickle.py in save_global(self, obj, name, pack)
840 self._save_parametrized_type_hint(obj)
841 elif name is not None:
--> 842 Pickler.save_global(self, obj, name=name)
843 elif not _is_importable_by_name(obj, name=name):
844 self.save_dynamic_class(obj)
~/.pyenv/versions/3.7.6/lib/python3.7/pickle.py in save_global(self, obj, name)
958 raise PicklingError(
959 "Can't pickle %r: it's not found as %s.%s" %
--> 960 (obj, module_name, name)) from None
961 else:
962 if obj2 is not obj:
PicklingError: Can't pickle <cyfunction int_validator at 0x116503460>: it's not found as pydantic.validators.lambda11
Note, I can pickle this function just fine:
>>> import pickle
>>> validate_data
<function __main__.validate_data(data)>
>>> pickled = pickle.dumps(validate_data)
>>> unpickled = pickle.loads(pickled)
>>> unpickled
<function __main__.validate_data(data)>
>>> unpickled({'id': 5, 'name': 'Foo'})
MyData(id=5, name='Foo')
Any ideas or tips on how to fix? (I'm not sure if this is an issue with dask or pydantic, so I have tagged both)
Thanks in advance!
System/package info:
Dask Version: 2.19.0
Pydantic Version: 1.5.1
❯ python -c "import pydantic.utils; print(pydantic.utils.version_info())"
pydantic version: 1.5.1
pydantic compiled: True
install path: /Users/ianwhitestone/Library/Caches/pypoetry/virtualenvs/domi-IWOYYLRr-py3.7/lib/python3.7/site-packages/pydantic
python version: 3.7.6 (default, Mar 7 2020, 14:34:51) [Clang 11.0.0 (clang-1100.0.33.17)]
platform: Darwin-19.5.0-x86_64-i386-64bit
optional deps. installed: ['typing-extensions']

Answer 1

Moving the pydantic model definition to a separate file solved this for me:
# my_data.py
from pydantic import BaseModel
class MyData(BaseModel):
id: int
name: str
# main.py
import dask.bag as db
from my_data import MyData
def validate_data(data):
return MyData(**data)
data = [
{'id': 1, 'name': 'Foo'},
{'id': 2, 'name': 'Bar'}
]
bag = db.from_sequence(data)
bag.map(validate_data).compute()
Marking this as the answer for now, if someone has an explanation for why this is the case I will mark that as the answer!