reduction.py 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277
  1. #
  2. # Module which deals with pickling of objects.
  3. #
  4. # multiprocessing/reduction.py
  5. #
  6. # Copyright (c) 2006-2008, R Oudkerk
  7. # Licensed to PSF under a Contributor Agreement.
  8. #
  9. from abc import ABCMeta
  10. import copyreg
  11. import functools
  12. import io
  13. import os
  14. import pickle
  15. import socket
  16. import sys
  17. from . import context
  18. __all__ = ['send_handle', 'recv_handle', 'ForkingPickler', 'register', 'dump']
  19. HAVE_SEND_HANDLE = (sys.platform == 'win32' or
  20. (hasattr(socket, 'CMSG_LEN') and
  21. hasattr(socket, 'SCM_RIGHTS') and
  22. hasattr(socket.socket, 'sendmsg')))
  23. #
  24. # Pickler subclass
  25. #
  26. class ForkingPickler(pickle.Pickler):
  27. '''Pickler subclass used by multiprocessing.'''
  28. _extra_reducers = {}
  29. _copyreg_dispatch_table = copyreg.dispatch_table
  30. def __init__(self, *args):
  31. super().__init__(*args)
  32. self.dispatch_table = self._copyreg_dispatch_table.copy()
  33. self.dispatch_table.update(self._extra_reducers)
  34. @classmethod
  35. def register(cls, type, reduce):
  36. '''Register a reduce function for a type.'''
  37. cls._extra_reducers[type] = reduce
  38. @classmethod
  39. def dumps(cls, obj, protocol=None):
  40. buf = io.BytesIO()
  41. cls(buf, protocol).dump(obj)
  42. return buf.getbuffer()
  43. loads = pickle.loads
  44. register = ForkingPickler.register
  45. def dump(obj, file, protocol=None):
  46. '''Replacement for pickle.dump() using ForkingPickler.'''
  47. ForkingPickler(file, protocol).dump(obj)
  48. #
  49. # Platform specific definitions
  50. #
  51. if sys.platform == 'win32':
  52. # Windows
  53. __all__ += ['DupHandle', 'duplicate', 'steal_handle']
  54. import _winapi
  55. def duplicate(handle, target_process=None, inheritable=False):
  56. '''Duplicate a handle. (target_process is a handle not a pid!)'''
  57. if target_process is None:
  58. target_process = _winapi.GetCurrentProcess()
  59. return _winapi.DuplicateHandle(
  60. _winapi.GetCurrentProcess(), handle, target_process,
  61. 0, inheritable, _winapi.DUPLICATE_SAME_ACCESS)
  62. def steal_handle(source_pid, handle):
  63. '''Steal a handle from process identified by source_pid.'''
  64. source_process_handle = _winapi.OpenProcess(
  65. _winapi.PROCESS_DUP_HANDLE, False, source_pid)
  66. try:
  67. return _winapi.DuplicateHandle(
  68. source_process_handle, handle,
  69. _winapi.GetCurrentProcess(), 0, False,
  70. _winapi.DUPLICATE_SAME_ACCESS | _winapi.DUPLICATE_CLOSE_SOURCE)
  71. finally:
  72. _winapi.CloseHandle(source_process_handle)
  73. def send_handle(conn, handle, destination_pid):
  74. '''Send a handle over a local connection.'''
  75. dh = DupHandle(handle, _winapi.DUPLICATE_SAME_ACCESS, destination_pid)
  76. conn.send(dh)
  77. def recv_handle(conn):
  78. '''Receive a handle over a local connection.'''
  79. return conn.recv().detach()
  80. class DupHandle(object):
  81. '''Picklable wrapper for a handle.'''
  82. def __init__(self, handle, access, pid=None):
  83. if pid is None:
  84. # We just duplicate the handle in the current process and
  85. # let the receiving process steal the handle.
  86. pid = os.getpid()
  87. proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False, pid)
  88. try:
  89. self._handle = _winapi.DuplicateHandle(
  90. _winapi.GetCurrentProcess(),
  91. handle, proc, access, False, 0)
  92. finally:
  93. _winapi.CloseHandle(proc)
  94. self._access = access
  95. self._pid = pid
  96. def detach(self):
  97. '''Get the handle. This should only be called once.'''
  98. # retrieve handle from process which currently owns it
  99. if self._pid == os.getpid():
  100. # The handle has already been duplicated for this process.
  101. return self._handle
  102. # We must steal the handle from the process whose pid is self._pid.
  103. proc = _winapi.OpenProcess(_winapi.PROCESS_DUP_HANDLE, False,
  104. self._pid)
  105. try:
  106. return _winapi.DuplicateHandle(
  107. proc, self._handle, _winapi.GetCurrentProcess(),
  108. self._access, False, _winapi.DUPLICATE_CLOSE_SOURCE)
  109. finally:
  110. _winapi.CloseHandle(proc)
  111. else:
  112. # Unix
  113. __all__ += ['DupFd', 'sendfds', 'recvfds']
  114. import array
  115. # On MacOSX we should acknowledge receipt of fds -- see Issue14669
  116. ACKNOWLEDGE = sys.platform == 'darwin'
  117. def sendfds(sock, fds):
  118. '''Send an array of fds over an AF_UNIX socket.'''
  119. fds = array.array('i', fds)
  120. msg = bytes([len(fds) % 256])
  121. sock.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)])
  122. if ACKNOWLEDGE and sock.recv(1) != b'A':
  123. raise RuntimeError('did not receive acknowledgement of fd')
  124. def recvfds(sock, size):
  125. '''Receive an array of fds over an AF_UNIX socket.'''
  126. a = array.array('i')
  127. bytes_size = a.itemsize * size
  128. msg, ancdata, flags, addr = sock.recvmsg(1, socket.CMSG_SPACE(bytes_size))
  129. if not msg and not ancdata:
  130. raise EOFError
  131. try:
  132. if ACKNOWLEDGE:
  133. sock.send(b'A')
  134. if len(ancdata) != 1:
  135. raise RuntimeError('received %d items of ancdata' %
  136. len(ancdata))
  137. cmsg_level, cmsg_type, cmsg_data = ancdata[0]
  138. if (cmsg_level == socket.SOL_SOCKET and
  139. cmsg_type == socket.SCM_RIGHTS):
  140. if len(cmsg_data) % a.itemsize != 0:
  141. raise ValueError
  142. a.frombytes(cmsg_data)
  143. if len(a) % 256 != msg[0]:
  144. raise AssertionError(
  145. "Len is {0:n} but msg[0] is {1!r}".format(
  146. len(a), msg[0]))
  147. return list(a)
  148. except (ValueError, IndexError):
  149. pass
  150. raise RuntimeError('Invalid data received')
  151. def send_handle(conn, handle, destination_pid):
  152. '''Send a handle over a local connection.'''
  153. with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
  154. sendfds(s, [handle])
  155. def recv_handle(conn):
  156. '''Receive a handle over a local connection.'''
  157. with socket.fromfd(conn.fileno(), socket.AF_UNIX, socket.SOCK_STREAM) as s:
  158. return recvfds(s, 1)[0]
  159. def DupFd(fd):
  160. '''Return a wrapper for an fd.'''
  161. popen_obj = context.get_spawning_popen()
  162. if popen_obj is not None:
  163. return popen_obj.DupFd(popen_obj.duplicate_for_child(fd))
  164. elif HAVE_SEND_HANDLE:
  165. from . import resource_sharer
  166. return resource_sharer.DupFd(fd)
  167. else:
  168. raise ValueError('SCM_RIGHTS appears not to be available')
  169. #
  170. # Try making some callable types picklable
  171. #
  172. def _reduce_method(m):
  173. if m.__self__ is None:
  174. return getattr, (m.__class__, m.__func__.__name__)
  175. else:
  176. return getattr, (m.__self__, m.__func__.__name__)
  177. class _C:
  178. def f(self):
  179. pass
  180. register(type(_C().f), _reduce_method)
  181. def _reduce_method_descriptor(m):
  182. return getattr, (m.__objclass__, m.__name__)
  183. register(type(list.append), _reduce_method_descriptor)
  184. register(type(int.__add__), _reduce_method_descriptor)
  185. def _reduce_partial(p):
  186. return _rebuild_partial, (p.func, p.args, p.keywords or {})
  187. def _rebuild_partial(func, args, keywords):
  188. return functools.partial(func, *args, **keywords)
  189. register(functools.partial, _reduce_partial)
  190. #
  191. # Make sockets picklable
  192. #
  193. if sys.platform == 'win32':
  194. def _reduce_socket(s):
  195. from .resource_sharer import DupSocket
  196. return _rebuild_socket, (DupSocket(s),)
  197. def _rebuild_socket(ds):
  198. return ds.detach()
  199. register(socket.socket, _reduce_socket)
  200. else:
  201. def _reduce_socket(s):
  202. df = DupFd(s.fileno())
  203. return _rebuild_socket, (df, s.family, s.type, s.proto)
  204. def _rebuild_socket(df, family, type, proto):
  205. fd = df.detach()
  206. return socket.socket(family, type, proto, fileno=fd)
  207. register(socket.socket, _reduce_socket)
  208. class AbstractReducer(metaclass=ABCMeta):
  209. '''Abstract base class for use in implementing a Reduction class
  210. suitable for use in replacing the standard reduction mechanism
  211. used in multiprocessing.'''
  212. ForkingPickler = ForkingPickler
  213. register = register
  214. dump = dump
  215. send_handle = send_handle
  216. recv_handle = recv_handle
  217. if sys.platform == 'win32':
  218. steal_handle = steal_handle
  219. duplicate = duplicate
  220. DupHandle = DupHandle
  221. else:
  222. sendfds = sendfds
  223. recvfds = recvfds
  224. DupFd = DupFd
  225. _reduce_method = _reduce_method
  226. _reduce_method_descriptor = _reduce_method_descriptor
  227. _rebuild_partial = _rebuild_partial
  228. _reduce_socket = _reduce_socket
  229. _rebuild_socket = _rebuild_socket
  230. def __init__(self, *args):
  231. register(type(_C().f), _reduce_method)
  232. register(type(list.append), _reduce_method_descriptor)
  233. register(type(int.__add__), _reduce_method_descriptor)
  234. register(functools.partial, _reduce_partial)
  235. register(socket.socket, _reduce_socket)