handle failed kernel restart in the notebook

IPython/frontend/html/notebook/handlers.py

@@ -514,12 +514,13 @@ class IOPubHandler(AuthenticatedZMQStreamHandler):
         # from the WebSocket close event. If the WebSocket connection is
         # closed before the ZMQ streams are setup, they could be None.
         km = self.application.kernel_manager
-        km.remove_restart_callback(
-            self.kernel_id, self.on_kernel_restarted,
-        )
-        km.remove_restart_callback(
-            self.kernel_id, self.on_restart_failed, 'dead',
-        )
+        if self.kernel_id in km:
+            km.remove_restart_callback(
+                self.kernel_id, self.on_kernel_restarted,
+            )
+            km.remove_restart_callback(
+                self.kernel_id, self.on_restart_failed, 'dead',
+            )
         if self.iopub_stream is not None and not self.iopub_stream.closed():
             self.iopub_stream.on_recv(None)
             self.iopub_stream.close()

IPython/frontend/html/notebook/kernelmanager.py

@@ -58,11 +58,10 @@ class MappingKernelManager(MultiKernelManager):
 
     def notebook_for_kernel(self, kernel_id):
         """Return the notebook_id for a kernel_id or None."""
-        notebook_ids = [k for k, v in self._notebook_mapping.iteritems() if v == kernel_id]
-        if len(notebook_ids) == 1:
-            return notebook_ids[0]
-        else:
-            return None
+        for notebook_id, kid in self._notebook_mapping.iteritems():
+            if kernel_id == kid:
+                return notebook_id
+        return None
 
     def delete_mapping_for_kernel(self, kernel_id):
         """Remove the kernel/notebook mapping for kernel_id."""
@@ -70,6 +69,12 @@ class MappingKernelManager(MultiKernelManager):
         if notebook_id is not None:
             del self._notebook_mapping[notebook_id]
 
+    def _handle_kernel_died(self, kernel_id):
+        """notice that a kernel died"""
+        self.log.warn("Kernel %s died, removing from map.", kernel_id)
+        self.delete_mapping_for_kernel(kernel_id)
+        self.remove_kernel(kernel_id, now=True)
+
     def start_kernel(self, notebook_id=None, **kwargs):
         """Start a kernel for a notebok an return its kernel_id.
 
@@ -87,8 +92,14 @@ class MappingKernelManager(MultiKernelManager):
             self.set_kernel_for_notebook(notebook_id, kernel_id)
             self.log.info("Kernel started: %s" % kernel_id)
             self.log.debug("Kernel args: %r" % kwargs)
+            # register callback for failed auto-restart
+            self.add_restart_callback(kernel_id,
+                lambda : self._handle_kernel_died(kernel_id),
+                'dead',
+            )
         else:
             self.log.info("Using existing kernel: %s" % kernel_id)
+
         return kernel_id
 
     # override _check_kernel_id to raise 404 instead of KeyError