SGE test related fixes

* allow iopub messages to arrive first on Hub
* SQLite no longer commits immediately
* parallelwave example
* typos in launcher.py

IPython/parallel/hub.py

@@ -1,1035 +1,1090
 #!/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
 import time
 from datetime import datetime
 
 import zmq
 from zmq.eventloop import ioloop
 from zmq.eventloop.zmqstream import ZMQStream
 
 # internal:
 from IPython.utils.importstring import import_item
 from IPython.utils.traitlets import HasTraits, Instance, Int, CStr, Str, Dict, Set, List, Bool
 
 from .entry_point import select_random_ports
 from .factory import RegistrationFactory, LoggingFactory
 
 from . import error
 from .heartmonitor import HeartMonitor
 from .util import validate_url_container, ISO8601
 
 #-----------------------------------------------------------------------------
 # Code
 #-----------------------------------------------------------------------------
 
 def _passer(*args, **kwargs):
     return
 
 def _printer(*args, **kwargs):
     print (args)
     print (kwargs)
 
+def empty_record():
+    """Return an empty dict with all record keys."""
+    return {
+        'msg_id' : None,
+        'header' : None,
+        'content': None,
+        'buffers': None,
+        'submitted': None,
+        'client_uuid' : None,
+        'engine_uuid' : None,
+        'started': None,
+        'completed': None,
+        'resubmitted': None,
+        'result_header' : None,
+        'result_content' : None,
+        'result_buffers' : None,
+        'queue' : None,
+        'pyin' : None,
+        'pyout': None,
+        'pyerr': None,
+        'stdout': '',
+        'stderr': '',
+    }
+    
 def init_record(msg):
     """Initialize a TaskRecord based on a request."""
     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,
         'pyin' : None,
         'pyout': None,
         'pyerr': None,
         'stdout': '',
         'stderr': '',
     }
 
 
 class EngineConnector(HasTraits):
     """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=Int(0)
     queue=Str()
     control=Str()
     registration=Str()
     heartbeat=Str()
     pending=Set()
 
 class HubFactory(RegistrationFactory):
     """The Configurable for setting up a Hub."""
     
     # name of a scheduler scheme
     scheme = Str('leastload', config=True)
     
     # port-pairs for monitoredqueues:
     hb = Instance(list, config=True)
     def _hb_default(self):
         return select_random_ports(2)
     
     mux = Instance(list, config=True)
     def _mux_default(self):
         return select_random_ports(2)
     
     task = Instance(list, config=True)
     def _task_default(self):
         return select_random_ports(2)
     
     control = Instance(list, config=True)
     def _control_default(self):
         return select_random_ports(2)
     
     iopub = Instance(list, config=True)
     def _iopub_default(self):
         return select_random_ports(2)
     
     # single ports:
     mon_port = Instance(int, config=True)
     def _mon_port_default(self):
         return select_random_ports(1)[0]
     
     notifier_port = Instance(int, config=True)
     def _notifier_port_default(self):
         return select_random_ports(1)[0]
     
     ping = Int(1000, config=True) # ping frequency
     
     engine_ip = CStr('127.0.0.1', config=True)
     engine_transport = CStr('tcp', config=True)
     
     client_ip = CStr('127.0.0.1', config=True)
     client_transport = CStr('tcp', config=True)
     
     monitor_ip = CStr('127.0.0.1', config=True)
     monitor_transport = CStr('tcp', config=True)
     
     monitor_url = CStr('')
     
     db_class = CStr('IPython.parallel.dictdb.DictDB', config=True)
     
     # not configurable
     db = Instance('IPython.parallel.dictdb.BaseDB')
     heartmonitor = Instance('IPython.parallel.heartmonitor.HeartMonitor')
     subconstructors = List()
     _constructed = Bool(False)
     
     def _ip_changed(self, name, old, new):
         self.engine_ip = new
         self.client_ip = new
         self.monitor_ip = new
         self._update_monitor_url()
     
     def _update_monitor_url(self):
         self.monitor_url = "%s://%s:%i"%(self.monitor_transport, self.monitor_ip, self.mon_port)
     
     def _transport_changed(self, name, old, new):
         self.engine_transport = new
         self.client_transport = new
         self.monitor_transport = new
         self._update_monitor_url()
         
     def __init__(self, **kwargs):
         super(HubFactory, self).__init__(**kwargs)
         self._update_monitor_url()
         # self.on_trait_change(self._sync_ips, 'ip')
         # self.on_trait_change(self._sync_transports, 'transport')
         self.subconstructors.append(self.construct_hub)
     
     
     def construct(self):
         assert not self._constructed, "already constructed!"
         
         for subc in self.subconstructors:
             subc()
         
         self._constructed = True
         
     
     def start(self):
         assert self._constructed, "must be constructed by self.construct() first!"
         self.heartmonitor.start()
         self.log.info("Heartmonitor started")
     
     def construct_hub(self):
         """construct"""
         client_iface = "%s://%s:"%(self.client_transport, self.client_ip) + "%i"
         engine_iface = "%s://%s:"%(self.engine_transport, self.engine_ip) + "%i"
         
         ctx = self.context
         loop = self.loop
         
         # Registrar socket
         q = ZMQStream(ctx.socket(zmq.XREP), loop)
         q.bind(client_iface % self.regport)
         self.log.info("Hub listening on %s for registration."%(client_iface%self.regport))
         if self.client_ip != self.engine_ip:
             q.bind(engine_iface % self.regport)
             self.log.info("Hub listening on %s for registration."%(engine_iface%self.regport))
         
         ### Engine connections ###
 
         # heartbeat
         hpub = ctx.socket(zmq.PUB)
         hpub.bind(engine_iface % self.hb[0])
         hrep = ctx.socket(zmq.XREP)
         hrep.bind(engine_iface % self.hb[1])
         self.heartmonitor = HeartMonitor(loop=loop, pingstream=ZMQStream(hpub,loop), pongstream=ZMQStream(hrep,loop), 
                                 period=self.ping, logname=self.log.name)
 
         ### Client connections ###
         # Notifier socket
         n = ZMQStream(ctx.socket(zmq.PUB), loop)
         n.bind(client_iface%self.notifier_port)
 
         ### build and launch the queues ###
 
         # monitor socket
         sub = ctx.socket(zmq.SUB)
         sub.setsockopt(zmq.SUBSCRIBE, "")
         sub.bind(self.monitor_url)
         sub.bind('inproc://monitor')
         sub = ZMQStream(sub, loop)
         
         # connect the db
         self.log.info('Hub using DB backend: %r'%(self.db_class.split()[-1]))
         # cdir = self.config.Global.cluster_dir
         self.db = import_item(self.db_class)(session=self.session.session, config=self.config)
         time.sleep(.25)
 
         # build connection dicts
         self.engine_info = {
             'control' : engine_iface%self.control[1],
             'mux': engine_iface%self.mux[1],
             'heartbeat': (engine_iface%self.hb[0], engine_iface%self.hb[1]),
             'task' : engine_iface%self.task[1],
             'iopub' : engine_iface%self.iopub[1],
             # 'monitor' : engine_iface%self.mon_port,
             }
 
         self.client_info = {
             'control' : client_iface%self.control[0],
             'mux': client_iface%self.mux[0],
             'task' : (self.scheme, client_iface%self.task[0]),
             'iopub' : client_iface%self.iopub[0],
             'notification': client_iface%self.notifier_port
             }
         self.log.debug("Hub engine addrs: %s"%self.engine_info)
         self.log.debug("Hub client addrs: %s"%self.client_info)
         self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor,
                 query=q, notifier=n, db=self.db,
                 engine_info=self.engine_info, client_info=self.client_info,
                 logname=self.log.name)
     
 
 class Hub(LoggingFactory):
     """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)
     query: ZMQStream for engine registration and client queries requests (XREP)
     heartbeat: HeartMonitor object checking the pulse of the engines
     notifier: ZMQStream for broadcasting engine registration changes (PUB)
     db: connection to db for out of memory logging of commands
                 NotImplemented
     engine_info: dict of zmq connection information for engines to connect
                 to the queues.
     client_info: dict of zmq connection information for engines to connect
                 to the queues.
     """
     # internal data structures:
     ids=Set() # engine IDs
     keytable=Dict()
     by_ident=Dict()
     engines=Dict()
     clients=Dict()
     hearts=Dict()
     pending=Set()
     queues=Dict()  # pending msg_ids keyed by engine_id
     tasks=Dict() # pending msg_ids submitted as tasks, keyed by client_id
     completed=Dict() # completed msg_ids keyed by engine_id
     all_completed=Set() # completed msg_ids keyed by engine_id
+    dead_engines=Set() # completed msg_ids keyed by engine_id
     # mia=None
     incoming_registrations=Dict()
     registration_timeout=Int()
     _idcounter=Int(0)
     
     # objects from constructor:
     loop=Instance(ioloop.IOLoop)
     query=Instance(ZMQStream)
     monitor=Instance(ZMQStream)
     heartmonitor=Instance(HeartMonitor)
     notifier=Instance(ZMQStream)
     db=Instance(object)
     client_info=Dict()
     engine_info=Dict()
     
     
     def __init__(self, **kwargs):
         """
         # universal:
         loop: IOLoop for creating future connections
         session: streamsession for sending serialized data
         # engine:
         queue: ZMQStream for monitoring queue messages
         query: ZMQStream for engine+client registration and client requests
         heartbeat: HeartMonitor object for tracking engines
         # extra:
         db: ZMQStream for db connection (NotImplemented)
         engine_info: zmq address/protocol dict for engine connections
         client_info: zmq address/protocol dict for client connections
         """
         
         super(Hub, self).__init__(**kwargs)
         self.registration_timeout = max(5000, 2*self.heartmonitor.period)
         
         # validate connection dicts:
         for k,v in self.client_info.iteritems():
             if k == 'task':
                 validate_url_container(v[1])
             else:
                 validate_url_container(v)
         # validate_url_container(self.client_info)
         validate_url_container(self.engine_info)
         
         # register our callbacks
         self.query.on_recv(self.dispatch_query)
         self.monitor.on_recv(self.dispatch_monitor_traffic)
         
         self.heartmonitor.add_heart_failure_handler(self.handle_heart_failure)
         self.heartmonitor.add_new_heart_handler(self.handle_new_heart)
         
         self.monitor_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,
                                 'iopub': self.save_iopub_message,
         }
         
         self.query_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,
                                 'registration_request' : self.register_engine,
                                 'unregistration_request' : self.unregister_engine,
                                 'connection_request': self.connection_request,
         }
         
         self.log.info("hub::created hub")
     
     @property
     def _next_id(self):
         """gemerate a new ID.
         
         No longer reuse old ids, just count from 0."""
         newid = self._idcounter
         self._idcounter += 1
         return newid
         # 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
     
     #-----------------------------------------------------------------------------
     # dispatch methods (1 per stream)
     #-----------------------------------------------------------------------------
     
     # def dispatch_registration_request(self, msg):
     #     """"""
     #     self.log.debug("registration::dispatch_register_request(%s)"%msg)
     #     idents,msg = self.session.feed_identities(msg)
     #     if not idents:
     #         self.log.error("Bad Query Message: %s"%msg, exc_info=True)
     #         return
     #     try:
     #         msg = self.session.unpack_message(msg,content=True)
     #     except:
     #         self.log.error("registration::got bad registration message: %s"%msg, exc_info=True)
     #         return
     #     
     #     msg_type = msg['msg_type']
     #     content = msg['content']
     #     
     #     handler = self.query_handlers.get(msg_type, None)
     #     if handler is None:
     #         self.log.error("registration::got bad registration message: %s"%msg)
     #     else:
     #         handler(idents, msg)
     
     def dispatch_monitor_traffic(self, msg):
         """all ME and Task queue messages come through here, as well as
         IOPub traffic."""
         self.log.debug("monitor traffic: %s"%msg[:2])
         switch = msg[0]
         idents, msg = self.session.feed_identities(msg[1:])
         if not idents:
             self.log.error("Bad Monitor Message: %s"%msg)
             return
         handler = self.monitor_handlers.get(switch, None)
         if handler is not None:
             handler(idents, msg)
         else:
             self.log.error("Invalid monitor topic: %s"%switch)
         
     
     def dispatch_query(self, msg):
         """Route registration requests and queries from clients."""
         idents, msg = self.session.feed_identities(msg)
         if not idents:
             self.log.error("Bad Query Message: %s"%msg)
             return
         client_id = idents[0]
         try:
             msg = self.session.unpack_message(msg, content=True)
         except:
             content = error.wrap_exception()
             self.log.error("Bad Query Message: %s"%msg, exc_info=True)
             self.session.send(self.query, "hub_error", ident=client_id, 
                     content=content)
             return
         
         # print client_id, header, parent, content
         #switch on message type:
         msg_type = msg['msg_type']
         self.log.info("client::client %s requested %s"%(client_id, msg_type))
         handler = self.query_handlers.get(msg_type, None)
         try:
             assert handler is not None, "Bad Message Type: %s"%msg_type
         except:
             content = error.wrap_exception()
             self.log.error("Bad Message Type: %s"%msg_type, exc_info=True)
             self.session.send(self.query, "hub_error", ident=client_id, 
                     content=content)
             return
         else:
             handler(idents, 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."""
         self.log.debug("heartbeat::handle_new_heart(%r)"%heart)
         if heart not in self.incoming_registrations:
             self.log.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"""
         self.log.debug("heartbeat::handle_heart_failure(%r)"%heart)
         eid = self.hearts.get(heart, None)
         queue = self.engines[eid].queue
         if eid is None:
             self.log.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:
             self.log.error("invalid identity prefix: %s"%idents)
             return
         queue_id, client_id = idents[:2]
         try:
             msg = self.session.unpack_message(msg, content=False)
         except:
             self.log.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:
             self.log.error("queue::target %r not registered"%queue_id)
             self.log.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'
 
+        try:
+            # it's posible iopub arrived first:
+            existing = self.db.get_record(msg_id)
+            for key,evalue in existing.iteritems():
+                rvalue = record[key]
+                if evalue and rvalue and evalue != rvalue:
+                    self.log.error("conflicting initial state for record: %s:%s <> %s"%(msg_id, rvalue, evalue))
+                elif evalue and not rvalue:
+                    record[key] = evalue
+            self.db.update_record(msg_id, record)
+        except KeyError:
+            self.db.add_record(msg_id, record)
+        
         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:
             self.log.error("invalid identity prefix: %s"%idents)
             return
             
         client_id, queue_id = idents[:2]
         try:
             msg = self.session.unpack_message(msg, content=False)
         except:
             self.log.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:
             self.log.error("queue::unknown engine %r is sending a reply: "%queue_id)
-            self.log.debug("queue::       %s"%msg[2:])
+            # self.log.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)
         elif msg_id not in self.all_completed:
             # it could be a result from a dead engine that died before delivering the
             # result
             self.log.warn("queue:: unknown msg finished %s"%msg_id)
             return
         # update record anyway, because the unregistration could have been premature
         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
         }
 
         result['result_buffers'] = msg['buffers']
         self.db.update_record(msg_id, result)
         
             
     #--------------------- 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:
             self.log.error("task::client %r sent invalid task message: %s"%(
                     client_id, msg), exc_info=True)
             return
         record = init_record(msg)
 
         record['client_uuid'] = client_id
         record['queue'] = 'task'
         header = msg['header']
         msg_id = header['msg_id']
         self.pending.add(msg_id)
+        try:
+            # it's posible iopub arrived first:
+            existing = self.db.get_record(msg_id)
+            for key,evalue in existing.iteritems():
+                rvalue = record[key]
+                if evalue and rvalue and evalue != rvalue:
+                    self.log.error("conflicting initial state for record: %s:%s <> %s"%(msg_id, rvalue, evalue))
+                elif evalue and not rvalue:
+                    record[key] = evalue
+            self.db.update_record(msg_id, record)
+        except KeyError:
         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:
             self.log.error("task::invalid task result message send to %r: %s"%(
                     client_id, msg), exc_info=True)
             raise
             return
         
         parent = msg['parent_header']
         if not parent:
             # print msg
             self.log.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
             }
 
             result['result_buffers'] = msg['buffers']
             self.db.update_record(msg_id, result)
             
         else:
             self.log.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:
             self.log.error("task::invalid task tracking message", exc_info=True)
             return
         content = msg['content']
         # print (content)
         msg_id = content['msg_id']
         engine_uuid = content['engine_id']
         eid = self.by_ident[engine_uuid]
         
         self.log.info("task::task %s arrived on %s"%(msg_id, eid))
         # if msg_id in self.mia:
         #     self.mia.remove(msg_id)
         # else:
         #     self.log.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)
         
         
     #--------------------- IOPub Traffic ------------------------------
     
     def save_iopub_message(self, topics, msg):
         """save an iopub message into the db"""
         # print (topics)
         try:
             msg = self.session.unpack_message(msg, content=True)
         except:
             self.log.error("iopub::invalid IOPub message", exc_info=True)
             return
         
         parent = msg['parent_header']
         if not parent:
             self.log.error("iopub::invalid IOPub message: %s"%msg)
             return
         msg_id = parent['msg_id']
         msg_type = msg['msg_type']
         content = msg['content']
         
         # ensure msg_id is in db
         try:
             rec = self.db.get_record(msg_id)
-        except:
-            self.log.error("iopub::IOPub message has invalid parent", exc_info=True)
-            return
+        except KeyError:
+            rec = empty_record()
+            rec['msg_id'] = msg_id
+            self.db.add_record(msg_id, rec)
         # stream
         d = {}
         if msg_type == 'stream':
             name = content['name']
             s = rec[name] or ''
             d[name] = s + content['data']
             
         elif msg_type == 'pyerr':
             d['pyerr'] = content
         else:
             d[msg_type] = content['data']
         
         self.db.update_record(msg_id, d)
         
     
             
     #-------------------------------------------------------------------------
     # Registration requests
     #-------------------------------------------------------------------------
         
     def connection_request(self, client_id, msg):
         """Reply with connection addresses for clients."""
         self.log.info("client::client %s connected"%client_id)
         content = dict(status='ok')
         content.update(self.client_info)
         jsonable = {}
         for k,v in self.keytable.iteritems():
+            if v not in self.dead_engines:
             jsonable[str(k)] = v
         content['engines'] = jsonable
         self.session.send(self.query, '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:
             self.log.error("registration::queue not specified", exc_info=True)
             return
         heart = content.get('heartbeat', None)
         """register a new engine, and create the socket(s) necessary"""
         eid = self._next_id
         # print (eid, queue, reg, heart)
         
         self.log.debug("registration::register_engine(%i, %r, %r, %r)"%(eid, queue, reg, heart))
         
         content = dict(id=eid,status='ok')
         content.update(self.engine_info)
         # check if requesting available IDs:
         if queue in self.by_ident:
             try:
                 raise KeyError("queue_id %r in use"%queue)
             except:
                 content = error.wrap_exception()
                 self.log.error("queue_id %r in use"%queue, exc_info=True)
         elif heart in self.hearts: # need to check unique hearts?
             try:
                 raise KeyError("heart_id %r in use"%heart)
             except:
                 self.log.error("heart_id %r in use"%heart, exc_info=True)
                 content = error.wrap_exception()
         else:
             for h, pack in self.incoming_registrations.iteritems():
                 if heart == h:
                     try:
                         raise KeyError("heart_id %r in use"%heart)
                     except:
                         self.log.error("heart_id %r in use"%heart, exc_info=True)
                         content = error.wrap_exception()
                     break
                 elif queue == pack[1]:
                     try:
                         raise KeyError("queue_id %r in use"%queue)
                     except:
                         self.log.error("queue_id %r in use"%queue, exc_info=True)
                         content = error.wrap_exception()
                     break
         
         msg = self.session.send(self.query, "registration_reply", 
                 content=content, 
                 ident=reg)
         
         if content['status'] == 'ok':
             if heart in self.heartmonitor.hearts:
                 # already beating
                 self.incoming_registrations[heart] = (eid,queue,reg[0],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[0],dc)
         else:
             self.log.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:
             self.log.error("registration::bad engine id for unregistration: %s"%ident, exc_info=True)
             return
         self.log.info("registration::unregister_engine(%s)"%eid)
         # print (eid)
-        content=dict(id=eid, queue=self.engines[eid].queue)
-        self.ids.remove(eid)
-        uuid = self.keytable.pop(eid)
-        ec = self.engines.pop(eid)
-        self.hearts.pop(ec.heartbeat)
-        self.by_ident.pop(ec.queue)
-        self.completed.pop(eid)
-        self._handle_stranded_msgs(eid, uuid)
+        uuid = self.keytable[eid]
+        content=dict(id=eid, queue=uuid)
+        self.dead_engines.add(uuid)
+        # self.ids.remove(eid)
+        # uuid = self.keytable.pop(eid)
+        # 
+        # ec = self.engines.pop(eid)
+        # self.hearts.pop(ec.heartbeat)
+        # self.by_ident.pop(ec.queue)
+        # self.completed.pop(eid)
+        handleit = lambda : self._handle_stranded_msgs(eid, uuid)
+        dc = ioloop.DelayedCallback(handleit, self.registration_timeout, self.loop)
+        dc.start()
             ############## TODO: HANDLE IT ################
         
         if self.notifier:
             self.session.send(self.notifier, "unregistration_notification", content=content)
     
     def _handle_stranded_msgs(self, eid, uuid):
         """Handle messages known to be on an engine when the engine unregisters.
         
         It is possible that this will fire prematurely - that is, an engine will
         go down after completing a result, and the client will be notified
         that the result failed and later receive the actual result.
         """
         
-        outstanding = self.queues.pop(eid)
+        outstanding = self.queues[eid]
         
         for msg_id in outstanding:
             self.pending.remove(msg_id)
             self.all_completed.add(msg_id)
             try:
                 raise error.EngineError("Engine %r died while running task %r"%(eid, msg_id))
             except:
                 content = error.wrap_exception()
             # build a fake header:
             header = {}
             header['engine'] = uuid
             header['date'] = datetime.now().strftime(ISO8601)
             rec = dict(result_content=content, result_header=header, result_buffers=[])
             rec['completed'] = header['date']
             rec['engine_uuid'] = uuid
             self.db.update_record(msg_id, rec)
     
     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:
             self.log.error("registration::tried to finish nonexistant registration", exc_info=True)
             return
         self.log.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(id=eid, queue=queue, registration=reg, 
                                     control=control, heartbeat=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)
         self.log.info("engine::Engine Connected: %i"%eid)
     
     def _purge_stalled_registration(self, heart):
         if heart in self.incoming_registrations:
             eid = self.incoming_registrations.pop(heart)[0]
             self.log.info("registration::purging stalled registration: %i"%eid)
         else:
             pass
             
     #-------------------------------------------------------------------------
     # Client Requests
     #-------------------------------------------------------------------------
     
     def shutdown_request(self, client_id, msg):
         """handle shutdown request."""
         self.session.send(self.query, 'shutdown_reply', content={'status': 'ok'}, ident=client_id)
         # also notify other clients of shutdown
         self.session.send(self.notifier, 'shutdown_notice', content={'status': 'ok'})
         dc = ioloop.DelayedCallback(lambda : self._shutdown(), 1000, self.loop)
         dc.start()
     
     def _shutdown(self):
         self.log.info("hub::hub 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 = error.wrap_exception()
             self.session.send(self.query, "hub_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.query, "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 = error.wrap_exception()
             self.session.send(self.query, "hub_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.query, "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 = error.wrap_exception()
                     else:
                         try:
                             raise IndexError("No such msg: %r"%msg_id)
                         except:
                             reply = error.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 = error.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.query, '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]
                     io_dict = {}
                     for key in 'pyin pyout pyerr stdout stderr'.split():
                             io_dict[key] = rec[key]
                     content[msg_id] = { 'result_content': rec['result_content'],
                                         'header': rec['header'],
                                         'result_header' : rec['result_header'],
                                         'io' : io_dict,
                                       }
                     if rec['result_buffers']:
                         buffers.extend(map(str, rec['result_buffers']))
             else:
                 try:
                     raise KeyError('No such message: '+msg_id)
                 except:
                     content = error.wrap_exception()
                 break
         self.session.send(self.query, "result_reply", content=content, 
                                             parent=msg, ident=client_id,
                                             buffers=buffers)
 

IPython/parallel/launcher.py

@@ -1,971 +1,971
 #!/usr/bin/env python
 # encoding: utf-8
 """
 Facilities for launching IPython processes asynchronously.
 """
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2008-2009  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 copy
 import logging
 import os
 import re
 import stat
 
 from signal import SIGINT, SIGTERM
 try:
     from signal import SIGKILL
 except ImportError:
     SIGKILL=SIGTERM
 
 from subprocess import Popen, PIPE, STDOUT
 try:
     from subprocess import check_output
 except ImportError:
     # pre-2.7, define check_output with Popen
     def check_output(*args, **kwargs):
         kwargs.update(dict(stdout=PIPE))
         p = Popen(*args, **kwargs)
         out,err = p.communicate()
         return out
 
 from zmq.eventloop import ioloop
 
 from IPython.external import Itpl
 # from IPython.config.configurable import Configurable
 from IPython.utils.traitlets import Any, Str, Int, List, Unicode, Dict, Instance, CUnicode
 from IPython.utils.path import get_ipython_module_path
 from IPython.utils.process import find_cmd, pycmd2argv, FindCmdError
 
 from .factory import LoggingFactory
 
 # load winhpcjob only on Windows
 try:
     from .winhpcjob import (
         IPControllerTask, IPEngineTask,
         IPControllerJob, IPEngineSetJob
     )
 except ImportError:
     pass
 
 
 #-----------------------------------------------------------------------------
 # Paths to the kernel apps
 #-----------------------------------------------------------------------------
 
 
 ipclusterz_cmd_argv = pycmd2argv(get_ipython_module_path(
     'IPython.parallel.ipclusterapp'
 ))
 
 ipenginez_cmd_argv = pycmd2argv(get_ipython_module_path(
     'IPython.parallel.ipengineapp'
 ))
 
 ipcontrollerz_cmd_argv = pycmd2argv(get_ipython_module_path(
     'IPython.parallel.ipcontrollerapp'
 ))
 
 #-----------------------------------------------------------------------------
 # Base launchers and errors
 #-----------------------------------------------------------------------------
 
 
 class LauncherError(Exception):
     pass
 
 
 class ProcessStateError(LauncherError):
     pass
 
 
 class UnknownStatus(LauncherError):
     pass
 
 
 class BaseLauncher(LoggingFactory):
     """An asbtraction for starting, stopping and signaling a process."""
 
     # In all of the launchers, the work_dir is where child processes will be
     # run. This will usually be the cluster_dir, but may not be. any work_dir
     # passed into the __init__ method will override the config value.
     # This should not be used to set the work_dir for the actual engine
     # and controller. Instead, use their own config files or the
     # controller_args, engine_args attributes of the launchers to add
     # the --work-dir option.
     work_dir = Unicode(u'.')
     loop = Instance('zmq.eventloop.ioloop.IOLoop')
     
     start_data = Any()
     stop_data = Any()
     
     def _loop_default(self):
         return ioloop.IOLoop.instance()
 
     def __init__(self, work_dir=u'.', config=None, **kwargs):
         super(BaseLauncher, self).__init__(work_dir=work_dir, config=config, **kwargs)
         self.state = 'before' # can be before, running, after
         self.stop_callbacks = []
         self.start_data = None
         self.stop_data = None
 
     @property
     def args(self):
         """A list of cmd and args that will be used to start the process.
 
         This is what is passed to :func:`spawnProcess` and the first element
         will be the process name.
         """
         return self.find_args()
 
     def find_args(self):
         """The ``.args`` property calls this to find the args list.
 
         Subcommand should implement this to construct the cmd and args.
         """
         raise NotImplementedError('find_args must be implemented in a subclass')
 
     @property
     def arg_str(self):
         """The string form of the program arguments."""
         return ' '.join(self.args)
 
     @property
     def running(self):
         """Am I running."""
         if self.state == 'running':
             return True
         else:
             return False
 
     def start(self):
         """Start the process.
 
         This must return a deferred that fires with information about the
         process starting (like a pid, job id, etc.).
         """
         raise NotImplementedError('start must be implemented in a subclass')
 
     def stop(self):
         """Stop the process and notify observers of stopping.
 
         This must return a deferred that fires with information about the
         processing stopping, like errors that occur while the process is
         attempting to be shut down. This deferred won't fire when the process
         actually stops. To observe the actual process stopping, see
         :func:`observe_stop`.
         """
         raise NotImplementedError('stop must be implemented in a subclass')
 
     def on_stop(self, f):
         """Get a deferred that will fire when the process stops.
 
         The deferred will fire with data that contains information about
         the exit status of the process.
         """
         if self.state=='after':
             return f(self.stop_data)
         else:
             self.stop_callbacks.append(f)
 
     def notify_start(self, data):
         """Call this to trigger startup actions.
 
         This logs the process startup and sets the state to 'running'.  It is
         a pass-through so it can be used as a callback.
         """
 
         self.log.info('Process %r started: %r' % (self.args[0], data))
         self.start_data = data
         self.state = 'running'
         return data
 
     def notify_stop(self, data):
         """Call this to trigger process stop actions.
 
         This logs the process stopping and sets the state to 'after'. Call
         this to trigger all the deferreds from :func:`observe_stop`."""
 
         self.log.info('Process %r stopped: %r' % (self.args[0], data))
         self.stop_data = data
         self.state = 'after'
         for i in range(len(self.stop_callbacks)):
             d = self.stop_callbacks.pop()
             d(data)
         return data
 
     def signal(self, sig):
         """Signal the process.
 
         Return a semi-meaningless deferred after signaling the process.
 
         Parameters
         ----------
         sig : str or int
             'KILL', 'INT', etc., or any signal number
         """
         raise NotImplementedError('signal must be implemented in a subclass')
 
 
 #-----------------------------------------------------------------------------
 # Local process launchers
 #-----------------------------------------------------------------------------
 
 
 class LocalProcessLauncher(BaseLauncher):
     """Start and stop an external process in an asynchronous manner.
 
     This will launch the external process with a working directory of
     ``self.work_dir``.
     """
 
     # This is used to to construct self.args, which is passed to 
     # spawnProcess.
     cmd_and_args = List([])
     poll_frequency = Int(100) # in ms
 
     def __init__(self, work_dir=u'.', config=None, **kwargs):
         super(LocalProcessLauncher, self).__init__(
             work_dir=work_dir, config=config, **kwargs
         )
         self.process = None
         self.start_deferred = None
         self.poller = None
 
     def find_args(self):
         return self.cmd_and_args
 
     def start(self):
         if self.state == 'before':
             self.process = Popen(self.args,
                 stdout=PIPE,stderr=PIPE,stdin=PIPE,
                 env=os.environ,
                 cwd=self.work_dir
             )
             
             self.loop.add_handler(self.process.stdout.fileno(), self.handle_stdout, self.loop.READ)
             self.loop.add_handler(self.process.stderr.fileno(), self.handle_stderr, self.loop.READ)
             self.poller = ioloop.PeriodicCallback(self.poll, self.poll_frequency, self.loop)
             self.poller.start()
             self.notify_start(self.process.pid)
         else:
             s = 'The process was already started and has state: %r' % self.state
             raise ProcessStateError(s)
 
     def stop(self):
         return self.interrupt_then_kill()
 
     def signal(self, sig):
         if self.state == 'running':
             self.process.send_signal(sig)
 
     def interrupt_then_kill(self, delay=2.0):
         """Send INT, wait a delay and then send KILL."""
         self.signal(SIGINT)
         self.killer  = ioloop.DelayedCallback(lambda : self.signal(SIGKILL), delay*1000, self.loop)
         self.killer.start()
 
     # callbacks, etc:
     
     def handle_stdout(self, fd, events):
         line = self.process.stdout.readline()
         # a stopped process will be readable but return empty strings
         if line:
             self.log.info(line[:-1])
         else:
             self.poll()
     
     def handle_stderr(self, fd, events):
         line = self.process.stderr.readline()
         # a stopped process will be readable but return empty strings
         if line:
             self.log.error(line[:-1])
         else:
             self.poll()
     
     def poll(self):
         status = self.process.poll()
         if status is not None:
             self.poller.stop()
             self.loop.remove_handler(self.process.stdout.fileno())
             self.loop.remove_handler(self.process.stderr.fileno())
             self.notify_stop(dict(exit_code=status, pid=self.process.pid))
         return status
 
 class LocalControllerLauncher(LocalProcessLauncher):
     """Launch a controller as a regular external process."""
 
     controller_cmd = List(ipcontrollerz_cmd_argv, config=True)
     # Command line arguments to ipcontroller.
     controller_args = List(['--log-to-file','--log-level', str(logging.INFO)], config=True)
 
     def find_args(self):
         return self.controller_cmd + self.controller_args
 
     def start(self, cluster_dir):
         """Start the controller by cluster_dir."""
         self.controller_args.extend(['--cluster-dir', cluster_dir])
         self.cluster_dir = unicode(cluster_dir)
         self.log.info("Starting LocalControllerLauncher: %r" % self.args)
         return super(LocalControllerLauncher, self).start()
 
 
 class LocalEngineLauncher(LocalProcessLauncher):
     """Launch a single engine as a regular externall process."""
 
     engine_cmd = List(ipenginez_cmd_argv, config=True)
     # Command line arguments for ipengine.
     engine_args = List(
         ['--log-to-file','--log-level', str(logging.INFO)], config=True
     )
 
     def find_args(self):
         return self.engine_cmd + self.engine_args
 
     def start(self, cluster_dir):
         """Start the engine by cluster_dir."""
         self.engine_args.extend(['--cluster-dir', cluster_dir])
         self.cluster_dir = unicode(cluster_dir)
         return super(LocalEngineLauncher, self).start()
 
 
 class LocalEngineSetLauncher(BaseLauncher):
     """Launch a set of engines as regular external processes."""
 
     # Command line arguments for ipengine.
     engine_args = List(
         ['--log-to-file','--log-level', str(logging.INFO)], config=True
     )
     # launcher class
     launcher_class = LocalEngineLauncher
     
     launchers = Dict()
     stop_data = Dict()
     
     def __init__(self, work_dir=u'.', config=None, **kwargs):
         super(LocalEngineSetLauncher, self).__init__(
             work_dir=work_dir, config=config, **kwargs
         )
         self.stop_data = {}
 
     def start(self, n, cluster_dir):
         """Start n engines by profile or cluster_dir."""
         self.cluster_dir = unicode(cluster_dir)
         dlist = []
         for i in range(n):
             el = self.launcher_class(work_dir=self.work_dir, config=self.config, logname=self.log.name)
             # Copy the engine args over to each engine launcher.
             el.engine_args = copy.deepcopy(self.engine_args)
             el.on_stop(self._notice_engine_stopped)
             d = el.start(cluster_dir)
             if i==0:
                 self.log.info("Starting LocalEngineSetLauncher: %r" % el.args)
             self.launchers[i] = el
             dlist.append(d)
         self.notify_start(dlist)
         # The consumeErrors here could be dangerous
         # dfinal = gatherBoth(dlist, consumeErrors=True)
         # dfinal.addCallback(self.notify_start)
         return dlist
 
     def find_args(self):
         return ['engine set']
 
     def signal(self, sig):
         dlist = []
         for el in self.launchers.itervalues():
             d = el.signal(sig)
             dlist.append(d)
         # dfinal = gatherBoth(dlist, consumeErrors=True)
         return dlist
 
     def interrupt_then_kill(self, delay=1.0):
         dlist = []
         for el in self.launchers.itervalues():
             d = el.interrupt_then_kill(delay)
             dlist.append(d)
         # dfinal = gatherBoth(dlist, consumeErrors=True)
         return dlist
 
     def stop(self):
         return self.interrupt_then_kill()
     
     def _notice_engine_stopped(self, data):
         pid = data['pid']
         for idx,el in self.launchers.iteritems():
             if el.process.pid == pid:
                 break
         self.launchers.pop(idx)
         self.stop_data[idx] = data
         if not self.launchers:
             self.notify_stop(self.stop_data)
 
 
 #-----------------------------------------------------------------------------
 # MPIExec launchers
 #-----------------------------------------------------------------------------
 
 
 class MPIExecLauncher(LocalProcessLauncher):
     """Launch an external process using mpiexec."""
 
     # The mpiexec command to use in starting the process.
     mpi_cmd = List(['mpiexec'], config=True)
     # The command line arguments to pass to mpiexec.
     mpi_args = List([], config=True)
     # The program to start using mpiexec.
     program = List(['date'], config=True)
     # The command line argument to the program.
     program_args = List([], config=True)
     # The number of instances of the program to start.
     n = Int(1, config=True)
 
     def find_args(self):
         """Build self.args using all the fields."""
         return self.mpi_cmd + ['-n', str(self.n)] + self.mpi_args + \
                self.program + self.program_args
 
     def start(self, n):
         """Start n instances of the program using mpiexec."""
         self.n = n
         return super(MPIExecLauncher, self).start()
 
 
 class MPIExecControllerLauncher(MPIExecLauncher):
     """Launch a controller using mpiexec."""
 
     controller_cmd = List(ipcontrollerz_cmd_argv, config=True)
     # Command line arguments to ipcontroller.
     controller_args = List(['--log-to-file','--log-level', str(logging.INFO)], config=True)
     n = Int(1, config=False)
 
     def start(self, cluster_dir):
         """Start the controller by cluster_dir."""
         self.controller_args.extend(['--cluster-dir', cluster_dir])
         self.cluster_dir = unicode(cluster_dir)
         self.log.info("Starting MPIExecControllerLauncher: %r" % self.args)
         return super(MPIExecControllerLauncher, self).start(1)
 
     def find_args(self):
         return self.mpi_cmd + ['-n', self.n] + self.mpi_args + \
                self.controller_cmd + self.controller_args
 
 
 class MPIExecEngineSetLauncher(MPIExecLauncher):
 
     program = List(ipenginez_cmd_argv, config=True)
     # Command line arguments for ipengine.
     program_args = List(
         ['--log-to-file','--log-level', str(logging.INFO)], config=True
     )
     n = Int(1, config=True)
 
     def start(self, n, cluster_dir):
         """Start n engines by profile or cluster_dir."""
         self.program_args.extend(['--cluster-dir', cluster_dir])
         self.cluster_dir = unicode(cluster_dir)
         self.n = n
         self.log.info('Starting MPIExecEngineSetLauncher: %r' % self.args)
         return super(MPIExecEngineSetLauncher, self).start(n)
 
 #-----------------------------------------------------------------------------
 # SSH launchers
 #-----------------------------------------------------------------------------
 
 # TODO: Get SSH Launcher working again.
 
 class SSHLauncher(LocalProcessLauncher):
     """A minimal launcher for ssh.
 
     To be useful this will probably have to be extended to use the ``sshx``
     idea for environment variables.  There could be other things this needs
     as well.
     """
 
     ssh_cmd = List(['ssh'], config=True)
     ssh_args = List(['-tt'], config=True)
     program = List(['date'], config=True)
     program_args = List([], config=True)
     hostname = CUnicode('', config=True)
     user = CUnicode('', config=True)
     location = CUnicode('')
 
     def _hostname_changed(self, name, old, new):
         if self.user:
             self.location = u'%s@%s' % (self.user, new)
         else:
             self.location = new
 
     def _user_changed(self, name, old, new):
         self.location = u'%s@%s' % (new, self.hostname)
 
     def find_args(self):
         return self.ssh_cmd + self.ssh_args + [self.location] + \
                self.program + self.program_args
 
     def start(self, cluster_dir, hostname=None, user=None):
         self.cluster_dir = unicode(cluster_dir)
         if hostname is not None:
             self.hostname = hostname
         if user is not None:
             self.user = user
         
         return super(SSHLauncher, self).start()
     
     def signal(self, sig):
         if self.state == 'running':
             # send escaped ssh connection-closer
             self.process.stdin.write('~.')
             self.process.stdin.flush()
         
 
 
 class SSHControllerLauncher(SSHLauncher):
 
     program = List(ipcontrollerz_cmd_argv, config=True)
     # Command line arguments to ipcontroller.
     program_args = List(['-r', '--log-to-file','--log-level', str(logging.INFO)], config=True)
 
 
 class SSHEngineLauncher(SSHLauncher):
     program = List(ipenginez_cmd_argv, config=True)
     # Command line arguments for ipengine.
     program_args = List(
         ['--log-to-file','--log-level', str(logging.INFO)], config=True
     )
     
 class SSHEngineSetLauncher(LocalEngineSetLauncher):
     launcher_class = SSHEngineLauncher
     engines = Dict(config=True)
     
     def start(self, n, cluster_dir):
         """Start engines by profile or cluster_dir.
         `n` is ignored, and the `engines` config property is used instead.
         """
         
         self.cluster_dir = unicode(cluster_dir)
         dlist = []
         for host, n in self.engines.iteritems():
             if isinstance(n, (tuple, list)):
                 n, args = n
             else:
                 args = copy.deepcopy(self.engine_args)
             
             if '@' in host:
                 user,host = host.split('@',1)
             else:
                 user=None
             for i in range(n):
                 el = self.launcher_class(work_dir=self.work_dir, config=self.config, logname=self.log.name)
                 
                 # Copy the engine args over to each engine launcher.
                 i
                 el.program_args = args
                 el.on_stop(self._notice_engine_stopped)
                 d = el.start(cluster_dir, user=user, hostname=host)
                 if i==0:
                     self.log.info("Starting SSHEngineSetLauncher: %r" % el.args)
                 self.launchers[host+str(i)] = el
                 dlist.append(d)
         self.notify_start(dlist)
         return dlist
     
 
 
 #-----------------------------------------------------------------------------
 # Windows HPC Server 2008 scheduler launchers
 #-----------------------------------------------------------------------------
 
 
 # This is only used on Windows.
 def find_job_cmd():
     if os.name=='nt':
         try:
             return find_cmd('job')
         except FindCmdError:
             return 'job'
     else:
         return 'job'
 
 
 class WindowsHPCLauncher(BaseLauncher):
 
     # A regular expression used to get the job id from the output of the 
     # submit_command.
     job_id_regexp = Str(r'\d+', config=True)
     # The filename of the instantiated job script.
     job_file_name = CUnicode(u'ipython_job.xml', config=True)
     # The full path to the instantiated job script. This gets made dynamically
     # by combining the work_dir with the job_file_name.
     job_file = CUnicode(u'')
     # The hostname of the scheduler to submit the job to
     scheduler = CUnicode('', config=True)
     job_cmd = CUnicode(find_job_cmd(), config=True)
 
     def __init__(self, work_dir=u'.', config=None, **kwargs):
         super(WindowsHPCLauncher, self).__init__(
             work_dir=work_dir, config=config, **kwargs
         )
 
     @property
     def job_file(self):
         return os.path.join(self.work_dir, self.job_file_name)
 
     def write_job_file(self, n):
         raise NotImplementedError("Implement write_job_file in a subclass.")
 
     def find_args(self):
         return [u'job.exe']
         
     def parse_job_id(self, output):
         """Take the output of the submit command and return the job id."""
         m = re.search(self.job_id_regexp, output)
         if m is not None:
             job_id = m.group()
         else:
             raise LauncherError("Job id couldn't be determined: %s" % output)
         self.job_id = job_id
         self.log.info('Job started with job id: %r' % job_id)
         return job_id
 
     def start(self, n):
         """Start n copies of the process using the Win HPC job scheduler."""
         self.write_job_file(n)
         args = [
             'submit',
             '/jobfile:%s' % self.job_file,
             '/scheduler:%s' % self.scheduler
         ]
         self.log.info("Starting Win HPC Job: %s" % (self.job_cmd + ' ' + ' '.join(args),))
         # Twisted will raise DeprecationWarnings if we try to pass unicode to this
         output = check_output([self.job_cmd]+args,
             env=os.environ,
             cwd=self.work_dir,
             stderr=STDOUT
         )
         job_id = self.parse_job_id(output)
         self.notify_start(job_id)
         return job_id
 
     def stop(self):
         args = [
             'cancel',
             self.job_id,
             '/scheduler:%s' % self.scheduler
         ]
         self.log.info("Stopping Win HPC Job: %s" % (self.job_cmd + ' ' + ' '.join(args),))
         try:
             output = check_output([self.job_cmd]+args,
                 env=os.environ,
                 cwd=self.work_dir,
                 stderr=STDOUT
             )
         except:
             output = 'The job already appears to be stoppped: %r' % self.job_id
         self.notify_stop(dict(job_id=self.job_id, output=output))  # Pass the output of the kill cmd
         return output
 
 
 class WindowsHPCControllerLauncher(WindowsHPCLauncher):
 
     job_file_name = CUnicode(u'ipcontroller_job.xml', config=True)
     extra_args = List([], config=False)
 
     def write_job_file(self, n):
         job = IPControllerJob(config=self.config)
 
         t = IPControllerTask(config=self.config)
         # The tasks work directory is *not* the actual work directory of 
         # the controller. It is used as the base path for the stdout/stderr
         # files that the scheduler redirects to.
         t.work_directory = self.cluster_dir
         # Add the --cluster-dir and from self.start().
         t.controller_args.extend(self.extra_args)
         job.add_task(t)
 
         self.log.info("Writing job description file: %s" % self.job_file)
         job.write(self.job_file)
 
     @property
     def job_file(self):
         return os.path.join(self.cluster_dir, self.job_file_name)
 
     def start(self, cluster_dir):
         """Start the controller by cluster_dir."""
         self.extra_args = ['--cluster-dir', cluster_dir]
         self.cluster_dir = unicode(cluster_dir)
         return super(WindowsHPCControllerLauncher, self).start(1)
 
 
 class WindowsHPCEngineSetLauncher(WindowsHPCLauncher):
 
     job_file_name = CUnicode(u'ipengineset_job.xml', config=True)
     extra_args = List([], config=False)
 
     def write_job_file(self, n):
         job = IPEngineSetJob(config=self.config)
 
         for i in range(n):
             t = IPEngineTask(config=self.config)
             # The tasks work directory is *not* the actual work directory of 
             # the engine. It is used as the base path for the stdout/stderr
             # files that the scheduler redirects to.
             t.work_directory = self.cluster_dir
             # Add the --cluster-dir and from self.start().
             t.engine_args.extend(self.extra_args)
             job.add_task(t)
 
         self.log.info("Writing job description file: %s" % self.job_file)
         job.write(self.job_file)
 
     @property
     def job_file(self):
         return os.path.join(self.cluster_dir, self.job_file_name)
 
     def start(self, n, cluster_dir):
         """Start the controller by cluster_dir."""
         self.extra_args = ['--cluster-dir', cluster_dir]
         self.cluster_dir = unicode(cluster_dir)
         return super(WindowsHPCEngineSetLauncher, self).start(n)
 
 
 #-----------------------------------------------------------------------------
 # Batch (PBS) system launchers
 #-----------------------------------------------------------------------------
 
 class BatchSystemLauncher(BaseLauncher):
     """Launch an external process using a batch system.
 
     This class is designed to work with UNIX batch systems like PBS, LSF,
     GridEngine, etc.  The overall model is that there are different commands
     like qsub, qdel, etc. that handle the starting and stopping of the process.
 
     This class also has the notion of a batch script. The ``batch_template``
     attribute can be set to a string that is a template for the batch script.
     This template is instantiated using Itpl. Thus the template can use
     ${n} fot the number of instances. Subclasses can add additional variables
     to the template dict.
     """
 
     # Subclasses must fill these in.  See PBSEngineSet
     # The name of the command line program used to submit jobs.
     submit_command = List([''], config=True)
     # The name of the command line program used to delete jobs.
     delete_command = List([''], config=True)
     # A regular expression used to get the job id from the output of the 
     # submit_command.
     job_id_regexp = CUnicode('', config=True)
     # The string that is the batch script template itself.
     batch_template = CUnicode('', config=True)
     # The file that contains the batch template
     batch_template_file = CUnicode(u'', config=True)
     # The filename of the instantiated batch script.
     batch_file_name = CUnicode(u'batch_script', config=True)
     # The PBS Queue
     queue = CUnicode(u'', config=True)
     
     # not configurable, override in subclasses
     # PBS Job Array regex
     job_array_regexp = CUnicode('')
     job_array_template = CUnicode('')
     # PBS Queue regex
     queue_regexp = CUnicode('')
     queue_template = CUnicode('')
     # The default batch template, override in subclasses
     default_template = CUnicode('')
     # The full path to the instantiated batch script.
     batch_file = CUnicode(u'')
     # the format dict used with batch_template:
     context = Dict()
 
     
     def find_args(self):
         return self.submit_command + [self.batch_file]
     
     def __init__(self, work_dir=u'.', config=None, **kwargs):
         super(BatchSystemLauncher, self).__init__(
             work_dir=work_dir, config=config, **kwargs
         )
         self.batch_file = os.path.join(self.work_dir, self.batch_file_name)
 
     def parse_job_id(self, output):
         """Take the output of the submit command and return the job id."""
         m = re.search(self.job_id_regexp, output)
         if m is not None:
             job_id = m.group()
         else:
             raise LauncherError("Job id couldn't be determined: %s" % output)
         self.job_id = job_id
         self.log.info('Job submitted with job id: %r' % job_id)
         return job_id
 
     def write_batch_script(self, n):
         """Instantiate and write the batch script to the work_dir."""
         self.context['n'] = n
         self.context['queue'] = self.queue
         print self.context
         # first priority is batch_template if set
         if self.batch_template_file and not self.batch_template:
             # second priority is batch_template_file
             with open(self.batch_template_file) as f:
                 self.batch_template = f.read()
         if not self.batch_template:
             # third (last) priority is default_template
             self.batch_template = self.default_template
         
         regex = re.compile(self.job_array_regexp)
         # print regex.search(self.batch_template)
         if not regex.search(self.batch_template):
             self.log.info("adding job array settings to batch script")
             firstline, rest = self.batch_template.split('\n',1)
             self.batch_template = u'\n'.join([firstline, self.job_array_template, rest])
         
         regex = re.compile(self.queue_regexp)
         # print regex.search(self.batch_template)
         if self.queue and not regex.search(self.batch_template):
             self.log.info("adding PBS queue settings to batch script")
             firstline, rest = self.batch_template.split('\n',1)
             self.batch_template = u'\n'.join([firstline, self.queue_template, rest])
             
         script_as_string = Itpl.itplns(self.batch_template, self.context)
         self.log.info('Writing instantiated batch script: %s' % self.batch_file)
     
         with open(self.batch_file, 'w') as f:
             f.write(script_as_string)
         os.chmod(self.batch_file, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR)
 
     def start(self, n, cluster_dir):
         """Start n copies of the process using a batch system."""
         # Here we save profile and cluster_dir in the context so they
         # can be used in the batch script template as ${profile} and
         # ${cluster_dir}
         self.context['cluster_dir'] = cluster_dir
         self.cluster_dir = unicode(cluster_dir)
         self.write_batch_script(n)
         output = check_output(self.args, env=os.environ)
         
         job_id = self.parse_job_id(output)
         self.notify_start(job_id)
         return job_id
 
     def stop(self):
         output = check_output(self.delete_command+[self.job_id], env=os.environ)
         self.notify_stop(dict(job_id=self.job_id, output=output)) # Pass the output of the kill cmd
         return output
 
 
 class PBSLauncher(BatchSystemLauncher):
     """A BatchSystemLauncher subclass for PBS."""
 
     submit_command = List(['qsub'], config=True)
     delete_command = List(['qdel'], config=True)
     job_id_regexp = CUnicode(r'\d+', config=True)
     
     batch_file = CUnicode(u'')
     job_array_regexp = CUnicode('#PBS\W+-t\W+[\w\d\-\$]+')
     job_array_template = CUnicode('#PBS -t 1-$n')
     queue_regexp = CUnicode('#PBS\W+-q\W+\$?\w+')
     queue_template = CUnicode('#PBS -q $queue')
 
 
 class PBSControllerLauncher(PBSLauncher):
     """Launch a controller using PBS."""
 
     batch_file_name = CUnicode(u'pbs_controller', config=True)
     default_template= CUnicode("""#!/bin/sh
 #PBS -V
 #PBS -N ipcontrollerz
 %s --log-to-file --cluster-dir $cluster_dir
 """%(' '.join(ipcontrollerz_cmd_argv)))
 
     def start(self, cluster_dir):
         """Start the controller by profile or cluster_dir."""
         self.log.info("Starting PBSControllerLauncher: %r" % self.args)
         return super(PBSControllerLauncher, self).start(1, cluster_dir)
 
 
 class PBSEngineSetLauncher(PBSLauncher):
     """Launch Engines using PBS"""
     batch_file_name = CUnicode(u'pbs_engines', config=True)
     default_template= CUnicode(u"""#!/bin/sh
 #PBS -V
 #PBS -N ipenginez
 %s --cluster-dir $cluster_dir
 """%(' '.join(ipenginez_cmd_argv)))
 
     def start(self, n, cluster_dir):
         """Start n engines by profile or cluster_dir."""
-        self.log.info('Starting %n engines with PBSEngineSetLauncher: %r' % (n, self.args))
+        self.log.info('Starting %i engines with PBSEngineSetLauncher: %r' % (n, self.args))
         return super(PBSEngineSetLauncher, self).start(n, cluster_dir)
 
 #SGE is very similar to PBS
 
 class SGELauncher(PBSLauncher):
     """Sun GridEngine is a PBS clone with slightly different syntax"""
     job_array_regexp = CUnicode('#$$\W+-t\W+[\w\d\-\$]+')
     job_array_template = CUnicode('#$$ -t 1-$n')
     queue_regexp = CUnicode('#$$\W+-q\W+\$?\w+')
     queue_template = CUnicode('#$$ -q $queue')
 
 class SGEControllerLauncher(SGELauncher):
     """Launch a controller using SGE."""
 
     batch_file_name = CUnicode(u'sge_controller', config=True)
     default_template= CUnicode(u"""#$$ -V
 #$$ -S /bin/sh
 #$$ -N ipcontrollerz
 %s --log-to-file --cluster-dir $cluster_dir
 """%(' '.join(ipcontrollerz_cmd_argv)))
 
     def start(self, cluster_dir):
         """Start the controller by profile or cluster_dir."""
         self.log.info("Starting PBSControllerLauncher: %r" % self.args)
         return super(PBSControllerLauncher, self).start(1, cluster_dir)
 
 class SGEEngineSetLauncher(SGELauncher):
     """Launch Engines with SGE"""
     batch_file_name = CUnicode(u'sge_engines', config=True)
     default_template = CUnicode("""#$$ -V
 #$$ -S /bin/sh
 #$$ -N ipenginez
 %s --cluster-dir $cluster_dir
 """%(' '.join(ipenginez_cmd_argv)))
 
     def start(self, n, cluster_dir):
         """Start n engines by profile or cluster_dir."""
-        self.log.info('Starting %n engines with SGEEngineSetLauncher: %r' % (n, self.args))
+        self.log.info('Starting %i engines with SGEEngineSetLauncher: %r' % (n, self.args))
         return super(SGEEngineSetLauncher, self).start(n, cluster_dir)
 
 
 #-----------------------------------------------------------------------------
 # A launcher for ipcluster itself!
 #-----------------------------------------------------------------------------
 
 
 class IPClusterLauncher(LocalProcessLauncher):
     """Launch the ipcluster program in an external process."""
 
     ipcluster_cmd = List(ipclusterz_cmd_argv, config=True)
     # Command line arguments to pass to ipcluster.
     ipcluster_args = List(
         ['--clean-logs', '--log-to-file', '--log-level', str(logging.INFO)], config=True)
     ipcluster_subcommand = Str('start')
     ipcluster_n = Int(2)
 
     def find_args(self):
         return self.ipcluster_cmd + [self.ipcluster_subcommand] + \
             ['-n', repr(self.ipcluster_n)] + self.ipcluster_args
 
     def start(self):
         self.log.info("Starting ipcluster: %r" % self.args)
         return super(IPClusterLauncher, self).start()
 

IPython/parallel/sqlitedb.py

@@ -1,273 +1,284
 """A TaskRecord backend using sqlite3"""
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 import json
 import os
 import cPickle as pickle
 from datetime import datetime
 
 import sqlite3
 
+from zmq.eventloop import ioloop
+
 from IPython.utils.traitlets import CUnicode, CStr, Instance, List
 from .dictdb import BaseDB
 from .util import ISO8601
 
 #-----------------------------------------------------------------------------
 # SQLite operators, adapters, and converters
 #-----------------------------------------------------------------------------
 
 operators = {
  '$lt' : lambda a,b: "%s < ?",
  '$gt' : ">",
  # null is handled weird with ==,!=
  '$eq' : "IS",
  '$ne' : "IS NOT",
  '$lte': "<=",
  '$gte': ">=",
  '$in' : ('IS', ' OR '),
  '$nin': ('IS NOT', ' AND '),
  # '$all': None,
  # '$mod': None,
  # '$exists' : None
 }
 
 def _adapt_datetime(dt):
     return dt.strftime(ISO8601)
 
 def _convert_datetime(ds):
     if ds is None:
         return ds
     else:
         return datetime.strptime(ds, ISO8601)
 
 def _adapt_dict(d):
     return json.dumps(d)
 
 def _convert_dict(ds):
     if ds is None:
         return ds
     else:
         return json.loads(ds)
 
 def _adapt_bufs(bufs):
     # this is *horrible*
     # copy buffers into single list and pickle it:
     if bufs and isinstance(bufs[0], (bytes, buffer)):
         return sqlite3.Binary(pickle.dumps(map(bytes, bufs),-1))
     elif bufs:
         return bufs
     else:
         return None
 
 def _convert_bufs(bs):
     if bs is None:
         return []
     else:
         return pickle.loads(bytes(bs))
 
 #-----------------------------------------------------------------------------
 # SQLiteDB class
 #-----------------------------------------------------------------------------
 
 class SQLiteDB(BaseDB):
     """SQLite3 TaskRecord backend."""
     
     filename = CUnicode('tasks.db', config=True)
     location = CUnicode('', config=True)
     table = CUnicode("", config=True)
     
     _db = Instance('sqlite3.Connection')
     _keys = List(['msg_id' ,
             'header' ,
             'content',
             'buffers',
             'submitted',
             'client_uuid' ,
             'engine_uuid' ,
             'started',
             'completed',
             'resubmitted',
             'result_header' ,
             'result_content' ,
             'result_buffers' ,
             'queue' ,
             'pyin' ,
             'pyout',
             'pyerr',
             'stdout',
             'stderr',
         ])
     
     def __init__(self, **kwargs):
         super(SQLiteDB, self).__init__(**kwargs)
         if not self.table:
             # use session, and prefix _, since starting with # is illegal
             self.table = '_'+self.session.replace('-','_')
         if not self.location:
             if hasattr(self.config.Global, 'cluster_dir'):
                 self.location = self.config.Global.cluster_dir
             else:
                 self.location = '.'
         self._init_db()
     
+        # register db commit as 2s periodic callback
+        # to prevent clogging pipes
+        # assumes we are being run in a zmq ioloop app
+        loop = ioloop.IOLoop.instance()
+        pc = ioloop.PeriodicCallback(self._db.commit, 2000, loop)
+        pc.start()
+    
     def _defaults(self):
         """create an empty record"""
         d = {}
         for key in self._keys:
             d[key] = None
         return d
     
     def _init_db(self):
         """Connect to the database and get new session number."""
         # register adapters
         sqlite3.register_adapter(datetime, _adapt_datetime)
         sqlite3.register_converter('datetime', _convert_datetime)
         sqlite3.register_adapter(dict, _adapt_dict)
         sqlite3.register_converter('dict', _convert_dict)
         sqlite3.register_adapter(list, _adapt_bufs)
         sqlite3.register_converter('bufs', _convert_bufs)
         # connect to the db
         dbfile = os.path.join(self.location, self.filename)
-        self._db = sqlite3.connect(dbfile, detect_types=sqlite3.PARSE_DECLTYPES, cached_statements=16)
+        self._db = sqlite3.connect(dbfile, detect_types=sqlite3.PARSE_DECLTYPES, 
+            # isolation_level = None)#,
+             cached_statements=64)
         # print dir(self._db)
         
         self._db.execute("""CREATE TABLE IF NOT EXISTS %s 
                 (msg_id text PRIMARY KEY,
                 header dict text,
                 content dict text,
                 buffers bufs blob,
                 submitted datetime text,
                 client_uuid text,
                 engine_uuid text,
                 started datetime text,
                 completed datetime text,
                 resubmitted datetime text,
                 result_header dict text,
                 result_content dict text,
                 result_buffers bufs blob,
                 queue text,
                 pyin text,
                 pyout text,
                 pyerr text,
                 stdout text,
                 stderr text)
                 """%self.table)
         # self._db.execute("""CREATE TABLE IF NOT EXISTS %s_buffers 
         #         (msg_id text, result integer, buffer blob)
         #         """%self.table)
         self._db.commit()
     
     def _dict_to_list(self, d):
         """turn a mongodb-style record dict into a list."""
         
         return [ d[key] for key in self._keys ]
     
     def _list_to_dict(self, line):
         """Inverse of dict_to_list"""
         d = self._defaults()
         for key,value in zip(self._keys, line):
             d[key] = value
         
         return d
     
     def _render_expression(self, check):
         """Turn a mongodb-style search dict into an SQL query."""
         expressions = []
         args = []
         
         skeys = set(check.keys())
         skeys.difference_update(set(self._keys))
         skeys.difference_update(set(['buffers', 'result_buffers']))
         if skeys:
             raise KeyError("Illegal testing key(s): %s"%skeys)
         
         for name,sub_check in check.iteritems():
             if isinstance(sub_check, dict):
                 for test,value in sub_check.iteritems():
                     try:
                         op = operators[test]
                     except KeyError:
                         raise KeyError("Unsupported operator: %r"%test)
                     if isinstance(op, tuple):
                         op, join = op
                     expr = "%s %s ?"%(name, op)
                     if isinstance(value, (tuple,list)):
                         expr = '( %s )'%( join.join([expr]*len(value)) )
                         args.extend(value)
                     else:
                         args.append(value)
                     expressions.append(expr)
             else:
                 # it's an equality check
                 expressions.append("%s IS ?"%name)
                 args.append(sub_check)
         
         expr = " AND ".join(expressions)
         return expr, args
     
     def add_record(self, msg_id, rec):
         """Add a new Task Record, by msg_id."""
         d = self._defaults()
         d.update(rec)
         d['msg_id'] = msg_id
         line = self._dict_to_list(d)
         tups = '(%s)'%(','.join(['?']*len(line)))
         self._db.execute("INSERT INTO %s VALUES %s"%(self.table, tups), line)
