Merge pull request #5764 from takluyver/parallel-kernel-shutdown

Shut down kernels in parallel

IPython/kernel/manager.py

@@ -246,12 +246,40 b' 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.
+        """
         content = dict(restart=restart)
         msg = self.session.msg("shutdown_request", content=content)
         self.session.send(self._control_socket, msg)
 
+    def finish_shutdown(self, waittime=1, pollinterval=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(waittime/pollinterval)):
+            if self.is_alive():
+                time.sleep(pollinterval)
+            else:
+                break
+        else:
+            # OK, we've waited long enough.
+            if self.has_kernel:
+                self._kill_kernel()
+
+    def cleanup(self, connection_file=True):
+        """Clean up resources when the kernel is shut down"""
+        if connection_file:
+            self.cleanup_connection_file()
+
+        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 +301,16 b' 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:
+        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.finish_shutdown()
 
-        if not restart:
-            self.cleanup_connection_file()
-            self.cleanup_ipc_files()
-        else:
-            self.cleanup_ipc_files()
-        
-        self._close_control_socket()
+        self.cleanup(connection_file=not restart)
 
     def restart_kernel(self, now=False, **kw):
         """Restarts a kernel with the arguments that were used to launch it.
@@ -331,11 +343,6 b' class KernelManager(LoggingConfigurable, ConnectionFileMixin):'
             self._launch_args.update(kw)
             self.start_kernel(**self._launch_args)
 
-            # FIXME: Messages get dropped in Windows due to probable ZMQ bug
-            # unless there is some delay here.
-            if sys.platform == 'win32':
-                time.sleep(0.2)
-
     @property
     def has_kernel(self):
         """Has a kernel been started that we are managing."""

IPython/kernel/multikernelmanager.py

@@ -118,7 +118,7 b' class MultiKernelManager(LoggingConfigurable):'
         return kernel_id
 
     @kernel_method
-    def shutdown_kernel(self, kernel_id, now=False):
+    def shutdown_kernel(self, kernel_id, now=False, restart=False):
         """Shutdown a kernel by its kernel uuid.
 
         Parameters
@@ -127,10 +127,26 b' class MultiKernelManager(LoggingConfigurable):'
             The id of the kernel to shutdown.
         now : bool
             Should the kernel be shutdown forcibly using a signal.
+        restart : bool
+            Will the kernel be restarted?
         """
         self.log.info("Kernel shutdown: %s" % kernel_id)
         self.remove_kernel(kernel_id)
 
+    @kernel_method
+    def request_shutdown(self, kernel_id, restart=False):
+        """Ask a kernel to shut down by its kernel uuid"""
+
+    @kernel_method
+    def finish_shutdown(self, kernel_id, waittime=1, pollinterval=0.1):
+        """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, connection_file=True):
+        """Clean up a kernel's resources"""
+
     def remove_kernel(self, kernel_id):
         """remove a kernel from our mapping.
 
@@ -143,8 +159,12 b' 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.finish_shutdown(kid)
+            self.cleanup(kid)
 
     @kernel_method
     def interrupt_kernel(self, kernel_id):