added dependency decorator

IPython/zmq/parallel/dependency.py

@@ -0,0 +1,66
+"""Dependency utilities"""
+
+from IPython.external.decorator import decorator
+
+# flags
+ALL = 1 << 0
+ANY = 1 << 1
+HERE = 1 << 2
+ANYWHERE = 1 << 3
+
+class UnmetDependency(Exception):
+    pass
+
+class depend2(object):
+    """dependency decorator"""
+    def __init__(self, f, *args, **kwargs):
+        print "Inside __init__()"
+        self.dependency = (f,args,kwargs)
+    
+    def __call__(self, f, *args, **kwargs):
+        f._dependency = self.dependency
+        return decorator(_depend_wrapper, f)
+    
+class depend(object):
+    """Dependency decorator, for use with tasks."""
+    def __init__(self, f, *args, **kwargs):
+        print "Inside __init__()"
+        self.f = f
+        self.args = args
+        self.kwargs = kwargs
+    
+    def __call__(self, f):
+        return dependent(f, self.f, *self.args, **self.kwargs)
+
+class dependent(object):
+    """A function that depends on another function.
+    This is an object to prevent the closure used
+    in traditional decorators, which are not picklable.
+    """
+
+    def __init__(self, f, df, *dargs, **dkwargs):
+        self.f = f
+        self.func_name = self.f.func_name
+        self.df = df
+        self.dargs = dargs
+        self.dkwargs = dkwargs
+
+    def __call__(self, *args, **kwargs):
+        if self.df(*self.dargs, **self.dkwargs) is False:
+            raise UnmetDependency()
+        return self.f(*args, **kwargs)
+
+def evaluate_dependency(deps):
+    """Evaluate wheter dependencies are met.
+    
+    Parameters
+    ----------
+    deps : dict
+    """
+    pass
+
+def _check_dependency(flag):
+    pass
+
+
+__all__ = ['UnmetDependency', 'depend', 'evaluate_dependencies']
  No newline at end of file

IPython/zmq/parallel/streamkernel.py

@@ -1,510 +1,507
 #!/usr/bin/env python
 """
 Kernel adapted from kernel.py to use ZMQ Streams
 """
 
 import __builtin__
 import os
 import sys
 import time
 import traceback
 from signal import SIGTERM, SIGKILL
 from pprint import pprint
 
 from code import CommandCompiler
 
 import zmq
 from zmq.eventloop import ioloop, zmqstream
 
+from IPython.zmq.completer import KernelCompleter
+
 from streamsession import StreamSession, Message, extract_header, serialize_object,\
                 unpack_apply_message
