diff --git a/Doc/library/multiprocessing.rst b/Doc/library/multiprocessing.rst index e80fdfda5bea36..5549d587e9a69b 100644 --- a/Doc/library/multiprocessing.rst +++ b/Doc/library/multiprocessing.rst @@ -1436,6 +1436,103 @@ For example: the data in the pipe is likely to become corrupted, because it may become impossible to be sure where the message boundaries lie. +Custom Reduction +~~~~~~~~~~~~~~~~ + +.. currentmodule:: multiprocessing + +Several primitives of the :mod:`multiprocessing` module such as +:class:`multiprocessing.Queue`, :class:`multiprocessing.connection.Listener` or +:class:`multiprocessing.connection.Server` need to serialize and deserialize Python +objects to communicate between processes. Sometimes it is useful to control what +serialization is to be used for the transport of data in order to support communication +with different versions of Python, use more performant third party tools or custom +strategies. + +For this purpose a set of hooks is available to provide alternate implementations of +the reduction mechanism: + +.. currentmodule:: multiprocessing.reduction + +.. class:: AbstractPickler(file, protocol) + + Base class that can be implemented in order to override + serialization methods of the reduction machinery used by multiprocessing. + + .. method:: dump(obj) + + Write a pickled representation of obj to the open file. + + Defaults to :meth:`pickle.Pickler.dump` + +.. class:: AbstractUnpickler(bytes_object, *, fix_imports=True, encoding="ASCII", errors="strict") + + Base class that can be implemented in order to override + multiprocessing's default unserialization mechanism. + + .. method:: load() + + Read a pickled object hierarchy from the open file and return the + reconstituted object hierarchy specified therein. + + Defaults to :meth:`pickle.Unpickler.load` + +.. class:: AbstractReducer() + + Base class providing access to custom ``Pickler`` and ``Unpickler`` + classes to be used by ``multiprocessing`` when serializing objects. + + .. method:: get_pickler_class(): + + This method must return an subclass of :class:`AbstractPickler` to be used + by ``multiprocessing`` when serializing objects. + + .. method:: get_unpickler_class(): + + This method must return an subclass of :class:`AbstractUnpickler` to be used + by ``multiprocessing`` when unserializing objects. + + +Note that both methods of the :class:`AbstractReducer` are optional. If +:meth:`get_pickler_class` (resp. :meth:`get_unpickler_class`) is not +implemented, multiprocessing will fallback to the :class:`pickle.Pickler` +(resp. :class:`pickle.Unpickler`) class to serialize objects. + +.. currentmodule:: multiprocessing + +.. method:: set_reducer(reduction) + + Sets a reduction instance to be used for serialization and deserialization + by the module primitive internals. **reduction** must be an instance of a + subclass of :class:`AbstractReducer`. + +.. method:: get_reducer() + + Gets the current reduction class in use by the module's primitive internals. + +For example, substituting the reducer class to use the :mod:`pickle` protocol +version 2 to be able to communicate with a Python 2.x programs.:: + + import multiprocessing + from multiprocessing.reduction import AbstractReducer, AbstractPickler + + class PicklerProtocol2(AbstractPickler): + def __init__(self, file, protocol=2, **kwargs): + super().__init__(file, protocol, **kwargs) + + def dump(self, obj): + return super().dump(obj) + + + class PickleProtocol2Reducer(AbstractReducer): + def get_pickler_class(self): + return PicklerProtocol2 + + multiprocessing.set_reducer(PickleProtocol2Reducer()) + +Notice that using :meth:`multiprocessing.set_reducer` changes the reducer globally. If +changing this setting globally is undesirable you could call :meth:`context.set_reducer`, +where *context* is a context object obtained by calling :func:`get_context`. Synchronization primitives ^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/Lib/multiprocessing/__init__.py b/Lib/multiprocessing/__init__.py index 8336f381decaac..2b70c044177e51 100644 --- a/Lib/multiprocessing/__init__.py +++ b/Lib/multiprocessing/__init__.py @@ -20,7 +20,11 @@ # __all__ = [x for x in dir(context._default_context) if not x.startswith('_')] -globals().update((name, getattr(context._default_context, name)) for name in __all__) + + +def __getattr__(name): + return getattr(context._default_context, name) + # # XXX These should not really be documented or public. diff --git a/Lib/multiprocessing/connection.py b/Lib/multiprocessing/connection.py index e33d1ed800f319..4bb4b7e7b5d54f 100644 --- a/Lib/multiprocessing/connection.py +++ b/Lib/multiprocessing/connection.py @@ -22,8 +22,8 @@ from . import util from . import AuthenticationError, BufferTooShort -from .context import reduction -_ForkingPickler = reduction.ForkingPickler +from . import context +from . import get_context try: import _multiprocessing @@ -122,7 +122,8 @@ def address_type(address): class _ConnectionBase: _handle = None - def __init__(self, handle, readable=True, writable=True): + def __init__(self, handle, readable=True, writable=True, ctx=None): + self._ctx = ctx or get_context() handle = handle.__index__() if handle < 0: raise ValueError("invalid handle") @@ -214,7 +215,7 @@ def send(self, obj): """Send a (picklable) object""" self._check_closed() self._check_writable() - self._send_bytes(_ForkingPickler.dumps(obj)) + self._send_bytes(self._ctx.get_reducer().dumps(obj)) def recv_bytes(self, maxlength=None): """ @@ -259,7 +260,7 @@ def recv(self): self._check_closed() self._check_readable() buf = self._recv_bytes() - return _ForkingPickler.loads(buf.getbuffer()) + return self._ctx.get_reducer().loads(buf.getbuffer()) def poll(self, timeout=0.0): """Whether there is any input available to be read""" @@ -474,7 +475,7 @@ class Listener(object): This is a wrapper for a bound socket which is 'listening' for connections, or for a Windows named pipe. ''' - def __init__(self, address=None, family=None, backlog=1, authkey=None): + def __init__(self, address=None, family=None, backlog=1, authkey=None, *, ctx=None): family = family or (address and address_type(address)) \ or default_family _validate_family(family) @@ -483,12 +484,12 @@ def __init__(self, address=None, family=None, backlog=1, authkey=None): if family == 'AF_PIPE': if address: - self._listener = PipeListener(address, backlog) + self._listener = PipeListener(address, backlog, ctx=ctx) else: for attempts in itertools.count(): address = arbitrary_address(family) try: - self._listener = PipeListener(address, backlog) + self._listener = PipeListener(address, backlog, ctx=ctx) break except OSError as e: if attempts >= _MAX_PIPE_ATTEMPTS: @@ -498,7 +499,7 @@ def __init__(self, address=None, family=None, backlog=1, authkey=None): raise else: address = address or arbitrary_address(family) - self._listener = SocketListener(address, family, backlog) + self._listener = SocketListener(address, family, backlog, ctx=ctx) self._authkey = authkey @@ -541,16 +542,16 @@ def __exit__(self, exc_type, exc_value, exc_tb): self.close() -def Client(address, family=None, authkey=None): +def Client(address, family=None, authkey=None, *, ctx=None): ''' Returns a connection to the address of a `Listener` ''' family = family or address_type(address) _validate_family(family) if family == 'AF_PIPE': - c = PipeClient(address) + c = PipeClient(address, ctx=ctx) else: - c = SocketClient(address) + c = SocketClient(address, ctx=ctx) if authkey is not None and not isinstance(authkey, bytes): raise TypeError('authkey should be a byte string') @@ -640,7 +641,7 @@ class SocketListener(object): ''' Representation of a socket which is bound to an address and listening ''' - def __init__(self, address, family, backlog=1): + def __init__(self, address, family, backlog=1, *, ctx=None): self._socket = socket.socket(getattr(socket, family)) try: # SO_REUSEADDR has different semantics on Windows (issue #2550). @@ -664,11 +665,12 @@ def __init__(self, address, family, backlog=1): ) else: self._unlink = None + self._ctx = ctx def accept(self): s, self._last_accepted = self._socket.accept() s.setblocking(True) - return Connection(s.detach()) + return Connection(s.detach(), ctx=self._ctx) def close(self): try: @@ -680,7 +682,7 @@ def close(self): unlink() -def SocketClient(address): +def SocketClient(address, *, ctx=None): ''' Return a connection object connected to the socket given by `address` ''' @@ -688,7 +690,7 @@ def SocketClient(address): with socket.socket( getattr(socket, family) ) as s: s.setblocking(True) s.connect(address) - return Connection(s.detach()) + return Connection(s.detach(), ctx=ctx) # # Definitions for connections based on named pipes @@ -700,7 +702,7 @@ class PipeListener(object): ''' Representation of a named pipe ''' - def __init__(self, address, backlog=None): + def __init__(self, address, backlog=None, *, ctx=None): self._address = address self._handle_queue = [self._new_handle(first=True)] @@ -710,6 +712,7 @@ def __init__(self, address, backlog=None): self, PipeListener._finalize_pipe_listener, args=(self._handle_queue, self._address), exitpriority=0 ) + self._ctx = ctx def _new_handle(self, first=False): flags = _winapi.PIPE_ACCESS_DUPLEX | _winapi.FILE_FLAG_OVERLAPPED @@ -743,7 +746,7 @@ def accept(self): finally: _, err = ov.GetOverlappedResult(True) assert err == 0 - return PipeConnection(handle) + return PipeConnection(handle, ctx=self._ctx) @staticmethod def _finalize_pipe_listener(queue, address): @@ -751,7 +754,7 @@ def _finalize_pipe_listener(queue, address): for handle in queue: _winapi.CloseHandle(handle) - def PipeClient(address): + def PipeClient(address, *, ctx=None): ''' Return a connection object connected to the pipe given by `address` ''' @@ -776,7 +779,7 @@ def PipeClient(address): _winapi.SetNamedPipeHandleState( h, _winapi.PIPE_READMODE_MESSAGE, None, None ) - return PipeConnection(h) + return PipeConnection(h, ctx=ctx) # # Authentication stuff @@ -1218,23 +1221,23 @@ def reduce_connection(conn): def rebuild_connection(ds, readable, writable): sock = ds.detach() return Connection(sock.detach(), readable, writable) - reduction.register(Connection, reduce_connection) + context.reduction.register(Connection, reduce_connection) def reduce_pipe_connection(conn): access = ((_winapi.FILE_GENERIC_READ if conn.readable else 0) | (_winapi.FILE_GENERIC_WRITE if conn.writable else 0)) - dh = reduction.DupHandle(conn.fileno(), access) + dh = context.reduction.DupHandle(conn.fileno(), access) return rebuild_pipe_connection, (dh, conn.readable, conn.writable) def rebuild_pipe_connection(dh, readable, writable): handle = dh.detach() return PipeConnection(handle, readable, writable) - reduction.register(PipeConnection, reduce_pipe_connection) + context.reduction.register(PipeConnection, reduce_pipe_connection) else: def reduce_connection(conn): - df = reduction.DupFd(conn.fileno()) + df = context.reduction.DupFd(conn.fileno()) return rebuild_connection, (df, conn.readable, conn.writable) def rebuild_connection(df, readable, writable): fd = df.detach() return Connection(fd, readable, writable) - reduction.register(Connection, reduce_connection) + context.reduction.register(Connection, reduce_connection) diff --git a/Lib/multiprocessing/context.py b/Lib/multiprocessing/context.py index e9f8aadceb9ced..933dccb17df375 100644 --- a/Lib/multiprocessing/context.py +++ b/Lib/multiprocessing/context.py @@ -6,6 +6,8 @@ from . import process from . import reduction +original_reducer = reduction + __all__ = () # @@ -38,6 +40,7 @@ class BaseContext(object): current_process = staticmethod(process.current_process) parent_process = staticmethod(process.parent_process) active_children = staticmethod(process.active_children) + _reducer = None def cpu_count(self): '''Returns the number of CPUs in the system''' @@ -47,6 +50,18 @@ def cpu_count(self): else: return num + @property + def Process(self): + if self._reducer is None: + # Ensure backward compatibility by returning a class when no + # custom reducer was specified + return self._Process + else: + return self._process + + def _process(self, *args, **kwargs): + return self._Process(*args, reducer=self.get_reducer(), **kwargs) + def Manager(self): '''Returns a manager associated with a running server process @@ -63,6 +78,16 @@ def Pipe(self, duplex=True): from .connection import Pipe return Pipe(duplex) + def Listener(self, *args, **kwargs): + '''Returns a connection Listener object''' + from .connection import Listener + return Listener(*args, ctx=self.get_context(), **kwargs) + + def Client(self, *args, **kwargs): + '''Returns a connection Client object''' + from .connection import Client + return Client(*args, ctx=self.get_context(), **kwargs) + def Lock(self): '''Returns a non-recursive lock object''' from .synchronize import Lock @@ -210,15 +235,34 @@ def get_start_method(self, allow_none=False): def set_start_method(self, method, force=False): raise ValueError('cannot set start method of concrete context') - @property - def reducer(self): + def get_reducer(self): '''Controls how objects will be reduced to a form that can be shared with other processes.''' - return globals().get('reduction') + reducer = self._reducer + if reducer is None: + return reduction + return reducer + + def set_reducer(self, reduction): + if reduction is original_reducer: + self._reducer = None + return - @reducer.setter - def reducer(self, reduction): - globals()['reduction'] = reduction + if not isinstance(reduction, original_reducer.AbstractReducer): + raise TypeError("Custom reducers must be instances of a " + "subclass of multiprocessing.reduction.AbstractReducer") + if not issubclass(reduction.get_pickler_class(), original_reducer.AbstractPickler): + raise TypeError("Custom reducers must return an subclass of " + "multiprocessing.reduction.AbstractPickler " + "in the get_pickler_class() method") + if not issubclass(reduction.get_unpickler_class(), + original_reducer.AbstractUnpickler): + raise TypeError( + "Custom reducers must return an subclass of " + "multiprocessing.reduction.AbstractUnpickler in the " + "get_unpickler_class() method") + + self._reducer = reduction def _check_available(self): pass @@ -231,14 +275,14 @@ class Process(process.BaseProcess): _start_method = None @staticmethod def _Popen(process_obj): - return _default_context.get_context().Process._Popen(process_obj) + return _default_context.get_context()._Process._Popen(process_obj) @staticmethod def _after_fork(): return _default_context.get_context().Process._after_fork() class DefaultContext(BaseContext): - Process = Process + _Process = Process def __init__(self, context=None): if context is not None: @@ -320,15 +364,15 @@ def _Popen(process_obj): class ForkContext(BaseContext): _name = 'fork' - Process = ForkProcess + _Process = ForkProcess class SpawnContext(BaseContext): _name = 'spawn' - Process = SpawnProcess + _Process = SpawnProcess class ForkServerContext(BaseContext): _name = 'forkserver' - Process = ForkServerProcess + _Process = ForkServerProcess def _check_available(self): if not reduction.HAVE_SEND_HANDLE: raise ValueError('forkserver start method not available') @@ -372,7 +416,7 @@ def _after_fork(): class SpawnContext(BaseContext): _name = 'spawn' - Process = SpawnProcess + _Process = SpawnProcess _concrete_contexts = { 'spawn': SpawnContext(), diff --git a/Lib/multiprocessing/forkserver.py b/Lib/multiprocessing/forkserver.py index b204a3ad0f1f01..f5b979a0cd6f47 100644 --- a/Lib/multiprocessing/forkserver.py +++ b/Lib/multiprocessing/forkserver.py @@ -12,8 +12,8 @@ from . import AuthenticationError from . import connection from . import process -from .context import reduction from . import resource_tracker +from . import context from . import spawn from . import util @@ -130,7 +130,7 @@ def connect_to_new_process(self, fds): finally: wrapped_client._detach() del wrapped_client - reduction.sendfds(client, allfds) + context.reduction.sendfds(client, allfds) return parent_r, parent_w except: os.close(parent_r) @@ -393,7 +393,7 @@ def sigchld_handler(*_unused): wrapped_s._detach() del wrapped_s # Receive fds from client - fds = reduction.recvfds(s, MAXFDS_TO_SEND + 1) + fds = context.reduction.recvfds(s, MAXFDS_TO_SEND + 1) except (EOFError, BrokenPipeError, AuthenticationError): s.close() continue diff --git a/Lib/multiprocessing/heap.py b/Lib/multiprocessing/heap.py index 5c835648395f79..8573a6927282a8 100644 --- a/Lib/multiprocessing/heap.py +++ b/Lib/multiprocessing/heap.py @@ -15,7 +15,8 @@ import tempfile import threading -from .context import reduction, assert_spawning +from . import context +from .context import assert_spawning from . import util __all__ = ['BufferWrapper'] @@ -101,12 +102,12 @@ def reduce_arena(a): if a.fd == -1: raise ValueError('Arena is unpicklable because ' 'forking was enabled when it was created') - return rebuild_arena, (a.size, reduction.DupFd(a.fd)) + return rebuild_arena, (a.size, context.reduction.DupFd(a.fd)) def rebuild_arena(size, dupfd): return Arena(size, dupfd.detach()) - reduction.register(Arena, reduce_arena) + context.reduction.register(Arena, reduce_arena) # # Class allowing allocation of chunks of memory from arenas diff --git a/Lib/multiprocessing/managers.py b/Lib/multiprocessing/managers.py index 11c6e21de90e69..06b9e7452c0e1b 100644 --- a/Lib/multiprocessing/managers.py +++ b/Lib/multiprocessing/managers.py @@ -28,7 +28,8 @@ from traceback import format_exc from . import connection -from .context import reduction, get_spawning_popen, ProcessError +from . import context +from .context import get_spawning_popen, ProcessError from . import pool from . import process from . import util @@ -47,13 +48,13 @@ def reduce_array(a): return array.array, (a.typecode, a.tobytes()) -reduction.register(array.array, reduce_array) +context.reduction.register(array.array, reduce_array) view_types = [type(getattr({}, name)()) for name in ('items','keys','values')] def rebuild_as_list(obj): return list, (list(obj),) for view_type in view_types: - reduction.register(view_type, rebuild_as_list) + context.reduction.register(view_type, rebuild_as_list) del view_type, view_types # diff --git a/Lib/multiprocessing/popen_forkserver.py b/Lib/multiprocessing/popen_forkserver.py index a56eb9bf11080b..9f748314371fc4 100644 --- a/Lib/multiprocessing/popen_forkserver.py +++ b/Lib/multiprocessing/popen_forkserver.py @@ -1,8 +1,9 @@ import io import os -from .context import reduction, set_spawning_popen -if not reduction.HAVE_SEND_HANDLE: +from . import context +from .context import set_spawning_popen +if not context.reduction.HAVE_SEND_HANDLE: raise ImportError('No support for sending fds between processes') from . import forkserver from . import popen_fork @@ -43,8 +44,8 @@ def _launch(self, process_obj): buf = io.BytesIO() set_spawning_popen(self) try: - reduction.dump(prep_data, buf) - reduction.dump(process_obj, buf) + process_obj.reducer.dump(prep_data, buf) + process_obj.reducer.dump(process_obj, buf) finally: set_spawning_popen(None) diff --git a/Lib/multiprocessing/popen_spawn_posix.py b/Lib/multiprocessing/popen_spawn_posix.py index cccd659ae77637..ff22f4244a8b37 100644 --- a/Lib/multiprocessing/popen_spawn_posix.py +++ b/Lib/multiprocessing/popen_spawn_posix.py @@ -1,7 +1,8 @@ import io import os -from .context import reduction, set_spawning_popen +import multiprocessing +from .context import set_spawning_popen from . import popen_fork from . import spawn from . import util @@ -43,8 +44,8 @@ def _launch(self, process_obj): fp = io.BytesIO() set_spawning_popen(self) try: - reduction.dump(prep_data, fp) - reduction.dump(process_obj, fp) + process_obj.reducer.dump(prep_data, fp) + process_obj.reducer.dump(process_obj, fp) finally: set_spawning_popen(None) diff --git a/Lib/multiprocessing/popen_spawn_win32.py b/Lib/multiprocessing/popen_spawn_win32.py index 62fb0ddbf91a5d..f7eefdeb063ee7 100644 --- a/Lib/multiprocessing/popen_spawn_win32.py +++ b/Lib/multiprocessing/popen_spawn_win32.py @@ -5,7 +5,9 @@ import _winapi from subprocess import STARTUPINFO, STARTF_FORCEOFFFEEDBACK -from .context import reduction, get_spawning_popen, set_spawning_popen +import multiprocessing +from . import context +from .context import get_spawning_popen, set_spawning_popen from . import spawn from . import util @@ -93,14 +95,14 @@ def __init__(self, process_obj): # send information to child set_spawning_popen(self) try: - reduction.dump(prep_data, to_child) - reduction.dump(process_obj, to_child) + process_obj.reducer.dump(prep_data, to_child) + process_obj.reducer.dump(process_obj, to_child) finally: set_spawning_popen(None) def duplicate_for_child(self, handle): assert self is get_spawning_popen() - return reduction.duplicate(handle, self.sentinel) + return context.reduction.duplicate(handle, self.sentinel) def wait(self, timeout=None): if self.returncode is not None: diff --git a/Lib/multiprocessing/process.py b/Lib/multiprocessing/process.py index 262513f295fde5..ddec2139c9c44e 100644 --- a/Lib/multiprocessing/process.py +++ b/Lib/multiprocessing/process.py @@ -21,6 +21,7 @@ import threading from _weakrefset import WeakSet +from . import context # # # @@ -78,7 +79,7 @@ def _Popen(self): raise NotImplementedError def __init__(self, group=None, target=None, name=None, args=(), kwargs=None, - *, daemon=None): + *, daemon=None, reducer=None): assert group is None, 'group argument must be None for now' count = next(_process_counter) self._identity = _current_process._identity + (count,) @@ -95,11 +96,16 @@ def __init__(self, group=None, target=None, name=None, args=(), kwargs=None, if daemon is not None: self.daemon = daemon _dangling.add(self) + self._reducer = reducer def _check_closed(self): if self._closed: raise ValueError("process object is closed") + @property + def reducer(self): + return self._reducer or context.reduction + def run(self): ''' Method to be run in sub-process; can be overridden in sub-class diff --git a/Lib/multiprocessing/queues.py b/Lib/multiprocessing/queues.py index 981599acf5ef26..bd79d836668005 100644 --- a/Lib/multiprocessing/queues.py +++ b/Lib/multiprocessing/queues.py @@ -22,7 +22,6 @@ from . import connection from . import context -_ForkingPickler = context.reduction.ForkingPickler from .util import debug, info, Finalize, register_after_fork, is_exiting @@ -47,6 +46,7 @@ def __init__(self, maxsize=0, *, ctx): self._sem = ctx.BoundedSemaphore(maxsize) # For use by concurrent.futures self._ignore_epipe = False + self._ctx = ctx self._reset() if sys.platform != 'win32': @@ -55,11 +55,11 @@ def __init__(self, maxsize=0, *, ctx): def __getstate__(self): context.assert_spawning(self) return (self._ignore_epipe, self._maxsize, self._reader, self._writer, - self._rlock, self._wlock, self._sem, self._opid) + self._rlock, self._wlock, self._sem, self._opid, self._ctx) def __setstate__(self, state): (self._ignore_epipe, self._maxsize, self._reader, self._writer, - self._rlock, self._wlock, self._sem, self._opid) = state + self._rlock, self._wlock, self._sem, self._opid, self._ctx) = state self._reset() def _after_fork(self): @@ -117,7 +117,7 @@ def get(self, block=True, timeout=None): finally: self._rlock.release() # unserialize the data after having released the lock - return _ForkingPickler.loads(res) + return self._ctx.get_reducer().loads(res) def qsize(self): # Raises NotImplementedError on Mac OSX because of broken sem_getvalue() @@ -180,7 +180,7 @@ def _start_thread(self): args=(self._buffer, self._notempty, self._send_bytes, self._wlock, self._reader.close, self._writer.close, self._ignore_epipe, self._on_queue_feeder_error, - self._sem), + self._sem, self._ctx), name='QueueFeederThread', daemon=True, ) @@ -228,7 +228,7 @@ def _finalize_close(buffer, notempty): @staticmethod def _feed(buffer, notempty, send_bytes, writelock, reader_close, - writer_close, ignore_epipe, onerror, queue_sem): + writer_close, ignore_epipe, onerror, queue_sem, ctx): debug('starting thread to feed data to pipe') nacquire = notempty.acquire nrelease = notempty.release @@ -259,7 +259,7 @@ def _feed(buffer, notempty, send_bytes, writelock, reader_close, return # serialize the data before acquiring the lock - obj = _ForkingPickler.dumps(obj) + obj = ctx.get_reducer().dumps(obj) if wacquire is None: send_bytes(obj) else: @@ -357,6 +357,7 @@ def join(self): class SimpleQueue(object): def __init__(self, *, ctx): + self._ctx = ctx self._reader, self._writer = connection.Pipe(duplex=False) self._rlock = ctx.Lock() self._poll = self._reader.poll @@ -374,21 +375,23 @@ def empty(self): def __getstate__(self): context.assert_spawning(self) - return (self._reader, self._writer, self._rlock, self._wlock) + return (self._reader, self._writer, self._rlock, self._wlock, + self._ctx) def __setstate__(self, state): - (self._reader, self._writer, self._rlock, self._wlock) = state + (self._reader, self._writer, self._rlock, self._wlock, + self._ctx) = state self._poll = self._reader.poll def get(self): with self._rlock: res = self._reader.recv_bytes() # unserialize the data after having released the lock - return _ForkingPickler.loads(res) + return self._ctx.get_reducer().loads(res) def put(self, obj): # serialize the data before acquiring the lock - obj = _ForkingPickler.dumps(obj) + obj = self._ctx.get_reducer().dumps(obj) if self._wlock is None: # writes to a message oriented win32 pipe are atomic self._writer.send_bytes(obj) diff --git a/Lib/multiprocessing/reduction.py b/Lib/multiprocessing/reduction.py index fcccd3eef86cc7..0473750f04408e 100644 --- a/Lib/multiprocessing/reduction.py +++ b/Lib/multiprocessing/reduction.py @@ -7,7 +7,7 @@ # Licensed to PSF under a Contributor Agreement. # -from abc import ABCMeta +import abc import copyreg import functools import io @@ -51,14 +51,39 @@ def dumps(cls, obj, protocol=None): cls(buf, protocol).dump(obj) return buf.getbuffer() - loads = pickle.loads + @classmethod + def loads(cls, bytes_object, *, fix_imports=True, + encoding="ASCII", errors="strict"): + return pickle.loads(bytes_object, fix_imports=fix_imports, + encoding=encoding, errors=errors) + + +class AbstractPickler(ForkingPickler): + pass + + +class AbstractUnpickler(pickle.Unpickler): + pass + + +def register(type_, reduce_): + return ForkingPickler.register(type_, reduce_) -register = ForkingPickler.register def dump(obj, file, protocol=None): '''Replacement for pickle.dump() using ForkingPickler.''' ForkingPickler(file, protocol).dump(obj) + +def dumps(obj, protocol=None): + return ForkingPickler.dumps(obj, protocol=protocol) + + +def loads(s, *, fix_imports=True, encoding="ASCII", errors="strict"): + return ForkingPickler.loads(s, fix_imports=fix_imports, + encoding=encoding, errors=errors) + + # # Platform specific definitions # @@ -248,30 +273,27 @@ def _rebuild_socket(df, family, type, proto): register(socket.socket, _reduce_socket) -class AbstractReducer(metaclass=ABCMeta): +class AbstractReducer(metaclass=abc.ABCMeta): '''Abstract base class for use in implementing a Reduction class suitable for use in replacing the standard reduction mechanism used in multiprocessing.''' - ForkingPickler = ForkingPickler - register = register - dump = dump send_handle = send_handle recv_handle = recv_handle if sys.platform == 'win32': - steal_handle = steal_handle - duplicate = duplicate + steal_handle = staticmethod(steal_handle) + duplicate = staticmethod(duplicate) DupHandle = DupHandle else: - sendfds = sendfds - recvfds = recvfds - DupFd = DupFd + sendfds = staticmethod(sendfds) + recvfds = staticmethod(recvfds) + DupFd = staticmethod(DupFd) - _reduce_method = _reduce_method - _reduce_method_descriptor = _reduce_method_descriptor - _rebuild_partial = _rebuild_partial - _reduce_socket = _reduce_socket - _rebuild_socket = _rebuild_socket + _reduce_method = staticmethod(_reduce_method) + _reduce_method_descriptor = staticmethod(_reduce_method_descriptor) + _rebuild_partial = staticmethod(_rebuild_partial) + _reduce_socket = staticmethod(_reduce_socket) + _rebuild_socket = staticmethod(_rebuild_socket) def __init__(self, *args): register(type(_C().f), _reduce_method) @@ -279,3 +301,33 @@ def __init__(self, *args): register(type(int.__add__), _reduce_method_descriptor) register(functools.partial, _reduce_partial) register(socket.socket, _reduce_socket) + + @property + def HAVE_SEND_HANDLE(self): + import sys + return (sys.platform == 'win32' or + (hasattr(socket, 'CMSG_LEN') and + hasattr(socket, 'SCM_RIGHTS') and + hasattr(socket.socket, 'sendmsg'))) + + def get_pickler_class(self): + return AbstractPickler + + def get_unpickler_class(self): + return AbstractUnpickler + + def register(self, type_, reduce_): + return self.get_pickler_class().register(type_, reduce_) + + def dump(self, obj, file, protocol=None): + self.get_pickler_class()(file, protocol).dump(obj) + + def dumps(self, obj, protocol=None): + return self.get_pickler_class().dumps(obj, protocol=protocol) + + def loads(self, s, *, fix_imports=True, encoding="ASCII", errors="strict"): + unpickler_cls = self.get_unpickler_class() + file = io.BytesIO(s) + unpickler = unpickler_cls( + file, fix_imports=fix_imports, encoding=encoding, errors=errors) + return unpickler.load() diff --git a/Lib/multiprocessing/resource_sharer.py b/Lib/multiprocessing/resource_sharer.py index b8afb0fbed3a3c..17331eed16e9d6 100644 --- a/Lib/multiprocessing/resource_sharer.py +++ b/Lib/multiprocessing/resource_sharer.py @@ -15,7 +15,7 @@ import threading from . import process -from .context import reduction +from . import context from . import util __all__ = ['stop'] @@ -47,7 +47,7 @@ class DupFd(object): def __init__(self, fd): new_fd = os.dup(fd) def send(conn, pid): - reduction.send_handle(conn, new_fd, pid) + context.reduction.send_handle(conn, new_fd, pid) def close(): os.close(new_fd) self._id = _resource_sharer.register(send, close) @@ -55,7 +55,7 @@ def close(): def detach(self): '''Get the fd. This should only be called once.''' with _resource_sharer.get_connection(self._id) as conn: - return reduction.recv_handle(conn) + return context.reduction.recv_handle(conn) class _ResourceSharer(object): diff --git a/Lib/multiprocessing/spawn.py b/Lib/multiprocessing/spawn.py index d43864c939cb63..f4f74acda96680 100644 --- a/Lib/multiprocessing/spawn.py +++ b/Lib/multiprocessing/spawn.py @@ -15,7 +15,7 @@ from . import get_start_method, set_start_method from . import process -from .context import reduction +from . import context from . import util __all__ = ['_main', 'freeze_support', 'set_executable', 'get_executable', @@ -110,8 +110,8 @@ def spawn_main(pipe_handle, parent_pid=None, tracker_fd=None): False, parent_pid) else: source_process = None - new_handle = reduction.duplicate(pipe_handle, - source_process=source_process) + new_handle = context.reduction.duplicate(pipe_handle, + source_process=source_process) fd = msvcrt.open_osfhandle(new_handle, os.O_RDONLY) parent_sentinel = source_process else: @@ -127,9 +127,9 @@ def _main(fd, parent_sentinel): with os.fdopen(fd, 'rb', closefd=True) as from_parent: process.current_process()._inheriting = True try: - preparation_data = reduction.pickle.load(from_parent) + preparation_data = context.reduction.pickle.load(from_parent) prepare(preparation_data) - self = reduction.pickle.load(from_parent) + self = context.reduction.pickle.load(from_parent) finally: del process.current_process()._inheriting return self._bootstrap(parent_sentinel) diff --git a/Lib/test/_test_multiprocessing.py b/Lib/test/_test_multiprocessing.py index 00de9d0ac921d2..778d3eaa487ed3 100644 --- a/Lib/test/_test_multiprocessing.py +++ b/Lib/test/_test_multiprocessing.py @@ -49,6 +49,7 @@ support.skip_if_broken_multiprocessing_synchronize() import threading +import multiprocessing import multiprocessing.connection import multiprocessing.dummy import multiprocessing.heap @@ -7765,6 +7766,342 @@ def test_preload_main_large_sys_argv(self): '', ]) + +class SpyReducerBase(reduction.AbstractReducer): + def __init__(self): + self.spy = {'load': 0, 'dump': 0} + + +class SpyReducerWithPickler(SpyReducerBase): + """Custom reducer that records the call to the Pickler + class it returns. Mocks are not used because they do not + play good with pickling""" + def get_pickler_class(self): + spy = self.spy + class Pickler(multiprocessing.reduction.AbstractPickler): + def dump(self, obj, protocol=None): + spy["dump"] = 1 + super().dump(obj) + return Pickler + + +class SpyReducerWithUnpickler(SpyReducerBase): + """Custom reducer that records the call to the Unpickler class it returns. + """ + def get_unpickler_class(self): + spy = self.spy + + class Unpickler(multiprocessing.reduction.AbstractUnpickler): + def load(self): + spy["load"] = 1 + return super().load() + return Unpickler + +class SpyReducerWithPicklerAndUnpickler(SpyReducerWithUnpickler, + SpyReducerWithPickler): + pass + + +class _TestCustomReducer(BaseTestCase): + """Test custom reducer usage. + + The goal of these tests is to make sure that multiprocessing actually uses + the Reducer given as input of the set_reducer() methods of multiprocessing + and of context objects. This is done by counting each call of load/dump + methods of the custom pickler/unpickler classes returned by the reducer, + and making sure, when expected, that these counts are greater than 0 at the + end. + """ + + ALLOWED_TYPES = ('processes',) + + def setUp(self): + self.context = multiprocessing.get_context() + self.original_reducer = multiprocessing.get_reducer() + + def tearDown(self): + self.context.set_reducer(self.original_reducer) + + def _check_pickler_hits(self, reducer): + custom_pickler = ( + reducer.get_pickler_class() is not reduction.AbstractPickler) + custom_unpickler = ( + reducer.get_unpickler_class() is not reduction.AbstractUnpickler) + + self.assertEqual(reducer.spy['dump'], int(custom_pickler)) + self.assertEqual(reducer.spy['load'], int(custom_unpickler)) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_connection_custom_reducer(self): + for reducer_cls in [ + SpyReducerWithPickler, + SpyReducerWithUnpickler, + SpyReducerWithPicklerAndUnpickler + ]: + with self.subTest(reducer_cls=reducer_cls): + with self.connection.Listener() as l: + reducer = reducer_cls() + self.context.set_reducer(reducer) + with self.connection.Client(l.address) as c: + with l.accept() as d: + c.send(1729) + self.assertEqual(d.recv(), 1729) + + if self.TYPE == 'processes': + self.assertRaises(OSError, l.accept) + + self._check_pickler_hits(reducer) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_process_custom_reducer(self): + sm = multiprocessing.get_start_method() + if sm == 'fork': + # The fork method does send data to new Process objects + self.skipTest('test not appropriate for {}'.format(sm)) + + # It is not possible to customize unpickling of Process objects data + # -- we would need to send to the new process which Unpickler to use, + # which itself would require to unpickle this information before + # unpickling anyting else... using which Unpickler? + # Thus, just test the SpyReducer that customizes pickling. + reducer = SpyReducerWithPickler() + self.context.set_reducer(reducer) + p = self.context.Process(target=id, args=(2, )) + p.start() + p.join() + self.assertEqual(p.exitcode, 0) + self._check_pickler_hits(reducer) + + @classmethod + def _put_and_get_in_queue(cls, queue,): + queue.put("Something") + queue.get(timeout=support.SHORT_TIMEOUT) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_queue_custom_reducer(self): + for reducer_cls in [ + SpyReducerWithPickler, + SpyReducerWithUnpickler, + SpyReducerWithPicklerAndUnpickler + ]: + with self.subTest(reducer_cls=reducer_cls): + reducer = reducer_cls() + self.context.set_reducer(reducer) + queue = self.Queue() + p = self.Process(target=self._put_and_get_in_queue, + args=(queue,)) + p.start() + element = queue.get(timeout=support.SHORT_TIMEOUT) + self.assertEqual(element, "Something") + queue.put("Other_Something") + close_queue(queue) + p.join(timeout=support.SHORT_TIMEOUT) + p.terminate() + self.assertEqual(element, "Something") + self.assertEqual(p.exitcode, 0) + self._check_pickler_hits(reducer) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_listener_custom_reduction(self): + for reducer_cls in [ + SpyReducerWithPickler, + SpyReducerWithUnpickler, + SpyReducerWithPicklerAndUnpickler + ]: + with self.subTest(reducer_cls=reducer_cls): + reducer = reducer_cls() + self.context.set_reducer(reducer) + with self.connection.Listener() as l: + with self.connection.Client(l.address) as c: + with l.accept() as d: + c.send(1729) + self.assertEqual(d.recv(), 1729) + + if self.TYPE == 'processes': + self.assertRaises(OSError, l.accept) + + self._check_pickler_hits(reducer) + + +class CustomContext(multiprocessing.context.BaseContext): + _name = "custom" + _Process = multiprocessing.Process + + +class _TestCustomReducerWithContext(BaseTestCase): + """Test case for per-context reducers""" + + ALLOWED_TYPES = ('processes',) + + def setUp(self): + self.custom_ctx = CustomContext() + self.default_ctx = multiprocessing.get_context() + self.original_custom_reducer = self.custom_ctx.get_reducer() + self.original_default_reducer = self.default_ctx.get_reducer() + + def tearDown(self): + self.custom_ctx.set_reducer(self.original_custom_reducer) + self.default_ctx.set_reducer(self.original_default_reducer) + + def _check_pickler_hits(self, reducer, alternative_reducer): + custom_pickler = ( + reducer.get_pickler_class() is not reduction.AbstractPickler) + custom_unpickler = ( + reducer.get_unpickler_class() is not reduction.AbstractUnpickler) + + self.assertEqual(reducer.spy['dump'], int(custom_pickler)) + self.assertEqual(reducer.spy['load'], int(custom_unpickler)) + + self.assertEqual(alternative_reducer.spy['dump'], 0) + self.assertEqual(alternative_reducer.spy['load'], 0) + + @classmethod + def _put_and_get_in_queue(cls, queue): + queue.put("Something") + queue.get(timeout=support.SHORT_TIMEOUT) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_process_custom_reducer_over_custom_context(self): + sm = multiprocessing.get_start_method() + if sm == 'fork': + # The fork method does send data to new Process objects + self.skipTest('test not appropriate for {}'.format(sm)) + + default_reducer = SpyReducerWithPickler() + custom_reducer = SpyReducerWithPickler() + + self.default_ctx.set_reducer(default_reducer) + self.custom_ctx.set_reducer(custom_reducer) + p = self.custom_ctx.Process(target=id, args=(2, )) + p.start() + p.join() + self.assertEqual(p.exitcode, 0) + self._check_pickler_hits(custom_reducer, default_reducer) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_queue_custom_reducer_over_custom_context(self): + for reducer_cls in [ + SpyReducerWithPickler, + SpyReducerWithUnpickler, + SpyReducerWithPicklerAndUnpickler + ]: + with self.subTest(reducer_cls=reducer_cls): + default_reducer = reducer_cls() + custom_reducer = reducer_cls() + + self.default_ctx.set_reducer(default_reducer) + self.custom_ctx.set_reducer(custom_reducer) + + queue = self.custom_ctx.Queue() + p = self.custom_ctx.Process(target=self._put_and_get_in_queue, + args=(queue,)) + p.start() + element = queue.get(timeout=support.SHORT_TIMEOUT) + self.assertEqual(element, "Something") + queue.put("Other_Something") + close_queue(queue) + p.join(timeout=support.SHORT_TIMEOUT) + p.terminate() + self.assertEqual(p.exitcode, 0) + self.assertEqual(element, "Something") + self._check_pickler_hits(custom_reducer, default_reducer) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_listener_custom_reduction_custom_context(self): + for reducer_cls in [ + SpyReducerWithPickler, + SpyReducerWithUnpickler, + SpyReducerWithPicklerAndUnpickler + ]: + with self.subTest(reducer_cls=reducer_cls): + custom_reducer = reducer_cls() + default_reducer = reducer_cls() + self.default_ctx.set_reducer(default_reducer) + self.custom_ctx.set_reducer(custom_reducer) + + with self.custom_ctx.Listener() as l: + with self.custom_ctx.Client(l.address) as c: + with l.accept() as d: + c.send(1729) + self.assertEqual(d.recv(), 1729) + + if self.TYPE == 'processes': + self.assertRaises(OSError, l.accept) + self._check_pickler_hits(custom_reducer, default_reducer) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_queue_custom_reducer_over_default_context(self): + for reducer_cls in [ + SpyReducerWithPickler, + SpyReducerWithUnpickler, + SpyReducerWithPicklerAndUnpickler + ]: + with self.subTest(reducer_cls=reducer_cls): + default_reducer = reducer_cls() + custom_reducer = reducer_cls() + + self.default_ctx.set_reducer(default_reducer) + self.custom_ctx.set_reducer(custom_reducer) + + queue = self.default_ctx.Queue() + p = self.default_ctx.Process(target=self._put_and_get_in_queue, + args=(queue, )) + p.start() + element = queue.get(timeout=support.SHORT_TIMEOUT) + self.assertEqual(element, "Something") + queue.put("Other_Something") + p.join(timeout=support.SHORT_TIMEOUT) + self.assertEqual(p.exitcode, 0) + self.assertEqual(element, "Something") + close_queue(queue) + self._check_pickler_hits(default_reducer, custom_reducer) + + @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction") + def test_process_custom_reducer_over_default_context(self): + sm = multiprocessing.get_start_method() + if sm == 'fork': + # The fork method does send data to new Process objects + self.skipTest('test not appropriate for {}'.format(sm)) + + default_reducer = SpyReducerWithPickler() + custom_reducer = SpyReducerWithPickler() + self.default_ctx.set_reducer(default_reducer) + p = self.default_ctx.Process(target=id, args=(2, )) + p.start() + p.join() + self.assertEqual(p.exitcode, 0) + self._check_pickler_hits(default_reducer, custom_reducer) + +class _TestCustomReducerInvalidParameters(BaseTestCase): + def test_setting_invalid_reducer(self): + class BadReducer(): + def get_pickler_class(self): + return + + class BadReducerUnpickler(): + def get_unpickler_class(self): + return + + with self.assertRaises(TypeError): + multiprocessing.set_reducer(BadReducer()) + + with self.assertRaises(TypeError): + multiprocessing.set_reducer(BadReducerUnpickler()) + + + def test_setting_valid_reducer_that_returns_invalid_pickler(self): + + class PicklerClass: + pass + + class BadReducer(reduction.AbstractReducer): + def get_pickler_class(self): + return PicklerClass + + with self.assertRaises(TypeError): + multiprocessing.set_reducer(BadReducer()) + + # # Mixins # diff --git a/Misc/NEWS.d/next/Library/2018-12-09-19-20-12.bpo-28053.JiOeSY.rst b/Misc/NEWS.d/next/Library/2018-12-09-19-20-12.bpo-28053.JiOeSY.rst new file mode 100644 index 00000000000000..c672a205d135cf --- /dev/null +++ b/Misc/NEWS.d/next/Library/2018-12-09-19-20-12.bpo-28053.JiOeSY.rst @@ -0,0 +1 @@ +Fix the implementation of custom reducers for in `multiprocessing`.