Work on the kernel manager.

IPython/zmq/kernelmanager.py

@@ -0,0 +1,271 b''
+"""Kernel frontend classes.
+
+To do:
+
+1. Create custom channel subclasses for Qt.
+2. Create logger to handle debugging and console messages.
+
+"""
+
+from Queue import Queue, Empty
+from threading import Thread
+import traceback
+
+import zmq
+from zmq import POLLIN, POLLOUT, POLLERR
+from zmq.eventloop import ioloop
+from session import Session
+
+
+class MissingHandlerError(Exception):
+    pass
+
+
+class KernelManager(object):
+
+    def __init__(self, xreq_addr, sub_addr, rep_addr, 
+                 context=None, session=None):
+        self.context = zmq.Context() if context is None else context
+        self.session = Session() if session is None else session
+        self.xreq_addr = xreq_addr
+        self.sub_addr = sub_addr
+        self.rep_addr = rep_addr
+
+    def start_kernel(self):
+        """Start a localhost kernel on ip and port.
+        
+        The SUB channel is for the frontend to receive messages published by
+        the kernel.
+        
+        The REQ channel is for the frontend to make requests of the kernel.
+        
+        The REP channel is for the kernel to request stdin (raw_input) from 
+        the frontend.
+        """
+
+    def kill_kernel(self):
+        """Kill the running kernel"""
+
+    def is_alive(self):
+        """Is the kernel alive?"""
+        return True
+
+    def signal_kernel(self, signum):
+        """Send signum to the kernel."""
+
+    def get_sub_channel(self):
+        """Get the SUB socket channel object."""
+        return SubSocketChannel(self.context, self.session, self.sub_addr)
+    
+    def get_xreq_channel(self):
+        """Get the REQ socket channel object to make requests of the kernel."""
+        return XReqSocketChannel(self.context, self.session, self.xreq_addr)
+    
+    def get_rep_channel(self):
+        """Get the REP socket channel object to handle stdin (raw_input)."""
+        return RepSocketChannel(self.context, self.session, self.rep_addr)
+
+
+class ZmqSocketChannel(Thread):
+
+    socket = None
+
+    def __init__(self, context, session, addr):
+        self.context = context
+        self.session = session
+        self.addr = addr
+        super(ZmqSocketChannel, self).__init__()
+        self.daemon = True
+
+
+class SubSocketChannel(ZmqSocketChannel):
+
+    handlers = None
+    _overriden_call_handler = None
+
+    def __init__(self, context, session, addr):
+        self.handlers = {}
+        super(SubSocketChannel, self).__init__(context, session, addr)
+
+    def run(self):
+        self.socket = self.context.socket(zmq.SUB)
+        self.socket.setsockopt(zmq.SUBSCRIBE,'')
+        self.socket.setsockopt(zmq.IDENTITY, self.session.session)
+        self.socket.connect('tcp://%s:%i' % self.addr)
+        self.ioloop = ioloop.IOLoop()
+        self.ioloop.add_handler(self.socket, self._handle_events, 
+                                POLLIN|POLLERR)
+        self.ioloop.start()
+
+    def _handle_events(self, socket, events):
+        # Turn on and off POLLOUT depending on if we have made a request
+        if events & POLLERR:
+            self._handle_err()
+        if events & POLLIN:
+            self._handle_recv()
+
+    def _handle_err(self):
+        raise zmq.ZmqError()
+
+    def _handle_recv(self):
+        msg = self.socket.recv_json()
+        self.call_handlers(msg)
+
+    def override_call_handler(self, func):
+        """Permanently override the call_handler.
+    
+        The function func will be called as::
+
+            func(handler, msg)
+
+        And must call::
+        
+            handler(msg)
+
+        in the main thread.
+        """
+        assert callable(func), "not a callable: %r" % func
+        self._overriden_call_handler = func
+
+    def call_handlers(self, msg):
+        handler = self.handlers.get(msg['msg_type'], None)
+        if handler is not None:
+            try:
+                self.call_handler(handler, msg)
+            except:
+                # XXX: This should be logged at least
+                traceback.print_last()
+
+    def call_handler(self, handler, msg):
+        if self._overriden_call_handler is not None:
+            self._overriden_call_handler(handler, msg)
+        elif hasattr(self, '_call_handler'):
+           call_handler = getattr(self, '_call_handler')
+           call_handler(handler, msg)
+        else:
+            raise RuntimeError('no handler!')
+
+    def add_handler(self, callback, msg_type):
+        """Register a callback for msg type."""
+        self.handlers[msg_type] = callback
+
+    def remove_handler(self, msg_type):
+        self.handlers.pop(msg_type, None)
+
+
+class XReqSocketChannel(ZmqSocketChannel):
+
+    handler_queue = None
+    command_queue = None
+    handlers = None
+    _overriden_call_handler = None
+
+    def __init__(self, context, session, addr):
+        self.handlers = {}
+        self.handler_queue = Queue()
+        self.command_queue = Queue()
+        super(XReqSocketChannel, self).__init__(context, session, addr)
+
+    def run(self):
+        self.socket = self.context.socket(zmq.XREQ)
+        self.socket.setsockopt(zmq.IDENTITY, self.session.session)
+        self.socket.connect('tcp://%s:%i' % self.addr)
+        self.ioloop = ioloop.IOLoop()
+        self.ioloop.add_handler(self.socket, self._handle_events, 
+                                POLLIN|POLLOUT|POLLERR)
+        self.ioloop.start()
+
+    def _handle_events(self, socket, events):
+        # Turn on and off POLLOUT depending on if we have made a request
+        if events & POLLERR:
+            self._handle_err()
+        if events & POLLOUT:
+            self._handle_send()
+        if events & POLLIN:
+            self._handle_recv()
+
+    def _handle_recv(self):
+        msg = self.socket.recv_json()
+        print "Got reply:", msg
+        try:
+            handler = self.handler_queue.get(False)
+        except Empty:
+            print "Message received with no handler!!!"
+            print msg
+        else:
+            self.call_handler(handler, msg)
+
+    def _handle_send(self):
+        try:
+            msg = self.command_queue.get(False)
+        except Empty:
+            pass
+        else:
+            self.socket.send_json(msg)
+
+    def _handle_err(self):
+        raise zmq.ZmqError()
+
+    def _queue_request(self, msg, callback):
+        handler = self._find_handler(msg['msg_type'], callback)
+        self.handler_queue.put(handler)
+        self.command_queue.put(msg)
+
+    def execute(self, code, callback=None):
+        # Create class for content/msg creation. Related to, but possibly
+        # not in Session.
+        content = dict(code=code)
+        msg = self.session.msg('execute_request', content)
+        self._queue_request(msg, callback)
+        return msg['header']['msg_id']
+
+    def complete(self, text, line, block=None, callback=None):
+        content = dict(text=text, line=line)
+        msg = self.session.msg('complete_request', content)
+        return self._queue_request(msg, callback)
+        return msg['header']['msg_id']
+
+    def object_info(self, oname, callback=None):
+        content = dict(oname=oname)
+        msg = self.session.msg('object_info_request', content)
+        return self._queue_request(msg, callback)
+        return msg['header']['msg_id']
+
+    def _find_handler(self, name, callback):
+        if callback is not None:
+            return callback
+        handler = self.handlers.get(name)
+        if handler is None:
+            raise MissingHandlerError('No handler defined for method: %s' % name)
+        return handler
+
+    def override_call_handler(self, func):
+        """Permanently override the call_handler.
+    
+        The function func will be called as::
+
+            func(handler, msg)
+
+        And must call::
+        
+            handler(msg)
+
+        in the main thread.
+        """
+        assert callable(func), "not a callable: %r" % func
+        self._overriden_call_handler = func
+
+    def call_handler(self, handler, msg):
+        if self._overriden_call_handler is not None:
+            self._overriden_call_handler(handler, msg)
+        elif hasattr(self, '_call_handler'):
+           call_handler = getattr(self, '_call_handler')
+           call_handler(handler, msg)
+        else:
+            raise RuntimeError('no handler!')
+
+
+class RepSocketChannel(ZmqSocketChannel):
+
+    def on_raw_input():
+        pass

