Use reworked restarter callback system

jupyter · Oct 31, 2018 · 2d9e021 · 2d9e021
1 parent c76e6bf
commit 2d9e021
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/notebook/services/kernels/kernelmanager.py b/notebook/services/kernels/kernelmanager.py
@@ -44,7 +44,7 @@ def __init__(self, kernel_type, kernel_finder):
 
         self.restarter = TornadoKernelRestarter(self.manager, kernel_type,
                                            kernel_finder=self.kernel_finder)
-        self.restarter.add_callback(self._handle_kernel_restarted, 'restart')
+        self.restarter.add_callback('restarted', self._handle_kernel_restarted)
         self.restarter.start()
 
         self.buffer_for_key = None
@@ -76,10 +76,11 @@ def shutdown(self, now=False):
     def interrupt(self):
         self.manager.interrupt()
 
-    def _handle_kernel_restarted(self):
-        self.manager = self.restarter.kernel_manager
-        # TODO: connection_info
-        self.connect_client()
+    def _handle_kernel_restarted(self, data):
+        if data['new_manager']:
+            self.manager = data['manager']
+            self.connection_info = data['connection_info']
+            self.connect_client()
 
     def start_buffering(self, session_key):
         # record the session key because only one session can buffer
@@ -303,8 +304,8 @@ def _start_kernel(self, kernel_id, kernel_type):
         self.log.info("Kernel started: %s" % kernel_id)
 
         kernel.restarter.add_callback(
-            lambda: self._handle_kernel_died(kernel_id),
-            'dead'
+            'failed',
+            lambda data: self._handle_kernel_died(kernel_id),
         )
 
     def start_buffering(self, kernel_id, session_key, channels):