JezK
Edit File: reduction.py
# # Module which deals with pickling of objects. # # multiprocessing/reduction.py # # Copyright (c) 2006-2008, R Oudkerk # Licensed to PSF under a Contributor Agreement. # from abc import ABCMeta import copyreg import functools import io import os import pickle import socket import sys from . import context __all__ = ['send_handle', 'recv_handle', 'ForkingPickler', 'register', 'dump'] HAVE_SEND_HANDLE = (sys.platform == 'win32' or (hasattr(socket, 'CMSG_LEN') and hasattr(socket, 'SCM_RIGHTS') and hasattr(socket.socket, 'sendmsg'))) # # Pickler subclass # class ForkingPickler(pickle.Pickler): '''Pickler subclass used by multiprocessing.''' _extra_reducers = {} _copyreg_dispatch_table = copyreg.dispatch_table def __init__(self, *args): super().__init__(*args) self.dispatch_table = self._copyreg_dispatch_table.copy() self.dispatch_table.update(self._extra_reducers) @classmethod def register(cls, type, reduce): '''Register a reduce function for a type.''' cls._extra_reducers[type] = reduce @classmethod def dumps(cls, obj, protocol=None): buf = io.BytesIO() cls(buf, protocol).dump(obj) return buf.getbuffer() loads = pickle.loads register = ForkingPickler.register def dump(obj, file, protocol=None): '''Replacement for pickle.dump() using ForkingPickler.''' ForkingPickler(file, protocol).dump(obj) # # Platform specific definitions # if sys.platform == 'win32': # Windows __all__ += ['DupHandle', 'duplicate', 'steal_handle'] import _winapi def duplicate(handle, target_process=None, inheritable=False, *, source_process=None): '''Duplicate a handle. (target_process is a handle not a pid!)''' current_process = _winapi.GetCurrentProcess() if source_process is None: source_process = current_process if target_process is None: target_process = current_process return _winapi.DuplicateHandle( source_process, handle, target_process, 0, inheritable, _winapi.DUPLICATE_SAME_ACCESS) def steal_handle(source_pid, handle): '''Steal a handle from process identified by source_pid.''' source_process_handle = _winapi.OpenProcess( _winapi.PROCESS_DUP_HANDLE, False, source_pid) try: return _winapi.DuplicateHandle( source_process_handle, handle, _winapi.GetCurrentProcess(), 0, False, _winapi.DUPLICATE_SAME_ACCESS | _winapi.DUPLICATE_CLOSE_SOURCE) finally: _winapi.CloseHandle(source_process_handle) def send_handle(conn, handle, destination_pid): '''Send a handle over a local connection.''' dh = DupHandle(handle, _winapi.DUPLICATE_SAME_ACCESS, destination_pid) conn.send(dh) def recv_handle(conn): '''Receive a handle over a local connection.''' return conn.recv().detach() class DupHandle(object): '''Picklable wrapper for a handle.''' def __init__(self, handle, access, pid=None): if pid is None: # We just duplicate the handle in the current process and # let the receiving process steal the handle. pid = os.getpid() proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False, pid) try: self._handle = _winapi.DuplicateHandle( _winapi.GetCurrentProcess(), handle, proc, access, False, 0) finally: _winapi.CloseHandle(proc) self._access = access self._pid = pid def detach(self): '''Get the handle. This should only be called once.''' # retrieve handle from process which currently owns it if self._pid == os.getpid(): # The handle has already been duplicated for this process. return self._handle # We must steal the handle from the process whose pid is self._pid. proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False, self._pid) try: return _winapi.DuplicateHandle( proc, self._handle, _winapi.GetCurrentProcess(), self._access, False, _winapi.DUPLICATE_CLOSE_SOURCE) finally: _winapi.CloseHandle(proc) else: # Unix __all__ += ['DupFd', 'sendfds', 'recvfds'] import array # On MacOSX we should acknowledge receipt of fds -- see Issue14669 ACKNOWLEDGE = sys.platform == 'darwin' def sendfds(sock, fds): '''Send an array of fds over an AF_UNIX socket.''' fds = array.array('i', fds) msg = bytes([len(fds) % 256]) sock.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)]) if ACKNOWLEDGE and sock.recv(1) != b'A': raise RuntimeError('did not receive acknowledgement of fd') def recvfds(sock, size): '''Receive an array of fds over an AF_UNIX socket.''' a = array.array('i') bytes_size = a.itemsize * size msg, ancdata, flags, addr = sock.recvmsg(1, socket.CMSG_SPACE(bytes_size)) if not msg and not ancdata: raise EOFError try: if ACKNOWLEDGE: sock.send(b'A') if len(ancdata) != 1: raise RuntimeError('received %d items of ancdata' % len(ancdata)) cmsg_level, cmsg_type, cmsg_data = ancdata[0] if (cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS): if len(cmsg_data) % a.itemsize != 0: raise ValueError a.frombytes(cmsg_data) if len(a) % 256 != msg[0]: raise AssertionError( "Len is {0:n} but msg[0] is {1!r}".format( len(a), msg[0])) return list(a) except (ValueError, IndexError): pass raise RuntimeError('Invalid data received') def send_handle(conn, handle, destination_pid): '''Send a handle over a local connection.''' with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s: sendfds(s, [handle]) def recv_handle(conn): '''Receive a handle over a local connection.''' with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s: return recvfds(s, 1)[0] def DupFd(fd): '''Return a wrapper for an fd.''' popen_obj = context.get_spawning_popen() if popen_obj is not None: return popen_obj.DupFd(popen_obj.duplicate_for_child(fd)) elif HAVE_SEND_HANDLE: from . import resource_sharer return resource_sharer.DupFd(fd) else: raise ValueError('SCM_RIGHTS appears not to be available') # # Try making some callable types 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 register(type(_C().f), _reduce_method) 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) 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) # # Make sockets picklable # if sys.platform == 'win32': def _reduce_socket(s): from .resource_sharer import DupSocket return _rebuild_socket, (DupSocket(s),) def _rebuild_socket(ds): return ds.detach() register(socket.socket, _reduce_socket) else: def _reduce_socket(s): df = DupFd(s.fileno()) return _rebuild_socket, (df, s.family, s.type, s.proto) def _rebuild_socket(df, family, type, proto): fd = df.detach() return socket.socket(family, type, proto, fileno=fd) register(socket.socket, _reduce_socket) class AbstractReducer(metaclass=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 DupHandle = DupHandle else: sendfds = sendfds recvfds = recvfds DupFd = DupFd _reduce_method = _reduce_method _reduce_method_descriptor = _reduce_method_descriptor _rebuild_partial = _rebuild_partial _reduce_socket = _reduce_socket _rebuild_socket = _rebuild_socket def __init__(self, *args): register(type(_C().f), _reduce_method) register(type(list.append), _reduce_method_descriptor) register(type(int.__add__), _reduce_method_descriptor) register(functools.partial, _reduce_partial) register(socket.socket, _reduce_socket)