IPython/zmq/ktest.py

@@ -0,0 +1,52 b''
+from Queue import Queue, Empty
+import time
+
+from kernelmanager import KernelManager
+
+xreq_addr = ('127.0.0.1',5575)
+sub_addr = ('127.0.0.1', 5576)
+rep_addr = ('127.0.0.1', 5577)
+
+
+km = KernelManager(xreq_addr, sub_addr, rep_addr)
+# xreq_channel = km.get_xreq_channel()
+sub_channel = km.get_sub_channel()
+
+# xreq_channel.start()
+sub_channel.start()
+
+print "Channels are started"
+
+def printer(msg):
+    print
+    print msg
+
+class CallHandler(object):
+
+    def __init__(self):
+        self.queue = Queue()
+
+    def __call__(self, handler, msg):
+        self.queue.put((handler, msg))
+
+    def handle(self):
+        try:
+            handler, msg = self.queue.get(block=False)
+        except Empty:
+            pass
+        else:
+            handler(msg)
+
+call_handler = CallHandler()
+sub_channel.override_call_handler(call_handler)
+sub_channel.add_handler(printer, 'pyin')
+sub_channel.add_handler(printer, 'pyout')
+sub_channel.add_handler(printer, 'stdout')
+sub_channel.add_handler(printer, 'stderr')
+
+for i in range(100):
+    call_handler.handle()
+    time.sleep(1)
+
+# xreq_channel.join()
+sub_channel.join()
  No newline at end of file

