Merge pull request #2069 from minrk/betterserial

start improving serialization in parallel code

* newserialized is now totally unused
* never double-pickle anything
* canning in pickleutil is extensible via can_map dict
* Session has configurable zero-copy threshold

IPython/zmq/tests/test_serialize.py

@@ -0,0 +1,168 b''
+"""test serialization tools"""
+
+#-------------------------------------------------------------------------------
+#  Copyright (C) 2011  The IPython Development Team
+#
+#  Distributed under the terms of the BSD License.  The full license is in
+#  the file COPYING, distributed as part of this software.
+#-------------------------------------------------------------------------------
+
+#-------------------------------------------------------------------------------
+# Imports
+#-------------------------------------------------------------------------------
+
+import pickle
+
+import nose.tools as nt
+
+# from unittest import TestCaes
+from IPython.zmq.serialize import serialize_object, unserialize_object
+from IPython.testing import decorators as dec
+from IPython.utils.pickleutil import CannedArray
+
+#-------------------------------------------------------------------------------
+# Globals and Utilities
+#-------------------------------------------------------------------------------
+
+def roundtrip(obj):
+    """roundtrip an object through serialization"""
+    bufs = serialize_object(obj)
+    obj2, remainder = unserialize_object(bufs)
+    nt.assert_equals(remainder, [])
+    return obj2
+
+class C(object):
+    """dummy class for """
+    
+    def __init__(self, **kwargs):
+        for key,value in kwargs.iteritems():
+            setattr(self, key, value)
+
+SHAPES = ((100,), (1024,10), (10,8,6,5), (), (0,))
+DTYPES = ('uint8', 'float64', 'int32', [('g', 'float32')], '|S10')
+#-------------------------------------------------------------------------------
+# Tests
+#-------------------------------------------------------------------------------
+
[e-mail]
+def test_roundtrip_simple():
+    for obj in [
+        'hello',
+        dict(a='b', b=10),
+        [1,2,'hi'],
+        (b'123', 'hello'),
+    ]:
+        obj2 = roundtrip(obj)
+        yield nt.assert_equals(obj, obj2)
+
[e-mail]
+def test_roundtrip_nested():
+    for obj in [
+        dict(a=range(5), b={1:b'hello'}),
+        [range(5),[range(3),(1,[b'whoda'])]],
+    ]:
+        obj2 = roundtrip(obj)
+        yield nt.assert_equals(obj, obj2)
+
[e-mail]
+def test_roundtrip_buffered():
+    for obj in [
+        dict(a=b"x"*1025),
+        b"hello"*500,
+        [b"hello"*501, 1,2,3]
+    ]:
+        bufs = serialize_object(obj)
+        yield nt.assert_equals(len(bufs), 2)
+        obj2, remainder = unserialize_object(bufs)
+        yield nt.assert_equals(remainder, [])
+        yield nt.assert_equals(obj, obj2)
+
+def _scrub_nan(A):
+    """scrub nans out of empty arrays
+    
+    since nan != nan
+    """
+    import numpy
+    if A.dtype.fields and A.shape:
+        for field in A.dtype.fields.keys():
+            try:
+                A[field][numpy.isnan(A[field])] = 0
+            except (TypeError, NotImplementedError):
+                # e.g. str dtype
+                pass
+
[e-mail]
[e-mail]_without('numpy')
+def test_numpy():
+    import numpy
+    from numpy.testing.utils import assert_array_equal
+    for shape in SHAPES:
+        for dtype in DTYPES:
+            A = numpy.empty(shape, dtype=dtype)
+            _scrub_nan(A)
+            bufs = serialize_object(A)
+            B, r = unserialize_object(bufs)
+            yield nt.assert_equals(r, [])
+            yield nt.assert_equals(A.shape, B.shape)
+            yield nt.assert_equals(A.dtype, B.dtype)
+            yield assert_array_equal(A,B)
+
[e-mail]
[e-mail]_without('numpy')
+def test_recarray():
+    import numpy
+    from numpy.testing.utils import assert_array_equal
+    for shape in SHAPES:
+        for dtype in [
+            [('f', float), ('s', '|S10')],
+            [('n', int), ('s', '|S1'), ('u', 'uint32')],
+        ]:
+            A = numpy.empty(shape, dtype=dtype)
+            _scrub_nan(A)
+            
+            bufs = serialize_object(A)
+            B, r = unserialize_object(bufs)
+            yield nt.assert_equals(r, [])
+            yield nt.assert_equals(A.shape, B.shape)
+            yield nt.assert_equals(A.dtype, B.dtype)
+            yield assert_array_equal(A,B)
+
[e-mail]
[e-mail]_without('numpy')
+def test_numpy_in_seq():
+    import numpy
+    from numpy.testing.utils import assert_array_equal
+    for shape in SHAPES:
+        for dtype in DTYPES:
+            A = numpy.empty(shape, dtype=dtype)
+            _scrub_nan(A)
+            bufs = serialize_object((A,1,2,b'hello'))
+            canned = pickle.loads(bufs[0])
+            yield nt.assert_true(canned[0], CannedArray)
+            tup, r = unserialize_object(bufs)
+            B = tup[0]
+            yield nt.assert_equals(r, [])
+            yield nt.assert_equals(A.shape, B.shape)
+            yield nt.assert_equals(A.dtype, B.dtype)
+            yield assert_array_equal(A,B)
+
[e-mail]
[e-mail]_without('numpy')
+def test_numpy_in_dict():
+    import numpy
+    from numpy.testing.utils import assert_array_equal
+    for shape in SHAPES:
+        for dtype in DTYPES:
+            A = numpy.empty(shape, dtype=dtype)
+            _scrub_nan(A)
+            bufs = serialize_object(dict(a=A,b=1,c=range(20)))
+            canned = pickle.loads(bufs[0])
+            yield nt.assert_true(canned['a'], CannedArray)
+            d, r = unserialize_object(bufs)
+            B = d['a']
+            yield nt.assert_equals(r, [])
+            yield nt.assert_equals(A.shape, B.shape)
+            yield nt.assert_equals(A.dtype, B.dtype)
+            yield assert_array_equal(A,B)
+    
+