-from IPython.zmq.completer import KernelCompleter
+from dependency import UnmetDependency
 
 def printer(*args):
     pprint(args)
 
 class OutStream(object):
     """A file like object that publishes the stream to a 0MQ PUB socket."""
 
     def __init__(self, session, pub_socket, name, max_buffer=200):
         self.session = session
         self.pub_socket = pub_socket
         self.name = name
         self._buffer = []
         self._buffer_len = 0
         self.max_buffer = max_buffer
         self.parent_header = {}
 
     def set_parent(self, parent):
         self.parent_header = extract_header(parent)
 
     def close(self):
         self.pub_socket = None
 
     def flush(self):
         if self.pub_socket is None:
             raise ValueError(u'I/O operation on closed file')
         else:
             if self._buffer:
                 data = ''.join(self._buffer)
                 content = {u'name':self.name, u'data':data}
                 # msg = self.session.msg(u'stream', content=content,
                 #                        parent=self.parent_header)
                 msg = self.session.send(self.pub_socket, u'stream', content=content, parent=self.parent_header)
                 # print>>sys.__stdout__, Message(msg)
                 # self.pub_socket.send_json(msg)
                 self._buffer_len = 0
                 self._buffer = []
 
     def isattr(self):
         return False
 
     def next(self):
         raise IOError('Read not supported on a write only stream.')
 
     def read(self, size=None):
         raise IOError('Read not supported on a write only stream.')
 
     readline=read
 
     def write(self, s):
         if self.pub_socket is None:
             raise ValueError('I/O operation on closed file')
         else:
             self._buffer.append(s)
             self._buffer_len += len(s)
             self._maybe_send()
 
     def _maybe_send(self):
         if '\n' in self._buffer[-1]:
             self.flush()
         if self._buffer_len > self.max_buffer:
             self.flush()
 
     def writelines(self, sequence):
         if self.pub_socket is None:
             raise ValueError('I/O operation on closed file')
         else:
             for s in sequence:
                 self.write(s)
 
 
 class DisplayHook(object):
 
     def __init__(self, session, pub_socket):
         self.session = session
         self.pub_socket = pub_socket
         self.parent_header = {}
 
     def __call__(self, obj):
         if obj is None:
             return
 
         __builtin__._ = obj
         # msg = self.session.msg(u'pyout', {u'data':repr(obj)},
         #                        parent=self.parent_header)
         # self.pub_socket.send_json(msg)
         self.session.send(self.pub_socket, u'pyout', content={u'data':repr(obj)}, parent=self.parent_header)
 
     def set_parent(self, parent):
         self.parent_header = extract_header(parent)
 
 
 class RawInput(object):
 
     def __init__(self, session, socket):
         self.session = session
         self.socket = socket
 
     def __call__(self, prompt=None):
         msg = self.session.msg(u'raw_input')
         self.socket.send_json(msg)
         while True:
             try:
                 reply = self.socket.recv_json(zmq.NOBLOCK)
             except zmq.ZMQError, e:
                 if e.errno == zmq.EAGAIN:
                     pass
                 else:
                     raise
             else:
                 break
         return reply[u'content'][u'data']
 
 
 class Kernel(object):
 
     def __init__(self, session, control_stream, reply_stream, pub_stream, 
                                             task_stream=None, client=None):
         self.session = session
         self.control_stream = control_stream
         # self.control_socket = control_stream.socket
         self.reply_stream = reply_stream
         self.task_stream = task_stream
         self.pub_stream = pub_stream
         self.client = client
         self.user_ns = {}
         self.history = []
         self.compiler = CommandCompiler()
         self.completer = KernelCompleter(self.user_ns)
         self.aborted = set()
         
         # Build dict of handlers for message types
         self.queue_handlers = {}
         self.control_handlers = {}
         for msg_type in ['execute_request', 'complete_request', 'apply_request']:
             self.queue_handlers[msg_type] = getattr(self, msg_type)
         
-        for msg_type in ['kill_request', 'abort_request']:
+        for msg_type in ['kill_request', 'abort_request']+self.queue_handlers.keys():
             self.control_handlers[msg_type] = getattr(self, msg_type)
 
     #-------------------- control handlers -----------------------------
     def abort_queues(self):
         for stream in (self.task_stream, self.reply_stream):
             if stream:
                 self.abort_queue(stream)
     
     def abort_queue(self, stream):
         while True:
             try:
                 msg = self.session.recv(stream, zmq.NOBLOCK,content=True)
             except zmq.ZMQError, e:
                 if e.errno == zmq.EAGAIN:
                     break
                 else:
                     return
             else:
                 if msg is None:
                     return
                 else:
                     idents,msg = msg
                 
                 # assert self.reply_socketly_socket.rcvmore(), "Unexpected missing message part."
                 # msg = self.reply_socket.recv_json()
             print>>sys.__stdout__, "Aborting:"
             print>>sys.__stdout__, Message(msg)
             msg_type = msg['msg_type']
             reply_type = msg_type.split('_')[0] + '_reply'
             # reply_msg = self.session.msg(reply_type, {'status' : 'aborted'}, msg)
             # self.reply_socket.send(ident,zmq.SNDMORE)
             # self.reply_socket.send_json(reply_msg)
             reply_msg = self.session.send(stream, reply_type, 
                         content={'status' : 'aborted'}, parent=msg, ident=idents)
             print>>sys.__stdout__, Message(reply_msg)
             # We need to wait a bit for requests to come in. This can probably
             # be set shorter for true asynchronous clients.
             time.sleep(0.05)
     
     def abort_request(self, stream, ident, parent):
         """abort a specifig msg by id"""
         msg_ids = parent['content'].get('msg_ids', None)
         if isinstance(msg_ids, basestring):
             msg_ids = [msg_ids]
         if not msg_ids:
             self.abort_queues()
         for mid in msg_ids:
             self.aborted.add(str(mid))
         
         content = dict(status='ok')
         reply_msg = self.session.send(stream, 'abort_reply', content=content, parent=parent,
                                                                     ident=ident)
         print>>sys.__stdout__, Message(reply_msg)
     
     def kill_request(self, stream, idents, parent):
         """kill ourselves.  This should really be handled in an external process"""
         self.abort_queues()
         msg = self.session.send(stream, 'kill_reply', ident=idents, parent=parent, 
                 content = dict(status='ok'))
         # we can know that a message is done if we *don't* use streams, but 
         # use a socket directly with MessageTracker
         time.sleep(.5)
         os.kill(os.getpid(), SIGTERM)
         time.sleep(1)
         os.kill(os.getpid(), SIGKILL)
     
     def dispatch_control(self, msg):
         idents,msg = self.session.feed_identities(msg, copy=False)
         msg = self.session.unpack_message(msg, content=True, copy=False)
         
         header = msg['header']
         msg_id = header['msg_id']
         
         handler = self.control_handlers.get(msg['msg_type'], None)
         if handler is None:
             print >> sys.__stderr__, "UNKNOWN CONTROL MESSAGE TYPE:", msg
         else:
             handler(self.control_stream, idents, msg)
     
     # def flush_control(self):
     #     while any(zmq.select([self.control_socket],[],[],1e-4)):
     #         try:
     #             msg = self.control_socket.recv_multipart(zmq.NOBLOCK, copy=False)
     #         except zmq.ZMQError, e:
     #             if e.errno != zmq.EAGAIN:
     #                 raise e
     #             return
     #         else:
     #             self.dispatch_control(msg)
     
 
     #-------------------- queue helpers ------------------------------
     
     def check_dependencies(self, dependencies):
         if not dependencies:
             return True
         if len(dependencies) == 2 and dependencies[0] in 'any all'.split():
             anyorall = dependencies[0]
             dependencies = dependencies[1]
         else:
             anyorall = 'all'
         results = self.client.get_results(dependencies,status_only=True)
         if results['status'] != 'ok':
             return False
         
         if anyorall == 'any':
             if not results['completed']:
                 return False
         else:
             if results['pending']:
                 return False
         
         return True
     
     def check_aborted(self, msg_id):
         return msg_id in self.aborted
     
