cleanup boundaries of MultiKernelManager and KernelRestarter classes

- MultiKernelManager no longer has IOLoop affiliations
- move connect_[channel] method to base KernelManager
- connect_[channel] wraps in ZMQStream in IOLoopKM subclass
- add KernelRestarter base class
- use decorators to minimize code in wrapped KM methods

IPython/kernel/restarter.py

@@ -0,0 +1,67 b''
+"""A basic kernel monitor with autorestarting.
+
+This watches a kernel's state using KernelManager.is_alive and auto
+restarts the kernel if it dies.
+
+It is an incomplete base class, and must be subclassed.
+"""
+
+#-----------------------------------------------------------------------------
+#  Copyright (C) 2013  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 IPython.config.configurable import LoggingConfigurable
+from IPython.utils.traitlets import (
+    Instance, Float, List,
+)
+
+#-----------------------------------------------------------------------------
+# Code
+#-----------------------------------------------------------------------------
+
+class KernelRestarter(LoggingConfigurable):
+    """Monitor and autorestart a kernel."""
+
+    kernel_manager = Instance('IPython.kernel.KernelManager')
+
+    time_to_dead = Float(3.0, config=True,
+        help="""Kernel heartbeat interval in seconds."""
+    )
+    _callbacks = List()
+
+    def start(self):
+        """Start the polling of the kernel."""
+        raise NotImplementedError("Must be implemented in a subclass")
+
+    def stop(self):
+        """Stop the kernel polling."""
+        raise NotImplementedError("Must be implemented in a subclass")
+
+    def register_callback(self, f):
+        """register a callback to fire"""
+        self.callbacks.append(f)
+
+    def unregister_callback(self, f):
+        try:
+            self.callbacks.remove(f)
+        except ValueError:
+            pass
+
+    def poll(self):
+        self.log.debug('Polling kernel...')
+        if not self.kernel_manager.is_alive():
+            self.log.info('KernelRestarter: restarting kernel')
+            for callback in self.callbacks:
+                try:
+                    callback()
+                except Exception as e:
+                    self.log.error("Kernel restart callback %r failed", callback, exc_info=True)
+            self.kernel_manager.restart_kernel(now=True)

IPython/frontend/html/notebook/handlers.py

@@ -480,7 +480,7 b' class IOPubHandler(AuthenticatedZMQStreamHandler):'
         km = self.application.kernel_manager
         kernel_id = self.kernel_id
         try:
-            self.iopub_stream = km.create_iopub_stream(kernel_id)
+            self.iopub_stream = km.connect_iopub(kernel_id)
         except web.HTTPError:
             # WebSockets don't response to traditional error codes so we
             # close the connection.
@@ -517,7 +517,7 b' class ShellHandler(AuthenticatedZMQStreamHandler):'
         self.max_msg_size = km.max_msg_size
         kernel_id = self.kernel_id
         try:
-            self.shell_stream = km.create_shell_stream(kernel_id)
+            self.shell_stream = km.connect_shell(kernel_id)
         except web.HTTPError:
             # WebSockets don't response to traditional error codes so we
             # close the connection.

IPython/frontend/html/notebook/kernelmanager.py

@@ -31,6 +31,9 b' from IPython.utils.traitlets import ('
 class MappingKernelManager(MultiKernelManager):
     """A KernelManager that handles notebok mapping and HTTP error handling"""
 
