You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

186 lines
6.0 KiB

4 years ago
  1. """serialization utilities for apply messages"""
  2. # Copyright (c) IPython Development Team.
  3. # Distributed under the terms of the Modified BSD License.
  4. import warnings
  5. warnings.warn("ipykernel.serialize is deprecated. It has moved to ipyparallel.serialize", DeprecationWarning)
  6. try:
  7. import cPickle
  8. pickle = cPickle
  9. except:
  10. cPickle = None
  11. import pickle
  12. from itertools import chain
  13. from ipython_genutils.py3compat import PY3, buffer_to_bytes_py2
  14. from ipykernel.pickleutil import (
  15. can, uncan, can_sequence, uncan_sequence, CannedObject,
  16. istype, sequence_types, PICKLE_PROTOCOL,
  17. )
  18. from jupyter_client.session import MAX_ITEMS, MAX_BYTES
  19. if PY3:
  20. buffer = memoryview
  21. #-----------------------------------------------------------------------------
  22. # Serialization Functions
  23. #-----------------------------------------------------------------------------
  24. def _extract_buffers(obj, threshold=MAX_BYTES):
  25. """extract buffers larger than a certain threshold"""
  26. buffers = []
  27. if isinstance(obj, CannedObject) and obj.buffers:
  28. for i,buf in enumerate(obj.buffers):
  29. if len(buf) > threshold:
  30. # buffer larger than threshold, prevent pickling
  31. obj.buffers[i] = None
  32. buffers.append(buf)
  33. # buffer too small for separate send, coerce to bytes
  34. # because pickling buffer objects just results in broken pointers
  35. elif isinstance(buf, memoryview):
  36. obj.buffers[i] = buf.tobytes()
  37. elif isinstance(buf, buffer):
  38. obj.buffers[i] = bytes(buf)
  39. return buffers
  40. def _restore_buffers(obj, buffers):
  41. """restore buffers extracted by """
  42. if isinstance(obj, CannedObject) and obj.buffers:
  43. for i,buf in enumerate(obj.buffers):
  44. if buf is None:
  45. obj.buffers[i] = buffers.pop(0)
  46. def serialize_object(obj, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS):
  47. """Serialize an object into a list of sendable buffers.
  48. Parameters
  49. ----------
  50. obj : object
  51. The object to be serialized
  52. buffer_threshold : int
  53. The threshold (in bytes) for pulling out data buffers
  54. to avoid pickling them.
  55. item_threshold : int
  56. The maximum number of items over which canning will iterate.
  57. Containers (lists, dicts) larger than this will be pickled without
  58. introspection.
  59. Returns
  60. -------
  61. [bufs] : list of buffers representing the serialized object.
  62. """
  63. buffers = []
  64. if istype(obj, sequence_types) and len(obj) < item_threshold:
  65. cobj = can_sequence(obj)
  66. for c in cobj:
  67. buffers.extend(_extract_buffers(c, buffer_threshold))
  68. elif istype(obj, dict) and len(obj) < item_threshold:
  69. cobj = {}
  70. for k in sorted(obj):
  71. c = can(obj[k])
  72. buffers.extend(_extract_buffers(c, buffer_threshold))
  73. cobj[k] = c
  74. else:
  75. cobj = can(obj)
  76. buffers.extend(_extract_buffers(cobj, buffer_threshold))
  77. buffers.insert(0, pickle.dumps(cobj, PICKLE_PROTOCOL))
  78. return buffers
  79. def deserialize_object(buffers, g=None):
  80. """reconstruct an object serialized by serialize_object from data buffers.
  81. Parameters
  82. ----------
  83. bufs : list of buffers/bytes
  84. g : globals to be used when uncanning
  85. Returns
  86. -------
  87. (newobj, bufs) : unpacked object, and the list of remaining unused buffers.
  88. """
  89. bufs = list(buffers)
  90. pobj = buffer_to_bytes_py2(bufs.pop(0))
  91. canned = pickle.loads(pobj)
  92. if istype(canned, sequence_types) and len(canned) < MAX_ITEMS:
  93. for c in canned:
  94. _restore_buffers(c, bufs)
  95. newobj = uncan_sequence(canned, g)
  96. elif istype(canned, dict) and len(canned) < MAX_ITEMS:
  97. newobj = {}
  98. for k in sorted(canned):
  99. c = canned[k]
  100. _restore_buffers(c, bufs)
  101. newobj[k] = uncan(c, g)
  102. else:
  103. _restore_buffers(canned, bufs)
  104. newobj = uncan(canned, g)
  105. return newobj, bufs
  106. def pack_apply_message(f, args, kwargs, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS):
  107. """pack up a function, args, and kwargs to be sent over the wire
  108. Each element of args/kwargs will be canned for special treatment,
  109. but inspection will not go any deeper than that.
  110. Any object whose data is larger than `threshold` will not have their data copied
  111. (only numpy arrays and bytes/buffers support zero-copy)
  112. Message will be a list of bytes/buffers of the format:
  113. [ cf, pinfo, <arg_bufs>, <kwarg_bufs> ]
  114. With length at least two + len(args) + len(kwargs)
  115. """
  116. arg_bufs = list(chain.from_iterable(
  117. serialize_object(arg, buffer_threshold, item_threshold) for arg in args))
  118. kw_keys = sorted(kwargs.keys())
  119. kwarg_bufs = list(chain.from_iterable(
  120. serialize_object(kwargs[key], buffer_threshold, item_threshold) for key in kw_keys))
  121. info = dict(nargs=len(args), narg_bufs=len(arg_bufs), kw_keys=kw_keys)
  122. msg = [pickle.dumps(can(f), PICKLE_PROTOCOL)]
  123. msg.append(pickle.dumps(info, PICKLE_PROTOCOL))
  124. msg.extend(arg_bufs)
  125. msg.extend(kwarg_bufs)
  126. return msg
  127. def unpack_apply_message(bufs, g=None, copy=True):
  128. """unpack f,args,kwargs from buffers packed by pack_apply_message()
  129. Returns: original f,args,kwargs"""
  130. bufs = list(bufs) # allow us to pop
  131. assert len(bufs) >= 2, "not enough buffers!"
  132. pf = buffer_to_bytes_py2(bufs.pop(0))
  133. f = uncan(pickle.loads(pf), g)
  134. pinfo = buffer_to_bytes_py2(bufs.pop(0))
  135. info = pickle.loads(pinfo)
  136. arg_bufs, kwarg_bufs = bufs[:info['narg_bufs']], bufs[info['narg_bufs']:]
  137. args = []
  138. for i in range(info['nargs']):
  139. arg, arg_bufs = deserialize_object(arg_bufs, g)
  140. args.append(arg)
  141. args = tuple(args)
  142. assert not arg_bufs, "Shouldn't be any arg bufs left over"
  143. kwargs = {}
  144. for key in info['kw_keys']:
  145. kwarg, kwarg_bufs = deserialize_object(kwarg_bufs, g)
  146. kwargs[key] = kwarg
  147. assert not kwarg_bufs, "Shouldn't be any kwarg bufs left over"
  148. return f,args,kwargs