| """serialization utilities for apply messages""" |
|
|
| |
| |
|
|
| import pickle |
| import warnings |
| from itertools import chain |
|
|
| try: |
| |
| from ipyparallel.serialize.canning import ( |
| CannedObject, |
| can, |
| can_sequence, |
| istype, |
| sequence_types, |
| uncan, |
| uncan_sequence, |
| ) |
| from ipyparallel.serialize.serialize import PICKLE_PROTOCOL |
| except ImportError: |
| |
| from ipykernel.pickleutil import ( |
| PICKLE_PROTOCOL, |
| CannedObject, |
| can, |
| can_sequence, |
| istype, |
| sequence_types, |
| uncan, |
| uncan_sequence, |
| ) |
|
|
| from jupyter_client.session import MAX_BYTES, MAX_ITEMS |
|
|
| warnings.warn( |
| "ipykernel.serialize is deprecated. It has moved to ipyparallel.serialize", |
| DeprecationWarning, |
| stacklevel=2, |
| ) |
|
|
| |
| |
| |
|
|
|
|
| def _extract_buffers(obj, threshold=MAX_BYTES): |
| """extract buffers larger than a certain threshold""" |
| buffers = [] |
| if isinstance(obj, CannedObject) and obj.buffers: |
| for i, buf in enumerate(obj.buffers): |
| if len(buf) > threshold: |
| |
| obj.buffers[i] = None |
| buffers.append(buf) |
| |
| |
| elif isinstance(buf, memoryview): |
| obj.buffers[i] = buf.tobytes() |
| return buffers |
|
|
|
|
| def _restore_buffers(obj, buffers): |
| """restore buffers extracted by""" |
| if isinstance(obj, CannedObject) and obj.buffers: |
| for i, buf in enumerate(obj.buffers): |
| if buf is None: |
| obj.buffers[i] = buffers.pop(0) |
|
|
|
|
| def serialize_object(obj, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): |
| """Serialize an object into a list of sendable buffers. |
| |
| Parameters |
| ---------- |
| obj : object |
| The object to be serialized |
| buffer_threshold : int |
| The threshold (in bytes) for pulling out data buffers |
| to avoid pickling them. |
| item_threshold : int |
| The maximum number of items over which canning will iterate. |
| Containers (lists, dicts) larger than this will be pickled without |
| introspection. |
| |
| Returns |
| ------- |
| [bufs] : list of buffers representing the serialized object. |
| """ |
| buffers = [] |
| if istype(obj, sequence_types) and len(obj) < item_threshold: |
| cobj = can_sequence(obj) |
| for c in cobj: |
| buffers.extend(_extract_buffers(c, buffer_threshold)) |
| elif istype(obj, dict) and len(obj) < item_threshold: |
| cobj = {} |
| for k in sorted(obj): |
| c = can(obj[k]) |
| buffers.extend(_extract_buffers(c, buffer_threshold)) |
| cobj[k] = c |
| else: |
| cobj = can(obj) |
| buffers.extend(_extract_buffers(cobj, buffer_threshold)) |
|
|
| buffers.insert(0, pickle.dumps(cobj, PICKLE_PROTOCOL)) |
| return buffers |
|
|
|
|
| def deserialize_object(buffers, g=None): |
| """reconstruct an object serialized by serialize_object from data buffers. |
| |
| Parameters |
| ---------- |
| buffers : list of buffers/bytes |
| g : globals to be used when uncanning |
| |
| Returns |
| ------- |
| (newobj, bufs) : unpacked object, and the list of remaining unused buffers. |
| """ |
| bufs = list(buffers) |
| pobj = bufs.pop(0) |
| canned = pickle.loads(pobj) |
| if istype(canned, sequence_types) and len(canned) < MAX_ITEMS: |
| for c in canned: |
| _restore_buffers(c, bufs) |
| newobj = uncan_sequence(canned, g) |
| elif istype(canned, dict) and len(canned) < MAX_ITEMS: |
| newobj = {} |
| for k in sorted(canned): |
| c = canned[k] |
| _restore_buffers(c, bufs) |
| newobj[k] = uncan(c, g) |
| else: |
| _restore_buffers(canned, bufs) |
| newobj = uncan(canned, g) |
|
|
| return newobj, bufs |
|
|
|
|
| def pack_apply_message(f, args, kwargs, buffer_threshold=MAX_BYTES, item_threshold=MAX_ITEMS): |
| """pack up a function, args, and kwargs to be sent over the wire |
| |
| Each element of args/kwargs will be canned for special treatment, |
| but inspection will not go any deeper than that. |
| |
| Any object whose data is larger than `threshold` will not have their data copied |
| (only numpy arrays and bytes/buffers support zero-copy) |
| |
| Message will be a list of bytes/buffers of the format: |
| |
| [ cf, pinfo, <arg_bufs>, <kwarg_bufs> ] |
| |
| With length at least two + len(args) + len(kwargs) |
| """ |
|
|
| arg_bufs = list( |
| chain.from_iterable(serialize_object(arg, buffer_threshold, item_threshold) for arg in args) |
| ) |
|
|
| kw_keys = sorted(kwargs.keys()) |
| kwarg_bufs = list( |
| chain.from_iterable( |
| serialize_object(kwargs[key], buffer_threshold, item_threshold) for key in kw_keys |
| ) |
| ) |
|
|
| info = dict(nargs=len(args), narg_bufs=len(arg_bufs), kw_keys=kw_keys) |
|
|
| msg = [pickle.dumps(can(f), PICKLE_PROTOCOL)] |
| msg.append(pickle.dumps(info, PICKLE_PROTOCOL)) |
| msg.extend(arg_bufs) |
| msg.extend(kwarg_bufs) |
|
|
| return msg |
|
|
|
|
| def unpack_apply_message(bufs, g=None, copy=True): |
| """unpack f,args,kwargs from buffers packed by pack_apply_message() |
| Returns: original f,args,kwargs""" |
| bufs = list(bufs) |
| assert len(bufs) >= 2, "not enough buffers!" |
| pf = bufs.pop(0) |
| f = uncan(pickle.loads(pf), g) |
| pinfo = bufs.pop(0) |
| info = pickle.loads(pinfo) |
| arg_bufs, kwarg_bufs = bufs[: info["narg_bufs"]], bufs[info["narg_bufs"] :] |
|
|
| args_list = [] |
| for _ in range(info["nargs"]): |
| arg, arg_bufs = deserialize_object(arg_bufs, g) |
| args_list.append(arg) |
| args = tuple(args_list) |
| assert not arg_bufs, "Shouldn't be any arg bufs left over" |
|
|
| kwargs = {} |
| for key in info["kw_keys"]: |
| kwarg, kwarg_bufs = deserialize_object(kwarg_bufs, g) |
| kwargs[key] = kwarg |
| assert not kwarg_bufs, "Shouldn't be any kwarg bufs left over" |
|
|
| return f, args, kwargs |
|
|
