copyright statements

IPython/zmq/parallel/asyncresult.py

@@ -1,322 +1,322 b''
 """AsyncResult objects for the client"""
 #-----------------------------------------------------------------------------
-#  Copyright (C) 2010  The IPython Development Team
+#  Copyright (C) 2010-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 time
 
 from IPython.external.decorator import decorator
 from . import error
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 @decorator
 def check_ready(f, self, *args, **kwargs):
     """Call spin() to sync state prior to calling the method."""
     self.wait(0)
     if not self._ready:
         raise error.TimeoutError("result not ready")
     return f(self, *args, **kwargs)
 
 class AsyncResult(object):
     """Class for representing results of non-blocking calls.
     
     Provides the same interface as :py:class:`multiprocessing.pool.AsyncResult`.
     """
     
     msg_ids = None
     _targets = None
     _tracker = None
     
     def __init__(self, client, msg_ids, fname='unknown', targets=None, tracker=None):
         self._client = client
         if isinstance(msg_ids, basestring):
             msg_ids = [msg_ids]
         self.msg_ids = msg_ids
         self._fname=fname
         self._targets = targets
         self._tracker = tracker
         self._ready = False
         self._success = None
         self._single_result = len(msg_ids) == 1
     
     def __repr__(self):
         if self._ready:
             return "<%s: finished>"%(self.__class__.__name__)
         else:
             return "<%s: %s>"%(self.__class__.__name__,self._fname)
     
     
     def _reconstruct_result(self, res):
         """Reconstruct our result from actual result list (always a list)
         
         Override me in subclasses for turning a list of results
         into the expected form.
         """
         if self._single_result:
             return res[0]
         else:
             return res
         
     def get(self, timeout=-1):
         """Return the result when it arrives. 
         
         If `timeout` is not ``None`` and the result does not arrive within
         `timeout` seconds then ``TimeoutError`` is raised. If the
         remote call raised an exception then that exception will be reraised
         by get() inside a `RemoteError`.
         """
         if not self.ready():
             self.wait(timeout)
         
         if self._ready:
             if self._success:
                 return self._result
             else:
                 raise self._exception
         else:
             raise error.TimeoutError("Result not ready.")
     
     def ready(self):
         """Return whether the call has completed."""
         if not self._ready:
             self.wait(0)
         return self._ready
     
     def wait(self, timeout=-1):
         """Wait until the result is available or until `timeout` seconds pass.
         
         This method always returns None.
         """
         if self._ready:
             return
         self._ready = self._client.barrier(self.msg_ids, timeout)
         if self._ready:
             try:
                 results = map(self._client.results.get, self.msg_ids)
                 self._result = results
                 if self._single_result:
                     r = results[0]
                     if isinstance(r, Exception):
                         raise r
                 else:
                     results = error.collect_exceptions(results, self._fname)
                 self._result = self._reconstruct_result(results)
             except Exception, e:
                 self._exception = e
                 self._success = False
             else:
                 self._success = True
             finally:
                 self._metadata = map(self._client.metadata.get, self.msg_ids)
             
     
     def successful(self):
         """Return whether the call completed without raising an exception. 
         
         Will raise ``AssertionError`` if the result is not ready.
         """
         assert self.ready()
         return self._success
     
     #----------------------------------------------------------------
     # Extra methods not in mp.pool.AsyncResult
     #----------------------------------------------------------------
     
     def get_dict(self, timeout=-1):
         """Get the results as a dict, keyed by engine_id.
         
         timeout behavior is described in `get()`.
         """
         
         results = self.get(timeout)
         engine_ids = [ md['engine_id'] for md in self._metadata ]
         bycount = sorted(engine_ids, key=lambda k: engine_ids.count(k))
         maxcount = bycount.count(bycount[-1])
         if maxcount > 1:
             raise ValueError("Cannot build dict, %i jobs ran on engine #%i"%(
                     maxcount, bycount[-1]))
         
         return dict(zip(engine_ids,results))
     
     @property
     @check_ready
     def result(self):
         """result property wrapper for `get(timeout=0)`."""
         return self._result
     
     # abbreviated alias:
     r = result
     
     @property
     @check_ready
     def metadata(self):
         """property for accessing execution metadata."""
         if self._single_result:
             return self._metadata[0]
         else:
             return self._metadata
     
     @property
     def result_dict(self):
         """result property as a dict."""
         return self.get_dict(0)
     
     def __dict__(self):
         return self.get_dict(0)
     
     def abort(self):
         """abort my tasks."""
         assert not self.ready(), "Can't abort, I am already done!"
         return self.client.abort(self.msg_ids, targets=self._targets, block=True)
 
     @property
     def sent(self):
         """check whether my messages have been sent"""
         if self._tracker is None:
             return True
         else:
             return self._tracker.done
 
     #-------------------------------------
     # dict-access
     #-------------------------------------
     
     @check_ready
     def __getitem__(self, key):
         """getitem returns result value(s) if keyed by int/slice, or metadata if key is str.
         """
         if isinstance(key, int):
             return error.collect_exceptions([self._result[key]], self._fname)[0]
         elif isinstance(key, slice):
             return error.collect_exceptions(self._result[key], self._fname)
         elif isinstance(key, basestring):
             values = [ md[key] for md in self._metadata ]
             if self._single_result:
                 return values[0]
             else:
                 return values
         else:
             raise TypeError("Invalid key type %r, must be 'int','slice', or 'str'"%type(key))
     
     @check_ready
     def __getattr__(self, key):
         """getattr maps to getitem for convenient attr access to metadata."""
         if key not in self._metadata[0].keys():
             raise AttributeError("%r object has no attribute %r"%(
                     self.__class__.__name__, key))
         return self.__getitem__(key)
     
     # asynchronous iterator:
     def __iter__(self):
         if self._single_result:
             raise TypeError("AsyncResults with a single result are not iterable.")
         try:
             rlist = self.get(0)
         except error.TimeoutError:
             # wait for each result individually
             for msg_id in self.msg_ids:
                 ar = AsyncResult(self._client, msg_id, self._fname)
                 yield ar.get()
         else:
             # already done
             for r in rlist:
                 yield r
 
 
     
 class AsyncMapResult(AsyncResult):
     """Class for representing results of non-blocking gathers.
     
     This will properly reconstruct the gather.
     """
     
     def __init__(self, client, msg_ids, mapObject, fname=''):
         AsyncResult.__init__(self, client, msg_ids, fname=fname)
         self._mapObject = mapObject
         self._single_result = False
     
     def _reconstruct_result(self, res):
         """Perform the gather on the actual results."""
         return self._mapObject.joinPartitions(res)
     
     # asynchronous iterator:
     def __iter__(self):
         try:
             rlist = self.get(0)
         except error.TimeoutError:
             # wait for each result individually
             for msg_id in self.msg_ids:
                 ar = AsyncResult(self._client, msg_id, self._fname)
                 rlist = ar.get()
                 try:
                     for r in rlist:
                         yield r
                 except TypeError:
                     # flattened, not a list
                     # this could get broken by flattened data that returns iterables
                     # but most calls to map do not expose the `flatten` argument
                     yield rlist
         else:
             # already done
             for r in rlist:
                 yield r
 
 
 class AsyncHubResult(AsyncResult):
     """Class to wrap pending results that must be requested from the Hub.
     
     Note that waiting/polling on these objects requires polling the Hubover the network,
     so use `AsyncHubResult.wait()` sparingly.
     """
     
     def wait(self, timeout=-1):
         """wait for result to complete."""
         start = time.time()
         if self._ready:
             return
         local_ids = filter(lambda msg_id: msg_id in self._client.outstanding, self.msg_ids)
         local_ready = self._client.barrier(local_ids, timeout)
         if local_ready:
             remote_ids = filter(lambda msg_id: msg_id not in self._client.results, self.msg_ids)
             if not remote_ids:
                 self._ready = True
             else:
                 rdict = self._client.result_status(remote_ids, status_only=False)
                 pending = rdict['pending']
                 while pending and (timeout < 0 or time.time() < start+timeout):
                     rdict = self._client.result_status(remote_ids, status_only=False)
                     pending = rdict['pending']
                     if pending:
                         time.sleep(0.1)
                 if not pending:
                     self._ready = True
         if self._ready:
             try:
                 results = map(self._client.results.get, self.msg_ids)
                 self._result = results
                 if self._single_result:
                     r = results[0]
                     if isinstance(r, Exception):
                         raise r
                 else:
                     results = error.collect_exceptions(results, self._fname)
                 self._result = self._reconstruct_result(results)
             except Exception, e:
                 self._exception = e
                 self._success = False
             else:
                 self._success = True
             finally:
                 self._metadata = map(self._client.metadata.get, self.msg_ids)
         
 __all__ = ['AsyncResult', 'AsyncMapResult', 'AsyncHubResult']
  No newline at end of file

IPython/zmq/parallel/dependency.py

@@ -1,153 +1,159 b''
 """Dependency utilities"""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 from IPython.external.decorator import decorator
 
 from .asyncresult import AsyncResult
 from .error import UnmetDependency
 
 
 class depend(object):
     """Dependency decorator, for use with tasks.
     
     `@depend` lets you define a function for engine dependencies
     just like you use `apply` for tasks.
     
     
     Examples
     --------
     ::
     
         @depend(df, a,b, c=5)
         def f(m,n,p)
         
         view.apply(f, 1,2,3)
     
     will call df(a,b,c=5) on the engine, and if it returns False or
     raises an UnmetDependency error, then the task will not be run
     and another engine will be tried.
     """
     def __init__(self, f, *args, **kwargs):
         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 = getattr(f, '__name__', 'f')
         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)
     
     @property
     def __name__(self):
         return self.func_name
 
 def _require(*names):
     """Helper for @require decorator."""
     for name in names:
         try:
             __import__(name)
         except ImportError:
             return False
     return True
 
 def require(*names):
     """Simple decorator for requiring names to be importable.
     
     Examples
     --------
     
     In [1]: @require('numpy')
        ...: def norm(a):
        ...:     import numpy
        ...:     return numpy.linalg.norm(a,2)
     """
     return depend(_require, *names)
 
 class Dependency(set):
     """An object for representing a set of msg_id dependencies.
     
     Subclassed from set().
     
     Parameters
     ----------
     dependencies: list/set of msg_ids or AsyncResult objects or output of Dependency.as_dict()
         The msg_ids to depend on
     all : bool [default True]
         Whether the dependency should be considered met when *all* depending tasks have completed
         or only when *any* have been completed.
     success_only : bool [default True]
         Whether to consider only successes for Dependencies, or consider failures as well.
         If `all=success_only=True`, then this task will fail with an ImpossibleDependency
         as soon as the first depended-upon task fails.
     """
     
     all=True
     success_only=True
     
     def __init__(self, dependencies=[], all=True, success_only=True):
         if isinstance(dependencies, dict):
             # load from dict
             all = dependencies.get('all', True)
             success_only = dependencies.get('success_only', success_only)
             dependencies = dependencies.get('dependencies', [])
         ids = []
         if isinstance(dependencies, AsyncResult):
             ids.extend(AsyncResult.msg_ids)
         else:
             for d in dependencies:
                 if isinstance(d, basestring):
                     ids.append(d)
                 elif isinstance(d, AsyncResult):
                     ids.extend(d.msg_ids)
                 else:
                     raise TypeError("invalid dependency type: %r"%type(d))
         set.__init__(self, ids)
         self.all = all
         self.success_only=success_only
     
     def check(self, completed, failed=None):
         if failed is not None and not self.success_only:
             completed = completed.union(failed)
         if len(self) == 0:
             return True
         if self.all:
             return self.issubset(completed)
         else:
             return not self.isdisjoint(completed)
     
     def unreachable(self, failed):
         if len(self) == 0 or len(failed) == 0 or not self.success_only:
             return False
         # print self, self.success_only, self.all, failed
         if self.all:
             return not self.isdisjoint(failed)
         else:
             return self.issubset(failed)
         
     
     def as_dict(self):
         """Represent this dependency as a dict. For json compatibility."""
         return dict(
             dependencies=list(self),
             all=self.all,
             success_only=self.success_only,
         )
     
 
 __all__ = ['depend', 'require', 'dependent', 'Dependency']
 