-    def unmet_dependencies(self, stream, idents, msg):
-        reply_type = msg['msg_type'].split('_')[0] + '_reply'
-        content = dict(status='resubmitted', reason='unmet dependencies')
-        reply_msg = self.session.send(stream, reply_type, 
-                    content=content, parent=msg, ident=idents)
-        ### TODO: actually resubmit it ###
-    
     #-------------------- queue handlers -----------------------------
     
     def execute_request(self, stream, ident, parent):
         try:
             code = parent[u'content'][u'code']
         except:
             print>>sys.__stderr__, "Got bad msg: "
             print>>sys.__stderr__, Message(parent)
             return
         # pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent)
         # self.pub_stream.send(pyin_msg)
         self.session.send(self.pub_stream, u'pyin', {u'code':code},parent=parent)
         try:
             comp_code = self.compiler(code, '<zmq-kernel>')
             # allow for not overriding displayhook
             if hasattr(sys.displayhook, 'set_parent'):
                 sys.displayhook.set_parent(parent)
             exec comp_code in self.user_ns, self.user_ns
         except:
             # result = u'error'
             etype, evalue, tb = sys.exc_info()
             tb = traceback.format_exception(etype, evalue, tb)
             exc_content = {
                 u'status' : u'error',
                 u'traceback' : tb,
                 u'etype' : unicode(etype),
                 u'evalue' : unicode(evalue)
             }
             # exc_msg = self.session.msg(u'pyerr', exc_content, parent)
             self.session.send(self.pub_stream, u'pyerr', exc_content, parent=parent)
             reply_content = exc_content
         else:
             reply_content = {'status' : 'ok'}
         # reply_msg = self.session.msg(u'execute_reply', reply_content, parent)
         # self.reply_socket.send(ident, zmq.SNDMORE)
         # self.reply_socket.send_json(reply_msg)
         reply_msg = self.session.send(stream, u'execute_reply', reply_content, parent=parent, ident=ident)
         # print>>sys.__stdout__, Message(reply_msg)
         if reply_msg['content']['status'] == u'error':
             self.abort_queues()
 
     def complete_request(self, stream, ident, parent):
         matches = {'matches' : self.complete(parent),
                    'status' : 'ok'}
         completion_msg = self.session.send(stream, 'complete_reply',
                                            matches, parent, ident)
         # print >> sys.__stdout__, completion_msg
 
     def complete(self, msg):
         return self.completer.complete(msg.content.line, msg.content.text)
     
     def apply_request(self, stream, ident, parent):
         try:
             content = parent[u'content']
             bufs = parent[u'buffers']
             msg_id = parent['header']['msg_id']
             bound = content.get('bound', False)
         except:
             print>>sys.__stderr__, "Got bad msg: "
             print>>sys.__stderr__, Message(parent)
             return
         # pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent)
         # self.pub_stream.send(pyin_msg)
         # self.session.send(self.pub_stream, u'pyin', {u'code':code},parent=parent)