-        self._db.commit()
+        # self._db.commit()
     
     def get_record(self, msg_id):
         """Get a specific Task Record, by msg_id."""
         cursor = self._db.execute("""SELECT * FROM %s WHERE msg_id==?"""%self.table, (msg_id,))
         line = cursor.fetchone()
         if line is None:
             raise KeyError("No such msg: %r"%msg_id)
         return self._list_to_dict(line)
     
     def update_record(self, msg_id, rec):
         """Update the data in an existing record."""
         query = "UPDATE %s SET "%self.table
         sets = []
         keys = sorted(rec.keys())
         values = []
         for key in keys:
             sets.append('%s = ?'%key)
             values.append(rec[key])
         query += ', '.join(sets)
         query += ' WHERE msg_id == %r'%msg_id
         self._db.execute(query, values)
-        self._db.commit()
+        # self._db.commit()
     
     def drop_record(self, msg_id):
         """Remove a record from the DB."""
         self._db.execute("""DELETE FROM %s WHERE mgs_id==?"""%self.table, (msg_id,))
-        self._db.commit()
+        # self._db.commit()
     
     def drop_matching_records(self, check):
         """Remove a record from the DB."""
         expr,args = self._render_expression(check)
         query = "DELETE FROM %s WHERE %s"%(self.table, expr)
         self._db.execute(query,args)