IPython/parallel/client/client.py

@@ -1209,7 +1209,10 b' class Client(HasTraits):'
         if not isinstance(metadata, dict):
             raise TypeError("metadata must be dict, not %s"%type(metadata))
 
-        bufs = util.pack_apply_message(f,args,kwargs)
+        bufs = util.pack_apply_message(f, args, kwargs,
+            buffer_threshold=self.session.buffer_threshold,
+            item_threshold=self.session.item_threshold,
+        )
 
         msg = self.session.send(socket, "apply_request", buffers=bufs, ident=ident,
                             metadata=metadata, track=track)

IPython/parallel/tests/test_view.py

@@ -239,7 +239,7 b' class TestView(ClusterTestCase, ParametricTestCase):'
         from numpy.testing.utils import assert_array_equal, assert_array_almost_equal
         view = self.client[:]
         a = numpy.arange(64)
-        view.scatter('a', a)
+        view.scatter('a', a, block=True)
         b = view.gather('a', block=True)
         assert_array_equal(b, a)
     
@@ -325,7 +325,7 b' class TestView(ClusterTestCase, ParametricTestCase):'
         r = view.map_sync(lambda x:x, arr)
         self.assertEqual(r, list(arr))
     
-    def test_scatterGatherNonblocking(self):
+    def test_scatter_gather_nonblocking(self):
         data = range(16)
         view = self.client[:]
         view.scatter('a', data, block=False)

IPython/parallel/util.py

@@ -43,17 +43,11 b' from IPython.external.decorator import decorator'
 
 # IPython imports
 from IPython.config.application import Application
-from IPython.utils import py3compat
-from IPython.utils.pickleutil import can, uncan, canSequence, uncanSequence
-from IPython.utils.newserialized import serialize, unserialize
 from IPython.zmq.log import EnginePUBHandler
 from IPython.zmq.serialize import (
     unserialize_object, serialize_object, pack_apply_message, unpack_apply_message
 )
 