+        sub = {'dependencies_met' : True}
         try:
             # allow for not overriding displayhook
             if hasattr(sys.displayhook, 'set_parent'):
                 sys.displayhook.set_parent(parent)
             # exec "f(*args,**kwargs)" in self.user_ns, self.user_ns
             if bound:
                 working = self.user_ns
                 suffix = str(msg_id).replace("-","")
                 prefix = "_"
                 
             else:
                 working = dict()
                 suffix = prefix = "_" # prevent keyword collisions with lambda
             f,args,kwargs = unpack_apply_message(bufs, working, copy=False)
             # if f.fun
             fname = prefix+f.func_name.strip('<>')+suffix
             argname = prefix+"args"+suffix
             kwargname = prefix+"kwargs"+suffix
             resultname = prefix+"result"+suffix
             
             ns = { fname : f, argname : args, kwargname : kwargs }
             # print ns
             working.update(ns)
             code = "%s=%s(*%s,**%s)"%(resultname, fname, argname, kwargname)
             exec code in working, working
             result = working.get(resultname)
             # clear the namespace
             if bound:
                 for key in ns.iterkeys():
                     self.user_ns.pop(key)
             else:
                 del working
             
             packed_result,buf = serialize_object(result)
             result_buf = [packed_result]+buf
         except:
             result = u'error'
             etype, evalue, tb = sys.exc_info()
             tb = traceback.format_exception(etype, evalue, tb)
             exc_content = {
                 u'status' : u'error',
                 u'traceback' : tb,
                 u'etype' : unicode(etype),
                 u'evalue' : unicode(evalue)
             }
             # exc_msg = self.session.msg(u'pyerr', exc_content, parent)
             self.session.send(self.pub_stream, u'pyerr', exc_content, parent=parent)
             reply_content = exc_content
             result_buf = []
+            
+            if etype is UnmetDependency:
+                sub = {'dependencies_met' : False}
         else:
             reply_content = {'status' : 'ok'}
         # reply_msg = self.session.msg(u'execute_reply', reply_content, parent)
         # self.reply_socket.send(ident, zmq.SNDMORE)
         # self.reply_socket.send_json(reply_msg)
-        reply_msg = self.session.send(stream, u'apply_reply', reply_content, parent=parent, ident=ident,buffers=result_buf)
+        reply_msg = self.session.send(stream, u'apply_reply', reply_content, 
+                    parent=parent, ident=ident,buffers=result_buf, subheader=sub)
         # print>>sys.__stdout__, Message(reply_msg)
-        if reply_msg['content']['status'] == u'error':
-            self.abort_queues()
+        # if reply_msg['content']['status'] == u'error':
+        #     self.abort_queues()
     
     def dispatch_queue(self, stream, msg):
         self.control_stream.flush()
         idents,msg = self.session.feed_identities(msg, copy=False)
         msg = self.session.unpack_message(msg, content=True, copy=False)
         
         header = msg['header']
         msg_id = header['msg_id']
-        dependencies = header.get('dependencies', [])
         if self.check_aborted(msg_id):
             self.aborted.remove(msg_id)
             # is it safe to assume a msg_id will not be resubmitted?
             reply_type = msg['msg_type'].split('_')[0] + '_reply'
             reply_msg = self.session.send(stream, reply_type, 
                         content={'status' : 'aborted'}, parent=msg, ident=idents)
             return