-        self._db.commit()
+        # self._db.commit()
         
     def find_records(self, check, id_only=False):
         """Find records matching a query dict."""
         req = 'msg_id' if id_only else '*'
         expr,args = self._render_expression(check)
         query = """SELECT %s FROM %s WHERE %s"""%(req, self.table, expr)
         cursor = self._db.execute(query, args)
         matches = cursor.fetchall()
         if id_only:
             return [ m[0] for m in matches ]
         else:
             records = {}
             for line in matches:
                 rec = self._list_to_dict(line)
                 records[rec['msg_id']] = rec
             return records
 
 __all__ = ['SQLiteDB']
  No newline at end of file

docs/examples/newparallel/wave2D/parallelwave-mpi.py

@@ -1,205 +1,205
 #!/usr/bin/env python
 """
 A simple python program of solving a 2D wave equation in parallel.
 Domain partitioning and inter-processor communication
 are done by an object of class MPIRectPartitioner2D
 (which is a subclass of RectPartitioner2D and uses MPI via mpi4py)
 
 An example of running the program is (8 processors, 4x2 partition,
 400x100 grid cells)::
 
    $ ipclusterz start --profile mpi -n 8 # start 8 engines (assuming mpi profile has been configured)
    $ ./parallelwave-mpi.py --grid 400 100 --partition 4 2 --profile mpi
 
 See also parallelwave-mpi, which runs the same program, but uses MPI 
 (via mpi4py) for the inter-engine communication.
 
 Authors
 -------
 
  * Xing Cai
  * Min Ragan-Kelley
 
 """
 
 import sys
 import time
 
 from numpy import exp, zeros, newaxis, sqrt
 
 from IPython.external import argparse