-if py3compat.PY3:
-    buffer = memoryview
-
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------

IPython/utils/pickleutil.py

@@ -16,10 +16,28 b' __docformat__ = "restructuredtext en"'
 #-------------------------------------------------------------------------------
 
 import copy
+import logging
 import sys
 from types import FunctionType
 
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
+
+try:
+    import numpy
+except:
+    numpy = None
+
 import codeutil
+import py3compat
+from importstring import import_item
+
+from IPython.config import Application
+
+if py3compat.PY3:
+    buffer = memoryview
 
 #-------------------------------------------------------------------------------
 # Classes
@@ -32,14 +50,16 b' class CannedObject(object):'
         self.obj = copy.copy(obj)
         for key in keys:
             setattr(self.obj, key, can(getattr(obj, key)))
+        
+        self.buffers = []
 
-
-    def getObject(self, g=None):
+    def get_object(self, g=None):
         if g is None:
-            g = globals()
+            g = {}
         for key in self.keys:
             setattr(self.obj, key, uncan(getattr(self.obj, key), g))
         return self.obj
+    
 
 class Reference(CannedObject):
     """object for wrapping a remote reference by name."""
@@ -47,13 +67,14 b' class Reference(CannedObject):'
         if not isinstance(name, basestring):
             raise TypeError("illegal name: %r"%name)
         self.name = name
+        self.buffers = []
 
     def __repr__(self):
         return "<Reference: %r>"%self.name
 
-    def getObject(self, g=None):
+    def get_object(self, g=None):
         if g is None:
-            g = globals()
+            g = {}
         
         return eval(self.name, g)
 
@@ -61,16 +82,17 b' class Reference(CannedObject):'
 class CannedFunction(CannedObject):
 
     def __init__(self, f):
-        self._checkType(f)
+        self._check_type(f)
         self.code = f.func_code
         self.defaults = f.func_defaults
         self.module = f.__module__ or '__main__'
         self.__name__ = f.__name__
+        self.buffers = []
 
-    def _checkType(self, obj):
+    def _check_type(self, obj):
         assert isinstance(obj, FunctionType), "Not a function type"
 
-    def getObject(self, g=None):
+    def get_object(self, g=None):
         # try to load function back into its module:
         if not self.module.startswith('__'):
             try:
@@ -81,30 +103,73 b' class CannedFunction(CannedObject):'
                 g = sys.modules[self.module].__dict__
 
         if g is None:
-            g = globals()
+            g = {}
         newFunc = FunctionType(self.code, g, self.__name__, self.defaults)
         return newFunc
 
+
+class CannedArray(CannedObject):
+    def __init__(self, obj):
+        self.shape = obj.shape
+        self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
+        if sum(obj.shape) == 0:
+            # just pickle it
+            self.buffers = [pickle.dumps(obj, -1)]
+        else:
+            # ensure contiguous
+            obj = numpy.ascontiguousarray(obj, dtype=None)
+            self.buffers = [buffer(obj)]
+    
+    def get_object(self, g=None):
+        data = self.buffers[0]
+        if sum(self.shape) == 0:
+            # no shape, we just pickled it
+            return pickle.loads(data)
+        else:
+            return numpy.frombuffer(data, dtype=self.dtype).reshape(self.shape)
+
+
+class CannedBytes(CannedObject):
+    wrap = bytes
+    def __init__(self, obj):
+        self.buffers = [obj]
+    
+    def get_object(self, g=None):
+        data = self.buffers[0]
+        return self.wrap(data)
+
+def CannedBuffer(CannedBytes):
+    wrap = buffer
+
 #-------------------------------------------------------------------------------
 # Functions
 #-------------------------------------------------------------------------------
 
-def can(obj):
-    # import here to prevent module-level circular imports
-    from IPython.parallel import dependent
-    if isinstance(obj, dependent):
-        keys = ('f','df')
-        return CannedObject(obj, keys=keys)
-    elif isinstance(obj, FunctionType):
-        return CannedFunction(obj)
-    elif isinstance(obj,dict):
-        return canDict(obj)
-    elif isinstance(obj, (list,tuple)):
-        return canSequence(obj)
+def _error(*args, **kwargs):
+    if Application.initialized():
+        logger = Application.instance().log
     else:
