migrate from dill to cloudpickle for advanced serialization

[jrwalk]

Description

Usage of dill for optional serialization in PythonVirtualenvOperator may be replaced with cloudpickle as its serialization library. This should be a mostly drop-in replacement.

Use case / motivation

Currently, the PythonVirtualenvOperator optionally uses dill in place of stock pickle to serialize advanced types. However, most major distributed compute frameworks have opted to shift to cloudpickle, meaning using dill for Airflow can introduce redundant dependencies for calling out to other distributed compute (e.g., farming compute-heavy tasks out to a remote dask cluster), and can interfere with serialization of tasks for those tools.

Since both dill and cloudpickle are largely drop-in replacements for pickle, the migration should be fairly minor.

Related Issues

kubeflow/pipelines#1387

dask/distributed#3606

piskvorky/gensim#558 (comment)

uqfoundation/multiprocess#22 (comment)

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template!

[ashb]

Is it worth making the pickle module a config setting, or just always using cloudpickle instead?

[jrwalk]

so dill is optional to use (set by a flag in the PythonVirtualenvOperator) but is always imported and treated as a required dependency for Airflow, right? Keeping that pattern it should be super easy to cut over to cloudpickle.

I'd generally advocate for allowing fewer dependencies though, so it could be nice to make cloudpickle a proper extra-dependency and put some checks around the import. That would let people keep minimal environments if they choose/need to. I've also run into situations where it wasn't being used (no virtualenv operators in my dag) but because it was imported dill was still getting captured by serializers like cloudpickle since they scoop up a fair amount of runtime environment information.

[jrwalk]

either way it would probably also necessitate reworking some of the tests, since (for example) cloudpickle could conceivably serialize lambda functions for the python_callable inputs

[sumeshpremraj]

Hi @jrwalk, can I work on this?
dill and cloudpickle seem to be fairly similar, is there anything else I should consider (except for the tests your mentioned)?

[potiuk]

Assigned you. The issue is pretty old - we also have ExternalPythonOperator using same approach now. I think just focusing on tests for those might be enough

[Felix-neko]

#35529 -- maybe using cloudpickle could cure this problem?

[potiuk]

#35529 -- maybe using cloudpickle could cure this problem?

If you would like to take a stab on it and attempt to try it - fee free to verify that hypothesis @Felix-neko - PRs are always most welcome.

[Felix-neko]

#35529 -- maybe using cloudpickle could cure this problem?

If you would like to take a stab on it and attempt to try it - fee free to verify that hypothesis @Felix-neko - PRs are always most welcome.

Heh, it looks like you're overestimating my power (and my current knowledge in cloudpickle and airflow)
; )

[Taragolis]

It might help. There is also limitations for cloudpickle exists:

Cloudpickle can only be used to send objects between the exact same version of Python.

[Felix-neko]

Cloudpickle can only be used to send objects between the exact same version of Python.

That won't be a big problem. Especially if simple dill will still be supported.

[VladaZakharova]

@sumeshpremraj
Hi there! May I ask how is the progress going on this issue?

[gyli]

With #39270 completed, argument serializer was added to support choosing from serializers. Looks like this issue can be closed.