+    def _kernel_manager_class_default(self):
+        return "IPython.kernel.ioloop.IOLoopKernelManager"
+
     kernel_argv = List(Unicode)
     
     max_msg_size = Integer(65536, config=True, help="""
@@ -88,42 +91,7 b' class MappingKernelManager(MultiKernelManager):'
             self.log.info("Using existing kernel: %s" % kernel_id)
         return kernel_id
 
-    def shutdown_kernel(self, kernel_id, now=False):
-        """Shutdown a kernel and remove its notebook association."""
-        self._check_kernel_id(kernel_id)
-        super(MappingKernelManager, self).shutdown_kernel(
-            kernel_id, now=now
-        )
-        self.delete_mapping_for_kernel(kernel_id)
-        self.log.info("Kernel shutdown: %s" % kernel_id)
-
-    def interrupt_kernel(self, kernel_id):
-        """Interrupt a kernel."""
-        self._check_kernel_id(kernel_id)
-        super(MappingKernelManager, self).interrupt_kernel(kernel_id)
-        self.log.info("Kernel interrupted: %s" % kernel_id)
-
-    def restart_kernel(self, kernel_id):
-        """Restart a kernel while keeping clients connected."""
-        self._check_kernel_id(kernel_id)
-        super(MappingKernelManager, self).restart_kernel(kernel_id)
-        self.log.info("Kernel restarted: %s" % kernel_id)
-
-    def create_iopub_stream(self, kernel_id):
-        """Create a new iopub stream."""
-        self._check_kernel_id(kernel_id)
-        return super(MappingKernelManager, self).create_iopub_stream(kernel_id)
-
-    def create_shell_stream(self, kernel_id):
-        """Create a new shell stream."""
-        self._check_kernel_id(kernel_id)
-        return super(MappingKernelManager, self).create_shell_stream(kernel_id)
-
-    def create_hb_stream(self, kernel_id):
-        """Create a new hb stream."""
-        self._check_kernel_id(kernel_id)
-        return super(MappingKernelManager, self).create_hb_stream(kernel_id)
-
+    # override _check_kernel_id to raise 404 instead of KeyError
     def _check_kernel_id(self, kernel_id):
         """Check a that a kernel_id exists and raise 404 if not."""
         if kernel_id not in self:

IPython/kernel/ioloop/manager.py

@@ -15,6 +15,7 b' from __future__ import absolute_import'
 
 import zmq
 from zmq.eventloop import ioloop
+from zmq.eventloop.zmqstream import ZMQStream
 
 from IPython.utils.traitlets import (
     Instance
@@ -27,6 +28,13 b' from .restarter import IOLoopKernelRestarter'
 # Code
 #-----------------------------------------------------------------------------
 
+
+def as_zmqstream(f):
+    def wrapped(self, *args, **kwargs):
+        socket = f(self, *args, **kwargs)
+        return ZMQStream(socket, self.loop)
+    return wrapped
+
 class IOLoopKernelManager(KernelManager):
 
     loop = Instance('zmq.eventloop.ioloop.IOLoop', allow_none=False)
@@ -48,3 +56,8 b' class IOLoopKernelManager(KernelManager):'
         if self.autorestart:
             if self._restarter is not None:
                 self._restarter.stop()
+
+    connect_shell = as_zmqstream(KernelManager.connect_shell)
+    connect_iopub = as_zmqstream(KernelManager.connect_iopub)
+    connect_stdin = as_zmqstream(KernelManager.connect_stdin)
+    connect_hb = as_zmqstream(KernelManager.connect_hb)

IPython/kernel/ioloop/restarter.py

@@ -21,24 +21,22 b' import zmq'
 from zmq.eventloop import ioloop
 
 
-from IPython.config.configurable import LoggingConfigurable
+from IPython.kernel.restarter import KernelRestarter
 from IPython.utils.traitlets import (
-    Instance, Float
+    Instance, Float, List,
 )
 
 #-----------------------------------------------------------------------------
 # Code
 #-----------------------------------------------------------------------------
 
-class IOLoopKernelRestarter(LoggingConfigurable):
+class IOLoopKernelRestarter(KernelRestarter):
     """Monitor and autorestart a kernel."""
 
     loop = Instance('zmq.eventloop.ioloop.IOLoop', allow_none=False)
     def _loop_default(self):
         return ioloop.IOLoop.instance()
 
-    kernel_manager = Instance('IPython.kernel.KernelManager')
-
     time_to_dead = Float(3.0, config=True,
         help="""Kernel heartbeat interval in seconds."""
     )
@@ -49,7 +47,7 b' class IOLoopKernelRestarter(LoggingConfigurable):'
         """Start the polling of the kernel."""
         if self._pcallback is None:
             self._pcallback = ioloop.PeriodicCallback(
-                self._poll, 1000*self.time_to_dead, self.loop
+                self.poll, 1000*self.time_to_dead, self.loop
             )
         self._pcallback.start()
 
@@ -61,14 +59,4 b' class IOLoopKernelRestarter(LoggingConfigurable):'
     def clear(self):
         """Clear the underlying PeriodicCallback."""
         self.stop()
-        if self._pcallback is not None:
-            self._pcallback = None
-
-    def _poll(self):
-        self.log.debug('Polling kernel...')
-        if not self.kernel_manager.is_alive():
-            # This restart event should leave the connection file in place so
-            # the ports are the same. Because this takes place below the
-            # MappingKernelManager, the kernel_id will also remain the same.
-            self.log.info('KernelRestarter: restarting kernel')
-            self.kernel_manager.restart_kernel(now=True)
+        self._pcallback = None

IPython/kernel/manager.py

@@ -41,6 +41,13 b' from .managerabc import ('
 # Main kernel manager class
 #-----------------------------------------------------------------------------
 
+_socket_types = {
+    'hb' : zmq.REQ,
+    'shell' : zmq.DEALER,
+    'iopub' : zmq.SUB,
+    'stdin' : zmq.DEALER,
+}
+
 class KernelManager(LoggingConfigurable, ConnectionFileMixin):
     """Manages a single kernel in a subprocess on this host.
 
@@ -93,6 +100,59 b' class KernelManager(LoggingConfigurable, ConnectionFileMixin):'
         pass
 
     #--------------------------------------------------------------------------
+    # Connection info
+    #--------------------------------------------------------------------------
+
+    def get_connection_info(self):
+        """return the connection info as a dict"""
+        return dict(
+            transport=self.transport,
+            ip=self.ip,
+            shell_port=self.shell_port,
+            iopub_port=self.iopub_port,
+            stdin_port=self.stdin_port,
+            hb_port=self.hb_port,
+        )
+
+    def _make_url(self, channel):
+        """Make a ZeroMQ URL for a given channel."""
+        transport = self.transport
+        ip = self.ip
+        port = getattr(self, '%s_port' % channel)
+
+        if transport == 'tcp':
+            return "tcp://%s:%i" % (ip, port)
+        else:
+            return "%s://%s-%s" % (transport, ip, port)
+
+    def _create_connected_socket(self, channel):
+        """Create a zmq Socket and connect it to the kernel."""
+        url = self._make_url(channel)
+        socket_type = _socket_types[channel]
+        sock = self.context.socket(socket_type)
+        self.log.info("Connecting to: %s" % url)
+        sock.connect(url)
+        return sock
+
+    def connect_iopub(self):
+        """return zmq Socket connected to the IOPub channel"""
+        sock = self._create_connected_socket('iopub')
+        sock.setsockopt(zmq.SUBSCRIBE, b'')
+        return sock
+
+    def connect_shell(self):
+        """return zmq Socket connected to the Shell channel"""
+        return self._create_connected_socket('shell')
+
+    def connect_stdin(self):
+        """return zmq Socket connected to the StdIn channel"""
+        return self._create_connected_socket('stdin')
+
+    def connect_hb(self):
+        """return zmq Socket connected to the Heartbeat channel"""
+        return self._create_connected_socket('hb')
+
+    #--------------------------------------------------------------------------
     # Kernel management
     #--------------------------------------------------------------------------
 

IPython/kernel/multikernelmanager.py

@@ -22,7 +22,6 b' import os'
 import uuid
 
 import zmq
-from zmq.eventloop.zmqstream import ZMQStream
 
 from IPython.config.configurable import LoggingConfigurable
 from IPython.utils.importstring import import_item
@@ -38,6 +37,23 b' class DuplicateKernelError(Exception):'
     pass
 
 
+
+def kernel_method(f):
+    """decorator for proxying MKM.method(kernel_id) to individual KMs by ID"""
+    def wrapped(self, kernel_id, *args, **kwargs):
+        # get the kernel
+        km = self.get_kernel(kernel_id)
+        method = getattr(km, f.__name__)
+        # call the kernel's method
+        r = method(*args, **kwargs)
+        # last thing, call anything defined in the actual class method
+        # such as logging messages
+        f(self, kernel_id, *args, **kwargs)
+        # return the method result
+        return r
+    return wrapped
+
+
 class MultiKernelManager(LoggingConfigurable):
     """A class for managing multiple kernels."""
     
@@ -100,6 +116,7 b' class MultiKernelManager(LoggingConfigurable):'
         self._kernels[kernel_id] = km
         return kernel_id
 
+    @kernel_method
     def shutdown_kernel(self, kernel_id, now=False):
         """Shutdown a kernel by its kernel uuid.
 
@@ -110,8 +127,7 b' class MultiKernelManager(LoggingConfigurable):'
         now : bool
             Should the kernel be shutdown forcibly using a signal.
         """
-        k = self.get_kernel(kernel_id)
-        k.shutdown_kernel(now=now)
+        self.log.info("Kernel shutdown: %s" % kernel_id)
         del self._kernels[kernel_id]
 
     def shutdown_all(self, now=False):
@@ -119,6 +135,7 b' class MultiKernelManager(LoggingConfigurable):'
         for kid in self.list_kernel_ids():
             self.shutdown_kernel(kid, now=now)
 
+    @kernel_method
     def interrupt_kernel(self, kernel_id):
         """Interrupt (SIGINT) the kernel by its uuid.
 
@@ -127,8 +144,9 b' class MultiKernelManager(LoggingConfigurable):'
         kernel_id : uuid
             The id of the kernel to interrupt.
         """
-        return self.get_kernel(kernel_id).interrupt_kernel()
+        self.log.info("Kernel interrupted: %s" % kernel_id)
 
+    @kernel_method
     def signal_kernel(self, kernel_id, signum):
         """Sends a signal to the kernel by its uuid.
 
@@ -140,8 +158,9 b' class MultiKernelManager(LoggingConfigurable):'
         kernel_id : uuid
             The id of the kernel to signal.
         """
-        return self.get_kernel(kernel_id).signal_kernel(signum)
+        self.log.info("Signaled Kernel %s with %s" % (kernel_id, signum))
 
+    @kernel_method
     def restart_kernel(self, kernel_id):
         """Restart a kernel by its uuid, keeping the same ports.
 
@@ -150,9 +169,9 b' class MultiKernelManager(LoggingConfigurable):'
         kernel_id : uuid
             The id of the kernel to interrupt.
         """
-        km = self.get_kernel(kernel_id)
-        km.restart_kernel()
+        self.log.info("Kernel restarted: %s" % kernel_id)
 
+    @kernel_method
     def is_alive(self, kernel_id):
         """Is the kernel alive.
 
@@ -164,7 +183,11 b' class MultiKernelManager(LoggingConfigurable):'
         kernel_id : uuid
             The id of the kernel.
         """
-        return self.get_kernel(kernel_id).is_alive()
+
+    def _check_kernel_id(self, kernel_id):
+        """check that a kernel id is valid"""
+        if kernel_id not in self:
+            raise KeyError("Kernel with id not found: %s" % kernel_id)
 
     def get_kernel(self, kernel_id):
         """Get the single KernelManager object for a kernel by its uuid.
@@ -174,12 +197,10 b' class MultiKernelManager(LoggingConfigurable):'
         kernel_id : uuid
             The id of the kernel.
         """
-        km = self._kernels.get(kernel_id)
-        if km is not None:
-            return km
-        else:
-            raise KeyError("Kernel with id not found: %s" % kernel_id)
+        self._check_kernel_id(kernel_id)
+        return self._kernels[kernel_id]
 
+    @kernel_method
     def get_connection_info(self, kernel_id):
         """Return a dictionary of connection data for a kernel.
 
@@ -196,35 +217,10 b' class MultiKernelManager(LoggingConfigurable):'
             numbers of the different channels (stdin_port, iopub_port,
             shell_port, hb_port).
         """
-        km = self.get_kernel(kernel_id)
-        return dict(transport=km.transport,
-                    ip=km.ip,
-                    shell_port=km.shell_port,
-                    iopub_port=km.iopub_port,
-                    stdin_port=km.stdin_port,
-                    hb_port=km.hb_port,
-        )  
-
-    def _make_url(self, transport, ip, port):
-        """Make a ZeroMQ URL for a given transport, ip and port."""
-        if transport == 'tcp':
-            return "tcp://%s:%i" % (ip, port)
-        else:
-            return "%s://%s-%s" % (transport, ip, port)
-
-    def _create_connected_stream(self, kernel_id, socket_type, channel):
-        """Create a connected ZMQStream for a kernel."""
-        cinfo = self.get_connection_info(kernel_id)
-        url = self._make_url(cinfo['transport'], cinfo['ip'],
-                cinfo['%s_port' % channel]
-        )
-        sock = self.context.socket(socket_type)
-        self.log.info("Connecting to: %s" % url)
-        sock.connect(url)
-        return ZMQStream(sock)
 
-    def create_iopub_stream(self, kernel_id):
-        """Return a ZMQStream object connected to the iopub channel.
+    @kernel_method
+    def connect_iopub(self, kernel_id):
+        """Return a zmq Socket connected to the iopub channel.
 
         Parameters
         ==========
@@ -233,14 +229,12 b' class MultiKernelManager(LoggingConfigurable):'
 
         Returns
         =======
-        stream : ZMQStream
+        stream : zmq Socket or ZMQStream
         """
-        iopub_stream = self._create_connected_stream(kernel_id, zmq.SUB, 'iopub')
-        iopub_stream.socket.setsockopt(zmq.SUBSCRIBE, b'')
-        return iopub_stream
 
-    def create_shell_stream(self, kernel_id):
-        """Return a ZMQStream object connected to the shell channel.
+    @kernel_method
+    def connect_shell(self, kernel_id):
+        """Return a zmq Socket connected to the shell channel.
 
         Parameters
         ==========
@@ -249,13 +243,12 b' class MultiKernelManager(LoggingConfigurable):'
 
         Returns
         =======
-        stream : ZMQStream
+        stream : zmq Socket or ZMQStream
         """
-        shell_stream = self._create_connected_stream(kernel_id, zmq.DEALER, 'shell')
-        return shell_stream
 
-    def create_hb_stream(self, kernel_id):
-        """Return a ZMQStream object connected to the hb channel.
+    @kernel_method
+    def connect_stdin(self, kernel_id):
+        """Return a zmq Socket connected to the stdin channel.
 
         Parameters
         ==========
@@ -264,8 +257,19 b' class MultiKernelManager(LoggingConfigurable):'
 
         Returns
         =======
-        stream : ZMQStream
+        stream : zmq Socket or ZMQStream
         """
-        hb_stream = self._create_connected_stream(kernel_id, zmq.REQ, 'hb')
-        return hb_stream
 
+    @kernel_method
+    def connect_hb(self, kernel_id):
+        """Return a zmq Socket connected to the hb channel.
+
+        Parameters
+        ==========
+        kernel_id : uuid
+            The id of the kernel.
+
+        Returns
+        =======
+        stream : zmq Socket or ZMQStream
+        """

IPython/kernel/tests/test_multikernelmanager.py

@@ -57,13 +57,13 b' class TestKernelManager(TestCase):'
         self.assertEqual(ip, cinfo['ip'])
         self.assertTrue('stdin_port' in cinfo)
         self.assertTrue('iopub_port' in cinfo)
-        stream = km.create_iopub_stream(kid)
+        stream = km.connect_iopub(kid)
         stream.close()
         self.assertTrue('shell_port' in cinfo)
-        stream = km.create_shell_stream(kid)
+        stream = km.connect_shell(kid)
         stream.close()
         self.assertTrue('hb_port' in cinfo)
-        stream = km.create_hb_stream(kid)
+        stream = km.connect_hb(kid)
         stream.close()
         km.shutdown_kernel(kid)