Add ability to register modules to be deeply serialized

[Samreay]

This PR is based on the work done by @kinghuang in PR391, but takes on the feedback provided by @ogrisel and adds testing.

Fixes #206

Issue Summary

To summarise the issue, in many cases cloudpickle is used to send code for remote execution. This is used in dask, prefect, mlflow and many libraries. For local functions, this works perfectly fine. But for any non-local function or class, cloudpickle assumes that external modules and packages are available at the location of deserialization. This may either not be the case, or the version of the package available at the end point may be different.

This PR adds the option to register modules for deep serialization by providing a register_deep_serialization function which takes either a name or a module. This is the original register_dynamic_module by @kinghuang.

import cloudpickle
from tests import external

cloudpickle.register_deep_serialization("tests.external")  # string name works
cloudpickle.register_deep_serialization(external)          # You can pass the module itself
cloudpickle.register_deep_serialization("tests")           # or the parent string/module

output = cloudpickle.dumps(external.an_external_function)

Original dumps:

b'\x80\x05\x95+\x00\x00\x00\x00\x00\x00\x00\x8c\x0etests.external\x94\x8c\x14an_external_function\x94\x93\x94.'

dumps after registering tests.external for deep serialization:

b'\x80\x05\x95<\x02\x00\x00\x00\x00\x00\x00\x8c\x17cloudpickle.cloudpickle\x94\x8c\r_builtin_type\x94\x93\x94\x8c\nLambdaType\x94\x85\x94R\x94(h\x02\x8c\x08CodeType\x94\x85\x94R\x94(K\x00K\x00K\x00K\x00K\x01KCC\x04d\x01S\x00\x94N\x8c\x11this is something\x94\x86\x94))\x8c5/Users/samreay/Projects/cloudpickle/tests/external.py\x94\x8c\x14an_external_function\x94K\x04C\x02\x00\x01\x94))t\x94R\x94}\x94(\x8c\x0b__package__\x94\x8c\x05tests\x94\x8c\x08__name__\x94\x8c\x0etests.external\x94\x8c\x08__file__\x94\x8c5/Users/samreay/Projects/cloudpickle/tests/external.py\x94uNNNt\x94R\x94\x8c\x1ccloudpickle.cloudpickle_fast\x94\x8c\x12_function_setstate\x94\x93\x94h\x19}\x94}\x94(h\x14h\r\x8c\x0c__qualname__\x94h\r\x8c\x0f__annotations__\x94}\x94\x8c\x0e__kwdefaults__\x94N\x8c\x0c__defaults__\x94N\x8c\n__module__\x94h\x15\x8c\x07__doc__\x94N\x8c\x0b__closure__\x94N\x8c\x17_cloudpickle_submodules\x94]\x94\x8c\x0b__globals__\x94}\x94u\x86\x94\x86R0.'

Modules can be unregistered via unregister_deep_serialization

Tests

One of the example tests with an explicit use case is shown above. On top of this, tests have been added to _lookup_module_and_qualname using the _cloudpickle_testpkg package, and also to the new _is_explicitly_serialized_module function.

[codecov]

Codecov Report

Merging #417 (1e9a48d) into master (0c62ae0) will increase coverage by 0.36%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #417      +/-   ##
==========================================
+ Coverage   91.99%   92.36%   +0.36%     
==========================================
  Files           4        4              
  Lines         687      720      +33     
  Branches      141      150       +9     
==========================================
+ Hits          632      665      +33     
  Misses         34       34              
  Partials       21       21              

Impacted Files	Coverage Δ
cloudpickle/cloudpickle.py	88.31% <100.00%> (+1.09%)	⬆️
cloudpickle/cloudpickle_fast.py	96.87% <100.00%> (ø)

---

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 0c62ae0...1e9a48d. Read the comment docs.

[kinghuang]

This PR is based on the work done by @kinghuang in PR391, but takes on the feedback provided by @ogrisel and adds testing.

Thanks for picking this up, @Samreay!

