From 633371e59ce8db949e43306195463ce32c798668 2014-04-30 23:04:55
From: Thomas Kluyver <takowl@gmail.com>
Date: 2014-04-30 23:04:55
Subject: [PATCH] Shut down kernels in parallel

When stopping the notebook server, it currently sends a shutdown request
to each kernel and then waits for the process to finish. This can be slow if
you have several kernels running. This makes it issues all the shutdown
requests before waiting on the processes, so shutdown happens in
parallel.

KernelManager (and MultiKernelManager) gain three new public API methods
to allow this:

* request_shutdown (promoted from a private method)
* wait_shutdown (refactored out of shutdown_kernel)
* cleanup (refactored out of shutdown_kernel)

---

diff --git a/IPython/kernel/manager.py b/IPython/kernel/manager.py
index ba85a8d..8ca6d57 100644
--- a/IPython/kernel/manager.py
+++ b/IPython/kernel/manager.py
@@ -246,12 +246,45 @@ class KernelManager(LoggingConfigurable, ConnectionFileMixin):
         self.start_restarter()
         self._connect_control_socket()
 
-    def _send_shutdown_request(self, restart=False):
-        """TODO: send a shutdown request via control channel"""
+    def request_shutdown(self, restart=False):
+        """Send a shutdown request via control channel
+        
+        On Windows, this just kills kernels instead, because the shutdown
+        messages don't work.
+        """
+        # FIXME: Shutdown does not work on Windows due to ZMQ errors!
+        if sys.platform == 'win32' and self.has_kernel:
+            return self._kill_kernel()
         content = dict(restart=restart)
         msg = self.session.msg("shutdown_request", content=content)
         self.session.send(self._control_socket, msg)
 
+    def wait_shutdown(self, totaltime=1, interval=0.1):
+        """Wait for kernel shutdown, then kill process if it doesn't shutdown.
+        
+        This does not send shutdown requests - use :meth:`request_shutdown`
+        first.
+        """
+        for i in range(int(totaltime/interval)):
+            if self.is_alive():
+                time.sleep(interval)
+            else:
+                break
+        else:
+            # OK, we've waited long enough.
+            if self.has_kernel:
+                self._kill_kernel()
+
+    def cleanup(self, restart=False):
+        """Clean up resources when the kernel is shut down"""
+        if not restart:
+            self.cleanup_connection_file()
+            self.cleanup_ipc_files()
+        else:
+            self.cleanup_ipc_files()
+
+        self._close_control_socket()
+
     def shutdown_kernel(self, now=False, restart=False):
         """Attempts to the stop the kernel process cleanly.
 
@@ -273,32 +306,16 @@ class KernelManager(LoggingConfigurable, ConnectionFileMixin):
         # Stop monitoring for restarting while we shutdown.
         self.stop_restarter()
 
-        # FIXME: Shutdown does not work on Windows due to ZMQ errors!
-        if now or sys.platform == 'win32':
-            if self.has_kernel:
-                self._kill_kernel()
+        if now:
+            self._kill_kernel()
         else:
+            self.request_shutdown(restart=restart)
             # Don't send any additional kernel kill messages immediately, to give
             # the kernel a chance to properly execute shutdown actions. Wait for at
             # most 1s, checking every 0.1s.
-            self._send_shutdown_request(restart=restart)
-            for i in range(10):
-                if self.is_alive():
-                    time.sleep(0.1)
-                else:
-                    break
-            else:
-                # OK, we've waited long enough.
-                if self.has_kernel:
-                    self._kill_kernel()
+            self.wait_shutdown()
 
-        if not restart:
-            self.cleanup_connection_file()
-            self.cleanup_ipc_files()
-        else:
-            self.cleanup_ipc_files()
-        
-        self._close_control_socket()
+        self.cleanup(restart=restart)
 
     def restart_kernel(self, now=False, **kw):
         """Restarts a kernel with the arguments that were used to launch it.
diff --git a/IPython/kernel/multikernelmanager.py b/IPython/kernel/multikernelmanager.py
index aa7854f..3d4fe26 100644
--- a/IPython/kernel/multikernelmanager.py
+++ b/IPython/kernel/multikernelmanager.py
@@ -131,6 +131,20 @@ class MultiKernelManager(LoggingConfigurable):
         self.log.info("Kernel shutdown: %s" % kernel_id)
         self.remove_kernel(kernel_id)
 
+    @kernel_method
+    def request_shutdown(self, kernel_id):
+        """Ask a kernel to shut down by its kernel uuid"""
+
+    @kernel_method
+    def wait_shutdown(self, kernel_id):
+        """Wait for a kernel to finish shutting down, and kill it if it doesn't
+        """
+        self.log.info("Kernel shutdown: %s" % kernel_id)
+
+    @kernel_method
+    def cleanup(self, kernel_id):
+        """Clean up a kernel's resources"""
+
     def remove_kernel(self, kernel_id):
         """remove a kernel from our mapping.
 
@@ -143,8 +157,12 @@ class MultiKernelManager(LoggingConfigurable):
 
     def shutdown_all(self, now=False):
         """Shutdown all kernels."""
-        for kid in self.list_kernel_ids():
-            self.shutdown_kernel(kid, now=now)
+        kids = self.list_kernel_ids()
+        for kid in kids:
+            self.request_shutdown(kid)
+        for kid in kids:
+            self.wait_shutdown(kid)
+            self.cleanup(kid)
 
     @kernel_method
     def interrupt_kernel(self, kernel_id):