IPython/zmq/kernel.py

@@ -1,270 +1,272 b''
 #!/usr/bin/env python
 """A simple interactive kernel that talks to a frontend over 0MQ.
 
 Things to do:
 
 * Finish implementing `raw_input`.
 * Implement `set_parent` logic. Right before doing exec, the Kernel should
   call set_parent on all the PUB objects with the message about to be executed.
 * Implement random port and security key logic.
 * Implement control messages.
 * Implement event loop and poll version.
 """
 
 import __builtin__
 import sys
 import time
 import traceback
 
 from code import CommandCompiler
 
 import zmq
 
 from session import Session, Message, extract_header
 from completer import KernelCompleter
 
 class OutStream(object):
     """A file like object that publishes the stream to a 0MQ PUB socket."""
 
     def __init__(self, session, pub_socket, name, max_buffer=200):
         self.session = session
         self.pub_socket = pub_socket
         self.name = name
         self._buffer = []
         self._buffer_len = 0
         self.max_buffer = max_buffer
         self.parent_header = {}
 
     def set_parent(self, parent):
         self.parent_header = extract_header(parent)
 
     def close(self):
         self.pub_socket = None
 
     def flush(self):
         if self.pub_socket is None:
             raise ValueError(u'I/O operation on closed file')
         else:
             if self._buffer:
                 data = ''.join(self._buffer)
                 content = {u'name':self.name, u'data':data}
                 msg = self.session.msg(u'stream', content=content,
                                        parent=self.parent_header)
                 print>>sys.__stdout__, Message(msg)
                 self.pub_socket.send_json(msg)
                 self._buffer_len = 0
                 self._buffer = []
 
     def isattr(self):
         return False
 
     def next(self):
         raise IOError('Read not supported on a write only stream.')
 
     def read(self, size=None):
         raise IOError('Read not supported on a write only stream.')
 
     readline=read
 
     def write(self, s):
         if self.pub_socket is None:
             raise ValueError('I/O operation on closed file')
         else:
             self._buffer.append(s)
             self._buffer_len += len(s)
             self._maybe_send()
 
     def _maybe_send(self):
         if '\n' in self._buffer[-1]:
             self.flush()
         if self._buffer_len > self.max_buffer:
             self.flush()
 
     def writelines(self, sequence):
         if self.pub_socket is None:
             raise ValueError('I/O operation on closed file')
         else:
             for s in sequence:
                 self.write(s)
 
 
 class DisplayHook(object):
 
     def __init__(self, session, pub_socket):
         self.session = session
         self.pub_socket = pub_socket
         self.parent_header = {}
 
     def __call__(self, obj):
         if obj is None:
             return
 
         __builtin__._ = obj
         msg = self.session.msg(u'pyout', {u'data':repr(obj)},
                                parent=self.parent_header)
         self.pub_socket.send_json(msg)
 
     def set_parent(self, parent):
         self.parent_header = extract_header(parent)
 
 
 class RawInput(object):
 
     def __init__(self, session, socket):
         self.session = session
         self.socket = socket
 
     def __call__(self, prompt=None):
         msg = self.session.msg(u'raw_input')
         self.socket.send_json(msg)
         while True:
             try:
                 reply = self.socket.recv_json(zmq.NOBLOCK)
             except zmq.ZMQError, e:
                 if e.errno == zmq.EAGAIN:
                     pass
                 else:
                     raise
             else:
                 break
         return reply[u'content'][u'data']
 
 
 class Kernel(object):
 
     def __init__(self, session, reply_socket, pub_socket):
         self.session = session
         self.reply_socket = reply_socket
         self.pub_socket = pub_socket
         self.user_ns = {}
         self.history = []
         self.compiler = CommandCompiler()
         self.completer = KernelCompleter(self.user_ns)
         
         # Build dict of handlers for message types
         self.handlers = {}
         for msg_type in ['execute_request', 'complete_request']:
             self.handlers[msg_type] = getattr(self, msg_type)
 
     def abort_queue(self):
         while True:
             try:
                 ident = self.reply_socket.recv(zmq.NOBLOCK)
             except zmq.ZMQError, e:
                 if e.errno == zmq.EAGAIN:
                     break
             else:
                 assert self.reply_socket.rcvmore(), "Unexpected missing message part."
                 msg = self.reply_socket.recv_json()
             print>>sys.__stdout__, "Aborting:"
             print>>sys.__stdout__, Message(msg)
             msg_type = msg['msg_type']
             reply_type = msg_type.split('_')[0] + '_reply'
             reply_msg = self.session.msg(reply_type, {'status' : 'aborted'}, msg)
             print>>sys.__stdout__, Message(reply_msg)
             self.reply_socket.send(ident,zmq.SNDMORE)
             self.reply_socket.send_json(reply_msg)
             # We need to wait a bit for requests to come in. This can probably
             # be set shorter for true asynchronous clients.
             time.sleep(0.1)
 
     def execute_request(self, ident, parent):
         try:
             code = parent[u'content'][u'code']
         except:
             print>>sys.__stderr__, "Got bad msg: "
             print>>sys.__stderr__, Message(parent)
             return
         pyin_msg = self.session.msg(u'pyin',{u'code':code}, parent=parent)
         self.pub_socket.send_json(pyin_msg)
         try:
             comp_code = self.compiler(code, '<zmq-kernel>')
             sys.displayhook.set_parent(parent)
             exec comp_code in self.user_ns, self.user_ns
         except:
             result = u'error'
             etype, evalue, tb = sys.exc_info()
             tb = traceback.format_exception(etype, evalue, tb)
             exc_content = {
                 u'status' : u'error',
                 u'traceback' : tb,
                 u'etype' : unicode(etype),
                 u'evalue' : unicode(evalue)
             }
             exc_msg = self.session.msg(u'pyerr', exc_content, parent)
             self.pub_socket.send_json(exc_msg)
             reply_content = exc_content
         else:
             reply_content = {'status' : 'ok'}
         reply_msg = self.session.msg(u'execute_reply', reply_content, parent)
         print>>sys.__stdout__, Message(reply_msg)
         self.reply_socket.send(ident, zmq.SNDMORE)
         self.reply_socket.send_json(reply_msg)
         if reply_msg['content']['status'] == u'error':
             self.abort_queue()
 
     def complete_request(self, ident, parent):
         matches = {'matches' : self.complete(parent),
                    'status' : 'ok'}
         completion_msg = self.session.send(self.reply_socket, 'complete_reply',
                                            matches, parent, ident)
         print >> sys.__stdout__, completion_msg
 
     def complete(self, msg):
         return self.completer.complete(msg.content.line, msg.content.text)
 
     def start(self):
         while True:
             ident = self.reply_socket.recv()
             assert self.reply_socket.rcvmore(), "Unexpected missing message part."
             msg = self.reply_socket.recv_json()
             omsg = Message(msg)