IPython/zmq/parallel/engine.py

@@ -1,148 +1,155 b''
 #!/usr/bin/env python
 """A simple engine that talks to a controller over 0MQ.
 it handles registration, etc. and launches a kernel
 connected to the Controller's Schedulers.
 """
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
+
 from __future__ import print_function
 
 import sys
 import time
 
 import zmq
 from zmq.eventloop import ioloop, zmqstream
 
 # internal
 from IPython.utils.traitlets import Instance, Str, Dict, Int, Type, CFloat
 # from IPython.utils.localinterfaces import LOCALHOST 
 
 from . import heartmonitor
 from .factory import RegistrationFactory
 from .streamkernel import Kernel
 from .streamsession import Message
 from .util import disambiguate_url
 
 class EngineFactory(RegistrationFactory):
     """IPython engine"""
     
     # configurables:
     user_ns=Dict(config=True)
     out_stream_factory=Type('IPython.zmq.iostream.OutStream', config=True)
     display_hook_factory=Type('IPython.zmq.displayhook.DisplayHook', config=True)
     location=Str(config=True)
     timeout=CFloat(2,config=True)
     
     # not configurable:
     id=Int(allow_none=True)
     registrar=Instance('zmq.eventloop.zmqstream.ZMQStream')
     kernel=Instance(Kernel)
     
     
     def __init__(self, **kwargs):
         super(EngineFactory, self).__init__(**kwargs)
         ctx = self.context
         
         reg = ctx.socket(zmq.XREQ)
         reg.setsockopt(zmq.IDENTITY, self.ident)
         reg.connect(self.url)
         self.registrar = zmqstream.ZMQStream(reg, self.loop)
         
     def register(self):
         """send the registration_request"""
         
         self.log.info("registering")
         content = dict(queue=self.ident, heartbeat=self.ident, control=self.ident)
         self.registrar.on_recv(self.complete_registration)
         # print (self.session.key)
         self.session.send(self.registrar, "registration_request",content=content)
     
     def complete_registration(self, msg):
         # print msg
         self._abort_dc.stop()
         ctx = self.context
         loop = self.loop
         identity = self.ident
         
         idents,msg = self.session.feed_identities(msg)
         msg = Message(self.session.unpack_message(msg))
         
         if msg.content.status == 'ok':
             self.id = int(msg.content.id)
             
             # create Shell Streams (MUX, Task, etc.):
             queue_addr = msg.content.mux
             shell_addrs = [ str(queue_addr) ]
             task_addr = msg.content.task
             if task_addr:
                 shell_addrs.append(str(task_addr))
             
             # Uncomment this to go back to two-socket model
             # shell_streams = []
             # for addr in shell_addrs:
             #     stream = zmqstream.ZMQStream(ctx.socket(zmq.XREP), loop)
             #     stream.setsockopt(zmq.IDENTITY, identity)
             #     stream.connect(disambiguate_url(addr, self.location))
             #     shell_streams.append(stream)
             
             # Now use only one shell stream for mux and tasks
             stream = zmqstream.ZMQStream(ctx.socket(zmq.XREP), loop)
             stream.setsockopt(zmq.IDENTITY, identity)
             shell_streams = [stream]
             for addr in shell_addrs:
                 stream.connect(disambiguate_url(addr, self.location))
             # end single stream-socket
             
             # control stream:
             control_addr = str(msg.content.control)
             control_stream = zmqstream.ZMQStream(ctx.socket(zmq.XREP), loop)
             control_stream.setsockopt(zmq.IDENTITY, identity)
             control_stream.connect(disambiguate_url(control_addr, self.location))
             
             # create iopub stream:
             iopub_addr = msg.content.iopub
             iopub_stream = zmqstream.ZMQStream(ctx.socket(zmq.PUB), loop)
             iopub_stream.setsockopt(zmq.IDENTITY, identity)
             iopub_stream.connect(disambiguate_url(iopub_addr, self.location))
             
             # launch heartbeat
             hb_addrs = msg.content.heartbeat
             # print (hb_addrs)
             
             # # Redirect input streams and set a display hook.
             if self.out_stream_factory:
                 sys.stdout = self.out_stream_factory(self.session, iopub_stream, u'stdout')
                 sys.stdout.topic = 'engine.%i.stdout'%self.id
                 sys.stderr = self.out_stream_factory(self.session, iopub_stream, u'stderr')
                 sys.stderr.topic = 'engine.%i.stderr'%self.id
             if self.display_hook_factory:
                 sys.displayhook = self.display_hook_factory(self.session, iopub_stream)
                 sys.displayhook.topic = 'engine.%i.pyout'%self.id
             
             self.kernel = Kernel(config=self.config, int_id=self.id, ident=self.ident, session=self.session, 
                     control_stream=control_stream, shell_streams=shell_streams, iopub_stream=iopub_stream, 
                     loop=loop, user_ns = self.user_ns, logname=self.log.name)
             self.kernel.start()
             hb_addrs = [ disambiguate_url(addr, self.location) for addr in hb_addrs ]
             heart = heartmonitor.Heart(*map(str, hb_addrs), heart_id=identity)
             # ioloop.DelayedCallback(heart.start, 1000, self.loop).start()
             heart.start()
             
             
         else:
             self.log.fatal("Registration Failed: %s"%msg)
             raise Exception("Registration Failed: %s"%msg)
         
         self.log.info("Completed registration with id %i"%self.id)
     
     
     def abort(self):
         self.log.fatal("Registration timed out")
         self.session.send(self.registrar, "unregistration_request", content=dict(id=self.id))
         time.sleep(1)
         sys.exit(255)
     
     def start(self):
         dc = ioloop.DelayedCallback(self.register, 0, self.loop)
         dc.start()
         self._abort_dc = ioloop.DelayedCallback(self.abort, self.timeout*1000, self.loop)
         self._abort_dc.start()
 

IPython/zmq/parallel/entry_point.py

@@ -1,119 +1,125 b''
 """ Defines helper functions for creating kernel entry points and process
 launchers.
 
 ************
 NOTE: Most of this module has been deprecated by moving to Configurables
 ************
 """
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 # Standard library imports.
 import atexit
 import logging
 import os
 import stat
 import socket
 import sys
 from signal import signal, SIGINT, SIGABRT, SIGTERM
 from subprocess import Popen, PIPE
 try:
     from signal import SIGKILL
 except ImportError:
     SIGKILL=None
 
 # System library imports.
 import zmq
 from zmq.log import handlers
 
 # Local imports.
 from IPython.core.ultratb import FormattedTB
 from IPython.external.argparse import ArgumentParser
 from IPython.zmq.log import EnginePUBHandler
 
 _random_ports = set()
 
 def select_random_ports(n):
     """Selects and return n random ports that are available."""
     ports = []
     for i in xrange(n):
         sock = socket.socket()
         sock.bind(('', 0))
         while sock.getsockname()[1] in _random_ports:
             sock.close()
             sock = socket.socket()
             sock.bind(('', 0))
         ports.append(sock)
     for i, sock in enumerate(ports):
         port = sock.getsockname()[1]
         sock.close()
         ports[i] = port
         _random_ports.add(port)
     return ports
 
 def signal_children(children):
     """Relay interupt/term signals to children, for more solid process cleanup."""
     def terminate_children(sig, frame):
         logging.critical("Got signal %i, terminating children..."%sig)
         for child in children:
             child.terminate()
         
         sys.exit(sig != SIGINT)
         # sys.exit(sig)
     for sig in (SIGINT, SIGABRT, SIGTERM):
         signal(sig, terminate_children)
 
 def generate_exec_key(keyfile):
     import uuid
     newkey = str(uuid.uuid4())
     with open(keyfile, 'w') as f:
         # f.write('ipython-key ')
         f.write(newkey+'\n')
     # set user-only RW permissions (0600)
     # this will have no effect on Windows
     os.chmod(keyfile, stat.S_IRUSR|stat.S_IWUSR)
 
 
 def integer_loglevel(loglevel):
     try:
         loglevel = int(loglevel)
     except ValueError:
         if isinstance(loglevel, str):
             loglevel = getattr(logging, loglevel)
     return loglevel
 
 def connect_logger(logname, context, iface, root="ip", loglevel=logging.DEBUG):
     logger = logging.getLogger(logname)
     if any([isinstance(h, handlers.PUBHandler) for h in logger.handlers]):
         # don't add a second PUBHandler
         return
     loglevel = integer_loglevel(loglevel)
     lsock = context.socket(zmq.PUB)
     lsock.connect(iface)
     handler = handlers.PUBHandler(lsock)
     handler.setLevel(loglevel)
     handler.root_topic = root
     logger.addHandler(handler)
     logger.setLevel(loglevel)
 
 def connect_engine_logger(context, iface, engine, loglevel=logging.DEBUG):
     logger = logging.getLogger()
     if any([isinstance(h, handlers.PUBHandler) for h in logger.handlers]):
         # don't add a second PUBHandler
         return
     loglevel = integer_loglevel(loglevel)
     lsock = context.socket(zmq.PUB)
     lsock.connect(iface)
     handler = EnginePUBHandler(engine, lsock)
     handler.setLevel(loglevel)
     logger.addHandler(handler)
     logger.setLevel(loglevel)
 
 def local_logger(logname, loglevel=logging.DEBUG):
     loglevel = integer_loglevel(loglevel)
     logger = logging.getLogger(logname)
     if any([isinstance(h, logging.StreamHandler) for h in logger.handlers]):
         # don't add a second StreamHandler
         return
     handler = logging.StreamHandler()
     handler.setLevel(loglevel)
     logger.addHandler(handler)
     logger.setLevel(loglevel)

IPython/zmq/parallel/heartmonitor.py

