Skip to content

Commit

Permalink
Add support for async kernel starts
Browse files Browse the repository at this point in the history
  • Loading branch information
kevin-bates committed Mar 7, 2019
1 parent 680fda9 commit 55483d7
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 15 deletions.
62 changes: 55 additions & 7 deletions jupyter_client/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import time

import zmq

from tornado import gen
from ipython_genutils.importstring import import_item
from .localinterfaces import is_local_ip, local_ips
from traitlets import (
Expand Down Expand Up @@ -192,6 +192,15 @@ def _launch_kernel(self, kernel_cmd, **kw):
"""
return launch_kernel(kernel_cmd, **kw)

@gen.coroutine
def launch_kernel_async(self, kernel_cmd, **kw):
"""actually launch the kernel
override in a subclass to launch kernel subprocesses differently
"""
res = yield gen.maybe_future(launch_kernel(kernel_cmd, **kw))
raise gen.Return(res)

# Control socket used for polite kernel shutdown

def _connect_control_socket(self):
Expand All @@ -205,8 +214,8 @@ def _close_control_socket(self):
self._control_socket.close()
self._control_socket = None

def start_kernel(self, **kw):
"""Starts a kernel on this host in a separate process.
def pre_start_kernel(self, **kw):
"""Prepares a kernel for startup in a separate process.
If random ports (port=0) are being used, this method must be called
before the channels are created.
Expand Down Expand Up @@ -243,14 +252,53 @@ def start_kernel(self, **kw):
env.update(self.kernel_spec.env or {})
elif self.extra_env:
env.update(self.extra_env)
kw['env'] = env

# launch the kernel subprocess
self.log.debug("Starting kernel: %s", kernel_cmd)
self.kernel = self._launch_kernel(kernel_cmd, env=env,
**kw)
return kernel_cmd, kw

def post_start_kernel(self, **kw):
self.start_restarter()
self._connect_control_socket()

def start_kernel(self, **kw):
"""Starts a kernel on this host in a separate process.
If random ports (port=0) are being used, this method must be called
before the channels are created.
Parameters
----------
`**kw` : optional
keyword arguments that are passed down to build the kernel_cmd
and launching the kernel (e.g. Popen kwargs).
"""
kernel_cmd, kw = self.pre_start_kernel(**kw)

# launch the kernel subprocess
self.log.debug("Starting kernel: %s", kernel_cmd)
self.kernel = self._launch_kernel(kernel_cmd, **kw)
self.post_start_kernel(**kw)

@gen.coroutine
def start_kernel_async(self, **kw):
"""Starts a kernel in a separate process in an asynchronous manner.
If random ports (port=0) are being used, this method must be called
before the channels are created.
Parameters
----------
`**kw` : optional
keyword arguments that are passed down to build the kernel_cmd
and launching the kernel (e.g. Popen kwargs).
"""
kernel_cmd, kw = self.pre_start_kernel(**kw)

# launch the kernel subprocess
self.log.debug("Starting kernel (async): %s", kernel_cmd)
self.kernel = yield self.launch_kernel_async(kernel_cmd, **kw)
self.post_start_kernel(**kw)

def request_shutdown(self, restart=False):
"""Send a shutdown request via control channel
"""
Expand Down
36 changes: 28 additions & 8 deletions jupyter_client/multikernelmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@

import zmq

from tornado import gen
from traitlets.config.configurable import LoggingConfigurable
from ipython_genutils.importstring import import_item
from traitlets import (
Expand Down Expand Up @@ -82,14 +83,8 @@ def __len__(self):
def __contains__(self, kernel_id):
return kernel_id in self._kernels

def start_kernel(self, kernel_name=None, **kwargs):
"""Start a new kernel.
The caller can pick a kernel_id by passing one in as a keyword arg,
otherwise one will be generated using new_kernel_id().
The kernel ID for the newly started kernel is returned.
"""
def pre_start_kernel(self, kernel_name=None, **kwargs):
"""Prepare kernel manager for startup """
kernel_id = kwargs.pop('kernel_id', self.new_kernel_id(**kwargs))
if kernel_id in self:
raise DuplicateKernelError('Kernel already exists: %s' % kernel_id)
Expand All @@ -107,10 +102,35 @@ def start_kernel(self, kernel_name=None, **kwargs):
parent=self, log=self.log, kernel_name=kernel_name,
**constructor_kwargs
)
return kernel_id, km

def start_kernel(self, kernel_name=None, **kwargs):
"""Start a new kernel.
The caller can pick a kernel_id by passing one in as a keyword arg,
otherwise one will be generated using new_kernel_id().
The kernel ID for the newly started kernel is returned.
"""
kernel_id, km = self.pre_start_kernel(kernel_name, **kwargs)
km.start_kernel(**kwargs)
self._kernels[kernel_id] = km
return kernel_id

@gen.coroutine
def start_kernel_async(self, kernel_name=None, **kwargs):
"""Start a new kernel asynchronously.
The caller can pick a kernel_id by passing one in as a keyword arg,
otherwise one will be generated using new_kernel_id().
The kernel ID for the newly started kernel is returned.
"""
kernel_id, km = self.pre_start_kernel(kernel_name, **kwargs)
yield km.start_kernel_async(**kwargs)
self._kernels[kernel_id] = km
raise gen.Return(kernel_id)

@kernel_method
def shutdown_kernel(self, kernel_id, now=False, restart=False):
"""Shutdown a kernel by its kernel uuid.
Expand Down

0 comments on commit 55483d7

Please sign in to comment.