+            print>>sys.__stdout__
             print>>sys.__stdout__, omsg
             handler = self.handlers.get(omsg.msg_type, None)
             if handler is None:
                 print >> sys.__stderr__, "UNKNOWN MESSAGE TYPE:", omsg
             else:
                 handler(ident, omsg)
 
 
 def main():
     c = zmq.Context()
 
     ip = '127.0.0.1'
     port_base = 5575
     connection = ('tcp://%s' % ip) + ':%i'
     rep_conn = connection % port_base
     pub_conn = connection % (port_base+1)
 
     print >>sys.__stdout__, "Starting the kernel..."
-    print >>sys.__stdout__, "On:",rep_conn, pub_conn
+    print >>sys.__stdout__, "XREP Channel:", rep_conn
+    print >>sys.__stdout__, "PUB Channel:", pub_conn
 
     session = Session(username=u'kernel')
 
     reply_socket = c.socket(zmq.XREP)
     reply_socket.bind(rep_conn)
 
     pub_socket = c.socket(zmq.PUB)
     pub_socket.bind(pub_conn)
 
     stdout = OutStream(session, pub_socket, u'stdout')
     stderr = OutStream(session, pub_socket, u'stderr')
     sys.stdout = stdout
     sys.stderr = stderr
 
     display_hook = DisplayHook(session, pub_socket)
     sys.displayhook = display_hook
 
     kernel = Kernel(session, reply_socket, pub_socket)
 
     # For debugging convenience, put sleep and a string in the namespace, so we
     # have them every time we start.
     kernel.user_ns['sleep'] = time.sleep
     kernel.user_ns['s'] = 'Test string'
     
     print >>sys.__stdout__, "Use Ctrl-\\ (NOT Ctrl-C!) to terminate."
     kernel.start()
 
 
 if __name__ == '__main__':
     main()