-        if not self.check_dependencies(dependencies):
-            return self.unmet_dependencies(stream, idents, msg)
         handler = self.queue_handlers.get(msg['msg_type'], None)
         if handler is None:
             print >> sys.__stderr__, "UNKNOWN MESSAGE TYPE:", msg
         else:
             handler(stream, idents, msg)
     
     def start(self):
         #### stream mode:
         if self.control_stream:
             self.control_stream.on_recv(self.dispatch_control, copy=False)
             self.control_stream.on_err(printer)
         if self.reply_stream:
             self.reply_stream.on_recv(lambda msg: 
                     self.dispatch_queue(self.reply_stream, msg), copy=False)
             self.reply_stream.on_err(printer)
         if self.task_stream:
             self.task_stream.on_recv(lambda msg: 
                     self.dispatch_queue(self.task_stream, msg), copy=False)
             self.task_stream.on_err(printer)
         
         #### while True mode:
         # while True:
         #     idle = True
         #     try:
         #         msg = self.reply_stream.socket.recv_multipart(
         #                     zmq.NOBLOCK, copy=False)
         #     except zmq.ZMQError, e:
         #         if e.errno != zmq.EAGAIN:
         #             raise e
         #     else:
         #         idle=False
         #         self.dispatch_queue(self.reply_stream, msg)
         #             
         #     if not self.task_stream.empty():
         #         idle=False
         #         msg = self.task_stream.recv_multipart()
         #         self.dispatch_queue(self.task_stream, msg)
         #     if idle:
         #         # don't busywait
         #         time.sleep(1e-3)
 
 
 def main():
     raise Exception("Don't run me anymore")
     loop = ioloop.IOLoop.instance()
     c = zmq.Context()
 
     ip = '127.0.0.1'
     port_base = 5575
     connection = ('tcp://%s' % ip) + ':%i'
     rep_conn = connection % port_base
     pub_conn = connection % (port_base+1)
 
     print >>sys.__stdout__, "Starting the kernel..."
     # print >>sys.__stdout__, "XREQ Channel:", rep_conn
     # print >>sys.__stdout__, "PUB Channel:", pub_conn
 
     session = StreamSession(username=u'kernel')
 
     reply_socket = c.socket(zmq.XREQ)
     reply_socket.connect(rep_conn)
 
     pub_socket = c.socket(zmq.PUB)
     pub_socket.connect(pub_conn)
 
     stdout = OutStream(session, pub_socket, u'stdout')
     stderr = OutStream(session, pub_socket, u'stderr')
     sys.stdout = stdout
     sys.stderr = stderr
 
     display_hook = DisplayHook(session, pub_socket)
     sys.displayhook = display_hook
     reply_stream = zmqstream.ZMQStream(reply_socket,loop)
     pub_stream = zmqstream.ZMQStream(pub_socket,loop)
     kernel = Kernel(session, reply_stream, pub_stream)
 
     # For debugging convenience, put sleep and a string in the namespace, so we
     # have them every time we start.
     kernel.user_ns['sleep'] = time.sleep
     kernel.user_ns['s'] = 'Test string'
     
     print >>sys.__stdout__, "Use Ctrl-\\ (NOT Ctrl-C!) to terminate."
     kernel.start()
     loop.start()
 
 
 if __name__ == '__main__':
     main()

IPython/zmq/parallel/streamsession.py