@@ -1,158 +1,164 b''
 #!/usr/bin/env python
 """
 A multi-heart Heartbeat system using PUB and XREP sockets. pings are sent out on the PUB,
 and hearts are tracked based on their XREQ identities.
 """
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 from __future__ import print_function
 import time
 import logging
 import uuid
 
 import zmq
 from zmq.devices import ProcessDevice,ThreadDevice
 from zmq.eventloop import ioloop, zmqstream
 
 from IPython.utils.traitlets import Set, Instance, CFloat, Bool
 from .factory import LoggingFactory
 
 class Heart(object):
     """A basic heart object for responding to a HeartMonitor.
     This is a simple wrapper with defaults for the most common
     Device model for responding to heartbeats.
     
     It simply builds a threadsafe zmq.FORWARDER Device, defaulting to using 
     SUB/XREQ for in/out.
     
     You can specify the XREQ's IDENTITY via the optional heart_id argument."""
     device=None
     id=None
     def __init__(self, in_addr, out_addr, in_type=zmq.SUB, out_type=zmq.XREQ, heart_id=None):
         self.device = ThreadDevice(zmq.FORWARDER, in_type, out_type)
         self.device.daemon=True
         self.device.connect_in(in_addr)
         self.device.connect_out(out_addr)
         if in_type == zmq.SUB:
             self.device.setsockopt_in(zmq.SUBSCRIBE, "")
         if heart_id is None:
             heart_id = str(uuid.uuid4())
         self.device.setsockopt_out(zmq.IDENTITY, heart_id)
         self.id = heart_id
     
     def start(self):
         return self.device.start()
         
 class HeartMonitor(LoggingFactory):
     """A basic HeartMonitor class
     pingstream: a PUB stream
     pongstream: an XREP stream
     period: the period of the heartbeat in milliseconds"""
     
     period=CFloat(1000, config=True) # in milliseconds
     
     pingstream=Instance('zmq.eventloop.zmqstream.ZMQStream')
     pongstream=Instance('zmq.eventloop.zmqstream.ZMQStream')
     loop = Instance('zmq.eventloop.ioloop.IOLoop')
     def _loop_default(self):
         return ioloop.IOLoop.instance()
     debug=Bool(False)
     
     # not settable:
     hearts=Set()
     responses=Set()
     on_probation=Set()
     last_ping=CFloat(0)
     _new_handlers = Set()
     _failure_handlers = Set()
     lifetime = CFloat(0)
     tic = CFloat(0)
     
     def __init__(self, **kwargs):
         super(HeartMonitor, self).__init__(**kwargs)
         
         self.pongstream.on_recv(self.handle_pong)
     
     def start(self):
         self.caller = ioloop.PeriodicCallback(self.beat, self.period, self.loop)
         self.caller.start()
     
     def add_new_heart_handler(self, handler):
         """add a new handler for new hearts"""
         self.log.debug("heartbeat::new_heart_handler: %s"%handler)
         self._new_handlers.add(handler)
         
     def add_heart_failure_handler(self, handler):
         """add a new handler for heart failure"""
         self.log.debug("heartbeat::new heart failure handler: %s"%handler)
         self._failure_handlers.add(handler)
             
     def beat(self):
         self.pongstream.flush() 
         self.last_ping = self.lifetime
         
         toc = time.time()
         self.lifetime += toc-self.tic
         self.tic = toc
         # self.log.debug("heartbeat::%s"%self.lifetime)
         goodhearts = self.hearts.intersection(self.responses)
         missed_beats = self.hearts.difference(goodhearts)
         heartfailures = self.on_probation.intersection(missed_beats)
         newhearts = self.responses.difference(goodhearts)
         map(self.handle_new_heart, newhearts)
         map(self.handle_heart_failure, heartfailures)
         self.on_probation = missed_beats.intersection(self.hearts)
         self.responses = set()
         # print self.on_probation, self.hearts
         # self.log.debug("heartbeat::beat %.3f, %i beating hearts"%(self.lifetime, len(self.hearts)))
         self.pingstream.send(str(self.lifetime))
     
     def handle_new_heart(self, heart):
         if self._new_handlers:
             for handler in self._new_handlers:
                 handler(heart)
         else:
             self.log.info("heartbeat::yay, got new heart %s!"%heart)
         self.hearts.add(heart)
     
     def handle_heart_failure(self, heart):
         if self._failure_handlers:
             for handler in self._failure_handlers:
                 try:
                     handler(heart)
                 except Exception as e:
                     self.log.error("heartbeat::Bad Handler! %s"%handler, exc_info=True)
                     pass
         else:
             self.log.info("heartbeat::Heart %s failed :("%heart)
         self.hearts.remove(heart)
         
     
     def handle_pong(self, msg):
         "a heart just beat"
         if msg[1] == str(self.lifetime):
             delta = time.time()-self.tic
             # self.log.debug("heartbeat::heart %r took %.2f ms to respond"%(msg[0], 1000*delta))
             self.responses.add(msg[0])
         elif msg[1] == str(self.last_ping):
             delta = time.time()-self.tic + (self.lifetime-self.last_ping)
             self.log.warn("heartbeat::heart %r missed a beat, and took %.2f ms to respond"%(msg[0], 1000*delta))
             self.responses.add(msg[0])
         else:
             self.log.warn("heartbeat::got bad heartbeat (possibly old?): %s (current=%.3f)"%
             (msg[1],self.lifetime))
 
 
 if __name__ == '__main__':
     loop = ioloop.IOLoop.instance()
     context = zmq.Context()
     pub = context.socket(zmq.PUB)
     pub.bind('tcp://127.0.0.1:5555')
     xrep = context.socket(zmq.XREP)
     xrep.bind('tcp://127.0.0.1:5556')
     
     outstream = zmqstream.ZMQStream(pub, loop)
     instream = zmqstream.ZMQStream(xrep, loop)
     
     hb = HeartMonitor(loop, outstream, instream)
     
     loop.start()

IPython/zmq/parallel/ipcluster.py

@@ -1,90 +1,97 b''
 #!/usr/bin/env python
+"""Old ipcluster script.  Possibly to be removed."""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 from __future__ import print_function
 
 import os
 import sys
 import time
 from subprocess import Popen, PIPE
 
 from IPython.external.argparse import ArgumentParser, SUPPRESS
 
 def _filter_arg(flag, args):
     filtered = []
     if flag in args:
         filtered.append(flag)
         idx = args.index(flag)
         if len(args) > idx+1:
             if not args[idx+1].startswith('-'):
                 filtered.append(args[idx+1])
     return filtered
 
 def filter_args(flags, args=sys.argv[1:]):
     filtered = []
     for flag in flags:
         if isinstance(flag, (list,tuple)):
             for f in flag:
                 filtered.extend(_filter_arg(f, args))
         else:
             filtered.extend(_filter_arg(flag, args))
     return filtered
 
 def _strip_arg(flag, args):
     while flag in args:
         idx = args.index(flag)
         args.pop(idx)
         if len(args) > idx:
             if not args[idx].startswith('-'):
                 args.pop(idx)
 
 def strip_args(flags, args=sys.argv[1:]):
     args = list(args)
     for flag in flags:
         if isinstance(flag, (list,tuple)):
             for f in flag:
                 _strip_arg(f, args)
         else:
             _strip_arg(flag, args)
     return args
     
 
 def launch_process(mod, args):
     """Launch a controller or engine in a subprocess."""
     code = "from IPython.zmq.parallel.%s import launch_new_instance;launch_new_instance()"%mod
     arguments = [ sys.executable, '-c', code ] + args
     blackholew = file(os.devnull, 'w')
     blackholer = file(os.devnull, 'r')
     
     proc = Popen(arguments, stdin=blackholer, stdout=blackholew, stderr=PIPE)
     return proc
 
 def main():
     parser = ArgumentParser(argument_default=SUPPRESS)
     parser.add_argument('--n', '-n', type=int, default=1,
                 help="The number of engines to start.")
     ns,args = parser.parse_known_args()
     n = ns.n
     
     controller = launch_process('ipcontrollerapp', args)
     for i in range(10):
         time.sleep(.1)
         if controller.poll() is not None:
             print("Controller failed to launch:")
             print (controller.stderr.read())
             sys.exit(255)
     
     print("Launched Controller")
     engines = [ launch_process('ipengineapp', args+['--ident', 'engine-%i'%i]) for i in range(n) ]
     print("%i Engines started"%n)
     
     def wait_quietly(p):
         try:
             p.wait()
         except KeyboardInterrupt:
             pass
     
     wait_quietly(controller)
     map(wait_quietly, engines)
     print ("Engines cleaned up.")
 
 if __name__ == '__main__':
     main()
  No newline at end of file

IPython/zmq/parallel/kernelstarter.py

