IPython/kernel/multikernelmanager.py

"""A kernel manager for multiple kernels

Authors:

* Brian Granger
"""

#-----------------------------------------------------------------------------
#  Copyright (C) 2013  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

from __future__ import absolute_import

import os
import uuid

import zmq
from zmq.eventloop.zmqstream import ZMQStream

from IPython.config.configurable import LoggingConfigurable
from IPython.utils.importstring import import_item
from IPython.utils.traitlets import (
    Instance, Dict, Unicode, Any, DottedObjectName, Bool
)

#-----------------------------------------------------------------------------
# Classes
#-----------------------------------------------------------------------------

class DuplicateKernelError(Exception):
    pass


class MultiKernelManager(LoggingConfigurable):
    """A class for managing multiple kernels."""
    
    kernel_manager_class = DottedObjectName(
        "IPython.kernel.ioloop.IOLoopKernelManager", config=True,
        help="""The kernel manager class.  This is configurable to allow
        subclassing of the KernelManager for customized behavior.
        """
    )
    def _kernel_manager_class_changed(self, name, old, new):
        self.kernel_manager_factory = import_item(new)
    
    kernel_manager_factory = Any(help="this is kernel_manager_class after import")
    def _kernel_manager_factory_default(self):
        return import_item(self.kernel_manager_class)
    
    context = Instance('zmq.Context')
    def _context_default(self):
        return zmq.Context.instance()

    connection_dir = Unicode('')

    _kernels = Dict()

    def list_kernel_ids(self):
        """Return a list of the kernel ids of the active kernels."""
        # Create a copy so we can iterate over kernels in operations
        # that delete keys.
        return list(self._kernels.keys())

    def __len__(self):
        """Return the number of running kernels."""
        return len(self.list_kernel_ids())

    def __contains__(self, kernel_id):
        return kernel_id in self._kernels

    def start_kernel(self, **kwargs):
        """Start a new kernel.

        The caller can pick a kernel_id by passing one in as a keyword arg,
        otherwise one will be picked using a uuid.

        To silence the kernel's stdout/stderr, call this using::

            km.start_kernel(stdout=PIPE, stderr=PIPE)

        """
        kernel_id = kwargs.pop('kernel_id', unicode(uuid.uuid4()))
        if kernel_id in self:
            raise DuplicateKernelError('Kernel already exists: %s' % kernel_id)
        # kernel_manager_factory is the constructor for the KernelManager
        # subclass we are using. It can be configured as any Configurable,
        # including things like its transport and ip.
        km = self.kernel_manager_factory(connection_file=os.path.join(
                    self.connection_dir, "kernel-%s.json" % kernel_id),
                    config=self.config, autorestart=True, log=self.log
        )
        km.start_kernel(**kwargs)
        self._kernels[kernel_id] = km
        return kernel_id

    def shutdown_kernel(self, kernel_id, now=False):
        """Shutdown a kernel by its kernel uuid.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to shutdown.
        now : bool
            Should the kernel be shutdown forcibly using a signal.
        """
        k = self.get_kernel(kernel_id)
        k.shutdown_kernel(now=now)
        del self._kernels[kernel_id]

    def shutdown_all(self, now=False):
        """Shutdown all kernels."""
        for kid in self.list_kernel_ids():
            self.shutdown_kernel(kid, now=now)

    def interrupt_kernel(self, kernel_id):
        """Interrupt (SIGINT) the kernel by its uuid.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to interrupt.
        """
        return self.get_kernel(kernel_id).interrupt_kernel()

    def signal_kernel(self, kernel_id, signum):
        """Sends a signal to the kernel by its uuid.

        Note that since only SIGTERM is supported on Windows, this function
        is only useful on Unix systems.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to signal.
        """
        return self.get_kernel(kernel_id).signal_kernel(signum)

    def restart_kernel(self, kernel_id):
        """Restart a kernel by its uuid, keeping the same ports.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel to interrupt.
        """
        km = self.get_kernel(kernel_id)
        km.restart_kernel()

    def is_alive(self, kernel_id):
        """Is the kernel alive.

        This calls KernelManager.is_alive() which calls Popen.poll on the
        actual kernel subprocess.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.
        """
        return self.get_kernel(kernel_id).is_alive()

    def get_kernel(self, kernel_id):
        """Get the single KernelManager object for a kernel by its uuid.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.
        """
        km = self._kernels.get(kernel_id)
        if km is not None:
            return km
        else:
            raise KeyError("Kernel with id not found: %s" % kernel_id)

    def get_connection_info(self, kernel_id):
        """Return a dictionary of connection data for a kernel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.

        Returns
        =======
        connection_dict : dict
            A dict of the information needed to connect to a kernel.
            This includes the ip address and the integer port
            numbers of the different channels (stdin_port, iopub_port,
            shell_port, hb_port).
        """
        km = self.get_kernel(kernel_id)
        return dict(transport=km.transport,
                    ip=km.ip,
                    shell_port=km.shell_port,
                    iopub_port=km.iopub_port,
                    stdin_port=km.stdin_port,
                    hb_port=km.hb_port,
        )  

    def _make_url(self, transport, ip, port):
        """Make a ZeroMQ URL for a given transport, ip and port."""
        if transport == 'tcp':
            return "tcp://%s:%i" % (ip, port)
        else:
            return "%s://%s-%s" % (transport, ip, port)

    def _create_connected_stream(self, kernel_id, socket_type, channel):
        """Create a connected ZMQStream for a kernel."""
        cinfo = self.get_connection_info(kernel_id)
        url = self._make_url(cinfo['transport'], cinfo['ip'],
                cinfo['%s_port' % channel]
        )
        sock = self.context.socket(socket_type)
        self.log.info("Connecting to: %s" % url)
        sock.connect(url)
        return ZMQStream(sock)

    def create_iopub_stream(self, kernel_id):
        """Return a ZMQStream object connected to the iopub channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.

        Returns
        =======
        stream : ZMQStream
        """
        iopub_stream = self._create_connected_stream(kernel_id, zmq.SUB, 'iopub')
        iopub_stream.socket.setsockopt(zmq.SUBSCRIBE, b'')
        return iopub_stream

    def create_shell_stream(self, kernel_id):
        """Return a ZMQStream object connected to the shell channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.

        Returns
        =======
        stream : ZMQStream
        """
        shell_stream = self._create_connected_stream(kernel_id, zmq.DEALER, 'shell')
        return shell_stream

    def create_hb_stream(self, kernel_id):
        """Return a ZMQStream object connected to the hb channel.

        Parameters
        ==========
        kernel_id : uuid
            The id of the kernel.

        Returns
        =======
        stream : ZMQStream
        """
        hb_stream = self._create_connected_stream(kernel_id, zmq.REQ, 'hb')
        return hb_stream