-        return obj
+        logger = logging.getLogger()
+        if not logger.handlers:
+            logging.basicConfig()
+    logger.error(*args, **kwargs)
 
-def canDict(obj):
+def can(obj):
+    """prepare an object for pickling"""
+    for cls,canner in can_map.iteritems():
+        if isinstance(cls, basestring):
+            try:
+                cls = import_item(cls)
+            except Exception:
+                _error("cannning class not importable: %r", cls, exc_info=True)
+                cls = None
+                continue
+        if isinstance(obj, cls):
+            return canner(obj)
+    return obj
+
+def can_dict(obj):
+    """can the *values* of a dict"""
     if isinstance(obj, dict):
         newobj = {}
         for k, v in obj.iteritems():
@@ -113,7 +178,8 b' def canDict(obj):'
     else:
         return obj
 
-def canSequence(obj):
+def can_sequence(obj):
+    """can the elements of a sequence"""
     if isinstance(obj, (list, tuple)):
         t = type(obj)
         return t([can(i) for i in obj])
@@ -121,16 +187,20 b' def canSequence(obj):'
         return obj
 
 def uncan(obj, g=None):
-    if isinstance(obj, CannedObject):
-        return obj.getObject(g)
-    elif isinstance(obj,dict):
-        return uncanDict(obj, g)
-    elif isinstance(obj, (list,tuple)):
-        return uncanSequence(obj, g)
-    else:
-        return obj
-
-def uncanDict(obj, g=None):
+    """invert canning"""
+    for cls,uncanner in uncan_map.iteritems():
+        if isinstance(cls, basestring):
+            try:
+                cls = import_item(cls)
+            except Exception:
+                _error("uncanning class not importable: %r", cls, exc_info=True)
+                cls = None
+                continue
+        if isinstance(obj, cls):
+            return uncanner(obj, g)
+    return obj
+
+def uncan_dict(obj, g=None):
     if isinstance(obj, dict):
         newobj = {}
         for k, v in obj.iteritems():
@@ -139,7 +209,7 b' def uncanDict(obj, g=None):'
     else:
         return obj
 
-def uncanSequence(obj, g=None):
+def uncan_sequence(obj, g=None):
     if isinstance(obj, (list, tuple)):
         t = type(obj)
         return t([uncan(i,g) for i in obj])
@@ -147,5 +217,21 b' def uncanSequence(obj, g=None):'
         return obj
 
 
-def rebindFunctionGlobals(f, glbls):
-    return FunctionType(f.func_code, glbls)
+#-------------------------------------------------------------------------------
+# API dictionary
+#-------------------------------------------------------------------------------
+
+# These dicts can be extended for custom serialization of new objects
+
+can_map = {
+    'IPython.parallel.dependent' : lambda obj: CannedObject(obj, keys=('f','df')),
+    'numpy.ndarray' : CannedArray,
+    FunctionType : CannedFunction,
+    bytes : CannedBytes,
+    buffer : CannedBuffer,
+}
+
+uncan_map = {
+    CannedObject : lambda obj, g: obj.get_object(g),
+}
+

IPython/zmq/ipkernel.py

@@ -570,8 +570,11 b' class Kernel(Configurable):'
                 for key in ns.iterkeys():
                     working.pop(key)
 
-            packed_result,buf = serialize_object(result)
-            result_buf = [packed_result]+buf
+            result_buf = serialize_object(result,
+                buffer_threshold=self.session.buffer_threshold,
+                item_threshold=self.session.item_threshold,
+            )
+        
         except:
             # invoke IPython traceback formatting
             shell.showtraceback()

IPython/zmq/serialize.py

@@ -32,8 +32,9 b' except:'
 
 # IPython imports
 from IPython.utils import py3compat
