225 lines
6.9 KiB
Python
225 lines
6.9 KiB
Python
###############################################################################
|
|
# Customizable Pickler with some basic reducers
|
|
#
|
|
# author: Thomas Moreau
|
|
#
|
|
# adapted from multiprocessing/reduction.py (17/02/2017)
|
|
# * Replace the ForkingPickler with a similar _LokyPickler,
|
|
# * Add CustomizableLokyPickler to allow customizing pickling process
|
|
# on the fly.
|
|
#
|
|
import copyreg
|
|
import io
|
|
import functools
|
|
import types
|
|
import sys
|
|
import os
|
|
|
|
from multiprocessing import util
|
|
from pickle import loads, HIGHEST_PROTOCOL
|
|
|
|
###############################################################################
|
|
# Enable custom pickling in Loky.
|
|
|
|
_dispatch_table = {}
|
|
|
|
|
|
def register(type_, reduce_function):
|
|
_dispatch_table[type_] = reduce_function
|
|
|
|
|
|
###############################################################################
|
|
# Registers extra pickling routines to improve picklization for loky
|
|
|
|
|
|
# make methods picklable
|
|
def _reduce_method(m):
|
|
if m.__self__ is None:
|
|
return getattr, (m.__class__, m.__func__.__name__)
|
|
else:
|
|
return getattr, (m.__self__, m.__func__.__name__)
|
|
|
|
|
|
class _C:
|
|
def f(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def h(cls):
|
|
pass
|
|
|
|
|
|
register(type(_C().f), _reduce_method)
|
|
register(type(_C.h), _reduce_method)
|
|
|
|
|
|
if not hasattr(sys, "pypy_version_info"):
|
|
# PyPy uses functions instead of method_descriptors and wrapper_descriptors
|
|
def _reduce_method_descriptor(m):
|
|
return getattr, (m.__objclass__, m.__name__)
|
|
|
|
register(type(list.append), _reduce_method_descriptor)
|
|
register(type(int.__add__), _reduce_method_descriptor)
|
|
|
|
|
|
# Make partial func pickable
|
|
def _reduce_partial(p):
|
|
return _rebuild_partial, (p.func, p.args, p.keywords or {})
|
|
|
|
|
|
def _rebuild_partial(func, args, keywords):
|
|
return functools.partial(func, *args, **keywords)
|
|
|
|
|
|
register(functools.partial, _reduce_partial)
|
|
|
|
if sys.platform != "win32":
|
|
from ._posix_reduction import _mk_inheritable # noqa: F401
|
|
else:
|
|
from . import _win_reduction # noqa: F401
|
|
|
|
# global variable to change the pickler behavior
|
|
try:
|
|
from joblib.externals import cloudpickle # noqa: F401
|
|
|
|
DEFAULT_ENV = "cloudpickle"
|
|
except ImportError:
|
|
# If cloudpickle is not present, fallback to pickle
|
|
DEFAULT_ENV = "pickle"
|
|
|
|
ENV_LOKY_PICKLER = os.environ.get("LOKY_PICKLER", DEFAULT_ENV)
|
|
_LokyPickler = None
|
|
_loky_pickler_name = None
|
|
|
|
|
|
def set_loky_pickler(loky_pickler=None):
|
|
global _LokyPickler, _loky_pickler_name
|
|
|
|
if loky_pickler is None:
|
|
loky_pickler = ENV_LOKY_PICKLER
|
|
|
|
loky_pickler_cls = None
|
|
|
|
# The default loky_pickler is cloudpickle
|
|
if loky_pickler in ["", None]:
|
|
loky_pickler = "cloudpickle"
|
|
|
|
if loky_pickler == _loky_pickler_name:
|
|
return
|
|
|
|
if loky_pickler == "cloudpickle":
|
|
from joblib.externals.cloudpickle import CloudPickler as loky_pickler_cls
|
|
else:
|
|
try:
|
|
from importlib import import_module
|
|
|
|
module_pickle = import_module(loky_pickler)
|
|
loky_pickler_cls = module_pickle.Pickler
|
|
except (ImportError, AttributeError) as e:
|
|
extra_info = (
|
|
"\nThis error occurred while setting loky_pickler to"
|
|
f" '{loky_pickler}', as required by the env variable "
|
|
"LOKY_PICKLER or the function set_loky_pickler."
|
|
)
|
|
e.args = (e.args[0] + extra_info,) + e.args[1:]
|
|
e.msg = e.args[0]
|
|
raise e
|
|
|
|
util.debug(
|
|
f"Using '{loky_pickler if loky_pickler else 'cloudpickle'}' for "
|
|
"serialization."
|
|
)
|
|
|
|
class CustomizablePickler(loky_pickler_cls):
|
|
_loky_pickler_cls = loky_pickler_cls
|
|
|
|
def _set_dispatch_table(self, dispatch_table):
|
|
for ancestor_class in self._loky_pickler_cls.mro():
|
|
dt_attribute = getattr(ancestor_class, "dispatch_table", None)
|
|
if isinstance(dt_attribute, types.MemberDescriptorType):
|
|
# Ancestor class (typically _pickle.Pickler) has a
|
|
# member_descriptor for its "dispatch_table" attribute. Use
|
|
# it to set the dispatch_table as a member instead of a
|
|
# dynamic attribute in the __dict__ of the instance,
|
|
# otherwise it will not be taken into account by the C
|
|
# implementation of the dump method if a subclass defines a
|
|
# class-level dispatch_table attribute as was done in
|
|
# cloudpickle 1.6.0:
|
|
# https://github.com/joblib/loky/pull/260
|
|
dt_attribute.__set__(self, dispatch_table)
|
|
break
|
|
|
|
# On top of member descriptor set, also use setattr such that code
|
|
# that directly access self.dispatch_table gets a consistent view
|
|
# of the same table.
|
|
self.dispatch_table = dispatch_table
|
|
|
|
def __init__(self, writer, reducers=None, protocol=HIGHEST_PROTOCOL):
|
|
loky_pickler_cls.__init__(self, writer, protocol=protocol)
|
|
if reducers is None:
|
|
reducers = {}
|
|
|
|
if hasattr(self, "dispatch_table"):
|
|
# Force a copy that we will update without mutating the
|
|
# any class level defined dispatch_table.
|
|
loky_dt = dict(self.dispatch_table)
|
|
else:
|
|
# Use standard reducers as bases
|
|
loky_dt = copyreg.dispatch_table.copy()
|
|
|
|
# Register loky specific reducers
|
|
loky_dt.update(_dispatch_table)
|
|
|
|
# Set the new dispatch table, taking care of the fact that we
|
|
# need to use the member_descriptor when we inherit from a
|
|
# subclass of the C implementation of the Pickler base class
|
|
# with an class level dispatch_table attribute.
|
|
self._set_dispatch_table(loky_dt)
|
|
|
|
# Register the reducers
|
|
for type, reduce_func in reducers.items():
|
|
self.register(type, reduce_func)
|
|
|
|
def register(self, type, reduce_func):
|
|
"""Attach a reducer function to a given type in the dispatch table."""
|
|
self.dispatch_table[type] = reduce_func
|
|
|
|
_LokyPickler = CustomizablePickler
|
|
_loky_pickler_name = loky_pickler
|
|
|
|
|
|
def get_loky_pickler_name():
|
|
global _loky_pickler_name
|
|
return _loky_pickler_name
|
|
|
|
|
|
def get_loky_pickler():
|
|
global _LokyPickler
|
|
return _LokyPickler
|
|
|
|
|
|
# Set it to its default value
|
|
set_loky_pickler()
|
|
|
|
|
|
def dump(obj, file, reducers=None, protocol=None):
|
|
"""Replacement for pickle.dump() using _LokyPickler."""
|
|
global _LokyPickler
|
|
_LokyPickler(file, reducers=reducers, protocol=protocol).dump(obj)
|
|
|
|
|
|
def dumps(obj, reducers=None, protocol=None):
|
|
global _LokyPickler
|
|
|
|
buf = io.BytesIO()
|
|
dump(obj, buf, reducers=reducers, protocol=protocol)
|
|
return buf.getbuffer()
|
|
|
|
|
|
__all__ = ["dump", "dumps", "loads", "register", "set_loky_pickler"]
|
|
|
|
if sys.platform == "win32":
|
|
from multiprocessing.reduction import duplicate
|
|
|
|
__all__ += ["duplicate"]
|