@@ -1,219 +1,225 b''
 """KernelStarter class that intercepts Control Queue messages, and handles process management."""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 from zmq.eventloop import ioloop
 
 from .streamsession import StreamSession
 
 class KernelStarter(object):
     """Object for resetting/killing the Kernel."""
     
     
     def __init__(self, session, upstream, downstream, *kernel_args, **kernel_kwargs):
         self.session = session
         self.upstream = upstream
         self.downstream = downstream
         self.kernel_args = kernel_args
         self.kernel_kwargs = kernel_kwargs
         self.handlers = {}
         for method in 'shutdown_request shutdown_reply'.split():
             self.handlers[method] = getattr(self, method)
     
     def start(self):
         self.upstream.on_recv(self.dispatch_request)
         self.downstream.on_recv(self.dispatch_reply)
         
     #--------------------------------------------------------------------------
     # Dispatch methods
     #--------------------------------------------------------------------------
     
     def dispatch_request(self, raw_msg):
         idents, msg = self.session.feed_identities()
         try:
             msg = self.session.unpack_message(msg, content=False)
         except:
             print ("bad msg: %s"%msg)
         
         msgtype = msg['msg_type']
         handler = self.handlers.get(msgtype, None)
         if handler is None:
             self.downstream.send_multipart(raw_msg, copy=False)
         else:
             handler(msg)
         
     def dispatch_reply(self, raw_msg):
         idents, msg = self.session.feed_identities()
         try:
             msg = self.session.unpack_message(msg, content=False)
         except:
             print ("bad msg: %s"%msg)
         
         msgtype = msg['msg_type']
         handler = self.handlers.get(msgtype, None)
         if handler is None:
             self.upstream.send_multipart(raw_msg, copy=False)
         else:
             handler(msg)
     
     #--------------------------------------------------------------------------
     # Handlers
     #--------------------------------------------------------------------------
     
     def shutdown_request(self, msg):
         """"""
         self.downstream.send_multipart(msg)
     
     #--------------------------------------------------------------------------
     # Kernel process management methods, from KernelManager:
     #--------------------------------------------------------------------------
     
     def _check_local(addr):
         if isinstance(addr, tuple):
             addr = addr[0]
         return addr in LOCAL_IPS
     
     def start_kernel(self, **kw):
         """Starts a kernel process and configures the manager to use it.
 
         If random ports (port=0) are being used, this method must be called
         before the channels are created.
 
         Parameters:
         -----------
         ipython : bool, optional (default True)
              Whether to use an IPython kernel instead of a plain Python kernel.
         """
         self.kernel = Process(target=make_kernel, args=self.kernel_args,
                                             kwargs=self.kernel_kwargs)
 
     def shutdown_kernel(self, restart=False):
         """ Attempts to the stop the kernel process cleanly. If the kernel
         cannot be stopped, it is killed, if possible.
         """
         # FIXME: Shutdown does not work on Windows due to ZMQ errors!
         if sys.platform == 'win32':
             self.kill_kernel()
             return
 
         # Don't send any additional kernel kill messages immediately, to give
         # the kernel a chance to properly execute shutdown actions. Wait for at
         # most 1s, checking every 0.1s.
         self.xreq_channel.shutdown(restart=restart)
         for i in range(10):
             if self.is_alive:
                 time.sleep(0.1)
             else:
                 break
         else:
             # OK, we've waited long enough.
             if self.has_kernel:
                 self.kill_kernel()
     
     def restart_kernel(self, now=False):
         """Restarts a kernel with the same arguments that were used to launch
         it. If the old kernel was launched with random ports, the same ports
         will be used for the new kernel.
 
         Parameters
         ----------
         now : bool, optional
           If True, the kernel is forcefully restarted *immediately*, without
           having a chance to do any cleanup action.  Otherwise the kernel is
           given 1s to clean up before a forceful restart is issued.
 
           In all cases the kernel is restarted, the only difference is whether
           it is given a chance to perform a clean shutdown or not.
         """
         if self._launch_args is None:
             raise RuntimeError("Cannot restart the kernel. "
                                "No previous call to 'start_kernel'.")
         else:
             if self.has_kernel:
                 if now:
                     self.kill_kernel()
                 else:
                     self.shutdown_kernel(restart=True)
             self.start_kernel(**self._launch_args)
 
             # FIXME: Messages get dropped in Windows due to probable ZMQ bug
             # unless there is some delay here.
             if sys.platform == 'win32':
                 time.sleep(0.2)
 
     @property
     def has_kernel(self):
         """Returns whether a kernel process has been specified for the kernel
         manager.
         """
         return self.kernel is not None
 
     def kill_kernel(self):
         """ Kill the running kernel. """
         if self.has_kernel:
             # Pause the heart beat channel if it exists.
             if self._hb_channel is not None:
                 self._hb_channel.pause()
 
             # Attempt to kill the kernel.
             try:
                 self.kernel.kill()
             except OSError, e:
                 # In Windows, we will get an Access Denied error if the process
                 # has already terminated. Ignore it.
                 if not (sys.platform == 'win32' and e.winerror == 5):
                     raise
             self.kernel = None
         else:
             raise RuntimeError("Cannot kill kernel. No kernel is running!")
 
     def interrupt_kernel(self):
         """ Interrupts the kernel. Unlike ``signal_kernel``, this operation is
         well supported on all platforms.
         """
         if self.has_kernel:
             if sys.platform == 'win32':
                 from parentpoller import ParentPollerWindows as Poller
                 Poller.send_interrupt(self.kernel.win32_interrupt_event)
             else:
                 self.kernel.send_signal(signal.SIGINT)
         else:
             raise RuntimeError("Cannot interrupt kernel. No kernel is running!")
 
     def signal_kernel(self, signum):
         """ Sends a signal to the kernel. Note that since only SIGTERM is
         supported on Windows, this function is only useful on Unix systems.
         """
         if self.has_kernel:
             self.kernel.send_signal(signum)
         else:
             raise RuntimeError("Cannot signal kernel. No kernel is running!")
 
     @property
     def is_alive(self):
         """Is the kernel process still running?"""
         # FIXME: not using a heartbeat means this method is broken for any
         # remote kernel, it's only capable of handling local kernels.
         if self.has_kernel:
             if self.kernel.poll() is None:
                 return True
             else:
                 return False
         else:
             # We didn't start the kernel with this KernelManager so we don't
             # know if it is running. We should use a heartbeat for this case.
             return True
 
 
 def make_starter(up_addr, down_addr, *args, **kwargs):
     """entry point function for launching a kernelstarter in a subprocess"""
     loop = ioloop.IOLoop.instance()
     ctx = zmq.Context()
     session = StreamSession()
     upstream = zmqstream.ZMQStream(ctx.socket(zmq.XREQ),loop)
     upstream.connect(up_addr)
     downstream = zmqstream.ZMQStream(ctx.socket(zmq.XREQ),loop)
     downstream.connect(down_addr)
     
     starter = KernelStarter(session, upstream, downstream, *args, **kwargs)
     starter.start()
     loop.start()
     
  No newline at end of file

IPython/zmq/parallel/remotenamespace.py

@@ -1,95 +1,101 b''
 """RemoteNamespace object, for dict style interaction with a remote
 execution kernel."""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 from functools import wraps
 from IPython.external.decorator import decorator
 
 def _clear():
     globals().clear()
 
 @decorator
 def spinfirst(f):
     @wraps(f)
     def spun_method(self, *args, **kwargs):
         self.spin()
         return f(self, *args, **kwargs)
     return spun_method
 
 @decorator
 def myblock(f, self, *args, **kwargs):
     block = self.client.block
     self.client.block = self.block
     ret = f(self, *args, **kwargs)
     self.client.block = block
     return ret
 
 class RemoteNamespace(object):
     """A RemoteNamespace object, providing dictionary 
     access to an engine via an IPython.zmq.client object.
     
     
     """
     client = None
     queue = None
     id = None
     block = False
     
     def __init__(self, client, id):
         self.client = client
         self.id = id
         self.block = client.block # initial state is same as client
     
     def __repr__(self):
         return "<RemoteNamespace[%i]>"%self.id
     
     @myblock
     def apply(self, f, *args, **kwargs):
         """call f(*args, **kwargs) in remote namespace
         
         This method has no access to the user namespace"""
         return self.client.apply_to(self.id, f, *args, **kwargs)
     
     @myblock
     def apply_bound(self, f, *args, **kwargs):
         """call `f(*args, **kwargs)` in remote namespace.
         
         `f` will have access to the user namespace as globals()."""
         return self.client.apply_bound_to(self.id, f, *args, **kwargs)
     
     @myblock
     def update(self, ns):
         """update remote namespace with dict `ns`"""
         return self.client.push(self.id, ns, self.block)
     
     def get(self, key_s):
         """get object(s) by `key_s` from remote namespace
         will return one object if it is a key.
         It also takes a list of keys, and will return a list of objects."""
         return self.client.pull(self.id, key_s, self.block)
     
     push = update
     pull = get
     
     def __getitem__(self, key):
         return self.get(key)
     
     def __setitem__(self,key,value):
         self.update({key:value})
     
     def clear(self):
         """clear the remote namespace"""
         return self.client.apply_bound_to(self.id, _clear)
     
     @decorator
     def withme(self, toapply):
         """for use as a decorator, this turns a function into
         one that executes remotely."""
         @wraps(toapply)
         def applied(self, *args, **kwargs):
             return self.apply_bound(self, toapply, *args, **kwargs)
         return applied
         
         
 
 
 

IPython/zmq/parallel/scheduler.py