-from IPython.utils.pickleutil import can, uncan, canSequence, uncanSequence
-from IPython.utils.newserialized import serialize, unserialize
+from IPython.utils.pickleutil import (
+    can, uncan, can_sequence, uncan_sequence, CannedObject
+)
 
 if py3compat.PY3:
     buffer = memoryview
@@ -42,7 +43,33 b' if py3compat.PY3:'
 # Serialization Functions
 #-----------------------------------------------------------------------------
 
-def serialize_object(obj, threshold=64e-6):
+# default values for the thresholds:
+MAX_ITEMS = 64
+MAX_BYTES = 1024
+
+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:
+                # buffer larger than threshold, prevent pickling
+                obj.buffers[i] = None
+                buffers.append(buf)
+            elif isinstance(buf, buffer):
+                # buffer too small for separate send, coerce to bytes
+                # because pickling buffer objects just results in broken pointers
+                obj.buffers[i] = bytes(buf)
+    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
@@ -50,76 +77,83 b' def serialize_object(obj, threshold=64e-6):'
     
     obj : object
         The object to be serialized
-    threshold : float
-        The threshold for not double-pickling the content.
-        
+    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
     -------
-    ('pmd', [bufs]) :
-        where pmd is the pickled metadata wrapper,
-        bufs is a list of data buffers
+    [bufs] : list of buffers representing the serialized object.
     """
-    databuffers = []
-    if isinstance(obj, (list, tuple)):
-        clist = canSequence(obj)
-        slist = map(serialize, clist)
-        for s in slist:
-            if s.typeDescriptor in ('buffer', 'ndarray') or s.getDataSize() > threshold:
-                databuffers.append(s.getData())
-                s.data = None
-        return pickle.dumps(slist,-1), databuffers
-    elif isinstance(obj, dict):
-        sobj = {}
+    buffers = []
+    if isinstance(obj, (list, tuple)) and len(obj) < item_threshold:
+        cobj = can_sequence(obj)
+        for c in cobj:
+            buffers.extend(_extract_buffers(c, buffer_threshold))
+    elif isinstance(obj, dict) and len(obj) < item_threshold:
+        cobj = {}
         for k in sorted(obj.iterkeys()):
-            s = serialize(can(obj[k]))
-            if s.typeDescriptor in ('buffer', 'ndarray') or s.getDataSize() > threshold:
-                databuffers.append(s.getData())
-                s.data = None
-            sobj[k] = s
-        return pickle.dumps(sobj,-1),databuffers
+            c = can(obj[k])
+            buffers.extend(_extract_buffers(c, buffer_threshold))
+            cobj[k] = c
     else:
-        s = serialize(can(obj))
-        if s.typeDescriptor in ('buffer', 'ndarray') or s.getDataSize() > threshold:
-            databuffers.append(s.getData())
-            s.data = None
-        return pickle.dumps(s,-1),databuffers
-            
-        
-def unserialize_object(bufs):
-    """reconstruct an object serialized by serialize_object from data buffers."""
-    bufs = list(bufs)
-    sobj = pickle.loads(bufs.pop(0))
-    if isinstance(sobj, (list, tuple)):
-        for s in sobj:
-            if s.data is None:
-                s.data = bufs.pop(0)
-        return uncanSequence(map(unserialize, sobj)), bufs
-    elif isinstance(sobj, dict):
+        cobj = can(obj)
+        buffers.extend(_extract_buffers(cobj, buffer_threshold))
+
+    buffers.insert(0, pickle.dumps(cobj,-1))
+    return buffers
+
+def unserialize_object(buffers, g=None):
+    """reconstruct an object serialized by serialize_object from data buffers.
+    
+    Parameters
+    ----------
+    
+    bufs : 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)
+    canned = pickle.loads(bufs.pop(0))
+    if isinstance(canned, (list, tuple)) and len(canned) < MAX_ITEMS:
+        for c in canned:
+            _restore_buffers(c, bufs)
+        newobj = uncan_sequence(canned, g)
+    elif isinstance(canned, dict) and len(canned) < MAX_ITEMS:
         newobj = {}
