Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow the execute preprocessor to make use of an existing kernel #852

Merged
merged 1 commit into from
Aug 9, 2018
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 50 additions & 24 deletions nbconvert/preprocessors/execute.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ class ExecutePreprocessor(Preprocessor):
"""
)
).tag(config=True)

@default('kernel_name')
def _kernel_name_default(self):
try:
Expand All @@ -167,7 +167,6 @@ def _kernel_name_default(self):
'the ExecutePreprocessor and you have not set '
'self.nb to be able to use that to infer the '
'kernel_name.')


raise_on_iopub_timeout = Bool(False,
help=dedent(
Expand Down Expand Up @@ -252,7 +251,7 @@ def start_new_kernel(self, **kwargs):
km = self.kernel_manager_class(kernel_name=self.kernel_name,
config=self.config)
km.start_kernel(extra_arguments=self.extra_arguments, **kwargs)

kc = km.client()
kc.start_channels()
try:
Expand All @@ -263,18 +262,18 @@ def start_new_kernel(self, **kwargs):
raise
kc.allow_stdin = False
return km, kc

@contextmanager
def setup_preprocessor(self, nb, resources):
def setup_preprocessor(self, nb, resources, km=None):
"""
Context manager for setting up the class to execute a notebook.

The assigns `nb` to `self.nb` where it will be modified in-place. It also creates
The assigns `nb` to `self.nb` where it will be modified in-place. It also creates
and assigns the Kernel Manager (`self.km`) and Kernel Client(`self.kc`).

It is intended to yield to a block that will execute codeself.
When control returns from the yield it stops the client's zmq channels, shuts

When control returns from the yield it stops the client's zmq channels, shuts
down the kernel, and removes the now unused attributes.

Parameters
Expand All @@ -285,6 +284,9 @@ def setup_preprocessor(self, nb, resources):
Additional resources used in the conversion process. For example,
passing ``{'metadata': {'path': run_path}}`` sets the
execution path to ``run_path``.
km : KernerlManager (optional)
Optional kernel manaher. If none is provided, a kernel manager will
be created.

Returns
-------
Expand All @@ -298,18 +300,37 @@ def setup_preprocessor(self, nb, resources):
# clear display_id map
self._display_id_map = {}

self.km, self.kc = self.start_new_kernel(cwd=path)
try:
# Yielding unbound args for more easier understanding and downstream consumption
yield nb, self.km, self.kc
finally:
self.kc.stop_channels()
self.km.shutdown_kernel(now=self.shutdown_kernel == 'immediate')

for attr in ['nb', 'km', 'kc']:
delattr(self, attr)

def preprocess(self, nb, resources):
if km is None:
self.km, self.kc = self.start_new_kernel(cwd=path)
try:
# Yielding unbound args for more easier understanding and downstream consumption
yield nb, self.km, self.kc
finally:
self.kc.stop_channels()
self.km.shutdown_kernel(now=self.shutdown_kernel == 'immediate')

for attr in ['nb', 'km', 'kc']:
delattr(self, attr)
else:
self.km = km
if not km.has_kernel:
km.start_kernel(extra_arguments=self.extra_arguments, **kwargs)
self.kc = km.client()
if not self.kc.channels_running:
self.kc.start_channels()
try:
self.kc.wait_for_ready(timeout=self.startup_timeout)
except RuntimeError:
self.kc.stop_channels()
raise
self.kc.allow_stdin = False
try:
yield nb, self.km, self.kc
finally:
for attr in ['nb', 'km', 'kc']:
delattr(self, attr)

def preprocess(self, nb, resources, km=None):
"""
Preprocess notebook executing each code cell.

Expand All @@ -323,6 +344,9 @@ def preprocess(self, nb, resources):
Additional resources used in the conversion process. For example,
passing ``{'metadata': {'path': run_path}}`` sets the
execution path to ``run_path``.
km: KernelManager (optional)
Optional kernel manager. If none is provided, a kernel manager will
be created.

Returns
-------
Expand All @@ -332,7 +356,7 @@ def preprocess(self, nb, resources):
Additional resources used in the conversion process.
"""

with self.setup_preprocessor(nb, resources):
with self.setup_preprocessor(nb, resources, km=km):
self.log.info("Executing notebook with kernel: %s" % self.kernel_name)
nb, resources = super(ExecutePreprocessor, self).preprocess(nb, resources)

Expand Down Expand Up @@ -491,7 +515,7 @@ def run_cell(self, cell, cell_index=0):
return exec_reply, outs


def executenb(nb, cwd=None, **kwargs):
def executenb(nb, cwd=None, km=None, **kwargs):
"""Execute a notebook's code, updating outputs within the notebook object.

This is a convenient wrapper around ExecutePreprocessor. It returns the
Expand All @@ -503,11 +527,13 @@ def executenb(nb, cwd=None, **kwargs):
The notebook object to be executed
cwd : str, optional
If supplied, the kernel will run in this directory
km : KernelManager, optional
If supplied, the specified kernel manager will be used for code execution.
kwargs :
Any other options for ExecutePreprocessor, e.g. timeout, kernel_name
"""
resources = {}
if cwd is not None:
resources['metadata'] = {'path': cwd}
ep = ExecutePreprocessor(**kwargs)
return ep.preprocess(nb, resources)[0]
return ep.preprocess(nb, resources, km=km)[0]