Controller renamed to Hub (keeping ipcontrollerz)

IPython/zmq/parallel/hub.py

@@ -0,0 +1,837 b''
+#!/usr/bin/env python
+"""The IPython Controller Hub with 0MQ
+This is the master object that handles connections from engines and clients,
+and monitors traffic through the various queues.
+"""
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2010  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 sys
+from datetime import datetime
+import time
+
+import zmq
+from zmq.eventloop import ioloop
+
+# internal:
+from IPython.zmq.log import logger # a Logger object
+
+from streamsession import Message, wrap_exception, ISO8601
+
+try:
+    from pymongo.binary import Binary
+except ImportError:
+    MongoDB=None
+else:
+    from mongodb import MongoDB
+    
+#-----------------------------------------------------------------------------
+# Code
+#-----------------------------------------------------------------------------
+
+def _passer(*args, **kwargs):
+    return
+
+def init_record(msg):
+    """return an empty TaskRecord dict, with all keys initialized with None."""
+    header = msg['header']
+    return {
+        'msg_id' : header['msg_id'],
+        'header' : header,
+        'content': msg['content'],
+        'buffers': msg['buffers'],
+        'submitted': datetime.strptime(header['date'], ISO8601),
+        'client_uuid' : None,
+        'engine_uuid' : None,
+        'started': None,
+        'completed': None,
+        'resubmitted': None,
+        'result_header' : None,
+        'result_content' : None,
+        'result_buffers' : None,
+        'queue' : None
+    }
+
+
+class EngineConnector(object):
+    """A simple object for accessing the various zmq connections of an object.
+    Attributes are:
+    id (int): engine ID
+    uuid (str): uuid (unused?)
+    queue (str): identity of queue's XREQ socket
+    registration (str): identity of registration XREQ socket
+    heartbeat (str): identity of heartbeat XREQ socket
+    """
+    id=0
+    queue=None
+    control=None
+    registration=None
+    heartbeat=None
+    pending=None
+    
+    def __init__(self, id, queue, registration, control, heartbeat=None):
+        logger.info("engine::Engine Connected: %i"%id)
+        self.id = id
+        self.queue = queue
+        self.registration = registration
+        self.control = control
+        self.heartbeat = heartbeat
+        
+class Hub(object):
+    """The IPython Controller Hub with 0MQ connections
+    
+    Parameters
+    ==========
+    loop: zmq IOLoop instance
+    session: StreamSession object
+    <removed> context: zmq context for creating new connections (?)
+    queue: ZMQStream for monitoring the command queue (SUB)
+    registrar: ZMQStream for engine registration requests (XREP)
+    heartbeat: HeartMonitor object checking the pulse of the engines
+    clientele: ZMQStream for client connections (XREP)
+                not used for jobs, only query/control commands
+    notifier: ZMQStream for broadcasting engine registration changes (PUB)
+    db: connection to db for out of memory logging of commands
+                NotImplemented
+    engine_addrs: dict of zmq connection information for engines to connect
+                to the queues.
+    client_addrs: dict of zmq connection information for engines to connect
+                to the queues.
+    """
+    # internal data structures:
+    ids=None # engine IDs
+    keytable=None
+    engines=None
+    clients=None
+    hearts=None
+    pending=None
+    results=None
+    tasks=None
+    completed=None
+    mia=None
+    incoming_registrations=None
+    registration_timeout=None
+    
+    #objects from constructor:
+    loop=None
+    registrar=None
+    clientelle=None
+    queue=None
+    heartbeat=None
+    notifier=None
+    db=None
+    client_addr=None
+    engine_addrs=None
+    
+    
+    def __init__(self, loop, session, queue, registrar, heartbeat, clientele, notifier, db, engine_addrs, client_addrs):
+        """
+        # universal:
+        loop: IOLoop for creating future connections
+        session: streamsession for sending serialized data
+        # engine:
+        queue: ZMQStream for monitoring queue messages
+        registrar: ZMQStream for engine registration
+        heartbeat: HeartMonitor object for tracking engines
+        # client:
+        clientele: ZMQStream for client connections
+        # extra:
+        db: ZMQStream for db connection (NotImplemented)
+        engine_addrs: zmq address/protocol dict for engine connections
+        client_addrs: zmq address/protocol dict for client connections
+        """
+        self.ids = set()
+        self.keytable={}
+        self.incoming_registrations={}
+        self.engines = {}
+        self.by_ident = {}
+        self.clients = {}
+        self.hearts = {}
+        # self.mia = set()
+        
+        # self.sockets = {}
+        self.loop = loop
+        self.session = session
+        self.registrar = registrar
+        self.clientele = clientele
+        self.queue = queue
+        self.heartbeat = heartbeat
+        self.notifier = notifier
+        self.db = db
+        
+        # validate connection dicts:
+        self.client_addrs = client_addrs
+        assert isinstance(client_addrs['queue'], str)
+        assert isinstance(client_addrs['control'], str)
+        # self.hb_addrs = hb_addrs
+        self.engine_addrs = engine_addrs
+        assert isinstance(engine_addrs['queue'], str)
+        assert isinstance(client_addrs['control'], str)
+        assert len(engine_addrs['heartbeat']) == 2
+        
+        # register our callbacks
+        self.registrar.on_recv(self.dispatch_register_request)
+        self.clientele.on_recv(self.dispatch_client_msg)
+        self.queue.on_recv(self.dispatch_queue_traffic)
+        
+        if heartbeat is not None:
+            heartbeat.add_heart_failure_handler(self.handle_heart_failure)
+            heartbeat.add_new_heart_handler(self.handle_new_heart)
+        
+        self.queue_handlers = { 'in' : self.save_queue_request,
+                                'out': self.save_queue_result,
+                                'intask': self.save_task_request,
+                                'outtask': self.save_task_result,
+                                'tracktask': self.save_task_destination,
+                                'incontrol': _passer,
+                                'outcontrol': _passer,
+        }
+        
+        self.client_handlers = {'queue_request': self.queue_status,
+                                'result_request': self.get_results,
+                                'purge_request': self.purge_results,
+                                'load_request': self.check_load,
+                                'resubmit_request': self.resubmit_task,
+                                'shutdown_request': self.shutdown_request,
+                                }
+        
+        self.registrar_handlers = {'registration_request' : self.register_engine,
+                                'unregistration_request' : self.unregister_engine,
+                                'connection_request': self.connection_request,
+        }
+        self.registration_timeout = max(5000, 2*self.heartbeat.period)
+        # this is the stuff that will move to DB:
+        # self.results = {} # completed results
+        self.pending = set() # pending messages, keyed by msg_id
+        self.queues = {} # pending msg_ids keyed by engine_id
+        self.tasks = {} # pending msg_ids submitted as tasks, keyed by client_id
+        self.completed = {} # completed msg_ids keyed by engine_id
+        self.all_completed = set()
+        
+        logger.info("controller::created controller")
+    
+    def _new_id(self):
+        """gemerate a new ID"""
+        newid = 0
+        incoming = [id[0] for id in self.incoming_registrations.itervalues()]
+        # print newid, self.ids, self.incoming_registrations
+        while newid in self.ids or newid in incoming:
+            newid += 1
+        return newid
+    
+    #-----------------------------------------------------------------------------
+    # message validation
+    #-----------------------------------------------------------------------------
+    
+    def _validate_targets(self, targets):
+        """turn any valid targets argument into a list of integer ids"""
+        if targets is None:
+            # default to all
+            targets = self.ids
+            
+        if isinstance(targets, (int,str,unicode)):
+            # only one target specified
+            targets = [targets]
+        _targets = []
+        for t in targets:
+            # map raw identities to ids
+            if isinstance(t, (str,unicode)):
+                t = self.by_ident.get(t, t)
+            _targets.append(t)
+        targets = _targets
+        bad_targets = [ t for t in targets if t not in self.ids ]
+        if bad_targets:
+            raise IndexError("No Such Engine: %r"%bad_targets)
+        if not targets:
+            raise IndexError("No Engines Registered")
+        return targets
+    
+    def _validate_client_msg(self, msg):
+        """validates and unpacks headers of a message. Returns False if invalid,
+        (ident, header, parent, content)"""
+        client_id = msg[0]
+        try:
+            msg = self.session.unpack_message(msg[1:], content=True)
+        except:
+            logger.error("client::Invalid Message %s"%msg)
+            return False
+        
+        msg_type = msg.get('msg_type', None)
+        if msg_type is None:
+            return False
+        header = msg.get('header')
+        # session doesn't handle split content for now:
+        return client_id, msg
+        
+    
+    #-----------------------------------------------------------------------------
+    # dispatch methods (1 per stream)
+    #-----------------------------------------------------------------------------
+    
+    def dispatch_register_request(self, msg):
+        """"""
+        logger.debug("registration::dispatch_register_request(%s)"%msg)
+        idents,msg = self.session.feed_identities(msg)
+        if not idents:
+            logger.error("Bad Queue Message: %s"%msg, exc_info=True)
+            return
+        try:
+            msg = self.session.unpack_message(msg,content=True)
+        except:
+            logger.error("registration::got bad registration message: %s"%msg, exc_info=True)
+            return
+        
+        msg_type = msg['msg_type']
+        content = msg['content']
+        
+        handler = self.registrar_handlers.get(msg_type, None)
+        if handler is None:
+            logger.error("registration::got bad registration message: %s"%msg)
+        else:
+            handler(idents, msg)
+    
+    def dispatch_queue_traffic(self, msg):
+        """all ME and Task queue messages come through here"""
+        logger.debug("queue traffic: %s"%msg[:2])
+        switch = msg[0]
+        idents, msg = self.session.feed_identities(msg[1:])
+        if not idents:
+            logger.error("Bad Queue Message: %s"%msg)
+            return
+        handler = self.queue_handlers.get(switch, None)
+        if handler is not None:
+            handler(idents, msg)
+        else:
+            logger.error("Invalid message topic: %s"%switch)
+        
+    
+    def dispatch_client_msg(self, msg):
+        """Route messages from clients"""
+        idents, msg = self.session.feed_identities(msg)
+        if not idents:
+            logger.error("Bad Client Message: %s"%msg)
+            return
+        client_id = idents[0]
+        try:
+            msg = self.session.unpack_message(msg, content=True)
+        except:
+            content = wrap_exception()
+            logger.error("Bad Client Message: %s"%msg, exc_info=True)
+            self.session.send(self.clientele, "controller_error", ident=client_id, 
+                    content=content)
+            return
+        
+        # print client_id, header, parent, content
+        #switch on message type:
+        msg_type = msg['msg_type']
+        logger.info("client:: client %s requested %s"%(client_id, msg_type))
+        handler = self.client_handlers.get(msg_type, None)
+        try:
+            assert handler is not None, "Bad Message Type: %s"%msg_type
+        except:
+            content = wrap_exception()
+            logger.error("Bad Message Type: %s"%msg_type, exc_info=True)
+            self.session.send(self.clientele, "controller_error", ident=client_id, 
+                    content=content)
+            return
+        else:
+            handler(client_id, msg)
+            
+    def dispatch_db(self, msg):
+        """"""
+        raise NotImplementedError
+    
+    #---------------------------------------------------------------------------
+    # handler methods (1 per event)
+    #---------------------------------------------------------------------------
+    
+    #----------------------- Heartbeat --------------------------------------
+    
+    def handle_new_heart(self, heart):
+        """handler to attach to heartbeater.
+        Called when a new heart starts to beat.
+        Triggers completion of registration."""
+        logger.debug("heartbeat::handle_new_heart(%r)"%heart)
+        if heart not in self.incoming_registrations:
+            logger.info("heartbeat::ignoring new heart: %r"%heart)
+        else:
+            self.finish_registration(heart)
+        
+    
+    def handle_heart_failure(self, heart):
+        """handler to attach to heartbeater.
+        called when a previously registered heart fails to respond to beat request.
+        triggers unregistration"""
+        logger.debug("heartbeat::handle_heart_failure(%r)"%heart)
+        eid = self.hearts.get(heart, None)
+        queue = self.engines[eid].queue
+        if eid is None:
+            logger.info("heartbeat::ignoring heart failure %r"%heart)
+        else:
+            self.unregister_engine(heart, dict(content=dict(id=eid, queue=queue)))
+    
+    #----------------------- MUX Queue Traffic ------------------------------
+    
+    def save_queue_request(self, idents, msg):
+        if len(idents) < 2:
+            logger.error("invalid identity prefix: %s"%idents)
+            return
+        queue_id, client_id = idents[:2]
+        try:
+            msg = self.session.unpack_message(msg, content=False)
+        except:
+            logger.error("queue::client %r sent invalid message to %r: %s"%(client_id, queue_id, msg), exc_info=True)
+            return
+        
+        eid = self.by_ident.get(queue_id, None)
+        if eid is None:
+            logger.error("queue::target %r not registered"%queue_id)
+            logger.debug("queue::    valid are: %s"%(self.by_ident.keys()))
+            return
+            
+        header = msg['header']
+        msg_id = header['msg_id']
+        record = init_record(msg)
+        record['engine_uuid'] = queue_id
+        record['client_uuid'] = client_id
+        record['queue'] = 'mux'
+        if MongoDB is not None and isinstance(self.db, MongoDB):
+            record['buffers'] = map(Binary, record['buffers'])
+        self.pending.add(msg_id)
+        self.queues[eid].append(msg_id)
+        self.db.add_record(msg_id, record)
+    
+    def save_queue_result(self, idents, msg):
+        if len(idents) < 2:
+            logger.error("invalid identity prefix: %s"%idents)
+            return
+            
+        client_id, queue_id = idents[:2]
+        try:
+            msg = self.session.unpack_message(msg, content=False)
+        except:
+            logger.error("queue::engine %r sent invalid message to %r: %s"%(
+                    queue_id,client_id, msg), exc_info=True)
+            return
+        
+        eid = self.by_ident.get(queue_id, None)
+        if eid is None:
+            logger.error("queue::unknown engine %r is sending a reply: "%queue_id)
+            logger.debug("queue::       %s"%msg[2:])
+            return
+        
+        parent = msg['parent_header']
+        if not parent:
+            return
+        msg_id = parent['msg_id']
+        if msg_id in self.pending:
+            self.pending.remove(msg_id)
+            self.all_completed.add(msg_id)
+            self.queues[eid].remove(msg_id)
+            self.completed[eid].append(msg_id)
+            rheader = msg['header']
+            completed = datetime.strptime(rheader['date'], ISO8601)
+            started = rheader.get('started', None)
+            if started is not None:
+                started = datetime.strptime(started, ISO8601)
+            result = {
+                'result_header' : rheader,
+                'result_content': msg['content'],
+                'started' : started,
+                'completed' : completed
+            }
+            if MongoDB is not None and isinstance(self.db, MongoDB):
+                result['result_buffers'] = map(Binary, msg['buffers'])
+            else:
+                result['result_buffers'] = msg['buffers']
+            self.db.update_record(msg_id, result)
+        else:
+            logger.debug("queue:: unknown msg finished %s"%msg_id)
+            
+    #--------------------- Task Queue Traffic ------------------------------
+    
+    def save_task_request(self, idents, msg):
+        """Save the submission of a task."""
+        client_id = idents[0]
+        
+        try:
+            msg = self.session.unpack_message(msg, content=False)
+        except:
+            logger.error("task::client %r sent invalid task message: %s"%(
+                    client_id, msg), exc_info=True)
+            return
+        record = init_record(msg)
+        if MongoDB is not None and isinstance(self.db, MongoDB):
+            record['buffers'] = map(Binary, record['buffers'])
+        record['client_uuid'] = client_id
+        record['queue'] = 'task'
+        header = msg['header']
+        msg_id = header['msg_id']
+        self.pending.add(msg_id)
+        self.db.add_record(msg_id, record)
+    
+    def save_task_result(self, idents, msg):
+        """save the result of a completed task."""
+        client_id = idents[0]
+        try:
+            msg = self.session.unpack_message(msg, content=False)
+        except:
+            logger.error("task::invalid task result message send to %r: %s"%(
+                    client_id, msg))
+            raise
+            return
+        
+        parent = msg['parent_header']
+        if not parent:
+            # print msg
+            logger.warn("Task %r had no parent!"%msg)
+            return
+        msg_id = parent['msg_id']
+        
+        header = msg['header']
+        engine_uuid = header.get('engine', None)
+        eid = self.by_ident.get(engine_uuid, None)
+        
+        if msg_id in self.pending:
+            self.pending.remove(msg_id)
+            self.all_completed.add(msg_id)
+            if eid is not None:
+                self.completed[eid].append(msg_id)
+                if msg_id in self.tasks[eid]:
+                    self.tasks[eid].remove(msg_id)
+            completed = datetime.strptime(header['date'], ISO8601)
+            started = header.get('started', None)
+            if started is not None:
+                started = datetime.strptime(started, ISO8601)
+            result = {
+                'result_header' : header,
+                'result_content': msg['content'],
+                'started' : started,
+                'completed' : completed,
+                'engine_uuid': engine_uuid
+            }
+            if MongoDB is not None and isinstance(self.db, MongoDB):
+                result['result_buffers'] = map(Binary, msg['buffers'])
+            else:
+                result['result_buffers'] = msg['buffers']
+            self.db.update_record(msg_id, result)
+            
+        else:
+            logger.debug("task::unknown task %s finished"%msg_id)
+    
+    def save_task_destination(self, idents, msg):
+        try:
+            msg = self.session.unpack_message(msg, content=True)
+        except:
+            logger.error("task::invalid task tracking message")
+            return
+        content = msg['content']
+        print (content)
+        msg_id = content['msg_id']
+        engine_uuid = content['engine_id']
+        eid = self.by_ident[engine_uuid]
+        
+        logger.info("task::task %s arrived on %s"%(msg_id, eid))
+        # if msg_id in self.mia:
+        #     self.mia.remove(msg_id)
+        # else:
+        #     logger.debug("task::task %s not listed as MIA?!"%(msg_id))
+        
+        self.tasks[eid].append(msg_id)
+        # self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
+        self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
+    
+    def mia_task_request(self, idents, msg):
+        raise NotImplementedError
+        client_id = idents[0]
+        # content = dict(mia=self.mia,status='ok')
+        # self.session.send('mia_reply', content=content, idents=client_id)
+        
+        
+            
+    #-------------------------------------------------------------------------
+    # Registration requests
+    #-------------------------------------------------------------------------
+        
+    def connection_request(self, client_id, msg):
+        """Reply with connection addresses for clients."""
+        logger.info("client::client %s connected"%client_id)
+        content = dict(status='ok')
+        content.update(self.client_addrs)
+        jsonable = {}
+        for k,v in self.keytable.iteritems():
+            jsonable[str(k)] = v
+        content['engines'] = jsonable
+        self.session.send(self.registrar, 'connection_reply', content, parent=msg, ident=client_id)
+    
+    def register_engine(self, reg, msg):
+        """Register a new engine."""
+        content = msg['content']
+        try:
+            queue = content['queue']
+        except KeyError:
+            logger.error("registration::queue not specified")
+            return
+        heart = content.get('heartbeat', None)
+        """register a new engine, and create the socket(s) necessary"""
+        eid = self._new_id()
+        # print (eid, queue, reg, heart)
+        
+        logger.debug("registration::register_engine(%i, %r, %r, %r)"%(eid, queue, reg, heart))
+        
+        content = dict(id=eid,status='ok')
+        content.update(self.engine_addrs)
+        # check if requesting available IDs:
+        if queue in self.by_ident:
+            try:
+                raise KeyError("queue_id %r in use"%queue)
+            except:
+                content = wrap_exception()
+        elif heart in self.hearts: # need to check unique hearts?
+            try:
+                raise KeyError("heart_id %r in use"%heart)
+            except:
+                content = wrap_exception()
+        else:
+            for h, pack in self.incoming_registrations.iteritems():
+                if heart == h:
+                    try:
+                        raise KeyError("heart_id %r in use"%heart)
+                    except:
+                        content = wrap_exception()
+                    break
+                elif queue == pack[1]:
+                    try:
+                        raise KeyError("queue_id %r in use"%queue)
+                    except:
+                        content = wrap_exception()
+                    break
+        
+        msg = self.session.send(self.registrar, "registration_reply", 
+                content=content, 
+                ident=reg)
+        
+        if content['status'] == 'ok':
+            if heart in self.heartbeat.hearts:
+                # already beating
+                self.incoming_registrations[heart] = (eid,queue,reg,None)
+                self.finish_registration(heart)
+            else:
+                purge = lambda : self._purge_stalled_registration(heart)
+                dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
+                dc.start()
+                self.incoming_registrations[heart] = (eid,queue,reg,dc)
+        else:
+            logger.error("registration::registration %i failed: %s"%(eid, content['evalue']))
+        return eid
+    
+    def unregister_engine(self, ident, msg):
+        """Unregister an engine that explicitly requested to leave."""
+        try:
+            eid = msg['content']['id']
+        except:
+            logger.error("registration::bad engine id for unregistration: %s"%ident)
+            return
+        logger.info("registration::unregister_engine(%s)"%eid)
+        content=dict(id=eid, queue=self.engines[eid].queue)
+        self.ids.remove(eid)
+        self.keytable.pop(eid)
+        ec = self.engines.pop(eid)
+        self.hearts.pop(ec.heartbeat)
+        self.by_ident.pop(ec.queue)
+        self.completed.pop(eid)
+        for msg_id in self.queues.pop(eid):
+            msg = self.pending.remove(msg_id)
+            ############## TODO: HANDLE IT ################
+        
+        if self.notifier:
+            self.session.send(self.notifier, "unregistration_notification", content=content)
+    
+    def finish_registration(self, heart):
+        """Second half of engine registration, called after our HeartMonitor
+        has received a beat from the Engine's Heart."""
+        try: 
+            (eid,queue,reg,purge) = self.incoming_registrations.pop(heart)
+        except KeyError:
+            logger.error("registration::tried to finish nonexistant registration")
+            return
+        logger.info("registration::finished registering engine %i:%r"%(eid,queue))
+        if purge is not None:
+            purge.stop()
+        control = queue
+        self.ids.add(eid)
+        self.keytable[eid] = queue
+        self.engines[eid] = EngineConnector(eid, queue, reg, control, heart)
+        self.by_ident[queue] = eid
+        self.queues[eid] = list()
+        self.tasks[eid] = list()
+        self.completed[eid] = list()
+        self.hearts[heart] = eid
+        content = dict(id=eid, queue=self.engines[eid].queue)
+        if self.notifier:
+            self.session.send(self.notifier, "registration_notification", content=content)
+    
+    def _purge_stalled_registration(self, heart):
+        if heart in self.incoming_registrations:
+            eid = self.incoming_registrations.pop(heart)[0]
+            logger.info("registration::purging stalled registration: %i"%eid)
+        else:
+            pass
+            
+    #-------------------------------------------------------------------------
+    # Client Requests
+    #-------------------------------------------------------------------------
+    
+    def shutdown_request(self, client_id, msg):
+        """handle shutdown request."""
+        # s = self.context.socket(zmq.XREQ)
+        # s.connect(self.client_connections['mux'])
+        # time.sleep(0.1)
+        # for eid,ec in self.engines.iteritems():
+        #     self.session.send(s, 'shutdown_request', content=dict(restart=False), ident=ec.queue)
+        # time.sleep(1)
+        self.session.send(self.clientele, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
+        dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
+        dc.start()
+    
+    def _shutdown(self):
+        logger.info("controller::controller shutting down.")
+        time.sleep(0.1)
+        sys.exit(0)
+        
+    
+    def check_load(self, client_id, msg):
+        content = msg['content']
+        try:
+            targets = content['targets']
+            targets = self._validate_targets(targets)
+        except:
+            content = wrap_exception()
+            self.session.send(self.clientele, "controller_error", 
+                    content=content, ident=client_id)
+            return
+        
+        content = dict(status='ok')
+        # loads = {}
+        for t in targets:
+            content[bytes(t)] = len(self.queues[t])+len(self.tasks[t])
+        self.session.send(self.clientele, "load_reply", content=content, ident=client_id)
+            
+    
+    def queue_status(self, client_id, msg):
+        """Return the Queue status of one or more targets.
+        if verbose: return the msg_ids
+        else: return len of each type.
+        keys: queue (pending MUX jobs)
+            tasks (pending Task jobs)
+            completed (finished jobs from both queues)"""
+        content = msg['content']
+        targets = content['targets']
+        try:
+            targets = self._validate_targets(targets)
+        except:
+            content = wrap_exception()
+            self.session.send(self.clientele, "controller_error", 
+                    content=content, ident=client_id)
+            return
+        verbose = content.get('verbose', False)
+        content = dict(status='ok')
+        for t in targets:
+            queue = self.queues[t]
+            completed = self.completed[t]
+            tasks = self.tasks[t]
+            if not verbose:
+                queue = len(queue)
+                completed = len(completed)
+                tasks = len(tasks)
+            content[bytes(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
+            # pending
+        self.session.send(self.clientele, "queue_reply", content=content, ident=client_id)
+    
+    def purge_results(self, client_id, msg):
+        """Purge results from memory. This method is more valuable before we move
+        to a DB based message storage mechanism."""
+        content = msg['content']
+        msg_ids = content.get('msg_ids', [])
+        reply = dict(status='ok')
+        if msg_ids == 'all':
+            self.db.drop_matching_records(dict(completed={'$ne':None}))
+        else:
+            for msg_id in msg_ids:
+                if msg_id in self.all_completed:
+                    self.db.drop_record(msg_id)
+                else:
+                    if msg_id in self.pending:
+                        try:
+                            raise IndexError("msg pending: %r"%msg_id)
+                        except:
+                            reply = wrap_exception()
+                    else:
+                        try:
+                            raise IndexError("No such msg: %r"%msg_id)
+                        except:
+                            reply = wrap_exception()
+                    break
+            eids = content.get('engine_ids', [])
+            for eid in eids:
+                if eid not in self.engines:
+                    try:
+                        raise IndexError("No such engine: %i"%eid)
+                    except:
+                        reply = wrap_exception()
+                    break
+                msg_ids = self.completed.pop(eid)
+                uid = self.engines[eid].queue
+                self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
+        
+        self.session.send(self.clientele, 'purge_reply', content=reply, ident=client_id)
+    
+    def resubmit_task(self, client_id, msg, buffers):
+        """Resubmit a task."""
+        raise NotImplementedError
+    
+    def get_results(self, client_id, msg):
+        """Get the result of 1 or more messages."""
+        content = msg['content']
+        msg_ids = sorted(set(content['msg_ids']))
+        statusonly = content.get('status_only', False)
+        pending = []
+        completed = []
+        content = dict(status='ok')
+        content['pending'] = pending
+        content['completed'] = completed
+        buffers = []
+        if not statusonly:
+            content['results'] = {}
+            records = self.db.find_records(dict(msg_id={'$in':msg_ids}))
+        for msg_id in msg_ids:
+            if msg_id in self.pending:
+                pending.append(msg_id)
+            elif msg_id in self.all_completed:
+                completed.append(msg_id)
+                if not statusonly:
+                    rec = records[msg_id]
+                    content[msg_id] = { 'result_content': rec['result_content'],
+                                        'header': rec['header'],
+                                        'result_header' : rec['result_header'],
+                                      }
+                    buffers.extend(map(str, rec['result_buffers']))
+            else:
+                try:
+                    raise KeyError('No such message: '+msg_id)
+                except:
+                    content = wrap_exception()
+                break
+        self.session.send(self.clientele, "result_reply", content=content, 
+                                            parent=msg, ident=client_id,
+                                            buffers=buffers)
+

IPython/zmq/parallel/controller.py

@@ -1,1035 +1,227 b''
 #!/usr/bin/env python
 """The IPython Controller with 0MQ
 This is the master object that handles connections from engines and clients,
 and monitors traffic through the various queues.
 """
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2010  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 sys
 import os
-from datetime import datetime
-import logging
 import time
-import uuid
+from multiprocessing import Process
 
 import zmq
-from zmq.eventloop import zmqstream, ioloop
+from zmq.eventloop import  ioloop
+from zmq.eventloop.zmqstream import ZMQStream
+from zmq.devices import ProcessMonitoredQueue
 
 # internal:
-from IPython.zmq.log import logger # a Logger object
 from IPython.zmq.entry_point import bind_port
 
-from streamsession import Message, wrap_exception, ISO8601
+from hub import Hub
 from entry_point import (make_base_argument_parser, select_random_ports, split_ports,
                         connect_logger, parse_url, signal_children, generate_exec_key)
+
+
+import streamsession as session
+import heartmonitor
+from scheduler import launch_scheduler
+
 from dictdb import DictDB
 try:
-    from pymongo.binary import Binary
+    import pymongo
 except ImportError:
     MongoDB=None
 else:
     from mongodb import MongoDB
     
-#-----------------------------------------------------------------------------
-# Code
-#-----------------------------------------------------------------------------
-
-def _passer(*args, **kwargs):
-    return
-
-def init_record(msg):
-    """return an empty TaskRecord dict, with all keys initialized with None."""
-    header = msg['header']
-    return {
-        'msg_id' : header['msg_id'],
-        'header' : header,
-        'content': msg['content'],
-        'buffers': msg['buffers'],
-        'submitted': datetime.strptime(header['date'], ISO8601),
-        'client_uuid' : None,
-        'engine_uuid' : None,
-        'started': None,
-        'completed': None,
-        'resubmitted': None,
-        'result_header' : None,
-        'result_content' : None,
-        'result_buffers' : None,
-        'queue' : None
-    }
-
-
-class EngineConnector(object):
-    """A simple object for accessing the various zmq connections of an object.
-    Attributes are:
-    id (int): engine ID
-    uuid (str): uuid (unused?)
-    queue (str): identity of queue's XREQ socket
-    registration (str): identity of registration XREQ socket
-    heartbeat (str): identity of heartbeat XREQ socket
-    """
-    id=0
-    queue=None
-    control=None
-    registration=None
-    heartbeat=None
-    pending=None
-    
-    def __init__(self, id, queue, registration, control, heartbeat=None):
-        logger.info("engine::Engine Connected: %i"%id)
-        self.id = id
-        self.queue = queue
-        self.registration = registration
-        self.control = control
-        self.heartbeat = heartbeat
-        
-class Controller(object):
-    """The IPython Controller with 0MQ connections
-    
-    Parameters
-    ==========
-    loop: zmq IOLoop instance
-    session: StreamSession object
-    <removed> context: zmq context for creating new connections (?)
-    queue: ZMQStream for monitoring the command queue (SUB)
-    registrar: ZMQStream for engine registration requests (XREP)
-    heartbeat: HeartMonitor object checking the pulse of the engines
-    clientele: ZMQStream for client connections (XREP)
-                not used for jobs, only query/control commands
-    notifier: ZMQStream for broadcasting engine registration changes (PUB)
-    db: connection to db for out of memory logging of commands
-                NotImplemented
-    engine_addrs: dict of zmq connection information for engines to connect
-                to the queues.
-    client_addrs: dict of zmq connection information for engines to connect
-                to the queues.
-    """
-    # internal data structures:
-    ids=None # engine IDs
-    keytable=None
-    engines=None
-    clients=None
-    hearts=None
-    pending=None
-    results=None
-    tasks=None
-    completed=None
-    mia=None
-    incoming_registrations=None
-    registration_timeout=None
-    
-    #objects from constructor:
-    loop=None
-    registrar=None
-    clientelle=None
-    queue=None
-    heartbeat=None
-    notifier=None
-    db=None
-    client_addr=None
-    engine_addrs=None
-    
-    
-    def __init__(self, loop, session, queue, registrar, heartbeat, clientele, notifier, db, engine_addrs, client_addrs):
-        """
-        # universal:
-        loop: IOLoop for creating future connections
-        session: streamsession for sending serialized data
-        # engine:
-        queue: ZMQStream for monitoring queue messages
-        registrar: ZMQStream for engine registration
-        heartbeat: HeartMonitor object for tracking engines
-        # client:
-        clientele: ZMQStream for client connections
-        # extra:
-        db: ZMQStream for db connection (NotImplemented)
-        engine_addrs: zmq address/protocol dict for engine connections
-        client_addrs: zmq address/protocol dict for client connections
-        """
-        self.ids = set()
-        self.keytable={}
-        self.incoming_registrations={}
-        self.engines = {}
-        self.by_ident = {}
-        self.clients = {}
-        self.hearts = {}
-        # self.mia = set()
-        
-        # self.sockets = {}
-        self.loop = loop
-        self.session = session
-        self.registrar = registrar
-        self.clientele = clientele
-        self.queue = queue
-        self.heartbeat = heartbeat
-        self.notifier = notifier
-        self.db = db
-        
-        # validate connection dicts:
-        self.client_addrs = client_addrs
-        assert isinstance(client_addrs['queue'], str)
-        assert isinstance(client_addrs['control'], str)
-        # self.hb_addrs = hb_addrs
-        self.engine_addrs = engine_addrs
-        assert isinstance(engine_addrs['queue'], str)
-        assert isinstance(client_addrs['control'], str)
-        assert len(engine_addrs['heartbeat']) == 2
-        
-        # register our callbacks
-        self.registrar.on_recv(self.dispatch_register_request)
-        self.clientele.on_recv(self.dispatch_client_msg)
-        self.queue.on_recv(self.dispatch_queue_traffic)
-        
-        if heartbeat is not None:
-            heartbeat.add_heart_failure_handler(self.handle_heart_failure)
-            heartbeat.add_new_heart_handler(self.handle_new_heart)
-        
-        self.queue_handlers = { 'in' : self.save_queue_request,
-                                'out': self.save_queue_result,
-                                'intask': self.save_task_request,
-                                'outtask': self.save_task_result,
-                                'tracktask': self.save_task_destination,
-                                'incontrol': _passer,
-                                'outcontrol': _passer,
-        }
-        
-        self.client_handlers = {'queue_request': self.queue_status,
-                                'result_request': self.get_results,
-                                'purge_request': self.purge_results,
-                                'load_request': self.check_load,
-                                'resubmit_request': self.resubmit_task,
-                                'shutdown_request': self.shutdown_request,
-                                }
-        
-        self.registrar_handlers = {'registration_request' : self.register_engine,
-                                'unregistration_request' : self.unregister_engine,
-                                'connection_request': self.connection_request,
-        }
-        self.registration_timeout = max(5000, 2*self.heartbeat.period)
-        # this is the stuff that will move to DB:
-        # self.results = {} # completed results
-        self.pending = set() # pending messages, keyed by msg_id
-        self.queues = {} # pending msg_ids keyed by engine_id
-        self.tasks = {} # pending msg_ids submitted as tasks, keyed by client_id
-        self.completed = {} # completed msg_ids keyed by engine_id
-        self.all_completed = set()
-        
-        logger.info("controller::created controller")
-    
-    def _new_id(self):
-        """gemerate a new ID"""
-        newid = 0
-        incoming = [id[0] for id in self.incoming_registrations.itervalues()]
-        # print newid, self.ids, self.incoming_registrations
-        while newid in self.ids or newid in incoming:
-            newid += 1
-        return newid
-    
-    #-----------------------------------------------------------------------------
-    # message validation
-    #-----------------------------------------------------------------------------
-    
-    def _validate_targets(self, targets):
-        """turn any valid targets argument into a list of integer ids"""
-        if targets is None:
-            # default to all
-            targets = self.ids
-            
-        if isinstance(targets, (int,str,unicode)):
-            # only one target specified
-            targets = [targets]
-        _targets = []
-        for t in targets:
-            # map raw identities to ids
-            if isinstance(t, (str,unicode)):
-                t = self.by_ident.get(t, t)
-            _targets.append(t)
-        targets = _targets
-        bad_targets = [ t for t in targets if t not in self.ids ]
-        if bad_targets:
-            raise IndexError("No Such Engine: %r"%bad_targets)
-        if not targets:
-            raise IndexError("No Engines Registered")
-        return targets
-    
-    def _validate_client_msg(self, msg):
-        """validates and unpacks headers of a message. Returns False if invalid,
-        (ident, header, parent, content)"""
-        client_id = msg[0]
-        try:
-            msg = self.session.unpack_message(msg[1:], content=True)
-        except:
-            logger.error("client::Invalid Message %s"%msg)
-            return False
-        
-        msg_type = msg.get('msg_type', None)
-        if msg_type is None:
-            return False
-        header = msg.get('header')
-        # session doesn't handle split content for now:
-        return client_id, msg
-        
-    
-    #-----------------------------------------------------------------------------
-    # dispatch methods (1 per stream)
-    #-----------------------------------------------------------------------------
-    
-    def dispatch_register_request(self, msg):
-        """"""
-        logger.debug("registration::dispatch_register_request(%s)"%msg)
-        idents,msg = self.session.feed_identities(msg)
-        if not idents:
-            logger.error("Bad Queue Message: %s"%msg, exc_info=True)
-            return
-        try:
-            msg = self.session.unpack_message(msg,content=True)
-        except:
-            logger.error("registration::got bad registration message: %s"%msg, exc_info=True)
-            return
-        
-        msg_type = msg['msg_type']
-        content = msg['content']
-        
-        handler = self.registrar_handlers.get(msg_type, None)
-        if handler is None:
-            logger.error("registration::got bad registration message: %s"%msg)
-        else:
-            handler(idents, msg)
-    
-    def dispatch_queue_traffic(self, msg):
-        """all ME and Task queue messages come through here"""
-        logger.debug("queue traffic: %s"%msg[:2])
-        switch = msg[0]
-        idents, msg = self.session.feed_identities(msg[1:])
-        if not idents:
-            logger.error("Bad Queue Message: %s"%msg)
-            return
-        handler = self.queue_handlers.get(switch, None)
-        if handler is not None:
-            handler(idents, msg)
-        else:
-            logger.error("Invalid message topic: %s"%switch)
-        
-    
-    def dispatch_client_msg(self, msg):
-        """Route messages from clients"""
-        idents, msg = self.session.feed_identities(msg)
-        if not idents:
-            logger.error("Bad Client Message: %s"%msg)
-            return
-        client_id = idents[0]
-        try:
-            msg = self.session.unpack_message(msg, content=True)
-        except:
-            content = wrap_exception()
-            logger.error("Bad Client Message: %s"%msg, exc_info=True)
-            self.session.send(self.clientele, "controller_error", ident=client_id, 
-                    content=content)
-            return
-        
-        # print client_id, header, parent, content
-        #switch on message type:
-        msg_type = msg['msg_type']
-        logger.info("client:: client %s requested %s"%(client_id, msg_type))
-        handler = self.client_handlers.get(msg_type, None)
-        try:
-            assert handler is not None, "Bad Message Type: %s"%msg_type
-        except:
-            content = wrap_exception()
-            logger.error("Bad Message Type: %s"%msg_type, exc_info=True)
-            self.session.send(self.clientele, "controller_error", ident=client_id, 
-                    content=content)
-            return
-        else:
-            handler(client_id, msg)
-            
-    def dispatch_db(self, msg):
-        """"""
-        raise NotImplementedError
-    
-    #---------------------------------------------------------------------------
-    # handler methods (1 per event)
-    #---------------------------------------------------------------------------
-    
-    #----------------------- Heartbeat --------------------------------------
-    
-    def handle_new_heart(self, heart):
-        """handler to attach to heartbeater.
-        Called when a new heart starts to beat.
-        Triggers completion of registration."""
-        logger.debug("heartbeat::handle_new_heart(%r)"%heart)
-        if heart not in self.incoming_registrations:
-            logger.info("heartbeat::ignoring new heart: %r"%heart)
-        else:
-            self.finish_registration(heart)
-        
-    
-    def handle_heart_failure(self, heart):
-        """handler to attach to heartbeater.
-        called when a previously registered heart fails to respond to beat request.
-        triggers unregistration"""
-        logger.debug("heartbeat::handle_heart_failure(%r)"%heart)
-        eid = self.hearts.get(heart, None)
-        queue = self.engines[eid].queue
-        if eid is None:
-            logger.info("heartbeat::ignoring heart failure %r"%heart)
-        else:
-            self.unregister_engine(heart, dict(content=dict(id=eid, queue=queue)))
-    
-    #----------------------- MUX Queue Traffic ------------------------------
-    
-    def save_queue_request(self, idents, msg):
-        if len(idents) < 2:
-            logger.error("invalid identity prefix: %s"%idents)
-            return
-        queue_id, client_id = idents[:2]
-        try:
-            msg = self.session.unpack_message(msg, content=False)
-        except:
-            logger.error("queue::client %r sent invalid message to %r: %s"%(client_id, queue_id, msg), exc_info=True)
-            return
-        
-        eid = self.by_ident.get(queue_id, None)
-        if eid is None:
-            logger.error("queue::target %r not registered"%queue_id)
-            logger.debug("queue::    valid are: %s"%(self.by_ident.keys()))
-            return
-            
-        header = msg['header']
-        msg_id = header['msg_id']
-        record = init_record(msg)
-        record['engine_uuid'] = queue_id
-        record['client_uuid'] = client_id
-        record['queue'] = 'mux'
-        if MongoDB is not None and isinstance(self.db, MongoDB):
-            record['buffers'] = map(Binary, record['buffers'])
-        self.pending.add(msg_id)
-        self.queues[eid].append(msg_id)
-        self.db.add_record(msg_id, record)
-    
-    def save_queue_result(self, idents, msg):
-        if len(idents) < 2:
-            logger.error("invalid identity prefix: %s"%idents)
-            return
-            
-        client_id, queue_id = idents[:2]
-        try:
-            msg = self.session.unpack_message(msg, content=False)
-        except:
-            logger.error("queue::engine %r sent invalid message to %r: %s"%(
-                    queue_id,client_id, msg), exc_info=True)
-            return
-        
-        eid = self.by_ident.get(queue_id, None)
-        if eid is None:
-            logger.error("queue::unknown engine %r is sending a reply: "%queue_id)
-            logger.debug("queue::       %s"%msg[2:])
-            return
-        
-        parent = msg['parent_header']
-        if not parent:
-            return
-        msg_id = parent['msg_id']
-        if msg_id in self.pending:
-            self.pending.remove(msg_id)
-            self.all_completed.add(msg_id)
-            self.queues[eid].remove(msg_id)
-            self.completed[eid].append(msg_id)
-            rheader = msg['header']
-            completed = datetime.strptime(rheader['date'], ISO8601)
-            started = rheader.get('started', None)
-            if started is not None:
-                started = datetime.strptime(started, ISO8601)
-            result = {
-                'result_header' : rheader,
-                'result_content': msg['content'],
-                'started' : started,
-                'completed' : completed
-            }
-            if MongoDB is not None and isinstance(self.db, MongoDB):
-                result['result_buffers'] = map(Binary, msg['buffers'])
-            else:
-                result['result_buffers'] = msg['buffers']
-            self.db.update_record(msg_id, result)
-        else:
-            logger.debug("queue:: unknown msg finished %s"%msg_id)
-            
-    #--------------------- Task Queue Traffic ------------------------------
-    
-    def save_task_request(self, idents, msg):
-        """Save the submission of a task."""
-        client_id = idents[0]
-        
-        try:
-            msg = self.session.unpack_message(msg, content=False)
-        except:
-            logger.error("task::client %r sent invalid task message: %s"%(
-                    client_id, msg), exc_info=True)
-            return
-        record = init_record(msg)
-        if MongoDB is not None and isinstance(self.db, MongoDB):
-            record['buffers'] = map(Binary, record['buffers'])
-        record['client_uuid'] = client_id
-        record['queue'] = 'task'
-        header = msg['header']
-        msg_id = header['msg_id']
-        self.pending.add(msg_id)
-        self.db.add_record(msg_id, record)
-    
-    def save_task_result(self, idents, msg):
-        """save the result of a completed task."""
-        client_id = idents[0]
-        try:
-            msg = self.session.unpack_message(msg, content=False)
-        except:
-            logger.error("task::invalid task result message send to %r: %s"%(
-                    client_id, msg))
-            raise
-            return
-        
-        parent = msg['parent_header']
-        if not parent:
-            # print msg
-            logger.warn("Task %r had no parent!"%msg)
-            return
-        msg_id = parent['msg_id']
-        
-        header = msg['header']
-        engine_uuid = header.get('engine', None)
-        eid = self.by_ident.get(engine_uuid, None)
-        
-        if msg_id in self.pending:
-            self.pending.remove(msg_id)
-            self.all_completed.add(msg_id)
-            if eid is not None:
-                self.completed[eid].append(msg_id)
-                if msg_id in self.tasks[eid]:
-                    self.tasks[eid].remove(msg_id)
-            completed = datetime.strptime(header['date'], ISO8601)
-            started = header.get('started', None)
-            if started is not None:
-                started = datetime.strptime(started, ISO8601)
-            result = {
-                'result_header' : header,
-                'result_content': msg['content'],
-                'started' : started,
-                'completed' : completed,
-                'engine_uuid': engine_uuid
-            }
-            if MongoDB is not None and isinstance(self.db, MongoDB):
-                result['result_buffers'] = map(Binary, msg['buffers'])
-            else:
-                result['result_buffers'] = msg['buffers']
-            self.db.update_record(msg_id, result)
-            
-        else:
-            logger.debug("task::unknown task %s finished"%msg_id)
-    
-    def save_task_destination(self, idents, msg):
-        try:
-            msg = self.session.unpack_message(msg, content=True)
-        except:
-            logger.error("task::invalid task tracking message")
-            return
-        content = msg['content']
-        print (content)
-        msg_id = content['msg_id']
-        engine_uuid = content['engine_id']
-        eid = self.by_ident[engine_uuid]
-        
-        logger.info("task::task %s arrived on %s"%(msg_id, eid))
-        # if msg_id in self.mia:
-        #     self.mia.remove(msg_id)
-        # else:
-        #     logger.debug("task::task %s not listed as MIA?!"%(msg_id))
-        
-        self.tasks[eid].append(msg_id)
-        # self.pending[msg_id][1].update(received=datetime.now(),engine=(eid,engine_uuid))
-        self.db.update_record(msg_id, dict(engine_uuid=engine_uuid))
-    
-    def mia_task_request(self, idents, msg):
-        raise NotImplementedError
-        client_id = idents[0]
-        # content = dict(mia=self.mia,status='ok')
-        # self.session.send('mia_reply', content=content, idents=client_id)
-        
-        
-            
-    #-------------------------------------------------------------------------
-    # Registration requests
-    #-------------------------------------------------------------------------
-        
-    def connection_request(self, client_id, msg):
-        """Reply with connection addresses for clients."""
-        logger.info("client::client %s connected"%client_id)
-        content = dict(status='ok')
-        content.update(self.client_addrs)
-        jsonable = {}
-        for k,v in self.keytable.iteritems():
-            jsonable[str(k)] = v
-        content['engines'] = jsonable
-        self.session.send(self.registrar, 'connection_reply', content, parent=msg, ident=client_id)
-    
-    def register_engine(self, reg, msg):
-        """Register a new engine."""
-        content = msg['content']
-        try:
-            queue = content['queue']
-        except KeyError:
-            logger.error("registration::queue not specified")
-            return
-        heart = content.get('heartbeat', None)
-        """register a new engine, and create the socket(s) necessary"""
-        eid = self._new_id()
-        # print (eid, queue, reg, heart)
-        
-        logger.debug("registration::register_engine(%i, %r, %r, %r)"%(eid, queue, reg, heart))
-        
-        content = dict(id=eid,status='ok')
-        content.update(self.engine_addrs)
-        # check if requesting available IDs:
-        if queue in self.by_ident:
-            try:
-                raise KeyError("queue_id %r in use"%queue)
-            except:
-                content = wrap_exception()
-        elif heart in self.hearts: # need to check unique hearts?
-            try:
-                raise KeyError("heart_id %r in use"%heart)
-            except:
-                content = wrap_exception()
-        else:
-            for h, pack in self.incoming_registrations.iteritems():
-                if heart == h:
-                    try:
-                        raise KeyError("heart_id %r in use"%heart)
-                    except:
-                        content = wrap_exception()
-                    break
-                elif queue == pack[1]:
-                    try:
-                        raise KeyError("queue_id %r in use"%queue)
-                    except:
-                        content = wrap_exception()
-                    break
-        
-        msg = self.session.send(self.registrar, "registration_reply", 
-                content=content, 
-                ident=reg)
-        
-        if content['status'] == 'ok':
-            if heart in self.heartbeat.hearts:
-                # already beating
-                self.incoming_registrations[heart] = (eid,queue,reg,None)
-                self.finish_registration(heart)
-            else:
-                purge = lambda : self._purge_stalled_registration(heart)
-                dc = ioloop.DelayedCallback(purge, self.registration_timeout, self.loop)
-                dc.start()
-                self.incoming_registrations[heart] = (eid,queue,reg,dc)
-        else:
-            logger.error("registration::registration %i failed: %s"%(eid, content['evalue']))
-        return eid
-    
-    def unregister_engine(self, ident, msg):
-        """Unregister an engine that explicitly requested to leave."""
-        try:
-            eid = msg['content']['id']
-        except:
-            logger.error("registration::bad engine id for unregistration: %s"%ident)
-            return
-        logger.info("registration::unregister_engine(%s)"%eid)
-        content=dict(id=eid, queue=self.engines[eid].queue)
-        self.ids.remove(eid)
-        self.keytable.pop(eid)
-        ec = self.engines.pop(eid)
-        self.hearts.pop(ec.heartbeat)
-        self.by_ident.pop(ec.queue)
-        self.completed.pop(eid)
-        for msg_id in self.queues.pop(eid):
-            msg = self.pending.remove(msg_id)
-            ############## TODO: HANDLE IT ################
-        
-        if self.notifier:
-            self.session.send(self.notifier, "unregistration_notification", content=content)
-    
-    def finish_registration(self, heart):
-        """Second half of engine registration, called after our HeartMonitor
-        has received a beat from the Engine's Heart."""
-        try: 
-            (eid,queue,reg,purge) = self.incoming_registrations.pop(heart)
-        except KeyError:
-            logger.error("registration::tried to finish nonexistant registration")
-            return
-        logger.info("registration::finished registering engine %i:%r"%(eid,queue))
-        if purge is not None:
-            purge.stop()
-        control = queue
-        self.ids.add(eid)
-        self.keytable[eid] = queue
-        self.engines[eid] = EngineConnector(eid, queue, reg, control, heart)
-        self.by_ident[queue] = eid
-        self.queues[eid] = list()
-        self.tasks[eid] = list()
-        self.completed[eid] = list()
-        self.hearts[heart] = eid
-        content = dict(id=eid, queue=self.engines[eid].queue)
-        if self.notifier:
-            self.session.send(self.notifier, "registration_notification", content=content)
-    
-    def _purge_stalled_registration(self, heart):
-        if heart in self.incoming_registrations:
-            eid = self.incoming_registrations.pop(heart)[0]
-            logger.info("registration::purging stalled registration: %i"%eid)
-        else:
-            pass
-            
-    #-------------------------------------------------------------------------
-    # Client Requests
-    #-------------------------------------------------------------------------
-    
-    def shutdown_request(self, client_id, msg):
-        """handle shutdown request."""
-        # s = self.context.socket(zmq.XREQ)
-        # s.connect(self.client_connections['mux'])
-        # time.sleep(0.1)
-        # for eid,ec in self.engines.iteritems():
-        #     self.session.send(s, 'shutdown_request', content=dict(restart=False), ident=ec.queue)
-        # time.sleep(1)
-        self.session.send(self.clientele, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
-        dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
-        dc.start()
-    
-    def _shutdown(self):
-        logger.info("controller::controller shutting down.")
-        time.sleep(0.1)
-        sys.exit(0)
-        
-    
-    def check_load(self, client_id, msg):
-        content = msg['content']
-        try:
-            targets = content['targets']
-            targets = self._validate_targets(targets)
-        except:
-            content = wrap_exception()
-            self.session.send(self.clientele, "controller_error", 
-                    content=content, ident=client_id)
-            return
-        
-        content = dict(status='ok')
-        # loads = {}
-        for t in targets:
-            content[bytes(t)] = len(self.queues[t])+len(self.tasks[t])
-        self.session.send(self.clientele, "load_reply", content=content, ident=client_id)
-            
-    
-    def queue_status(self, client_id, msg):
-        """Return the Queue status of one or more targets.
-        if verbose: return the msg_ids
-        else: return len of each type.
-        keys: queue (pending MUX jobs)
-            tasks (pending Task jobs)
-            completed (finished jobs from both queues)"""
-        content = msg['content']
-        targets = content['targets']
-        try:
-            targets = self._validate_targets(targets)
-        except:
-            content = wrap_exception()
-            self.session.send(self.clientele, "controller_error", 
-                    content=content, ident=client_id)
-            return
-        verbose = content.get('verbose', False)
-        content = dict(status='ok')
-        for t in targets:
-            queue = self.queues[t]
-            completed = self.completed[t]
-            tasks = self.tasks[t]
-            if not verbose:
-                queue = len(queue)
-                completed = len(completed)
-                tasks = len(tasks)
-            content[bytes(t)] = {'queue': queue, 'completed': completed , 'tasks': tasks}
-            # pending
-        self.session.send(self.clientele, "queue_reply", content=content, ident=client_id)
-    
-    def purge_results(self, client_id, msg):
-        """Purge results from memory. This method is more valuable before we move
-        to a DB based message storage mechanism."""
-        content = msg['content']
-        msg_ids = content.get('msg_ids', [])
-        reply = dict(status='ok')
-        if msg_ids == 'all':
-            self.db.drop_matching_records(dict(completed={'$ne':None}))
-        else:
-            for msg_id in msg_ids:
-                if msg_id in self.all_completed:
-                    self.db.drop_record(msg_id)
-                else:
-                    if msg_id in self.pending:
-                        try:
-                            raise IndexError("msg pending: %r"%msg_id)
-                        except:
-                            reply = wrap_exception()
-                    else:
-                        try:
-                            raise IndexError("No such msg: %r"%msg_id)
-                        except:
-                            reply = wrap_exception()
-                    break
-            eids = content.get('engine_ids', [])
-            for eid in eids:
-                if eid not in self.engines:
-                    try:
-                        raise IndexError("No such engine: %i"%eid)
-                    except:
-                        reply = wrap_exception()
-                    break
-                msg_ids = self.completed.pop(eid)
-                uid = self.engines[eid].queue
-                self.db.drop_matching_records(dict(engine_uuid=uid, completed={'$ne':None}))
-        
-        self.session.send(self.clientele, 'purge_reply', content=reply, ident=client_id)
-    
-    def resubmit_task(self, client_id, msg, buffers):
-        """Resubmit a task."""
-        raise NotImplementedError
-    
-    def get_results(self, client_id, msg):
-        """Get the result of 1 or more messages."""
-        content = msg['content']
-        msg_ids = sorted(set(content['msg_ids']))
-        statusonly = content.get('status_only', False)
-        pending = []
-        completed = []
-        content = dict(status='ok')
-        content['pending'] = pending
-        content['completed'] = completed
-        buffers = []
-        if not statusonly:
-            content['results'] = {}
-            records = self.db.find_records(dict(msg_id={'$in':msg_ids}))
-        for msg_id in msg_ids:
-            if msg_id in self.pending:
-                pending.append(msg_id)
-            elif msg_id in self.all_completed:
-                completed.append(msg_id)
-                if not statusonly:
-                    rec = records[msg_id]
-                    content[msg_id] = { 'result_content': rec['result_content'],
-                                        'header': rec['header'],
-                                        'result_header' : rec['result_header'],
-                                      }
-                    buffers.extend(map(str, rec['result_buffers']))
-            else:
-                try:
-                    raise KeyError('No such message: '+msg_id)
-                except:
-                    content = wrap_exception()
-                break
-        self.session.send(self.clientele, "result_reply", content=content, 
-                                            parent=msg, ident=client_id,
-                                            buffers=buffers)
-
-
 #-------------------------------------------------------------------------
 # Entry Point
 #-------------------------------------------------------------------------
 
 def make_argument_parser():
     """Make an argument parser"""
     parser = make_base_argument_parser()
     
     parser.add_argument('--client', type=int, metavar='PORT', default=0,
                         help='set the XREP port for clients [default: random]')
     parser.add_argument('--notice', type=int, metavar='PORT', default=0,
                         help='set the PUB socket for registration notification [default: random]')
     parser.add_argument('--hb', type=str, metavar='PORTS',
                         help='set the 2 ports for heartbeats [default: random]')
     parser.add_argument('--ping', type=int, default=3000,
                         help='set the heartbeat period in ms [default: 3000]')
     parser.add_argument('--monitor', type=int, metavar='PORT', default=0,
                         help='set the SUB port for queue monitoring [default: random]')
     parser.add_argument('--mux', type=str, metavar='PORTS',
                         help='set the XREP ports for the MUX queue [default: random]')
     parser.add_argument('--task', type=str, metavar='PORTS',
                         help='set the XREP/XREQ ports for the task queue [default: random]')
     parser.add_argument('--control', type=str, metavar='PORTS',
                         help='set the XREP ports for the control queue [default: random]')
     parser.add_argument('--scheduler', type=str, default='pure',
                         choices = ['pure', 'lru', 'plainrandom', 'weighted', 'twobin','leastload'],
                         help='select the task scheduler  [default: pure ZMQ]')
     parser.add_argument('--mongodb', action='store_true',
                         help='Use MongoDB task storage [default: in-memory]')
     
     return parser
     
 def main(argv=None):
-    import time
-    from multiprocessing import Process
-    
-    from zmq.eventloop.zmqstream import ZMQStream
-    from zmq.devices import ProcessMonitoredQueue
-    from zmq.log import handlers
-    
-    import streamsession as session
-    import heartmonitor
-    from scheduler import launch_scheduler
     
     parser = make_argument_parser()
     
     args = parser.parse_args(argv)
     parse_url(args)
     
     iface="%s://%s"%(args.transport,args.ip)+':%i'
     
     random_ports = 0
     if args.hb:
         hb = split_ports(args.hb, 2)
     else:
         hb = select_random_ports(2)
     if args.mux:
         mux = split_ports(args.mux, 2)
     else:
         mux = None
         random_ports += 2
     if args.task:
         task = split_ports(args.task, 2)
     else:
         task = None
         random_ports += 2
     if args.control:
         control = split_ports(args.control, 2)
     else:
         control = None
         random_ports += 2
     
     ctx = zmq.Context()
     loop = ioloop.IOLoop.instance()
     
     # setup logging
     connect_logger(ctx, iface%args.logport, root="controller", loglevel=args.loglevel)
     
     # Registrar socket
     reg = ZMQStream(ctx.socket(zmq.XREP), loop)
     regport = bind_port(reg, args.ip, args.regport)
     
     ### Engine connections ###
     
     # heartbeat
     hpub = ctx.socket(zmq.PUB)
     bind_port(hpub, args.ip, hb[0])
     hrep = ctx.socket(zmq.XREP)
     bind_port(hrep, args.ip, hb[1])
     
     hmon = heartmonitor.HeartMonitor(loop, ZMQStream(hpub,loop), ZMQStream(hrep,loop),args.ping)
     hmon.start()
     
     ### Client connections ###
     # Clientele socket
     c = ZMQStream(ctx.socket(zmq.XREP), loop)
     cport = bind_port(c, args.ip, args.client)
     # Notifier socket
     n = ZMQStream(ctx.socket(zmq.PUB), loop)
     nport = bind_port(n, args.ip, args.notice)
     
     ### Key File ###
     if args.execkey and not os.path.isfile(args.execkey):
             generate_exec_key(args.execkey)
     
     thesession = session.StreamSession(username=args.ident or "controller", keyfile=args.execkey)
     
     ### build and launch the queues ###
     
     # monitor socket
     sub = ctx.socket(zmq.SUB)
     sub.setsockopt(zmq.SUBSCRIBE, "")
     monport = bind_port(sub, args.ip, args.monitor)
     sub = ZMQStream(sub, loop)
     
     ports = select_random_ports(random_ports)
     children = []
     # Multiplexer Queue (in a Process)
     if not mux:
         mux = (ports.pop(),ports.pop())
     q = ProcessMonitoredQueue(zmq.XREP, zmq.XREP, zmq.PUB, 'in', 'out')
     q.bind_in(iface%mux[0])
     q.bind_out(iface%mux[1])
     q.connect_mon(iface%monport)
     q.daemon=True
     q.start()
     children.append(q.launcher)
     
     # Control Queue (in a Process)
     if not control:
         control = (ports.pop(),ports.pop())
     q = ProcessMonitoredQueue(zmq.XREP, zmq.XREP, zmq.PUB, 'incontrol', 'outcontrol')
     q.bind_in(iface%control[0])
     q.bind_out(iface%control[1])
     q.connect_mon(iface%monport)
     q.daemon=True
     q.start()
     children.append(q.launcher)
     # Task Queue (in a Process)
     if not task:
         task = (ports.pop(),ports.pop())
     if args.scheduler == 'pure':
         q = ProcessMonitoredQueue(zmq.XREP, zmq.XREQ, zmq.PUB, 'intask', 'outtask')
         q.bind_in(iface%task[0])
         q.bind_out(iface%task[1])
         q.connect_mon(iface%monport)
         q.daemon=True
         q.start()
         children.append(q.launcher)
     else:
         sargs = (iface%task[0],iface%task[1],iface%monport,iface%nport,args.scheduler)
         print (sargs)
         q = Process(target=launch_scheduler, args=sargs)
         q.daemon=True
         q.start()
         children.append(q)
     
     if args.mongodb:
         from mongodb import MongoDB
         db = MongoDB(thesession.session)
     else:
         db = DictDB()
     time.sleep(.25)
     
     # build connection dicts
     engine_addrs = {
         'control' : iface%control[1],
         'queue': iface%mux[1],
         'heartbeat': (iface%hb[0], iface%hb[1]),
         'task' : iface%task[1],
         'monitor' : iface%monport,
         }
     
     client_addrs = {
         'control' : iface%control[0],
         'query': iface%cport,
         'queue': iface%mux[0],
         'task' : iface%task[0],
         'notification': iface%nport
         }
     signal_children(children)
-    con = Controller(loop, thesession, sub, reg, hmon, c, n, db, engine_addrs, client_addrs)
+    hub = Hub(loop, thesession, sub, reg, hmon, c, n, db, engine_addrs, client_addrs)
     dc = ioloop.DelayedCallback(lambda : print("Controller started..."), 100, loop)
     dc.start()
     loop.start()
     
     
     
 
 if __name__ == '__main__':
     main()