-        for k in sorted(sobj.iterkeys()):
-            s = sobj[k]
-            if s.data is None:
-                s.data = bufs.pop(0)
-            newobj[k] = uncan(unserialize(s))
-        return newobj, bufs
+        for k in sorted(canned.iterkeys()):
+            c = canned[k]
+            _restore_buffers(c, bufs)
+            newobj[k] = uncan(c, g)
     else:
-        if sobj.data is None:
-            sobj.data = bufs.pop(0)
-        return uncan(unserialize(sobj)), bufs
+        _restore_buffers(canned, bufs)
+        newobj = uncan(canned, g)
+    
+    return newobj, bufs
 
-def pack_apply_message(f, args, kwargs, threshold=64e-6):
+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
+    
     as a series of buffers. Any object whose data is larger than `threshold`
-    will not have their data copied (currently only numpy arrays support zero-copy)"""
+    will not have their data copied (currently only numpy arrays support zero-copy)
+    """
     msg = [pickle.dumps(can(f),-1)]
     databuffers = [] # for large objects
-    sargs, bufs = serialize_object(args,threshold)
-    msg.append(sargs)
-    databuffers.extend(bufs)
-    skwargs, bufs = serialize_object(kwargs,threshold)
-    msg.append(skwargs)
-    databuffers.extend(bufs)
+    sargs = serialize_object(args, buffer_threshold, item_threshold)
+    msg.append(sargs[0])
+    databuffers.extend(sargs[1:])
+    skwargs = serialize_object(kwargs, buffer_threshold, item_threshold)
+    msg.append(skwargs[0])
+    databuffers.extend(skwargs[1:])
     msg.extend(databuffers)
     return msg
 
@@ -131,49 +165,16 b' def unpack_apply_message(bufs, g=None, copy=True):'
     if not copy:
         for i in range(3):
             bufs[i] = bufs[i].bytes
-    cf = pickle.loads(bufs.pop(0))
-    sargs = list(pickle.loads(bufs.pop(0)))
-    skwargs = dict(pickle.loads(bufs.pop(0)))
-    # print sargs, skwargs
-    f = uncan(cf, g)
-    for sa in sargs:
-        if sa.data is None:
-            m = bufs.pop(0)
-            if sa.getTypeDescriptor() in ('buffer', 'ndarray'):
-                # always use a buffer, until memoryviews get sorted out
-                sa.data = buffer(m)
-                # disable memoryview support
-                # if copy:
-                #     sa.data = buffer(m)
-                # else:
-                #     sa.data = m.buffer
-            else:
-                if copy:
-                    sa.data = m
-                else:
-                    sa.data = m.bytes
+    f = uncan(pickle.loads(bufs.pop(0)), g)
+    # sargs = bufs.pop(0)
+    # pop kwargs out, so first n-elements are args, serialized
+    skwargs = bufs.pop(1)
+    args, bufs = unserialize_object(bufs, g)
+    # put skwargs back in as the first element
+    bufs.insert(0, skwargs)
+    kwargs, bufs = unserialize_object(bufs, g)
     
-    args = uncanSequence(map(unserialize, sargs), g)
-    kwargs = {}
-    for k in sorted(skwargs.iterkeys()):
-        sa = skwargs[k]
-        if sa.data is None:
-            m = bufs.pop(0)
-            if sa.getTypeDescriptor() in ('buffer', 'ndarray'):
-                # always use a buffer, until memoryviews get sorted out
-                sa.data = buffer(m)
-                # disable memoryview support
-                # if copy:
-                #     sa.data = buffer(m)
-                # else:
-                #     sa.data = m.buffer
-            else:
-                if copy:
-                    sa.data = m
-                else:
-                    sa.data = m.bytes
-
-        kwargs[k] = uncan(unserialize(sa), g)
+    assert not bufs, "Shouldn't be any data left over"
     
     return f,args,kwargs
 

IPython/zmq/session.py

