IPython/zmq/parallel/kernelstarter.py

"""KernelStarter class that intercepts Control Queue messages, and handles process management."""

from zmq.eventloop import ioloop

from .streamsession import StreamSession

class KernelStarter(object):
    """Object for resetting/killing the Kernel."""
    
    
    def __init__(self, session, upstream, downstream, *kernel_args, **kernel_kwargs):
        self.session = session
        self.upstream = upstream
        self.downstream = downstream
        self.kernel_args = kernel_args
        self.kernel_kwargs = kernel_kwargs
        self.handlers = {}
        for method in 'shutdown_request shutdown_reply'.split():
            self.handlers[method] = getattr(self, method)
    
    def start(self):
        self.upstream.on_recv(self.dispatch_request)
        self.downstream.on_recv(self.dispatch_reply)
        
    #--------------------------------------------------------------------------
    # Dispatch methods
    #--------------------------------------------------------------------------
    
    def dispatch_request(self, raw_msg):
        idents, msg = self.session.feed_identities()
        try:
            msg = self.session.unpack_message(msg, content=False)
        except:
            print ("bad msg: %s"%msg)
        
        msgtype = msg['msg_type']
        handler = self.handlers.get(msgtype, None)
        if handler is None:
            self.downstream.send_multipart(raw_msg, copy=False)
        else:
            handler(msg)
        
    def dispatch_reply(self, raw_msg):
        idents, msg = self.session.feed_identities()
        try:
            msg = self.session.unpack_message(msg, content=False)
        except:
            print ("bad msg: %s"%msg)
        
        msgtype = msg['msg_type']
        handler = self.handlers.get(msgtype, None)
        if handler is None:
            self.upstream.send_multipart(raw_msg, copy=False)
        else:
            handler(msg)
    
    #--------------------------------------------------------------------------
    # Handlers
    #--------------------------------------------------------------------------
    
    def shutdown_request(self, msg):
        """"""
        self.downstream.send_multipart(msg)
    
    #--------------------------------------------------------------------------
    # Kernel process management methods, from KernelManager:
    #--------------------------------------------------------------------------
    
    def _check_local(addr):
        if isinstance(addr, tuple):
            addr = addr[0]
        return addr in LOCAL_IPS
    
    def start_kernel(self, **kw):
        """Starts a kernel process and configures the manager to use it.

        If random ports (port=0) are being used, this method must be called
        before the channels are created.

        Parameters:
        -----------
        ipython : bool, optional (default True)
             Whether to use an IPython kernel instead of a plain Python kernel.
        """
        self.kernel = Process(target=make_kernel, args=self.kernel_args,
                                            kwargs=self.kernel_kwargs)

    def shutdown_kernel(self, restart=False):
        """ Attempts to the stop the kernel process cleanly. If the kernel
        cannot be stopped, it is killed, if possible.
        """
        # FIXME: Shutdown does not work on Windows due to ZMQ errors!
        if sys.platform == 'win32':
            self.kill_kernel()
            return

        # 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.xreq_channel.shutdown(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()
    
    def restart_kernel(self, now=False):
        """Restarts a kernel with the same arguments that were used to launch
        it. If the old kernel was launched with random ports, the same ports
        will be used for the new kernel.

        Parameters
        ----------
        now : bool, optional
          If True, the kernel is forcefully restarted *immediately*, without
          having a chance to do any cleanup action.  Otherwise the kernel is
          given 1s to clean up before a forceful restart is issued.

          In all cases the kernel is restarted, the only difference is whether
          it is given a chance to perform a clean shutdown or not.
        """
        if self._launch_args is None:
            raise RuntimeError("Cannot restart the kernel. "
                               "No previous call to 'start_kernel'.")
        else:
            if self.has_kernel:
                if now:
                    self.kill_kernel()
                else:
                    self.shutdown_kernel(restart=True)
            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):
        """Returns whether a kernel process has been specified for the kernel
        manager.
        """
        return self.kernel is not None

    def kill_kernel(self):
        """ Kill the running kernel. """
        if self.has_kernel:
            # Pause the heart beat channel if it exists.
            if self._hb_channel is not None:
                self._hb_channel.pause()

            # Attempt to kill the kernel.
            try:
                self.kernel.kill()
            except OSError, e:
                # In Windows, we will get an Access Denied error if the process
                # has already terminated. Ignore it.
                if not (sys.platform == 'win32' and e.winerror == 5):
                    raise
            self.kernel = None
        else:
            raise RuntimeError("Cannot kill kernel. No kernel is running!")

    def interrupt_kernel(self):
        """ Interrupts the kernel. Unlike ``signal_kernel``, this operation is
        well supported on all platforms.
        """
        if self.has_kernel:
            if sys.platform == 'win32':
                from parentpoller import ParentPollerWindows as Poller
                Poller.send_interrupt(self.kernel.win32_interrupt_event)
            else:
                self.kernel.send_signal(signal.SIGINT)
        else:
            raise RuntimeError("Cannot interrupt kernel. No kernel is running!")

    def signal_kernel(self, signum):
        """ Sends a signal to the kernel. Note that since only SIGTERM is
        supported on Windows, this function is only useful on Unix systems.
        """
        if self.has_kernel:
            self.kernel.send_signal(signum)
        else:
            raise RuntimeError("Cannot signal kernel. No kernel is running!")

    @property
    def is_alive(self):
        """Is the kernel process still running?"""
        # FIXME: not using a heartbeat means this method is broken for any
        # remote kernel, it's only capable of handling local kernels.
        if self.has_kernel:
            if self.kernel.poll() is None:
                return True
            else:
                return False
        else:
            # We didn't start the kernel with this KernelManager so we don't
            # know if it is running. We should use a heartbeat for this case.
            return True


def make_starter(up_addr, down_addr, *args, **kwargs):
    """entry point function for launching a kernelstarter in a subprocess"""
    loop = ioloop.IOLoop.instance()
    ctx = zmq.Context()
    session = StreamSession()
    upstream = zmqstream.ZMQStream(ctx.socket(zmq.XREQ),loop)
    upstream.connect(up_addr)
    downstream = zmqstream.ZMQStream(ctx.socket(zmq.XREQ),loop)
    downstream.connect(down_addr)
    
    starter = KernelStarter(session, upstream, downstream, *args, **kwargs)
    starter.start()
    loop.start()