Add ability to register modules to be deeply serialized

CHANGES.md

@@ -6,9 +6,13 @@ dev
 
 - Python 3.5 is no longer supported.
 
+- Support for registering modules to be serialised by value. This will
+  allow for code defined in local modules to be serialised and executed
+  remotely without those local modules installed on the remote machine.
+  ([PR #417](https://github.com/cloudpipe/cloudpickle/pull/417))
+
 - Fix a side effect altering dynamic modules at pickling time.
   ([PR #426](https://github.com/cloudpipe/cloudpickle/pull/426))
-
 - Support for pickling type annotations on Python 3.10 as per [PEP 563](
   https://www.python.org/dev/peps/pep-0563/)
   ([PR #400](https://github.com/cloudpipe/cloudpickle/pull/400))

README.md

@@ -66,6 +66,50 @@ Pickling a function interactively defined in a Python shell session
 85
 ```
 
+
+Overriding pickle's serialization mechanism for importable constructs:
+----------------------------------------------------------------------
+
+An important difference between `cloudpickle` and `pickle` is that
+`cloudpickle` can serialize a function or class **by value**, whereas `pickle`
+can only serialize it **by reference**, e.g. by serializing its *module
+attribute path* (such as `my_module.my_function`).
+
+By default, `cloudpickle` only uses serialization by value in cases where
+serialization by reference is usually ineffective, for example when the
+function/class to be pickled was constructed in an interactive Python session.
+
+Since `cloudpickle 1.7.0`, it is possible to extend the use of serialization by
+value to functions or classes coming from **any pure Python module**. This feature
+is useful when the said module is unavailable in the unpickling environment
+(making traditional serialization by reference ineffective). To this end,
+`cloudpickle` exposes the
+`register_pickle_by_value`/`unregister_pickle_by_value` functions:
+
+```python
+>>> import cloudpickle
+>>> import my_module
+>>> cloudpickle.register_pickle_by_value(my_module)
+>>> cloudpickle.dumps(my_module.my_function)  # my_function is pickled by value
+>>> cloudpickle.unregister_pickle_by_value(my_module)
+>>> cloudpickle.dumps(my_module.my_function)  # my_function is pickled by reference
+```
+
+Note that this feature is still **experimental**, and may fail in the following
+situations:
+
+- If the body of a function/class pickled by value contains an `import` statement:
+  ```python
+  >>> def f():
+  >>> ... from another_module import g
+  >>> ... # calling f in the unpickling environment may fail if another_module
+  >>> ... # is unavailable
+  >>> ... return g() + 1
+  ```
+
+- If a function pickled by reference uses a function pickled by value during its execution.
+
+
 Running the tests
 -----------------
 

cloudpickle/cloudpickle.py

@@ -45,6 +45,7 @@
 import builtins
 import dis
 import opcode
+import inspect
 import platform
 import sys
 import types
@@ -88,6 +89,9 @@ def g():
 # communication speed over compatibility:
 DEFAULT_PROTOCOL = pickle.HIGHEST_PROTOCOL
 
+# Names of modules whose resources should be treated as dynamic.
+_PICKLE_BY_VALUE_MODULES = set()
+
 # Track the provenance of reconstructed dynamic classes to make it possible to
 # reconstruct instances from the matching singleton class definition when
 # appropriate and preserve the usual "isinstance" semantics of Python objects.
@@ -124,6 +128,60 @@ def _lookup_class_or_track(class_tracker_id, class_def):
     return class_def
 
 
+def register_pickle_by_value(module):
+    """Register that the input module should be pickled by value."""
+    if not isinstance(module, types.ModuleType):
+        raise ValueError(
+            f"Input should be a module object, got {type(module)} instead"
+        )
+    # In the future, cloudpickle may need a way to access any module registered
+    # for pickling by value in order to introspect relative imports inside
+    # functions pickled by value. (see
+    # https://github.com/cloudpipe/cloudpickle/pull/417#issuecomment-873684633).
+    # This access can be ensured by checking that module is present in
+    # sys.modules at registering time and assuming that it will still be in
+    # there when accessed during pickling. Another alternative would be to
+    # store a weakref to the module. Even though cloudpickle does not implement
+    # this introspection yet, in order to avoid a possible breaking change
+    # later, we still enforce the presence of module inside sys.modules.
+    if module.__name__ not in sys.modules:
+        raise ValueError(
+            f"{module} was not imported correctly, have you used an "
+            f"`import` statement to access it?"
+        )
+    _PICKLE_BY_VALUE_MODULES.add(module.__name__)
+
+
+def unregister_pickle_by_value(module):
+    """Unregister that the input module should be pickled by value."""
+    if not isinstance(module, types.ModuleType):
+        raise ValueError(
+            f"Input should be a module object, got {type(module)} instead"
+        )
+    if module.__name__ not in _PICKLE_BY_VALUE_MODULES:
+        raise ValueError(f"{module} is not registered for pickle by value")
+    else:
+        _PICKLE_BY_VALUE_MODULES.remove(module.__name__)
+
+
+def list_registry_pickle_by_value():
+    return _PICKLE_BY_VALUE_MODULES.copy()
+
+
+def _is_registered_pickle_by_value(module):
+    module_name = module.__name__ if inspect.ismodule(module) else module
+    if module_name in _PICKLE_BY_VALUE_MODULES:
+        return True
+    while True:
+        parent_name = module_name.rsplit(".", 1)[0]
+        if parent_name == module_name:
+            break
+        if parent_name in _PICKLE_BY_VALUE_MODULES:
+            return True
+        module_name = parent_name
+    return False
+
+
 def _whichmodule(obj, name):
     """Find the module an object belongs to.
 
@@ -170,18 +228,24 @@ def _whichmodule(obj, name):
     return None
 
 
-def _is_importable(obj, name=None):
-    """Dispatcher utility to test the importability of various constructs."""
-    if isinstance(obj, types.FunctionType):
-        return _lookup_module_and_qualname(obj, name=name) is not None
-    elif issubclass(type(obj), type):
-        return _lookup_module_and_qualname(obj, name=name) is not None
+def _should_pickle_by_reference(obj, name=None):
+    """Dispatcher utility to test whether an object should be serialised by
+    value or reference by using importability as a proxy."""
+    if isinstance(obj, types.FunctionType) or issubclass(type(obj), type):
+        module_and_name = _lookup_module_and_qualname(obj, name=name)
+        if module_and_name is None:
+            return False
+        module, name = module_and_name
+        return not _is_registered_pickle_by_value(module.__name__)
+
     elif isinstance(obj, types.ModuleType):
         # We assume that sys.modules is primarily used as a cache mechanism for
         # the Python import machinery. Checking if a module has been added in
-        # is sys.modules therefore a cheap and simple heuristic to tell us whether
-        # we can assume  that a given module could be imported by name in
-        # another Python process.
+        # is sys.modules therefore a cheap and simple heuristic to tell us
+        # whether we can assume  that a given module could be imported by name
+        # in another Python process.
+        if _is_registered_pickle_by_value(obj.__name__):
+            return False
         return obj.__name__ in sys.modules
     else:
         raise TypeError(
@@ -839,10 +903,15 @@ def _decompose_typevar(obj):
 
 
 def _typevar_reduce(obj):
-    # TypeVar instances have no __qualname__ hence we pass the name explicitly.
+    # TypeVar instances require the module information hence why we
+    # are not using the _should_pickle_by_reference directly
     module_and_name = _lookup_module_and_qualname(obj, name=obj.__name__)
+
     if module_and_name is None:
         return (_make_typevar, _decompose_typevar(obj))
+    elif _is_registered_pickle_by_value(module_and_name[0].__name__):
+        return (_make_typevar, _decompose_typevar(obj))
+
     return (getattr, module_and_name)
 
 

cloudpickle/cloudpickle_fast.py

@@ -28,7 +28,7 @@
 from .compat import pickle, Pickler
 from .cloudpickle import (
     _extract_code_globals, _BUILTIN_TYPE_NAMES, DEFAULT_PROTOCOL,
-    _find_imported_submodules, _get_cell_contents, _is_importable,
+    _find_imported_submodules, _get_cell_contents, _should_pickle_by_reference,
     _builtin_type, _get_or_create_tracker_id,  _make_skeleton_class,
     _make_skeleton_enum, _extract_class_dict, dynamic_subimport, subimport,
     _typevar_reduce, _get_bases, _make_cell, _make_empty_cell, CellType,
@@ -352,7 +352,7 @@ def _memoryview_reduce(obj):
 
 
 def _module_reduce(obj):
-    if _is_importable(obj):
+    if _should_pickle_by_reference(obj):
         return subimport, (obj.__name__,)
     else:
         # Some external libraries can populate the "__builtins__" entry of a
@@ -414,7 +414,7 @@ def _class_reduce(obj):
         return type, (NotImplemented,)
     elif obj in _BUILTIN_TYPE_NAMES:
         return _builtin_type, (_BUILTIN_TYPE_NAMES[obj],)
-    elif not _is_importable(obj):
+    elif not _should_pickle_by_reference(obj):
         return _dynamic_class_reduce(obj)
     return NotImplemented
 
@@ -559,7 +559,7 @@ def _function_reduce(self, obj):
         As opposed to cloudpickle.py, There no special handling for builtin
         pypy functions because cloudpickle_fast is CPython-specific.
         """
-        if _is_importable(obj):
+        if _should_pickle_by_reference(obj):
             return NotImplemented
         else:
             return self._dynamic_function_reduce(obj)
@@ -763,7 +763,7 @@ def save_global(self, obj, name=None, pack=struct.pack):
                 )
             elif name is not None:
                 Pickler.save_global(self, obj, name=name)
-            elif not _is_importable(obj, name=name):
+            elif not _should_pickle_by_reference(obj, name=name):
                 self._save_reduce_pickle5(*_dynamic_class_reduce(obj), obj=obj)
             else:
                 Pickler.save_global(self, obj, name=name)
@@ -775,7 +775,7 @@ def save_function(self, obj, name=None):
             Determines what kind of function obj is (e.g. lambda, defined at
             interactive prompt, etc) and handles the pickling appropriately.
             """
-            if _is_importable(obj, name=name):
+            if _should_pickle_by_reference(obj, name=name):
                 return Pickler.save_global(self, obj, name=name)
             elif PYPY and isinstance(obj.__code__, builtin_code_type):
                 return self.save_pypy_builtin_func(obj)