@@ -1,584 +1,590 b''
 """The Python scheduler for rich scheduling.
 
 The Pure ZMQ scheduler does not allow routing schemes other than LRU,
 nor does it check msg_id DAG dependencies. For those, a slightly slower
 Python Scheduler exists.
 """
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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
 #----------------------------------------------------------------------
 
 from __future__ import print_function
 
 import logging
 import sys
 
 from datetime import datetime, timedelta
 from random import randint, random
 from types import FunctionType
 
 try:
     import numpy
 except ImportError:
     numpy = None
 
 import zmq
 from zmq.eventloop import ioloop, zmqstream
 
 # local imports
 from IPython.external.decorator import decorator
 from IPython.utils.traitlets import Instance, Dict, List, Set
 
 from . import error
 from .dependency import Dependency
 from .entry_point import connect_logger, local_logger
 from .factory import SessionFactory
 
 
 @decorator
 def logged(f,self,*args,**kwargs):
     # print ("#--------------------")
     self.log.debug("scheduler::%s(*%s,**%s)"%(f.func_name, args, kwargs))
     # print ("#--")
     return f(self,*args, **kwargs)
 
 #----------------------------------------------------------------------
 # Chooser functions
 #----------------------------------------------------------------------
 
 def plainrandom(loads):
     """Plain random pick."""
     n = len(loads)
     return randint(0,n-1)
 
 def lru(loads):
     """Always pick the front of the line.
     
     The content of `loads` is ignored.
     
     Assumes LRU ordering of loads, with oldest first.
     """
     return 0
 
 def twobin(loads):
     """Pick two at random, use the LRU of the two.
     
     The content of loads is ignored.
     
     Assumes LRU ordering of loads, with oldest first.
     """
     n = len(loads)
     a = randint(0,n-1)
     b = randint(0,n-1)
     return min(a,b)
 
 def weighted(loads):
     """Pick two at random using inverse load as weight.
     
     Return the less loaded of the two.
     """
     # weight 0 a million times more than 1:
     weights = 1./(1e-6+numpy.array(loads))
     sums = weights.cumsum()
     t = sums[-1]
     x = random()*t
     y = random()*t
     idx = 0
     idy = 0
     while sums[idx] < x:
         idx += 1
     while sums[idy] < y:
         idy += 1
     if weights[idy] > weights[idx]:
         return idy
     else:
         return idx
 
 def leastload(loads):
     """Always choose the lowest load.
     
     If the lowest load occurs more than once, the first
     occurance will be used.  If loads has LRU ordering, this means
     the LRU of those with the lowest load is chosen.
     """
     return loads.index(min(loads))
 
 #---------------------------------------------------------------------
 # Classes
 #---------------------------------------------------------------------
 # store empty default dependency:
 MET = Dependency([])
 
 class TaskScheduler(SessionFactory):
     """Python TaskScheduler object.
     
     This is the simplest object that supports msg_id based
     DAG dependencies. *Only* task msg_ids are checked, not
     msg_ids of jobs submitted via the MUX queue.
     
     """
     
     # input arguments:
     scheme = Instance(FunctionType, default=leastload) # function for determining the destination
     client_stream = Instance(zmqstream.ZMQStream) # client-facing stream
     engine_stream = Instance(zmqstream.ZMQStream) # engine-facing stream
     notifier_stream = Instance(zmqstream.ZMQStream) # hub-facing sub stream
     mon_stream = Instance(zmqstream.ZMQStream) # hub-facing pub stream
     
     # internals:
     graph = Dict() # dict by msg_id of [ msg_ids that depend on key ]
     depending = Dict() # dict by msg_id of (msg_id, raw_msg, after, follow)
     pending = Dict() # dict by engine_uuid of submitted tasks
     completed = Dict() # dict by engine_uuid of completed tasks
     failed = Dict() # dict by engine_uuid of failed tasks
     destinations = Dict() # dict by msg_id of engine_uuids where jobs ran (reverse of completed+failed)
     clients = Dict() # dict by msg_id for who submitted the task
     targets = List() # list of target IDENTs
     loads = List() # list of engine loads
     all_completed = Set() # set of all completed tasks
     all_failed = Set() # set of all failed tasks
     all_done = Set() # set of all finished tasks=union(completed,failed)
     all_ids = Set() # set of all submitted task IDs
     blacklist = Dict() # dict by msg_id of locations where a job has encountered UnmetDependency
     auditor = Instance('zmq.eventloop.ioloop.PeriodicCallback')
     
     
     def start(self):
         self.engine_stream.on_recv(self.dispatch_result, copy=False)
         self._notification_handlers = dict(
             registration_notification = self._register_engine,
             unregistration_notification = self._unregister_engine
         )
         self.notifier_stream.on_recv(self.dispatch_notification)
         self.auditor = ioloop.PeriodicCallback(self.audit_timeouts, 2e3, self.loop) # 1 Hz
         self.auditor.start()
         self.log.info("Scheduler started...%r"%self)
     
     def resume_receiving(self):
         """Resume accepting jobs."""
         self.client_stream.on_recv(self.dispatch_submission, copy=False)
     
     def stop_receiving(self):
         """Stop accepting jobs while there are no engines.
         Leave them in the ZMQ queue."""
         self.client_stream.on_recv(None)
     
     #-----------------------------------------------------------------------
     # [Un]Registration Handling
     #-----------------------------------------------------------------------
     
     def dispatch_notification(self, msg):
         """dispatch register/unregister events."""
         idents,msg = self.session.feed_identities(msg)
         msg = self.session.unpack_message(msg)
         msg_type = msg['msg_type']
         handler = self._notification_handlers.get(msg_type, None)
         if handler is None:
             raise Exception("Unhandled message type: %s"%msg_type)
         else:
             try:
                 handler(str(msg['content']['queue']))
             except KeyError:
                 self.log.error("task::Invalid notification msg: %s"%msg)
     
     @logged
     def _register_engine(self, uid):
         """New engine with ident `uid` became available."""
         # head of the line:
         self.targets.insert(0,uid)
         self.loads.insert(0,0)
         # initialize sets
         self.completed[uid] = set()
         self.failed[uid] = set()
         self.pending[uid] = {}
         if len(self.targets) == 1:
             self.resume_receiving()
 
     def _unregister_engine(self, uid):
         """Existing engine with ident `uid` became unavailable."""
         if len(self.targets) == 1:
             # this was our only engine
             self.stop_receiving()
         
         # handle any potentially finished tasks:
         self.engine_stream.flush()
         
         self.completed.pop(uid)
         self.failed.pop(uid)
         # don't pop destinations, because it might be used later
         # map(self.destinations.pop, self.completed.pop(uid))
         # map(self.destinations.pop, self.failed.pop(uid))
         
         idx = self.targets.index(uid)
         self.targets.pop(idx)
         self.loads.pop(idx)
         
         # wait 5 seconds before cleaning up pending jobs, since the results might
         # still be incoming
         if self.pending[uid]:
             dc = ioloop.DelayedCallback(lambda : self.handle_stranded_tasks(uid), 5000, self.loop)
             dc.start()
     
     @logged
     def handle_stranded_tasks(self, engine):
         """Deal with jobs resident in an engine that died."""
         lost = self.pending.pop(engine)
         
         for msg_id, (raw_msg, targets, MET, follow, timeout) in lost.iteritems():
             self.all_failed.add(msg_id)
             self.all_done.add(msg_id)
             idents,msg = self.session.feed_identities(raw_msg, copy=False)
             msg = self.session.unpack_message(msg, copy=False, content=False)
             parent = msg['header']
             idents = [idents[0],engine]+idents[1:]
             print (idents)
             try:
                 raise error.EngineError("Engine %r died while running task %r"%(engine, msg_id))
             except:
                 content = error.wrap_exception()
             msg = self.session.send(self.client_stream, 'apply_reply', content, 
                                                     parent=parent, ident=idents)
             self.session.send(self.mon_stream, msg, ident=['outtask']+idents)
             self.update_graph(msg_id)
     
     
     #-----------------------------------------------------------------------
     # Job Submission
     #-----------------------------------------------------------------------
     @logged
     def dispatch_submission(self, raw_msg):
         """Dispatch job submission to appropriate handlers."""
         # ensure targets up to date:
         self.notifier_stream.flush()
         try:
             idents, msg = self.session.feed_identities(raw_msg, copy=False)
             msg = self.session.unpack_message(msg, content=False, copy=False)
         except:
             self.log.error("task::Invaid task: %s"%raw_msg, exc_info=True)
             return
         
         # send to monitor
         self.mon_stream.send_multipart(['intask']+raw_msg, copy=False)
         
         header = msg['header']
         msg_id = header['msg_id']
         self.all_ids.add(msg_id)
         
         # targets
         targets = set(header.get('targets', []))
         
         # time dependencies
         after = Dependency(header.get('after', []))
         if after.all:
             after.difference_update(self.all_completed)
             if not after.success_only:
                 after.difference_update(self.all_failed)
         if after.check(self.all_completed, self.all_failed):
             # recast as empty set, if `after` already met,
             # to prevent unnecessary set comparisons
             after = MET
         
         # location dependencies
         follow = Dependency(header.get('follow', []))
         
         # turn timeouts into datetime objects:
         timeout = header.get('timeout', None)
         if timeout:
             timeout = datetime.now() + timedelta(0,timeout,0)
         
         args = [raw_msg, targets, after, follow, timeout]
         
         # validate and reduce dependencies:
         for dep in after,follow:
             # check valid:
             if msg_id in dep or dep.difference(self.all_ids):
                 self.depending[msg_id] = args
                 return self.fail_unreachable(msg_id, error.InvalidDependency)
             # check if unreachable:
             if dep.unreachable(self.all_failed):
                 self.depending[msg_id] = args
                 return self.fail_unreachable(msg_id)
         
         if after.check(self.all_completed, self.all_failed):
             # time deps already met, try to run
             if not self.maybe_run(msg_id, *args):
                 # can't run yet
                 self.save_unmet(msg_id, *args)
         else:
             self.save_unmet(msg_id, *args)
     
     # @logged
     def audit_timeouts(self):
         """Audit all waiting tasks for expired timeouts."""
         now = datetime.now()
         for msg_id in self.depending.keys():
             # must recheck, in case one failure cascaded to another:
             if msg_id in self.depending:
                 raw,after,targets,follow,timeout = self.depending[msg_id]
                 if timeout and timeout < now:
                     self.fail_unreachable(msg_id, timeout=True)
                 
     @logged
     def fail_unreachable(self, msg_id, why=error.ImpossibleDependency):
         """a task has become unreachable, send a reply with an ImpossibleDependency
         error."""
         if msg_id not in self.depending:
             self.log.error("msg %r already failed!"%msg_id)
             return
         raw_msg,targets,after,follow,timeout = self.depending.pop(msg_id)
         for mid in follow.union(after):
             if mid in self.graph:
                 self.graph[mid].remove(msg_id)
         
         # FIXME: unpacking a message I've already unpacked, but didn't save:
         idents,msg = self.session.feed_identities(raw_msg, copy=False)
         msg = self.session.unpack_message(msg, copy=False, content=False)
         header = msg['header']
         
         try:
             raise why()
         except:
             content = error.wrap_exception()
         
         self.all_done.add(msg_id)
         self.all_failed.add(msg_id)
         
         msg = self.session.send(self.client_stream, 'apply_reply', content, 
                                                 parent=header, ident=idents)
         self.session.send(self.mon_stream, msg, ident=['outtask']+idents)
         
         self.update_graph(msg_id, success=False)
     
     @logged
     def maybe_run(self, msg_id, raw_msg, targets, after, follow, timeout):
         """check location dependencies, and run if they are met."""
         blacklist = self.blacklist.setdefault(msg_id, set())
         if follow or targets or blacklist:
             # we need a can_run filter
             def can_run(idx):
                 target = self.targets[idx]
                 # check targets
                 if targets and target not in targets:
                     return False
                 # check blacklist
                 if target in blacklist:
                     return False
                 # check follow
                 return follow.check(self.completed[target], self.failed[target])
             
             indices = filter(can_run, range(len(self.targets)))
             if not indices:
                 # couldn't run
                 if follow.all:
                     # check follow for impossibility
                     dests = set()
                     relevant = self.all_completed if follow.success_only else self.all_done
                     for m in follow.intersection(relevant):
                         dests.add(self.destinations[m])
                     if len(dests) > 1:
                         self.fail_unreachable(msg_id)
                         return False
                 if targets:
                     # check blacklist+targets for impossibility
                     targets.difference_update(blacklist)
                     if not targets or not targets.intersection(self.targets):
                         self.fail_unreachable(msg_id)
                         return False
                 return False
         else:
             indices = None
             
         self.submit_task(msg_id, raw_msg, targets, follow, timeout, indices)
         return True
             
     @logged
     def save_unmet(self, msg_id, raw_msg, targets, after, follow, timeout):
         """Save a message for later submission when its dependencies are met."""
         self.depending[msg_id] = [raw_msg,targets,after,follow,timeout]
         # track the ids in follow or after, but not those already finished
         for dep_id in after.union(follow).difference(self.all_done):
             if dep_id not in self.graph:
                 self.graph[dep_id] = set()
             self.graph[dep_id].add(msg_id)
     
     @logged
     def submit_task(self, msg_id, raw_msg, targets, follow, timeout, indices=None):
         """Submit a task to any of a subset of our targets."""
         if indices:
             loads = [self.loads[i] for i in indices]
         else:
             loads = self.loads
         idx = self.scheme(loads)
         if indices:
             idx = indices[idx]
         target = self.targets[idx]
         # print (target, map(str, msg[:3]))
         self.engine_stream.send(target, flags=zmq.SNDMORE, copy=False)
         self.engine_stream.send_multipart(raw_msg, copy=False)
         self.add_job(idx)
         self.pending[target][msg_id] = (raw_msg, targets, MET, follow, timeout)
         content = dict(msg_id=msg_id, engine_id=target)
         self.session.send(self.mon_stream, 'task_destination', content=content, 
                         ident=['tracktask',self.session.session])
     
     #-----------------------------------------------------------------------
     # Result Handling
     #-----------------------------------------------------------------------
     @logged
     def dispatch_result(self, raw_msg):
         """dispatch method for result replies"""
         try:
             idents,msg = self.session.feed_identities(raw_msg, copy=False)
             msg = self.session.unpack_message(msg, content=False, copy=False)
         except:
             self.log.error("task::Invaid result: %s"%raw_msg, exc_info=True)
             return
         
         header = msg['header']
         if header.get('dependencies_met', True):
             success = (header['status'] == 'ok')
             self.handle_result(idents, msg['parent_header'], raw_msg, success)
             # send to Hub monitor
             self.mon_stream.send_multipart(['outtask']+raw_msg, copy=False)
         else:
             self.handle_unmet_dependency(idents, msg['parent_header'])
         
     @logged
     def handle_result(self, idents, parent, raw_msg, success=True):
         """handle a real task result, either success or failure"""
         # first, relay result to client
         engine = idents[0]
         client = idents[1]
         # swap_ids for XREP-XREP mirror
         raw_msg[:2] = [client,engine]
         # print (map(str, raw_msg[:4]))
         self.client_stream.send_multipart(raw_msg, copy=False)
         # now, update our data structures
         msg_id = parent['msg_id']
         self.blacklist.pop(msg_id, None)
         self.pending[engine].pop(msg_id)
         if success:
             self.completed[engine].add(msg_id)
             self.all_completed.add(msg_id)
         else:
             self.failed[engine].add(msg_id)
             self.all_failed.add(msg_id)
         self.all_done.add(msg_id)
         self.destinations[msg_id] = engine
         
         self.update_graph(msg_id, success)
         
     @logged
     def handle_unmet_dependency(self, idents, parent):
         """handle an unmet dependency"""
         engine = idents[0]
         msg_id = parent['msg_id']
         
         if msg_id not in self.blacklist:
             self.blacklist[msg_id] = set()
         self.blacklist[msg_id].add(engine)
         
         args = self.pending[engine].pop(msg_id)
         raw,targets,after,follow,timeout = args
         
         if self.blacklist[msg_id] == targets:
             self.depending[msg_id] = args
             return self.fail_unreachable(msg_id)
         
         elif not self.maybe_run(msg_id, *args):
             # resubmit failed, put it back in our dependency tree
             self.save_unmet(msg_id, *args)
         
     
     @logged
     def update_graph(self, dep_id, success=True):
         """dep_id just finished. Update our dependency
         graph and submit any jobs that just became runable."""
         # print ("\n\n***********")
         # pprint (dep_id)
         # pprint (self.graph)
         # pprint (self.depending)
         # pprint (self.all_completed)
         # pprint (self.all_failed)
         # print ("\n\n***********\n\n")
         if dep_id not in self.graph:
             return
         jobs = self.graph.pop(dep_id)
         
         for msg_id in jobs:
             raw_msg, targets, after, follow, timeout = self.depending[msg_id]
             # if dep_id in after:
             #     if after.all and (success or not after.success_only):
             #         after.remove(dep_id)
             
             if after.unreachable(self.all_failed) or follow.unreachable(self.all_failed):
                 self.fail_unreachable(msg_id)
             
             elif after.check(self.all_completed, self.all_failed): # time deps met, maybe run
                 if self.maybe_run(msg_id, raw_msg, targets, MET, follow, timeout):
                     
                     self.depending.pop(msg_id)
                     for mid in follow.union(after):
                         if mid in self.graph:
                             self.graph[mid].remove(msg_id)
     
     #----------------------------------------------------------------------
     # methods to be overridden by subclasses
     #----------------------------------------------------------------------
     
     def add_job(self, idx):
         """Called after self.targets[idx] just got the job with header.
         Override with subclasses.  The default ordering is simple LRU.
         The default loads are the number of outstanding jobs."""
         self.loads[idx] += 1
         for lis in (self.targets, self.loads):
             lis.append(lis.pop(idx))
             
     
     def finish_job(self, idx):
         """Called after self.targets[idx] just finished a job.
         Override with subclasses."""
         self.loads[idx] -= 1
     
 
 
 def launch_scheduler(in_addr, out_addr, mon_addr, not_addr, config=None,logname='ZMQ', 
                             log_addr=None, loglevel=logging.DEBUG, scheme='lru',
                             identity=b'task'):
     from zmq.eventloop import ioloop
     from zmq.eventloop.zmqstream import ZMQStream
     
     ctx = zmq.Context()
     loop = ioloop.IOLoop()
     print (in_addr, out_addr, mon_addr, not_addr)
     ins = ZMQStream(ctx.socket(zmq.XREP),loop)
     ins.setsockopt(zmq.IDENTITY, identity)
     ins.bind(in_addr)
     
     outs = ZMQStream(ctx.socket(zmq.XREP),loop)
     outs.setsockopt(zmq.IDENTITY, identity)
     outs.bind(out_addr)
     mons = ZMQStream(ctx.socket(zmq.PUB),loop)
     mons.connect(mon_addr)
     nots = ZMQStream(ctx.socket(zmq.SUB),loop)
     nots.setsockopt(zmq.SUBSCRIBE, '')
     nots.connect(not_addr)
     
     scheme = globals().get(scheme, None)
     # setup logging
     if log_addr:
         connect_logger(logname, ctx, log_addr, root="scheduler", loglevel=loglevel)
     else:
         local_logger(logname, loglevel)
     
     scheduler = TaskScheduler(client_stream=ins, engine_stream=outs,
                             mon_stream=mons, notifier_stream=nots,
                             scheme=scheme, loop=loop, logname=logname,
                             config=config)
     scheduler.start()
     try:
         loop.start()
     except KeyboardInterrupt:
         print ("interrupted, exiting...", file=sys.__stderr__)
 

