Ability to delay shutdown

[rgbkrk]

We're having an issue where cleanup of subprocesses of our kernel is taking longer than shutdown_wait_time, a configurable from within the AsyncKernelManager. What are some ways we can delay shutdown or possibly some kernel messages we could send to let papermill (or nbclient) know we're still working on shutdown?

[MSeal]

So most nbclient configurations are passable from papermill's execute_notebook interface but the kernel_manager's config is not exposed on the nbclient __init__ interface so those nested config attributes are not modifiable directly from the python interface. What you can do is set the traitlet value for the class directly before running papermill:

import papermill
from jupyter_client import KernelManager
KernelManager.shutdown_wait_time = 60
# The shutdown_wait_time is now 60 seconds
papermill.execute_notebook('Untitled.ipynb', 'Untitled-out.ipynb');

I looked into if a config file could be used for the jupyter_client traits, but unfortunately this isn't loaded by default. A nice option for supporting jupyter's config system may be adding an option to nbclient for accepting a config_file path to load like the Jupyter application does for the notebook server.

[MSeal]

If there were done then we could send that along to nbclient:

papermill.execute_notebook('Untitled.ipynb', 'Untitled-out.ipynb', config_file='./path/to/my/config.py);

or from the cli something like:

papermill Untitled.ipynb Untitled-out.ipynb --config_file='./path/to/my/config.py'

[MSeal]

As to supporting kernel messages to delay shutdown, that could work but it'd probably be more at the jupyter protocol contract level to support a new message pattern to reset the shutdown timeout. It's possible this only applies to nbclient's use of headless execution but I'm not entirely sure. I'd be open to suggestions, but that should probably be a separate issue on jupyter_client where other folks working on that layer might be able to weigh in.

[rgbkrk]

As to supporting kernel messages to delay shutdown, that could work but it'd probably be more at the jupyter protocol contract level to support a new message pattern to reset the shutdown timeout. It's possible this only applies to nbclient's use of headless execution but I'm not entirely sure. I'd be open to suggestions, but that should probably be a separate issue on jupyter_client where other folks working on that layer might be able to weigh in.

I'm definitely open to a wider conversation on some messaging for this, so we don't have to keep pushing timeouts further along in the future when we could be sending some sort of keep alive in between.

[rgbkrk]

Looks like the other way to do this is to override execute_managed_notebook in a custom Engine.

class AsyncKernelManager(jupyter_client.AsyncKernelManager):
    shutdown_wait_time = 600

class CustomEngine(NBClientEngine):
    @classmethod
    def execute_managed_notebook(cls, *args, **kwargs):
        kwargs['kernel_manager_class'] = AsyncKernelManager
        return super().execute_managed_notebook(*args, **kwargs)

What do you think of this approach?

[MSeal]

That would also work. I started with posting this originally (which does the same thing without a new engine, but doesn't work from the CLI):

class LongerWaitAsyncKernelManager(jupyter_client.AsyncKernelManager):
    shutdown_wait_time = 600

papermill.execute_notebook('Untitled.ipynb', 'Untitled-out.ipynb', kernel_manager_class=LongerWaitAsyncKernelManager);

[rgbkrk]

The other thing I'm thinking of for future work is to use the pick kernel and create a magic that can be run in the last cell to do the cleanup.

[MSeal]

Registering a cleanup cell with a magic sounds useful