@@ -1,447 +1,447
 #!/usr/bin/env python
 """edited session.py to work with streams, and move msg_type to the header
 """
 
 
 import os
 import sys
 import traceback
 import pprint
 import uuid
 
 import zmq
 from zmq.utils import jsonapi
 from zmq.eventloop.zmqstream import ZMQStream
 
 from IPython.zmq.pickleutil import can, uncan, canSequence, uncanSequence
 from IPython.zmq.newserialized import serialize, unserialize
 
 try:
     import cPickle
     pickle = cPickle
 except:
     cPickle = None
     import pickle
 
 # packer priority: jsonlib[2], cPickle, simplejson/json, pickle
 json_name = '' if not jsonapi.jsonmod else jsonapi.jsonmod.__name__
 if json_name in ('jsonlib', 'jsonlib2'):
     use_json = True
 elif json_name:
     if cPickle is None:
         use_json = True
     else:
         use_json = False
 else:
     use_json = False
 
 if use_json:
     default_packer = jsonapi.dumps
     default_unpacker = jsonapi.loads
 else:
     default_packer = lambda o: pickle.dumps(o,-1)
     default_unpacker = pickle.loads
 
 
 DELIM="<IDS|MSG>"
 
 def wrap_exception():
     etype, evalue, tb = sys.exc_info()
     tb = traceback.format_exception(etype, evalue, tb)
     exc_content = {
         u'status' : u'error',
         u'traceback' : tb,
         u'etype' : unicode(etype),
         u'evalue' : unicode(evalue)
     }
     return exc_content
 
 class KernelError(Exception):
     pass
 
 def unwrap_exception(content):
     err = KernelError(content['etype'], content['evalue'])
     err.evalue = content['evalue']
     err.etype = content['etype']
     err.traceback = ''.join(content['traceback'])
     return err
     
 
 class Message(object):
     """A simple message object that maps dict keys to attributes.
 
     A Message can be created from a dict and a dict from a Message instance
     simply by calling dict(msg_obj)."""
     
     def __init__(self, msg_dict):
         dct = self.__dict__
         for k, v in dict(msg_dict).iteritems():
             if isinstance(v, dict):
                 v = Message(v)
             dct[k] = v
 
     # Having this iterator lets dict(msg_obj) work out of the box.
     def __iter__(self):
         return iter(self.__dict__.iteritems())
     
     def __repr__(self):
         return repr(self.__dict__)
 
     def __str__(self):
         return pprint.pformat(self.__dict__)
 
     def __contains__(self, k):
         return k in self.__dict__
 
     def __getitem__(self, k):
         return self.__dict__[k]
 
 
 def msg_header(msg_id, msg_type, username, session):
     return locals()
     # return {
     #     'msg_id' : msg_id,
     #     'msg_type': msg_type,
     #     'username' : username,
     #     'session' : session
     # }
 
 
 def extract_header(msg_or_header):
     """Given a message or header, return the header."""
     if not msg_or_header:
         return {}
     try:
         # See if msg_or_header is the entire message.
         h = msg_or_header['header']
     except KeyError:
         try:
             # See if msg_or_header is just the header
             h = msg_or_header['msg_id']
         except KeyError:
             raise
         else:
             h = msg_or_header
     if not isinstance(h, dict):
         h = dict(h)
     return h
 
 def rekey(dikt):
     """rekey a dict that has been forced to use str keys where there should be
     ints by json.  This belongs in the jsonutil added by fperez."""
     for k in dikt.iterkeys():
         if isinstance(k, str):
             ik=fk=None
             try:
                 ik = int(k)
             except ValueError:
                 try:
                     fk = float(k)
                 except ValueError:
                     continue
             if ik is not None:
                 nk = ik
             else:
                 nk = fk
             if nk in dikt:
                 raise KeyError("already have key %r"%nk)
             dikt[nk] = dikt.pop(k)
     return dikt
 
 def serialize_object(obj, threshold=64e-6):
     """serialize an object into a list of sendable buffers.
     
     Returns: (pmd, bufs)
         where pmd is the pickled metadata wrapper, and bufs
         is a list of data buffers"""
     # threshold is 100 B
     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 = {}
         for k in sorted(obj.iterkeys()):
             s = serialize(can(obj[k]))
             if s.getDataSize() > threshold:
                 databuffers.append(s.getData())
                 s.data = None
             sobj[k] = s
         return pickle.dumps(sobj,-1),databuffers
     else:
         s = serialize(can(obj))
         if 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))
     elif isinstance(sobj, dict):
         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
     else:
         if sobj.data is None:
             sobj.data = bufs.pop(0)
         return uncan(unserialize(sobj))
 
 def pack_apply_message(f, args, kwargs, threshold=64e-6):
     """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)"""
     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)
     msg.extend(databuffers)
     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) # allow us to pop
     assert len(bufs) >= 3, "not enough buffers!"
     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 = cf.getFunction(g)