IPython/zmq/parallel/streamkernel.py

@@ -1,487 +1,493 b''
 #!/usr/bin/env python
 """
 Kernel adapted from kernel.py to use ZMQ Streams
 """
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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
 #-----------------------------------------------------------------------------
 
 # Standard library imports.
 from __future__ import print_function
 
 import sys
 import time
 
 from code import CommandCompiler
 from datetime import datetime
 from pprint import pprint
 from signal import SIGTERM, SIGKILL
 
 # System library imports.
 import zmq
 from zmq.eventloop import ioloop, zmqstream
 
 # Local imports.
 from IPython.core import ultratb
 from IPython.utils.traitlets import Instance, List, Int, Dict, Set, Str
 from IPython.zmq.completer import KernelCompleter
 from IPython.zmq.iostream import OutStream
 from IPython.zmq.displayhook import DisplayHook
 
 from . import heartmonitor
 from .client import Client
 from .error import wrap_exception
 from .factory import SessionFactory
 from .streamsession import StreamSession
 from .util import serialize_object, unpack_apply_message, ISO8601, Namespace
 
 def printer(*args):
     pprint(args, stream=sys.__stdout__)
 
 
 class _Passer:
     """Empty class that implements `send()` that does nothing."""
     def send(self, *args, **kwargs):
         pass
     send_multipart = send
     
 
 #-----------------------------------------------------------------------------
 # Main kernel class
 #-----------------------------------------------------------------------------
 
 class Kernel(SessionFactory):
 
     #---------------------------------------------------------------------------
     # Kernel interface
     #---------------------------------------------------------------------------
     
     # kwargs:
     int_id = Int(-1, config=True)
     user_ns = Dict(config=True)
     exec_lines = List(config=True)
     
     control_stream = Instance(zmqstream.ZMQStream)
     task_stream = Instance(zmqstream.ZMQStream)
     iopub_stream = Instance(zmqstream.ZMQStream)
     client = Instance('IPython.zmq.parallel.client.Client')
     
     # internals
     shell_streams = List()
     compiler = Instance(CommandCompiler, (), {})
     completer = Instance(KernelCompleter)
     
     aborted = Set()
     shell_handlers = Dict()
     control_handlers = Dict()
     
     def _set_prefix(self):
         self.prefix = "engine.%s"%self.int_id
     
     def _connect_completer(self):
         self.completer = KernelCompleter(self.user_ns)
     
     def __init__(self, **kwargs):
         super(Kernel, self).__init__(**kwargs)
         self._set_prefix()
         self._connect_completer()
         
         self.on_trait_change(self._set_prefix, 'id')
         self.on_trait_change(self._connect_completer, 'user_ns')
         
         # Build dict of handlers for message types
         for msg_type in ['execute_request', 'complete_request', 'apply_request', 
                 'clear_request']:
             self.shell_handlers[msg_type] = getattr(self, msg_type)
         
         for msg_type in ['shutdown_request', 'abort_request']+self.shell_handlers.keys():
             self.control_handlers[msg_type] = getattr(self, msg_type)
         
         self._initial_exec_lines()
     
     def _wrap_exception(self, method=None):
         e_info = dict(engine_uuid=self.ident, engine_id=self.int_id, method=method)
         content=wrap_exception(e_info)
         return content
     
     def _initial_exec_lines(self):
         s = _Passer()
         content = dict(silent=True, user_variable=[],user_expressions=[])
         for line in self.exec_lines:
             self.log.debug("executing initialization: %s"%line)
             content.update({'code':line})
             msg = self.session.msg('execute_request', content)
             self.execute_request(s, [], msg)
         
         
     #-------------------- control handlers -----------------------------
     def abort_queues(self):
         for stream in self.shell_streams:
             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 as 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()
             self.log.info("Aborting:")
             self.log.info(str(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)[0]
             self.log.debug(str(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)[0]
         self.log.debug(str(reply_msg))
     
     def shutdown_request(self, stream, ident, parent):
         """kill ourself.  This should really be handled in an external process"""
         try:
             self.abort_queues()
         except:
             content = self._wrap_exception('shutdown')
         else:
             content = dict(parent['content'])
             content['status'] = 'ok'
         msg = self.session.send(stream, 'shutdown_reply',
                                 content=content, parent=parent, ident=ident)
         # msg = self.session.send(self.pub_socket, 'shutdown_reply',
         #                         content, parent, ident)
         # print >> sys.__stdout__, msg
         # time.sleep(0.2)
         dc = ioloop.DelayedCallback(lambda : sys.exit(0), 1000, self.loop)
         dc.start()
     
     def dispatch_control(self, msg):
         idents,msg = self.session.feed_identities(msg, copy=False)
         try:
             msg = self.session.unpack_message(msg, content=True, copy=False)
         except:
             self.log.error("Invalid Message", exc_info=True)
             return
         
         header = msg['header']
         msg_id = header['msg_id']
         
         handler = self.control_handlers.get(msg['msg_type'], None)
         if handler is None:
             self.log.error("UNKNOWN CONTROL MESSAGE TYPE: %r"%msg['msg_type'])
         else:
             handler(self.control_stream, idents, 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
     
     #-------------------- queue handlers -----------------------------
     
     def clear_request(self, stream, idents, parent):
         """Clear our namespace."""
         self.user_ns = {}
         msg = self.session.send(stream, 'clear_reply', ident=idents, parent=parent, 
                 content = dict(status='ok'))
         self._initial_exec_lines()
     
     def execute_request(self, stream, ident, parent):
         self.log.debug('execute request %s'%parent)
         try:
             code = parent[u'content'][u'code']
         except:
             self.log.error("Got bad msg: %s"%parent, exc_info=True)
             return
         self.session.send(self.iopub_stream, u'pyin', {u'code':code},parent=parent,
                             ident='%s.pyin'%self.prefix)
         started = datetime.now().strftime(ISO8601)
         try:
             comp_code = self.compiler(code, '<zmq-kernel>')
             # allow for not overriding displayhook
             if hasattr(sys.displayhook, 'set_parent'):
                 sys.displayhook.set_parent(parent)
                 sys.stdout.set_parent(parent)
                 sys.stderr.set_parent(parent)
             exec comp_code in self.user_ns, self.user_ns
         except:
             exc_content = self._wrap_exception('execute')
             # exc_msg = self.session.msg(u'pyerr', exc_content, parent)
             self.session.send(self.iopub_stream, u'pyerr', exc_content, parent=parent,
                             ident='%s.pyerr'%self.prefix)
             reply_content = exc_content
         else:
             reply_content = {'status' : 'ok'}
         
         reply_msg = self.session.send(stream, u'execute_reply', reply_content, parent=parent, 
                     ident=ident, subheader = dict(started=started))
         self.log.debug(str(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):
         # flush previous reply, so this request won't block it
         stream.flush(zmq.POLLOUT)
         
         try:
             content = parent[u'content']
             bufs = parent[u'buffers']
             msg_id = parent['header']['msg_id']
             bound = content.get('bound', False)
         except:
             self.log.error("Got bad msg: %s"%parent, exc_info=True)
             return
         # pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent)
         # self.iopub_stream.send(pyin_msg)
         # self.session.send(self.iopub_stream, u'pyin', {u'code':code},parent=parent)
         sub = {'dependencies_met' : True, 'engine' : self.ident,
                 'started': datetime.now().strftime(ISO8601)}
         try:
             # allow for not overriding displayhook
             if hasattr(sys.displayhook, 'set_parent'):
                 sys.displayhook.set_parent(parent)
                 sys.stdout.set_parent(parent)
                 sys.stderr.set_parent(parent)
             # exec "f(*args,**kwargs)" in self.user_ns, self.user_ns
             working = self.user_ns
             # suffix = 
             prefix = "_"+str(msg_id).replace("-","")+"_"
             # if bound:
             #     
             # else:
             #     working = dict()
             #     suffix = prefix = "_" # prevent keyword collisions with lambda
             f,args,kwargs = unpack_apply_message(bufs, working, copy=False)
             if bound:
                 bound_ns = Namespace(working)
                 args = [bound_ns]+list(args)
             # if f.fun
             fname = getattr(f, '__name__', 'f')
             
             fname = prefix+"f"
             argname = prefix+"args"
             kwargname = prefix+"kwargs"
             resultname = prefix+"result"
             
             ns = { fname : f, argname : args, kwargname : kwargs , resultname : None }
             # print ns
             working.update(ns)
             code = "%s=%s(*%s,**%s)"%(resultname, fname, argname, kwargname)
             try:
                 exec code in working,working
                 result = working.get(resultname)
             finally:
                 for key in ns.iterkeys():
                     working.pop(key)
             if bound:
                 working.update(bound_ns)
             
             packed_result,buf = serialize_object(result)
             result_buf = [packed_result]+buf
         except:
             exc_content = self._wrap_exception('apply')
             # exc_msg = self.session.msg(u'pyerr', exc_content, parent)
             self.session.send(self.iopub_stream, u'pyerr', exc_content, parent=parent,
                                 ident='%s.pyerr'%self.prefix)
             reply_content = exc_content
             result_buf = []
             
             if exc_content['ename'] == 'UnmetDependency':
                 sub['dependencies_met'] = False
         else:
             reply_content = {'status' : 'ok'}
         
         # put 'ok'/'error' status in header, for scheduler introspection:
         sub['status'] = reply_content['status']
         
         reply_msg = self.session.send(stream, u'apply_reply', reply_content, 
                     parent=parent, ident=ident,buffers=result_buf, subheader=sub)
 
         # 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)
         try:
             msg = self.session.unpack_message(msg, content=True, copy=False)
         except:
             self.log.error("Invalid Message", exc_info=True)
             return
             
         
         header = msg['header']
         msg_id = header['msg_id']
         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
         handler = self.shell_handlers.get(msg['msg_type'], None)
         if handler is None:
             self.log.error("UNKNOWN MESSAGE TYPE: %r"%msg['msg_type'])
         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)
         
         def make_dispatcher(stream):
             def dispatcher(msg):
                 return self.dispatch_queue(stream, msg)
             return dispatcher
         
         for s in self.shell_streams:
             s.on_recv(make_dispatcher(s), copy=False)
             s.on_err(printer)
         
         if self.iopub_stream:
             self.iopub_stream.on_err(printer)
         
         #### while True mode:
         # while True:
         #     idle = True
         #     try:
         #         msg = self.shell_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.shell_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 make_kernel(int_id, identity, control_addr, shell_addrs, iopub_addr, hb_addrs, 
                 client_addr=None, loop=None, context=None, key=None,
                 out_stream_factory=OutStream, display_hook_factory=DisplayHook):
     """NO LONGER IN USE"""
     # create loop, context, and session:
     if loop is None:
         loop = ioloop.IOLoop.instance()
     if context is None:
         context = zmq.Context()
     c = context
     session = StreamSession(key=key)
     # print (session.key)
     # print (control_addr, shell_addrs, iopub_addr, hb_addrs)
     
     # create Control Stream
     control_stream = zmqstream.ZMQStream(c.socket(zmq.PAIR), loop)
     control_stream.setsockopt(zmq.IDENTITY, identity)
     control_stream.connect(control_addr)
     
     # create Shell Streams (MUX, Task, etc.):
     shell_streams = []
     for addr in shell_addrs:
         stream = zmqstream.ZMQStream(c.socket(zmq.PAIR), loop)
         stream.setsockopt(zmq.IDENTITY, identity)
         stream.connect(addr)
         shell_streams.append(stream)
     
     # create iopub stream:
     iopub_stream = zmqstream.ZMQStream(c.socket(zmq.PUB), loop)
     iopub_stream.setsockopt(zmq.IDENTITY, identity)
     iopub_stream.connect(iopub_addr)
     
     # Redirect input streams and set a display hook.
     if out_stream_factory:
         sys.stdout = out_stream_factory(session, iopub_stream, u'stdout')
         sys.stdout.topic = 'engine.%i.stdout'%int_id
         sys.stderr = out_stream_factory(session, iopub_stream, u'stderr')
         sys.stderr.topic = 'engine.%i.stderr'%int_id
     if display_hook_factory:
         sys.displayhook = display_hook_factory(session, iopub_stream)
         sys.displayhook.topic = 'engine.%i.pyout'%int_id
     
     
     # launch heartbeat
     heart = heartmonitor.Heart(*map(str, hb_addrs), heart_id=identity)
     heart.start()
     
     # create (optional) Client
     if client_addr:
         client = Client(client_addr, username=identity)
     else:
         client = None
     
     kernel = Kernel(id=int_id, session=session, control_stream=control_stream,
             shell_streams=shell_streams, iopub_stream=iopub_stream, 
             client=client, loop=loop)
     kernel.start()
     return loop, c, kernel
 

IPython/zmq/parallel/streamsession.py

@@ -1,412 +1,418 b''
 #!/usr/bin/env python
 """edited session.py to work with streams, and move msg_type to the header
 """
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 
 import os
 import pprint
 import uuid
 from datetime import datetime
 
 try:
     import cPickle
     pickle = cPickle
 except:
     cPickle = None
     import pickle
 
 import zmq
 from zmq.utils import jsonapi
 from zmq.eventloop.zmqstream import ZMQStream
 
 from .util import ISO8601
 
 # 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
 
 def squash_unicode(obj):
     if isinstance(obj,dict):
         for key in obj.keys():
             obj[key] = squash_unicode(obj[key])
             if isinstance(key, unicode):
                 obj[squash_unicode(key)] = obj.pop(key)
     elif isinstance(obj, list):
         for i,v in enumerate(obj):
             obj[i] = squash_unicode(v)
     elif isinstance(obj, unicode):
         obj = obj.encode('utf8')
     return obj
 
 json_packer = jsonapi.dumps
 json_unpacker = lambda s: squash_unicode(jsonapi.loads(s))
 
 pickle_packer = lambda o: pickle.dumps(o,-1)
 pickle_unpacker = pickle.loads
 
 if use_json:
     default_packer = json_packer
     default_unpacker = json_unpacker
 else:
     default_packer = pickle_packer
     default_unpacker = pickle_unpacker
 
 
 DELIM="<IDS|MSG>"
 
 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):
     date=datetime.now().strftime(ISO8601)
     return locals()
 
 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
 
 class StreamSession(object):
     """tweaked version of IPython.zmq.session.Session, for development in Parallel"""
     debug=False
     key=None
     
     def __init__(self, username=None, session=None, packer=None, unpacker=None, key=None, keyfile=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
         
         if key is not None and keyfile is not None:
             raise TypeError("Must specify key OR keyfile, not both")
         if keyfile is not None:
             with open(keyfile) as f:
                 self.key = f.read().strip()
         else:
             self.key = key
         if isinstance(self.key, unicode):
             self.key = self.key.encode('utf8')
         # print key, keyfile, self.key
         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 check_key(self, msg_or_header):
         """Check that a message's header has the right key"""
         if self.key is None:
             return True
         header = extract_header(msg_or_header)
         return header.get('key', None) == self.key
             
         
     def send(self, stream, msg_or_type, content=None, buffers=None, parent=None, subheader=None, ident=None, track=False):
         """Build and send a message via stream or socket.
         
         Parameters
         ----------
         
         stream : zmq.Socket or ZMQStream
             the socket-like object used to send the data
         msg_or_type : str or Message/dict
             Normally, msg_or_type will be a msg_type unless a message is being sent more
             than once.
         
         content : dict or None
             the content of the message (ignored if msg_or_type is a message)
         buffers : list or None
             the already-serialized buffers to be appended to the message
         parent : Message or dict or None
             the parent or parent header describing the parent of this message
         subheader : dict or None
             extra header keys for this message's header
         ident : bytes or list of bytes
             the zmq.IDENTITY routing path
         track : bool
             whether to track.  Only for use with Sockets, because ZMQStream objects cannot track messages.
         
         Returns
         -------
         msg : message dict
             the constructed message
         (msg,tracker) : (message 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)):
             raise TypeError("stream must be Socket or ZMQStream, not %r"%type(stream))
         elif track and isinstance(stream, ZMQStream):
             raise TypeError("ZMQStream cannot track messages")
         
         if isinstance(msg_or_type, (Message, dict)):
             # we got a Message, not a msg_type
             # don't build a new Message
             msg = msg_or_type
             content = msg['content']
         else:
             msg = self.msg(msg_or_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)
         if self.key is not None:
             to_send.append(self.key)
         to_send.append(self.pack(msg['header']))
         to_send.append(self.pack(msg['parent_header']))
         
         if content is None:
             content = self.none
         elif isinstance(content, dict):
             content = self.pack(content)
         elif isinstance(content, bytes):
             # 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
             _track = False
         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)
                 
         # omsg = Message(msg)
         if self.debug:
             pprint.pprint(msg)
             pprint.pprint(to_send)
             pprint.pprint(buffers)
         
         msg['tracker'] = tracker
         
         return msg
     
     def send_raw(self, stream, msg, flags=0, copy=True, ident=None):
         """Send a raw message via ident path.
         
         Parameters
         ----------
         msg : list of sendable buffers"""
         to_send = []
         if isinstance(ident, bytes):
             ident = [ident]
         if ident is not None:
             to_send.extend(ident)
         to_send.append(DELIM)
         if self.key is not None:
             to_send.append(self.key)
         to_send.extend(msg)
         stream.send_multipart(msg, flags, copy=copy)
     
     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 as e:
             if e.errno == zmq.EAGAIN:
                 # We can convert EAGAIN to None as we know in this case
                 # recv_multipart 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 as e:
             print (idents, msg)
             # TODO: handle it
             raise e
     
     def feed_identities(self, msg, copy=True):
         """feed until DELIM is reached, then return the prefix as idents and remainder as
         msg. This is easily broken by setting an IDENT to DELIM, but that would be silly.
         
         Parameters
         ----------
         msg : a list of Message or bytes objects
             the message to be split
         copy : bool
             flag determining whether the arguments are bytes or Messages
         
         Returns
         -------
         (idents,msg) : two lists
             idents will always be a list of bytes - the indentity prefix
             msg will be a list of bytes or Messages, unchanged from input
             msg should be unpackable via self.unpack_message at this point.
         """
         ikey = int(self.key is not None)
         minlen = 3 + ikey
         msg = list(msg)
         idents = []
         while len(msg) > minlen:
             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)
         
         """
         ikey = int(self.key is not None)
         minlen = 3 + ikey
         message = {}
         if not copy:
             for i in range(minlen):
                 msg[i] = msg[i].bytes
         if ikey:
             if not self.key == msg[0]:
                 raise KeyError("Invalid Session Key: %s"%msg[0])
         if not len(msg) >= minlen:
             raise TypeError("malformed message, must have at least %i elements"%minlen)
         message['header'] = self.unpack(msg[ikey+0])
         message['msg_type'] = message['header']['msg_type']
         message['parent_header'] = self.unpack(msg[ikey+1])
         if content:
             message['content'] = self.unpack(msg[ikey+2])
         else:
             message['content'] = msg[ikey+2]
         
         message['buffers'] = msg[ikey+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/parallel/taskthread.py

@@ -1,100 +1,106 b''
 """Thread for popping Tasks from zmq to Python Queue"""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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.
+#-----------------------------------------------------------------------------
 
 
 import time
 from threading import Thread
 
 try:
     from queue import Queue
 except:
     from Queue import Queue
 
 import zmq
 from zmq.core.poll import _poll as poll
 from zmq.devices import ThreadDevice
 from IPython.zmq.parallel import streamsession as ss
 
 
 class QueueStream(object):
     def __init__(self, in_queue, out_queue):
         self.in_queue = in_queue
         self.out_queue = out_queue
     
     def send_multipart(self, *args, **kwargs):
         while self.out_queue.full():
             time.sleep(1e-3)
         self.out_queue.put(('send_multipart', args, kwargs))
     
     def send(self, *args, **kwargs):
         while self.out_queue.full():
             time.sleep(1e-3)
         self.out_queue.put(('send', args, kwargs))
     
     def recv_multipart(self):
         return self.in_queue.get()
     
     def empty(self):
         return self.in_queue.empty()
 
 class TaskThread(ThreadDevice):
     """Class for popping Tasks from C-ZMQ->Python Queue"""
     max_qsize = 100
     in_socket = None
     out_socket = None
     # queue = None
     
     def __init__(self, queue_type, mon_type, engine_id, max_qsize=100):
         ThreadDevice.__init__(self, 0, queue_type, mon_type)
         self.session = ss.StreamSession(username='TaskNotifier[%s]'%engine_id)
         self.engine_id = engine_id
         self.in_queue = Queue(max_qsize)
         self.out_queue = Queue(max_qsize)
         self.max_qsize = max_qsize
     
     @property
     def queues(self):
         return self.in_queue, self.out_queue
     
     @property
     def can_recv(self):
         # print self.in_queue.full(), poll((self.queue_socket, zmq.POLLIN),1e-3)
         return (not self.in_queue.full()) and poll([(self.queue_socket, zmq.POLLIN)], 1e-3 )
     
     @property
     def can_send(self):
         return not self.out_queue.empty()
     
     def run(self):
         print 'running'
         self.queue_socket,self.mon_socket = self._setup_sockets()
         print 'setup'
         
         while True:
             while not self.can_send and not self.can_recv:
                 # print 'idle'
                 # nothing to do, wait
                 time.sleep(1e-3)
             while self.can_send:
                 # flush out queue
                 print 'flushing...'
                 meth, args, kwargs = self.out_queue.get()
                 getattr(self.queue_socket, meth)(*args, **kwargs)
                 print 'flushed'
             
             if self.can_recv:
                 print 'recving'
                 # get another job from zmq
                 msg = self.queue_socket.recv_multipart(0, copy=False)
                 # put it in the Queue
                 self.in_queue.put(msg)
                 idents,msg = self.session.feed_identities(msg, copy=False)
                 msg = self.session.unpack_message(msg, content=False, copy=False)
                 # notify the Controller that we got it
                 self.mon_socket.send('tracktask', zmq.SNDMORE)
                 header = msg['header']
                 msg_id = header['msg_id']
                 content = dict(engine_id=self.engine_id, msg_id = msg_id)
                 self.session.send(self.mon_socket, 'task_receipt', content=content)
                 print 'recvd'
             
     
  No newline at end of file

IPython/zmq/parallel/util.py

@@ -1,299 +1,318 b''
 """some generic utilities for dealing with classes, urls, and serialization"""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010-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 re
 import socket
 
 try:
     import cPickle
     pickle = cPickle
 except:
     cPickle = None
     import pickle
 
 
 from IPython.utils.pickleutil import can, uncan, canSequence, uncanSequence
 from IPython.utils.newserialized import serialize, unserialize
 
 ISO8601="%Y-%m-%dT%H:%M:%S.%f"
 
+#-----------------------------------------------------------------------------
+# Classes
+#-----------------------------------------------------------------------------
+
 class Namespace(dict):
     """Subclass of dict for attribute access to keys."""
     
     def __getattr__(self, key):
         """getattr aliased to getitem"""
         if key in self.iterkeys():
             return self[key]
         else:
             raise NameError(key)
 
     def __setattr__(self, key, value):
         """setattr aliased to setitem, with strict"""
         if hasattr(dict, key):
             raise KeyError("Cannot override dict keys %r"%key)
         self[key] = value
     
 
 class ReverseDict(dict):
     """simple double-keyed subset of dict methods."""
     
     def __init__(self, *args, **kwargs):
         dict.__init__(self, *args, **kwargs)
         self._reverse = dict()
         for key, value in self.iteritems():
             self._reverse[value] = key
     
     def __getitem__(self, key):
         try:
             return dict.__getitem__(self, key)
         except KeyError:
             return self._reverse[key]
     
     def __setitem__(self, key, value):
         if key in self._reverse:
             raise KeyError("Can't have key %r on both sides!"%key)
         dict.__setitem__(self, key, value)
         self._reverse[value] = key
     
     def pop(self, key):
         value = dict.pop(self, key)
         self._reverse.pop(value)
         return value
     
     def get(self, key, default=None):
         try:
             return self[key]
         except KeyError:
             return default
 
+#-----------------------------------------------------------------------------
+# Functions
+#-----------------------------------------------------------------------------
+
 def validate_url(url):
     """validate a url for zeromq"""
     if not isinstance(url, basestring):
         raise TypeError("url must be a string, not %r"%type(url))
     url = url.lower()
     
     proto_addr = url.split('://')
     assert len(proto_addr) == 2, 'Invalid url: %r'%url
     proto, addr = proto_addr
     assert proto in ['tcp','pgm','epgm','ipc','inproc'], "Invalid protocol: %r"%proto
     
     # domain pattern adapted from http://www.regexlib.com/REDetails.aspx?regexp_id=391
     # author: Remi Sabourin
     pat = re.compile(r'^([\w\d]([\w\d\-]{0,61}[\w\d])?\.)*[\w\d]([\w\d\-]{0,61}[\w\d])?$')
     
     if proto == 'tcp':
         lis = addr.split(':')
         assert len(lis) == 2, 'Invalid url: %r'%url
         addr,s_port = lis
         try:
             port = int(s_port)
         except ValueError:
             raise AssertionError("Invalid port %r in url: %r"%(port, url))
         
         assert addr == '*' or pat.match(addr) is not None, 'Invalid url: %r'%url
         
     else:
         # only validate tcp urls currently
         pass
     
     return True
 
 
 def validate_url_container(container):
     """validate a potentially nested collection of urls."""
     if isinstance(container, basestring):
         url = container
         return validate_url(url)
     elif isinstance(container, dict):
         container = container.itervalues()
     
     for element in container:
         validate_url_container(element)
 
 
 def split_url(url):
     """split a zmq url (tcp://ip:port) into ('tcp','ip','port')."""
     proto_addr = url.split('://')
     assert len(proto_addr) == 2, 'Invalid url: %r'%url
     proto, addr = proto_addr
     lis = addr.split(':')
     assert len(lis) == 2, 'Invalid url: %r'%url
     addr,s_port = lis
     return proto,addr,s_port
     
 def disambiguate_ip_address(ip, location=None):
     """turn multi-ip interfaces '0.0.0.0' and '*' into connectable
     ones, based on the location (default interpretation of location is localhost)."""
     if ip in ('0.0.0.0', '*'):
         external_ips = socket.gethostbyname_ex(socket.gethostname())[2]
         if location is None or location in external_ips:
             ip='127.0.0.1'
         elif location:
             return location
     return ip
 
 def disambiguate_url(url, location=None):
     """turn multi-ip interfaces '0.0.0.0' and '*' into connectable
     ones, based on the location (default interpretation is localhost).
     
     This is for zeromq urls, such as tcp://*:10101."""
     try:
         proto,ip,port = split_url(url)
     except AssertionError:
         # probably not tcp url; could be ipc, etc.
         return url
     
     ip = disambiguate_ip_address(ip,location)
     
     return "%s://%s:%s"%(proto,ip,port)
 
 
 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.
     
     Parameters
     ----------
     
     obj : object
         The object to be serialized
     threshold : float
         The threshold for not double-pickling the content.
         
     
     Returns
     -------
     ('pmd', [bufs]) :
         where pmd is the pickled metadata wrapper,
         bufs is a list of data buffers
     """
     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.typeDescriptor in ('buffer', 'ndarray') or 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.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):
         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
     else:
         if sobj.data is None:
             sobj.data = bufs.pop(0)
         return uncan(unserialize(sobj)), bufs
 
 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 = 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:
             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
 
         kwargs[k] = uncan(unserialize(sa), g)
     
     return f,args,kwargs