@@ -49,7 +49,8 b' from IPython.utils.importstring import import_item'
 from IPython.utils.jsonutil import extract_dates, squash_dates, date_default
 from IPython.utils.py3compat import str_to_bytes
 from IPython.utils.traitlets import (CBytes, Unicode, Bool, Any, Instance, Set,
-                                        DottedObjectName, CUnicode, Dict)
+                                        DottedObjectName, CUnicode, Dict, Integer)
+from IPython.zmq.serialize import MAX_ITEMS, MAX_BYTES
 
 #-----------------------------------------------------------------------------
 # utility functions
@@ -84,8 +85,9 b' pickle_unpacker = pickle.loads'
 default_packer = json_packer
 default_unpacker = json_unpacker
 
-DELIM=b"<IDS|MSG>"
-
+DELIM = b"<IDS|MSG>"
+# singleton dummy tracker, which will always report as done
+DONE = zmq.MessageTracker()
 
 #-----------------------------------------------------------------------------
 # Mixin tools for apps that use Sessions
@@ -329,7 +331,18 b' class Session(Configurable):'
         # unpacker is not checked - it is assumed to be
         if not callable(new):
             raise TypeError("unpacker must be callable, not %s"%type(new))
-
+    
+    # thresholds:
+    copy_threshold = Integer(2**16, config=True,
+        help="Threshold (in bytes) beyond which a buffer should be sent without copying.")
+    buffer_threshold = Integer(MAX_BYTES, config=True,
+        help="Threshold (in bytes) beyond which an object's buffer should be extracted to avoid pickling.")
+    item_threshold = Integer(MAX_ITEMS, config=True,
+        help="""The maximum number of items for a container to be introspected for custom serialization.
+        Containers larger than this are pickled outright.
+        """
+    )
+    
     def __init__(self, **kwargs):
         """create a Session object
 
@@ -544,11 +557,6 b' class Session(Configurable):'
         -------
         msg : dict
             The constructed message.
-        (msg,tracker) : (dict, MessageTracker)
-            if track=True, then a 2-tuple will be returned,
-            the first element being the constructed
-            message, and the second being the MessageTracker
-
         """
 
         if not isinstance(stream, (zmq.Socket, ZMQStream)):
@@ -566,25 +574,18 b' class Session(Configurable):'
 
         buffers = [] if buffers is None else buffers
         to_send = self.serialize(msg, ident)
-        flag = 0
-        if buffers:
-            flag = zmq.SNDMORE
-            _track = False
+        to_send.extend(buffers)
+        longest = max([ len(s) for s in to_send ])
+        copy = (longest < self.copy_threshold)
+        
+        if buffers and track and not copy:
+            # only really track when we are doing zero-copy buffers
+            tracker = stream.send_multipart(to_send, copy=False, track=True)
         else:
-            _track=track
-        if track:
-            tracker = stream.send_multipart(to_send, flag, copy=False, track=_track)
-        else:
-            tracker = stream.send_multipart(to_send, flag, copy=False)
-        for b in buffers[:-1]:
-            stream.send(b, flag, copy=False)
-        if buffers:
-            if track:
-                tracker = stream.send(buffers[-1], copy=False, track=track)
-            else:
-                tracker = stream.send(buffers[-1], copy=False)
+            # use dummy tracker, which will be done immediately
+            tracker = DONE
+            stream.send_multipart(to_send, copy=copy)
 
-        # omsg = Message(msg)
         if self.debug:
             pprint.pprint(msg)
             pprint.pprint(to_send)

IPython/zmq/tests/test_session.py

@@ -146,9 +146,10 b' class TestSession(SessionTestCase):'
         """test tracking messages"""
         a,b = self.create_bound_pair(zmq.PAIR, zmq.PAIR)
         s = self.session
+        s.copy_threshold = 1
         stream = ZMQStream(a)
         msg = s.send(a, 'hello', track=False)
-        self.assertTrue(msg['tracker'] is None)
+        self.assertTrue(msg['tracker'] is ss.DONE)
         msg = s.send(a, 'hello', track=True)
         self.assertTrue(isinstance(msg['tracker'], zmq.MessageTracker))
         M = zmq.Message(b'hi there', track=True)