-from IPython.parallel.client import Client, Reference
+from IPython.parallel import Client, Reference
 
 def setup_partitioner(index, num_procs, gnum_cells, parts):
     """create a partitioner in the engine namespace"""
     global partitioner
     p = MPIRectPartitioner2D(my_id=index, num_procs=num_procs)
     p.redim(global_num_cells=gnum_cells, num_parts=parts)
     p.prepare_communication()
     # put the partitioner into the global namespace:
     partitioner=p
 
 def setup_solver(*args, **kwargs):
     """create a WaveSolver in the engine namespace"""
     global solver
     solver = WaveSolver(*args, **kwargs)
 
 def wave_saver(u, x, y, t):
     """save the wave log"""
     global u_hist
     global t_hist
     t_hist.append(t)
     u_hist.append(1.0*u)
     
 
 # main program:
 if __name__ == '__main__':
     
     parser = argparse.ArgumentParser()
     paa = parser.add_argument
     paa('--grid', '-g',
         type=int, nargs=2, default=[100,100], dest='grid',
         help="Cells in the grid, e.g. --grid 100 200")
     paa('--partition', '-p',
         type=int, nargs=2, default=None,
         help="Process partition grid, e.g. --partition 4 2 for 4x2")
     paa('-c',
         type=float, default=1.,
         help="Wave speed (I think)")
     paa('-Ly',
         type=float, default=1.,
         help="system size (in y)")
     paa('-Lx',
         type=float, default=1.,
         help="system size (in x)")
     paa('-t', '--tstop',
         type=float, default=1.,
         help="Time units to run")
     paa('--profile', 
         type=unicode, default=u'default',
         help="Specify the ipcluster profile for the client to connect to.")
     paa('--save', 
         action='store_true',
         help="Add this flag to save the time/wave history during the run.")
     paa('--scalar', 
         action='store_true',
         help="Also run with scalar interior implementation, to see vector speedup.")
     
     ns = parser.parse_args()
     # set up arguments
     grid = ns.grid
     partition = ns.partition
     Lx = ns.Lx
     Ly = ns.Ly
     c = ns.c
     tstop = ns.tstop
     if ns.save:
         user_action = wave_saver
     else:
         user_action = None
     
     num_cells = 1.0*(grid[0]-1)*(grid[1]-1)
     final_test = True
     
     # create the Client
     rc = Client(profile=ns.profile)
     num_procs = len(rc.ids)
     
     if partition is None:
         partition = [1,num_procs]
     
     assert partition[0]*partition[1] == num_procs, "can't map partition %s to %i engines"%(partition, num_procs)
     
     view = rc[:]
     print "Running %s system on %s processes until %f"%(grid, partition, tstop)
     
     # functions defining initial/boundary/source conditions
     def I(x,y):
         from numpy import exp
         return 1.5*exp(-100*((x-0.5)**2+(y-0.5)**2))
     def f(x,y,t):
         return 0.0
         # from numpy import exp,sin
         # return 10*exp(-(x - sin(100*t))**2)
     def bc(x,y,t):
         return 0.0
     
     # initial imports, setup rank
     view.execute('\n'.join([
     "from mpi4py import MPI",
     "import numpy",
     "mpi = MPI.COMM_WORLD",
     "my_id = MPI.COMM_WORLD.Get_rank()"]), block=True)
     
     # initialize t_hist/u_hist for saving the state at each step (optional)
     view['t_hist'] = []
     view['u_hist'] = []
     
     # set vector/scalar implementation details
     impl = {}
     impl['ic'] = 'vectorized'
     impl['inner'] = 'scalar'
     impl['bc'] = 'vectorized'
     
     # execute some files so that the classes we need will be defined on the engines:
     view.run('RectPartitioner.py')
     view.run('wavesolver.py')
     
     # setup remote partitioner
     # note that Reference means that the argument passed to setup_partitioner will be the
     # object named 'my_id' in the engine's namespace
     view.apply_sync(setup_partitioner, Reference('my_id'), num_procs, grid, partition)
     # wait for initial communication to complete
     view.execute('mpi.barrier()')
     # setup remote solvers
     view.apply_sync(setup_solver, I,f,c,bc,Lx,Ly,partitioner=Reference('partitioner'), dt=0,implementation=impl)
 
     # lambda for calling solver.solve:
     _solve = lambda *args, **kwargs: solver.solve(*args, **kwargs)
     
     if ns.scalar:
         impl['inner'] = 'scalar'
         # run first with element-wise Python operations for each cell
         t0 = time.time()
         ar = view.apply_async(_solve, tstop, dt=0, verbose=True, final_test=final_test, user_action=user_action)
         if final_test:
             # this sum is performed element-wise as results finish
             s = sum(ar)
             # the L2 norm (RMS) of the result:
             norm = sqrt(s/num_cells)
         else:
             norm = -1
         t1 = time.time()
         print 'scalar inner-version, Wtime=%g, norm=%g'%(t1-t0, norm)
     
     impl['inner'] = 'vectorized'
     # setup new solvers
     view.apply_sync(setup_solver, I,f,c,bc,Lx,Ly,partitioner=Reference('partitioner'), dt=0,implementation=impl)
     view.execute('mpi.barrier()')
     
     # run again with numpy vectorized inner-implementation
     t0 = time.time()
     ar = view.apply_async(_solve, tstop, dt=0, verbose=True, final_test=final_test)#, user_action=wave_saver)
     if final_test:
         # this sum is performed element-wise as results finish
         s = sum(ar)
         # the L2 norm (RMS) of the result:
         norm = sqrt(s/num_cells)
     else:
         norm = -1
     t1 = time.time()
     print 'vector inner-version, Wtime=%g, norm=%g'%(t1-t0, norm)
     
     # if ns.save is True, then u_hist stores the history of u as a list
     # If the partion scheme is Nx1, then u can be reconstructed via 'gather':
     if ns.save and partition[-1] == 1:
         import pylab
         view.execute('u_last=u_hist[-1]')
         # map mpi IDs to IPython IDs, which may not match
         ranks = view['my_id']
         targets = range(len(ranks))
         for idx in range(len(ranks)):
             targets[idx] = ranks.index(idx)
         u_last = rc[targets].gather('u_last', block=True)
         pylab.pcolor(u_last)
         pylab.show()