+    f = uncan(cf, g)
     for sa in sargs:
         if sa.data is None:
             m = bufs.pop(0)
             if sa.getTypeDescriptor() in ('buffer', 'ndarray'):
                 if copy:
                     sa.data = buffer(m)
                 else:
                     sa.data = m.buffer
             else:
                 if copy:
                     sa.data = m
                 else:
                     sa.data = m.bytes
     
     args = uncanSequence(map(unserialize, sargs), g)
     kwargs = {}
     for k in sorted(skwargs.iterkeys()):
         sa = skwargs[k]
         if sa.data is None:
             sa.data = bufs.pop(0)
         kwargs[k] = uncan(unserialize(sa), g)
     
     return f,args,kwargs
 
 class StreamSession(object):
     """tweaked version of IPython.zmq.session.Session, for development in Parallel"""
     debug=False
     def __init__(self, username=None, session=None, packer=None, unpacker=None):
         if username is None:
             username = os.environ.get('USER','username')
         self.username = username
         if session is None:
             self.session = str(uuid.uuid4())
         else:
             self.session = session
         self.msg_id = str(uuid.uuid4())
         if packer is None:
             self.pack = default_packer
         else:
             if not callable(packer):
                 raise TypeError("packer must be callable, not %s"%type(packer))
             self.pack = packer
         
         if unpacker is None:
             self.unpack = default_unpacker
         else:
             if not callable(unpacker):
                 raise TypeError("unpacker must be callable, not %s"%type(unpacker))
             self.unpack = unpacker
         
         self.none = self.pack({})
             
     def msg_header(self, msg_type):
         h = msg_header(self.msg_id, msg_type, self.username, self.session)
         self.msg_id = str(uuid.uuid4())
         return h
 
     def msg(self, msg_type, content=None, parent=None, subheader=None):
         msg = {}
         msg['header'] = self.msg_header(msg_type)
         msg['msg_id'] = msg['header']['msg_id']
         msg['parent_header'] = {} if parent is None else extract_header(parent)
         msg['msg_type'] = msg_type
         msg['content'] = {} if content is None else content
         sub = {} if subheader is None else subheader
         msg['header'].update(sub)
         return msg
 
     def send(self, stream, msg_type, content=None, buffers=None, parent=None, subheader=None, ident=None):
         """send a message via stream"""
         msg = self.msg(msg_type, content, parent, subheader)
         buffers = [] if buffers is None else buffers
         to_send = []
         if isinstance(ident, list):
             # accept list of idents
             to_send.extend(ident)
         elif ident is not None:
             to_send.append(ident)
         to_send.append(DELIM)
         to_send.append(self.pack(msg['header']))
         to_send.append(self.pack(msg['parent_header']))
         # if parent is None:
         #     to_send.append(self.none)
         # else:
         #     to_send.append(self.pack(dict(parent)))
         if content is None:
             content = self.none
         elif isinstance(content, dict):
             content = self.pack(content)
         elif isinstance(content, str):
             # content is already packed, as in a relayed message
             pass
         else:
             raise TypeError("Content incorrect type: %s"%type(content))
         to_send.append(content)
         flag = 0
         if buffers:
             flag = zmq.SNDMORE
         stream.send_multipart(to_send, flag, copy=False)
         for b in buffers[:-1]:
             stream.send(b, flag, copy=False)
         if buffers:
             stream.send(buffers[-1], copy=False)
         omsg = Message(msg)
         if self.debug:
             pprint.pprint(omsg)
             pprint.pprint(to_send)
             pprint.pprint(buffers)
         return omsg
 
     def recv(self, socket, mode=zmq.NOBLOCK, content=True, copy=True):
         """receives and unpacks a message
         returns [idents], msg"""
         if isinstance(socket, ZMQStream):
             socket = socket.socket
         try:
             msg = socket.recv_multipart(mode)
         except zmq.ZMQError, e:
             if e.errno == zmq.EAGAIN:
                 # We can convert EAGAIN to None as we know in this case
                 # recv_json won't return None.
                 return None
             else:
                 raise
         # return an actual Message object
         # determine the number of idents by trying to unpack them.
         # this is terrible:
         idents, msg = self.feed_identities(msg, copy)
         try:
             return idents, self.unpack_message(msg, content=content, copy=copy)
         except Exception, e:
             print idents, msg
             # TODO: handle it
             raise e
     
     def feed_identities(self, msg, copy=True):
         """This is a completely horrible thing, but it strips the zmq
         ident prefixes off of a message. It will break if any identities
         are unpackable by self.unpack."""
         msg = list(msg)
         idents = []
         while len(msg) > 3:
             if copy:
                 s = msg[0]
             else:
                 s = msg[0].bytes
             if s == DELIM:
                 msg.pop(0)
                 break
             else:
                 idents.append(s)
                 msg.pop(0)
                 
         return idents, msg
     
     def unpack_message(self, msg, content=True, copy=True):
         """return a message object from the format
         sent by self.send.
         
         parameters:
         
         content : bool (True)
             whether to unpack the content dict (True), 
             or leave it serialized (False)
         
         copy : bool (True)
             whether to return the bytes (True), 
             or the non-copying Message object in each place (False)
         
         """
         if not len(msg) >= 3:
             raise TypeError("malformed message, must have at least 3 elements")
         message = {}
         if not copy:
             for i in range(3):
                 msg[i] = msg[i].bytes
         message['header'] = self.unpack(msg[0])
         message['msg_type'] = message['header']['msg_type']
         message['parent_header'] = self.unpack(msg[1])
         if content:
             message['content'] = self.unpack(msg[2])
         else:
             message['content'] = msg[2]
     
         # message['buffers'] = msg[3:]
         # else:
         #     message['header'] = self.unpack(msg[0].bytes)
         #     message['msg_type'] = message['header']['msg_type']
         #     message['parent_header'] = self.unpack(msg[1].bytes)
         #     if content:
         #         message['content'] = self.unpack(msg[2].bytes)
         #     else:
         #         message['content'] = msg[2].bytes
         
         message['buffers'] = msg[3:]# [ m.buffer for m in msg[3:] ]
         return message
             
         
 
 def test_msg2obj():
     am = dict(x=1)
     ao = Message(am)
     assert ao.x == am['x']
 
     am['y'] = dict(z=1)
     ao = Message(am)
     assert ao.y.z == am['y']['z']
     
     k1, k2 = 'y', 'z'
     assert ao[k1][k2] == am[k1][k2]
     
     am2 = dict(ao)
     assert am['x'] == am2['x']
     assert am['y']['z'] == am2['y']['z']