IPython/zmq/session.py

@@ -1,119 +1,122 b''
 import os
 import uuid
 import pprint
 
 import zmq
 
 class Message(object):
     """A simple message object that maps dict keys to attributes.
 
     A Message can be created from a dict and a dict from a Message instance
     simply by calling dict(msg_obj)."""
     
     def __init__(self, msg_dict):
         dct = self.__dict__
         for k, v in msg_dict.iteritems():
             if isinstance(v, dict):
                 v = Message(v)
             dct[k] = v
 
     # Having this iterator lets dict(msg_obj) work out of the box.
     def __iter__(self):
         return iter(self.__dict__.iteritems())
     
     def __repr__(self):
         return repr(self.__dict__)
 
     def __str__(self):
         return pprint.pformat(self.__dict__)
 
     def __contains__(self, k):
         return k in self.__dict__
 
     def __getitem__(self, k):
         return self.__dict__[k]
 
 
 def msg_header(msg_id, username, session):
     return {
         'msg_id' : msg_id,
         'username' : username,
         'session' : session
     }
 
 
 def extract_header(msg_or_header):
     """Given a message or header, return the header."""
     if not msg_or_header:
         return {}
     try:
         # See if msg_or_header is the entire message.
         h = msg_or_header['header']
     except KeyError:
         try:
             # See if msg_or_header is just the header
             h = msg_or_header['msg_id']
         except KeyError:
             raise
         else:
             h = msg_or_header
     if not isinstance(h, dict):
         h = dict(h)
     return h
 
 
 class Session(object):
 
-    def __init__(self, username=os.environ.get('USER','username')):
+    def __init__(self, username=os.environ.get('USER','username'), session=None):
         self.username = username
-        self.session = str(uuid.uuid4())
+        if session is None:
+            self.session = str(uuid.uuid4())
+        else:
+            self.session = session
         self.msg_id = 0
 
     def msg_header(self):
         h = msg_header(self.msg_id, self.username, self.session)
         self.msg_id += 1
         return h
 
     def msg(self, msg_type, content=None, parent=None):
         msg = {}
         msg['header'] = self.msg_header()
         msg['parent_header'] = {} if parent is None else extract_header(parent)
         msg['msg_type'] = msg_type
         msg['content'] = {} if content is None else content
         return msg
 
     def send(self, socket, msg_type, content=None, parent=None, ident=None):
         msg = self.msg(msg_type, content, parent)
         if ident is not None:
             socket.send(ident, zmq.SNDMORE)
         socket.send_json(msg)
         omsg = Message(msg)
         return omsg
 
     def recv(self, socket, mode=zmq.NOBLOCK):
         try:
             msg = socket.recv_json(mode)
         except zmq.ZMQError, e:
             if e.errno == zmq.EAGAIN:
                 # We can convert EAGAIN to None as we know in this case
                 # recv_json won't return None.
                 return None
             else:
                 raise
         return Message(msg)
 
 def test_msg2obj():
     am = dict(x=1)
     ao = Message(am)
     assert ao.x == am['x']
 
     am['y'] = dict(z=1)
     ao = Message(am)
     assert ao.y.z == am['y']['z']
     
     k1, k2 = 'y', 'z'
     assert ao[k1][k2] == am[k1][k2]
     
     am2 = dict(ao)
     assert am['x'] == am2['x']
     assert am['y']['z'] == am2['y']['z']