docs/examples/newparallel/wave2D/parallelwave.py

@@ -1,209 +1,209
 #!/usr/bin/env python
 """
 A simple python program of solving a 2D wave equation in parallel.
 Domain partitioning and inter-processor communication
 are done by an object of class ZMQRectPartitioner2D
 (which is a subclass of RectPartitioner2D and uses 0MQ via pyzmq)
 
 An example of running the program is (8 processors, 4x2 partition,
 200x200 grid cells)::
 
    $ ipclusterz start -n 8 # start 8 engines
    $ ./parallelwave.py --grid 200 200 --partition 4 2
 
 See also parallelwave-mpi, which runs the same program, but uses MPI 
 (via mpi4py) for the inter-engine communication.
 
 Authors
 -------
 
  * Xing Cai
  * Min Ragan-Kelley
 
 """
 #
 import sys
 import time
 
 from numpy import exp, zeros, newaxis, sqrt
 
 from IPython.external import argparse
-from IPython.parallel.client import Client, Reference
+from IPython.parallel import Client, Reference
 
 def setup_partitioner(comm, addrs, index, num_procs, gnum_cells, parts):
     """create a partitioner in the engine namespace"""
     global partitioner
     p = ZMQRectPartitioner2D(comm, addrs, my_id=index, num_procs=num_procs)
     p.redim(global_num_cells=gnum_cells, num_parts=parts)
     p.prepare_communication()
     # put the partitioner into the global namespace:
     partitioner=p
 
 def setup_solver(*args, **kwargs):
     """create a WaveSolver in the engine namespace."""
     global solver
     solver = WaveSolver(*args, **kwargs)
 
 def wave_saver(u, x, y, t):
     """save the wave state for each timestep."""
     global u_hist
     global t_hist
     t_hist.append(t)
     u_hist.append(1.0*u)
     
 
 # main program:
 if __name__ == '__main__':
     
     parser = argparse.ArgumentParser()
     paa = parser.add_argument
     paa('--grid', '-g',
         type=int, nargs=2, default=[100,100], dest='grid',
         help="Cells in the grid, e.g. --grid 100 200")
     paa('--partition', '-p',
         type=int, nargs=2, default=None,
         help="Process partition grid, e.g. --partition 4 2 for 4x2")
     paa('-c',
         type=float, default=1.,
         help="Wave speed (I think)")
     paa('-Ly',
         type=float, default=1.,
         help="system size (in y)")
     paa('-Lx',
         type=float, default=1.,
         help="system size (in x)")
     paa('-t', '--tstop',
         type=float, default=1.,
         help="Time units to run")
     paa('--profile', 
         type=unicode, default=u'default',
         help="Specify the ipcluster profile for the client to connect to.")
     paa('--save', 
         action='store_true',
         help="Add this flag to save the time/wave history during the run.")
     paa('--scalar', 
         action='store_true',
         help="Also run with scalar interior implementation, to see vector speedup.")
         
     ns = parser.parse_args()
     # set up arguments
     grid = ns.grid
     partition = ns.partition
     Lx = ns.Lx
     Ly = ns.Ly
     c = ns.c
     tstop = ns.tstop
     if ns.save:
         user_action = wave_saver
     else:
         user_action = None
     
     num_cells = 1.0*(grid[0]-1)*(grid[1]-1)
     final_test = True
     
     # create the Client
     rc = Client(profile=ns.profile)
     num_procs = len(rc.ids)
     
     if partition is None:
         partition = [num_procs,1]
     else:
         num_procs = min(num_procs, partition[0]*partition[1])
     
     assert partition[0]*partition[1] == num_procs, "can't map partition %s to %i engines"%(partition, num_procs)
     
     # construct the View:
     view = rc[:num_procs]
     print "Running %s system on %s processes until %f"%(grid, partition, tstop)
     
     # functions defining initial/boundary/source conditions
     def I(x,y):
         from numpy import exp
         return 1.5*exp(-100*((x-0.5)**2+(y-0.5)**2))
     def f(x,y,t):
         return 0.0
         # from numpy import exp,sin
         # return 10*exp(-(x - sin(100*t))**2)
     def bc(x,y,t):
         return 0.0
     
     # initialize t_hist/u_hist for saving the state at each step (optional)
     view['t_hist'] = []
     view['u_hist'] = []
 
     # set vector/scalar implementation details
     impl = {}
     impl['ic'] = 'vectorized'
     impl['inner'] = 'scalar'
     impl['bc'] = 'vectorized'
     
     # execute some files so that the classes we need will be defined on the engines:
     view.execute('import numpy')
     view.run('communicator.py')
     view.run('RectPartitioner.py')
     view.run('wavesolver.py')
     
     # scatter engine IDs
     view.scatter('my_id', range(num_procs), flatten=True)
     
     # create the engine connectors
     view.execute('com = EngineCommunicator()')
 
     # gather the connection information into a single dict
     ar = view.apply_async(lambda : com.info)
     peers = ar.get_dict()
     # print peers
     # this is a dict, keyed by engine ID, of the connection info for the EngineCommunicators
     
     # setup remote partitioner
     # note that Reference means that the argument passed to setup_partitioner will be the
     # object named 'com' in the engine's namespace
     view.apply_sync(setup_partitioner, Reference('com'), peers, Reference('my_id'), num_procs, grid, partition)
     time.sleep(1)
     # convenience lambda to call solver.solve:
     _solve = lambda *args, **kwargs: solver.solve(*args, **kwargs)
 
     if ns.scalar:
         impl['inner'] = 'scalar'
         # setup remote solvers
         view.apply_sync(setup_solver, I,f,c,bc,Lx,Ly, partitioner=Reference('partitioner'), dt=0,implementation=impl)
 
         # run first with element-wise Python operations for each cell
         t0 = time.time()
         ar = view.apply_async(_solve, tstop, dt=0, verbose=True, final_test=final_test, user_action=user_action)
         if final_test:
             # this sum is performed element-wise as results finish
             s = sum(ar)
             # the L2 norm (RMS) of the result:
             norm = sqrt(s/num_cells)
         else:
             norm = -1
         t1 = time.time()
         print 'scalar inner-version, Wtime=%g, norm=%g'%(t1-t0, norm)
     
     # run again with faster numpy-vectorized inner implementation:
     impl['inner'] = 'vectorized'
     # setup remote solvers
     view.apply_sync(setup_solver, I,f,c,bc,Lx,Ly,partitioner=Reference('partitioner'), dt=0,implementation=impl)
 
     t0 = time.time()
     
     ar = view.apply_async(_solve, tstop, dt=0, verbose=True, final_test=final_test)#, user_action=wave_saver)
     if final_test:
         # this sum is performed element-wise as results finish
         s = sum(ar)
         # the L2 norm (RMS) of the result:
         norm = sqrt(s/num_cells)
     else:
         norm = -1
     t1 = time.time()
     print 'vector inner-version, Wtime=%g, norm=%g'%(t1-t0, norm)
     
     # if ns.save is True, then u_hist stores the history of u as a list
     # If the partion scheme is Nx1, then u can be reconstructed via 'gather':
     if ns.save and partition[-1] == 1:
         import pylab
         view.execute('u_last=u_hist[-1]')
         u_last = view.gather('u_last', block=True)
         pylab.pcolor(u_last)
         pylab.show()
  No newline at end of file