[Samreay]

@ogrisel, any thoughts on the PR? Is there someone I should direct it to?

[damienrj]

This is also of great interest to me! Very nice. Maybe @ogrisel would be the person to review?

[Ben-Epstein]

Hello, following up here, is there any plan on merging or moving forward? I see all but 1 test passing

[ogrisel]

This sounds like a legit need and the implementation looks fine from a quick look at it. However I am not a big fan of the "deep_serialization" naming in the API and comments. What about calling this "pickle_by_value" or "pickling_by_value" instead?

[Samreay]

@ogrisel Happy to update this to pickle_by_value() if we can lock that name in. The deep notation was taken from @pierreglaser's comment here which I liked because of the similarity to copy/deepcopy. However, if you believe there is still ambiguity, it's an easy change.

[Samreay]

Functions renamed to register_pickle_by_value and unregister_pickle_by_value

[ogrisel]

LGTM. Please document the change in the changelog.

@pierreglaser ok with you?

[Samreay]

Changes now documented :)

[ogrisel]

The failure of CloudPickleTest.test_locally_defined_class_with_type_hints on python nightly is unrelated. The memory usage on macos might be random but I am not 100% sure.

[Samreay]

Yeah, will be interesting to see if they roll back the string typehints given how many libraries its going to impact. Ive seen no memory issues when testing this locally on my macbook, can you point me to which test / build step is concern and I can try to dig a little deeper?

[ogrisel]

The macos failure looks random and is probably unrelated.

[pierreglaser]

Here is a first batch of comments - mostly nitpicks, but I raised a few more structural suggestions/questions.

As an important side point, I think that as of now, calling cloudpickle.dumps(module) on a module registered for pickling by value will still be pickled by reference: It is a behavior that we may want to change, cc @ogrisel

[pierreglaser]

Feel free to request my review again when all requested changes are made. I think also that a test making sure modules registered for pickling by value are actually pickled by value when doing a cloudpickle.dumps(module) call.

[Samreay]

@pierreglaser Ive added the module serialisation test in, however it is currently broken, with the main issue being the TypeVar being passed into _should_pickle_by_reference from somewhere (maybe the module reduce?). My lack of knowledge in how the fundamentals is making debugging difficult, would you be able to look over the changes and offer any suggestions on what is going wrong?

[ogrisel]

Thanks, I will try to review soon.

[ogrisel]

Very nice contrib and tests!

Just a few suggestions to further improve the doc + code nitpicks.

[ogrisel]

Suggested change

	assert w.run(
	lambda: LocalClass().method() == LocalClass().method()
	)
	assert (
	w.run(lambda: LocalClass().method())
	== LocalClass().method()
	)

[ogrisel]

Maybe the motivation (interactive dev on a distributed cluster with long running worker processes) could be made more explicit by reusing the text I suggest in the docstring of register_pickle_by_value above in this review.

[ogrisel]

@pierreglaser once my review comments are addressed, feel free to merge and release cloudpickle. But before releasing, please run the downstream ci tests in the release PR to try to avoid introducing any regression in this quite big release.

[pierreglaser]

@ogrisel thanks for the review! Should we bump cloudpickle's version to 2.0.0 for this release, or stick with 1.7.0?

[ogrisel]

+1 for 2.0.0 :)

[ogrisel]

Merged! Thank you very much @kinghuang @Samreay and @pierreglaser!

[pierreglaser]

Woot 🎉!! Thanks a lot @Samreay for the contribution!

[bonelli]

Thank you everybody, IMHO this is a great improvement towards general purpose clusters.
When do you think 2.0.0 will be tagged and released as a package? If possible I'd wait that moment to use this in my new docker images on my k8s cluster.

[pierreglaser]

Releasing will be done once #432 is addressed - hopefully within the next couple of weeks this should be done.

[Samreay]

Thanks @pierreglaser and @ogrisel for all the commits getting this ready. Fingers crossed #432 is a simple fix! :)