IPython/zmq/pickleutil.py

@@ -1,95 +1,115
 # encoding: utf-8
 
 """Pickle related utilities. Perhaps this should be called 'can'."""
 
 __docformat__ = "restructuredtext en"
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2008  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
 #-------------------------------------------------------------------------------
 
 from types import FunctionType
 
 # contents of codeutil should either be in here, or codeutil belongs in IPython/util
 from IPython.kernel import codeutil
+from IPython.zmq.parallel.dependency import dependent
 
 class CannedObject(object):
-    pass
+    def __init__(self, obj, keys=[]):
+        self.keys = keys
+        self.obj = obj
+        for key in keys:
+            setattr(obj, key, can(getattr(obj, key)))
+            
     
+    def getObject(self, g=None):
+        if g is None:
+            g = globals()
+        for key in self.keys:
+            setattr(self.obj, key, uncan(getattr(self.obj, key), g))
+        return self.obj
+
+        
+
 class CannedFunction(CannedObject):
     
     def __init__(self, f):
         self._checkType(f)    
         self.code = f.func_code
     
     def _checkType(self, obj):
         assert isinstance(obj, FunctionType), "Not a function type"
     
     def getFunction(self, g=None):
         if g is None:
             g = globals()
         newFunc = FunctionType(self.code, g)
         return newFunc
 
 def can(obj):
     if isinstance(obj, FunctionType):
         return CannedFunction(obj)
+    elif isinstance(obj, dependent):
+        keys = ('f','df')
+        return CannedObject(obj, keys=keys)
     elif isinstance(obj,dict):
         return canDict(obj)
     elif isinstance(obj, (list,tuple)):
         return canSequence(obj)
     else:
         return obj
 
 def canDict(obj):
     if isinstance(obj, dict):
         newobj = {}
         for k, v in obj.iteritems():
             newobj[k] = can(v)
         return newobj
     else:
         return obj
 
 def canSequence(obj):
     if isinstance(obj, (list, tuple)):
         t = type(obj)
         return t([can(i) for i in obj])
     else:
         return obj
 
 def uncan(obj, g=None):
     if isinstance(obj, CannedFunction):
         return obj.getFunction(g)
+    elif isinstance(obj, CannedObject):
+        return obj.getObject(g)
     elif isinstance(obj,dict):
         return uncanDict(obj)
     elif isinstance(obj, (list,tuple)):
         return uncanSequence(obj)
     else:
         return obj
 
 def uncanDict(obj, g=None):
     if isinstance(obj, dict):
         newobj = {}
         for k, v in obj.iteritems():
             newobj[k] = uncan(v,g)
         return newobj
     else:
         return obj
 
 def uncanSequence(obj, g=None):
     if isinstance(obj, (list, tuple)):
         t = type(obj)
         return t([uncan(i,g) for i in obj])
     else:
         return obj
 
 
 def rebindFunctionGlobals(f, glbls):
     return FunctionType(f.func_code, glbls)