first comit
This commit is contained in:
806
venv/lib/python3.10/site-packages/jupyter_client/manager.py
Normal file
806
venv/lib/python3.10/site-packages/jupyter_client/manager.py
Normal file
@@ -0,0 +1,806 @@
|
||||
"""Base class to manage a running kernel"""
|
||||
# Copyright (c) Jupyter Development Team.
|
||||
# Distributed under the terms of the Modified BSD License.
|
||||
import asyncio
|
||||
import functools
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import sys
|
||||
import typing as t
|
||||
import uuid
|
||||
import warnings
|
||||
from asyncio.futures import Future
|
||||
from concurrent.futures import Future as CFuture
|
||||
from contextlib import contextmanager
|
||||
from enum import Enum
|
||||
|
||||
import zmq
|
||||
from jupyter_core.utils import run_sync
|
||||
from traitlets import (
|
||||
Any,
|
||||
Bool,
|
||||
Dict,
|
||||
DottedObjectName,
|
||||
Float,
|
||||
Instance,
|
||||
Type,
|
||||
Unicode,
|
||||
default,
|
||||
observe,
|
||||
observe_compat,
|
||||
)
|
||||
from traitlets.utils.importstring import import_item
|
||||
|
||||
from . import kernelspec
|
||||
from .asynchronous import AsyncKernelClient
|
||||
from .blocking import BlockingKernelClient
|
||||
from .client import KernelClient
|
||||
from .connect import ConnectionFileMixin
|
||||
from .managerabc import KernelManagerABC
|
||||
from .provisioning import KernelProvisionerBase
|
||||
from .provisioning import KernelProvisionerFactory as KPF # noqa
|
||||
|
||||
|
||||
class _ShutdownStatus(Enum):
|
||||
"""
|
||||
|
||||
This is so far used only for testing in order to track the internal state of
|
||||
the shutdown logic, and verifying which path is taken for which
|
||||
missbehavior.
|
||||
|
||||
"""
|
||||
|
||||
Unset = None
|
||||
ShutdownRequest = "ShutdownRequest"
|
||||
SigtermRequest = "SigtermRequest"
|
||||
SigkillRequest = "SigkillRequest"
|
||||
|
||||
|
||||
F = t.TypeVar("F", bound=t.Callable[..., t.Any])
|
||||
|
||||
|
||||
def _get_future() -> t.Union[Future, CFuture]:
|
||||
"""Get an appropriate Future object"""
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
return Future()
|
||||
except RuntimeError:
|
||||
# No event loop running, use concurrent future
|
||||
return CFuture()
|
||||
|
||||
|
||||
def in_pending_state(method: F) -> F:
|
||||
"""Sets the kernel to a pending state by
|
||||
creating a fresh Future for the KernelManager's `ready`
|
||||
attribute. Once the method is finished, set the Future's results.
|
||||
"""
|
||||
|
||||
@t.no_type_check
|
||||
@functools.wraps(method)
|
||||
async def wrapper(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
|
||||
"""Create a future for the decorated method."""
|
||||
if self._attempted_start or not self._ready:
|
||||
self._ready = _get_future()
|
||||
try:
|
||||
# call wrapped method, await, and set the result or exception.
|
||||
out = await method(self, *args, **kwargs)
|
||||
# Add a small sleep to ensure tests can capture the state before done
|
||||
await asyncio.sleep(0.01)
|
||||
if self.owns_kernel:
|
||||
self._ready.set_result(None)
|
||||
return out
|
||||
except Exception as e:
|
||||
self._ready.set_exception(e)
|
||||
self.log.exception(self._ready.exception())
|
||||
raise e
|
||||
|
||||
return t.cast(F, wrapper)
|
||||
|
||||
|
||||
class KernelManager(ConnectionFileMixin):
|
||||
"""Manages a single kernel in a subprocess on this host.
|
||||
|
||||
This version starts kernels with Popen.
|
||||
"""
|
||||
|
||||
_ready: t.Optional[t.Union[Future, CFuture]]
|
||||
|
||||
def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
|
||||
"""Initialize a kernel manager."""
|
||||
if args:
|
||||
warnings.warn(
|
||||
"Passing positional only arguments to "
|
||||
"`KernelManager.__init__` is deprecated since jupyter_client"
|
||||
" 8.6, and will become an error on future versions. Positional "
|
||||
" arguments have been ignored since jupyter_client 7.0",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
self._owns_kernel = kwargs.pop("owns_kernel", True)
|
||||
super().__init__(**kwargs)
|
||||
self._shutdown_status = _ShutdownStatus.Unset
|
||||
self._attempted_start = False
|
||||
self._ready = None
|
||||
|
||||
_created_context: Bool = Bool(False)
|
||||
|
||||
# The PyZMQ Context to use for communication with the kernel.
|
||||
context: Instance = Instance(zmq.Context)
|
||||
|
||||
@default("context")
|
||||
def _context_default(self) -> zmq.Context:
|
||||
self._created_context = True
|
||||
return zmq.Context()
|
||||
|
||||
# the class to create with our `client` method
|
||||
client_class: DottedObjectName = DottedObjectName(
|
||||
"jupyter_client.blocking.BlockingKernelClient"
|
||||
)
|
||||
client_factory: Type = Type(klass=KernelClient)
|
||||
|
||||
@default("client_factory")
|
||||
def _client_factory_default(self) -> Type:
|
||||
return import_item(self.client_class)
|
||||
|
||||
@observe("client_class")
|
||||
def _client_class_changed(self, change: t.Dict[str, DottedObjectName]) -> None:
|
||||
self.client_factory = import_item(str(change["new"]))
|
||||
|
||||
kernel_id: t.Union[str, Unicode] = Unicode(None, allow_none=True)
|
||||
|
||||
# The kernel provisioner with which this KernelManager is communicating.
|
||||
# This will generally be a LocalProvisioner instance unless the kernelspec
|
||||
# indicates otherwise.
|
||||
provisioner: t.Optional[KernelProvisionerBase] = None
|
||||
|
||||
kernel_spec_manager: Instance = Instance(kernelspec.KernelSpecManager)
|
||||
|
||||
@default("kernel_spec_manager")
|
||||
def _kernel_spec_manager_default(self) -> kernelspec.KernelSpecManager:
|
||||
return kernelspec.KernelSpecManager(data_dir=self.data_dir)
|
||||
|
||||
@observe("kernel_spec_manager")
|
||||
@observe_compat
|
||||
def _kernel_spec_manager_changed(self, change: t.Dict[str, Instance]) -> None:
|
||||
self._kernel_spec = None
|
||||
|
||||
shutdown_wait_time: Float = Float(
|
||||
5.0,
|
||||
config=True,
|
||||
help="Time to wait for a kernel to terminate before killing it, "
|
||||
"in seconds. When a shutdown request is initiated, the kernel "
|
||||
"will be immediately sent an interrupt (SIGINT), followed"
|
||||
"by a shutdown_request message, after 1/2 of `shutdown_wait_time`"
|
||||
"it will be sent a terminate (SIGTERM) request, and finally at "
|
||||
"the end of `shutdown_wait_time` will be killed (SIGKILL). terminate "
|
||||
"and kill may be equivalent on windows. Note that this value can be"
|
||||
"overridden by the in-use kernel provisioner since shutdown times may"
|
||||
"vary by provisioned environment.",
|
||||
)
|
||||
|
||||
kernel_name: t.Union[str, Unicode] = Unicode(kernelspec.NATIVE_KERNEL_NAME)
|
||||
|
||||
@observe("kernel_name")
|
||||
def _kernel_name_changed(self, change: t.Dict[str, str]) -> None:
|
||||
self._kernel_spec = None
|
||||
if change["new"] == "python":
|
||||
self.kernel_name = kernelspec.NATIVE_KERNEL_NAME
|
||||
|
||||
_kernel_spec: t.Optional[kernelspec.KernelSpec] = None
|
||||
|
||||
@property
|
||||
def kernel_spec(self) -> t.Optional[kernelspec.KernelSpec]:
|
||||
if self._kernel_spec is None and self.kernel_name != "":
|
||||
self._kernel_spec = self.kernel_spec_manager.get_kernel_spec(self.kernel_name)
|
||||
return self._kernel_spec
|
||||
|
||||
cache_ports: Bool = Bool(
|
||||
False,
|
||||
config=True,
|
||||
help="True if the MultiKernelManager should cache ports for this KernelManager instance",
|
||||
)
|
||||
|
||||
@default("cache_ports")
|
||||
def _default_cache_ports(self) -> bool:
|
||||
return self.transport == "tcp"
|
||||
|
||||
@property
|
||||
def ready(self) -> t.Union[CFuture, Future]:
|
||||
"""A future that resolves when the kernel process has started for the first time"""
|
||||
if not self._ready:
|
||||
self._ready = _get_future()
|
||||
return self._ready
|
||||
|
||||
@property
|
||||
def ipykernel(self) -> bool:
|
||||
return self.kernel_name in {"python", "python2", "python3"}
|
||||
|
||||
# Protected traits
|
||||
_launch_args: t.Optional["Dict[str, Any]"] = Dict(allow_none=True)
|
||||
_control_socket: Any = Any()
|
||||
|
||||
_restarter: Any = Any()
|
||||
|
||||
autorestart: Bool = Bool(
|
||||
True, config=True, help="""Should we autorestart the kernel if it dies."""
|
||||
)
|
||||
|
||||
shutting_down: bool = False
|
||||
|
||||
def __del__(self) -> None:
|
||||
self._close_control_socket()
|
||||
self.cleanup_connection_file()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Kernel restarter
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
def start_restarter(self) -> None:
|
||||
"""Start the kernel restarter."""
|
||||
pass
|
||||
|
||||
def stop_restarter(self) -> None:
|
||||
"""Stop the kernel restarter."""
|
||||
pass
|
||||
|
||||
def add_restart_callback(self, callback: t.Callable, event: str = "restart") -> None:
|
||||
"""Register a callback to be called when a kernel is restarted"""
|
||||
if self._restarter is None:
|
||||
return
|
||||
self._restarter.add_callback(callback, event)
|
||||
|
||||
def remove_restart_callback(self, callback: t.Callable, event: str = "restart") -> None:
|
||||
"""Unregister a callback to be called when a kernel is restarted"""
|
||||
if self._restarter is None:
|
||||
return
|
||||
self._restarter.remove_callback(callback, event)
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# create a Client connected to our Kernel
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
def client(self, **kwargs: t.Any) -> BlockingKernelClient:
|
||||
"""Create a client configured to connect to our kernel"""
|
||||
kw: dict = {}
|
||||
kw.update(self.get_connection_info(session=True))
|
||||
kw.update(
|
||||
{
|
||||
"connection_file": self.connection_file,
|
||||
"parent": self,
|
||||
}
|
||||
)
|
||||
|
||||
# add kwargs last, for manual overrides
|
||||
kw.update(kwargs)
|
||||
return self.client_factory(**kw)
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# Kernel management
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
def update_env(self, *, env: t.Dict[str, str]) -> None:
|
||||
"""
|
||||
Allow to update the environment of a kernel manager.
|
||||
|
||||
This will take effect only after kernel restart when the new env is
|
||||
passed to the new kernel.
|
||||
|
||||
This is useful as some of the information of the current kernel reflect
|
||||
the state of the session that started it, and those session information
|
||||
(like the attach file path, or name), are mutable.
|
||||
|
||||
.. version-added: 8.5
|
||||
"""
|
||||
# Mypy think this is unreachable as it see _launch_args as Dict, not t.Dict
|
||||
if (
|
||||
isinstance(self._launch_args, dict)
|
||||
and "env" in self._launch_args
|
||||
and isinstance(self._launch_args["env"], dict) # type: ignore [unreachable]
|
||||
):
|
||||
self._launch_args["env"].update(env) # type: ignore [unreachable]
|
||||
|
||||
def format_kernel_cmd(self, extra_arguments: t.Optional[t.List[str]] = None) -> t.List[str]:
|
||||
"""Replace templated args (e.g. {connection_file})"""
|
||||
extra_arguments = extra_arguments or []
|
||||
assert self.kernel_spec is not None
|
||||
cmd = self.kernel_spec.argv + extra_arguments
|
||||
|
||||
if cmd and cmd[0] in {
|
||||
"python",
|
||||
"python%i" % sys.version_info[0],
|
||||
"python%i.%i" % sys.version_info[:2],
|
||||
}:
|
||||
# executable is 'python' or 'python3', use sys.executable.
|
||||
# These will typically be the same,
|
||||
# but if the current process is in an env
|
||||
# and has been launched by abspath without
|
||||
# activating the env, python on PATH may not be sys.executable,
|
||||
# but it should be.
|
||||
cmd[0] = sys.executable
|
||||
|
||||
# Make sure to use the realpath for the connection_file
|
||||
# On windows, when running with the store python, the connection_file path
|
||||
# is not usable by non python kernels because the path is being rerouted when
|
||||
# inside of a store app.
|
||||
# See this bug here: https://bugs.python.org/issue41196
|
||||
ns: t.Dict[str, t.Any] = {
|
||||
"connection_file": os.path.realpath(self.connection_file),
|
||||
"prefix": sys.prefix,
|
||||
}
|
||||
|
||||
if self.kernel_spec: # type:ignore[truthy-bool]
|
||||
ns["resource_dir"] = self.kernel_spec.resource_dir
|
||||
assert isinstance(self._launch_args, dict)
|
||||
|
||||
ns.update(self._launch_args)
|
||||
|
||||
pat = re.compile(r"\{([A-Za-z0-9_]+)\}")
|
||||
|
||||
def from_ns(match: t.Any) -> t.Any:
|
||||
"""Get the key out of ns if it's there, otherwise no change."""
|
||||
return ns.get(match.group(1), match.group())
|
||||
|
||||
return [pat.sub(from_ns, arg) for arg in cmd]
|
||||
|
||||
async def _async_launch_kernel(self, kernel_cmd: t.List[str], **kw: t.Any) -> None:
|
||||
"""actually launch the kernel
|
||||
|
||||
override in a subclass to launch kernel subprocesses differently
|
||||
Note that provisioners can now be used to customize kernel environments
|
||||
and
|
||||
"""
|
||||
assert self.provisioner is not None
|
||||
connection_info = await self.provisioner.launch_kernel(kernel_cmd, **kw)
|
||||
assert self.provisioner.has_process
|
||||
# Provisioner provides the connection information. Load into kernel manager
|
||||
# and write the connection file, if not already done.
|
||||
self._reconcile_connection_info(connection_info)
|
||||
|
||||
_launch_kernel = run_sync(_async_launch_kernel)
|
||||
|
||||
# Control socket used for polite kernel shutdown
|
||||
|
||||
def _connect_control_socket(self) -> None:
|
||||
if self._control_socket is None:
|
||||
self._control_socket = self._create_connected_socket("control")
|
||||
self._control_socket.linger = 100
|
||||
|
||||
def _close_control_socket(self) -> None:
|
||||
if self._control_socket is None:
|
||||
return
|
||||
self._control_socket.close()
|
||||
self._control_socket = None
|
||||
|
||||
async def _async_pre_start_kernel(
|
||||
self, **kw: t.Any
|
||||
) -> t.Tuple[t.List[str], t.Dict[str, t.Any]]:
|
||||
"""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.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
`**kw` : optional
|
||||
keyword arguments that are passed down to build the kernel_cmd
|
||||
and launching the kernel (e.g. Popen kwargs).
|
||||
"""
|
||||
self.shutting_down = False
|
||||
self.kernel_id = self.kernel_id or kw.pop("kernel_id", str(uuid.uuid4()))
|
||||
# save kwargs for use in restart
|
||||
# assigning Traitlets Dicts to Dict make mypy unhappy but is ok
|
||||
self._launch_args = kw.copy() # type:ignore [assignment]
|
||||
if self.provisioner is None: # will not be None on restarts
|
||||
self.provisioner = KPF.instance(parent=self.parent).create_provisioner_instance(
|
||||
self.kernel_id,
|
||||
self.kernel_spec,
|
||||
parent=self,
|
||||
)
|
||||
kw = await self.provisioner.pre_launch(**kw)
|
||||
kernel_cmd = kw.pop("cmd")
|
||||
return kernel_cmd, kw
|
||||
|
||||
pre_start_kernel = run_sync(_async_pre_start_kernel)
|
||||
|
||||
async def _async_post_start_kernel(self, **kw: t.Any) -> None:
|
||||
"""Performs any post startup tasks relative to the kernel.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
`**kw` : optional
|
||||
keyword arguments that were used in the kernel process's launch.
|
||||
"""
|
||||
self.start_restarter()
|
||||
self._connect_control_socket()
|
||||
assert self.provisioner is not None
|
||||
await self.provisioner.post_launch(**kw)
|
||||
|
||||
post_start_kernel = run_sync(_async_post_start_kernel)
|
||||
|
||||
@in_pending_state
|
||||
async def _async_start_kernel(self, **kw: t.Any) -> None:
|
||||
"""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).
|
||||
"""
|
||||
self._attempted_start = True
|
||||
kernel_cmd, kw = await self._async_pre_start_kernel(**kw)
|
||||
|
||||
# launch the kernel subprocess
|
||||
self.log.debug("Starting kernel: %s", kernel_cmd)
|
||||
await self._async_launch_kernel(kernel_cmd, **kw)
|
||||
await self._async_post_start_kernel(**kw)
|
||||
|
||||
start_kernel = run_sync(_async_start_kernel)
|
||||
|
||||
async def _async_request_shutdown(self, restart: bool = False) -> None:
|
||||
"""Send a shutdown request via control channel"""
|
||||
content = {"restart": restart}
|
||||
msg = self.session.msg("shutdown_request", content=content)
|
||||
# ensure control socket is connected
|
||||
self._connect_control_socket()
|
||||
self.session.send(self._control_socket, msg)
|
||||
assert self.provisioner is not None
|
||||
await self.provisioner.shutdown_requested(restart=restart)
|
||||
self._shutdown_status = _ShutdownStatus.ShutdownRequest
|
||||
|
||||
request_shutdown = run_sync(_async_request_shutdown)
|
||||
|
||||
async def _async_finish_shutdown(
|
||||
self,
|
||||
waittime: t.Optional[float] = None,
|
||||
pollinterval: float = 0.1,
|
||||
restart: bool = False,
|
||||
) -> None:
|
||||
"""Wait for kernel shutdown, then kill process if it doesn't shutdown.
|
||||
|
||||
This does not send shutdown requests - use :meth:`request_shutdown`
|
||||
first.
|
||||
"""
|
||||
if waittime is None:
|
||||
waittime = max(self.shutdown_wait_time, 0)
|
||||
if self.provisioner: # Allow provisioner to override
|
||||
waittime = self.provisioner.get_shutdown_wait_time(recommended=waittime)
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._async_wait(pollinterval=pollinterval), timeout=waittime / 2
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
self.log.debug("Kernel is taking too long to finish, terminating")
|
||||
self._shutdown_status = _ShutdownStatus.SigtermRequest
|
||||
await self._async_send_kernel_sigterm()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._async_wait(pollinterval=pollinterval), timeout=waittime / 2
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
self.log.debug("Kernel is taking too long to finish, killing")
|
||||
self._shutdown_status = _ShutdownStatus.SigkillRequest
|
||||
await self._async_kill_kernel(restart=restart)
|
||||
else:
|
||||
# Process is no longer alive, wait and clear
|
||||
if self.has_kernel:
|
||||
assert self.provisioner is not None
|
||||
await self.provisioner.wait()
|
||||
|
||||
finish_shutdown = run_sync(_async_finish_shutdown)
|
||||
|
||||
async def _async_cleanup_resources(self, restart: bool = False) -> None:
|
||||
"""Clean up resources when the kernel is shut down"""
|
||||
if not restart:
|
||||
self.cleanup_connection_file()
|
||||
|
||||
self.cleanup_ipc_files()
|
||||
self._close_control_socket()
|
||||
self.session.parent = None
|
||||
|
||||
if self._created_context and not restart:
|
||||
self.context.destroy(linger=100)
|
||||
|
||||
if self.provisioner:
|
||||
await self.provisioner.cleanup(restart=restart)
|
||||
|
||||
cleanup_resources = run_sync(_async_cleanup_resources)
|
||||
|
||||
@in_pending_state
|
||||
async def _async_shutdown_kernel(self, now: bool = False, restart: bool = False) -> None:
|
||||
"""Attempts to stop the kernel process cleanly.
|
||||
|
||||
This attempts to shutdown the kernels cleanly by:
|
||||
|
||||
1. Sending it a shutdown message over the control channel.
|
||||
2. If that fails, the kernel is shutdown forcibly by sending it
|
||||
a signal.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
now : bool
|
||||
Should the kernel be forcible killed *now*. This skips the
|
||||
first, nice shutdown attempt.
|
||||
restart: bool
|
||||
Will this kernel be restarted after it is shutdown. When this
|
||||
is True, connection files will not be cleaned up.
|
||||
"""
|
||||
if not self.owns_kernel:
|
||||
return
|
||||
|
||||
self.shutting_down = True # Used by restarter to prevent race condition
|
||||
# Stop monitoring for restarting while we shutdown.
|
||||
self.stop_restarter()
|
||||
|
||||
if self.has_kernel:
|
||||
await self._async_interrupt_kernel()
|
||||
|
||||
if now:
|
||||
await self._async_kill_kernel()
|
||||
else:
|
||||
await self._async_request_shutdown(restart=restart)
|
||||
# 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.
|
||||
await self._async_finish_shutdown(restart=restart)
|
||||
|
||||
await self._async_cleanup_resources(restart=restart)
|
||||
|
||||
shutdown_kernel = run_sync(_async_shutdown_kernel)
|
||||
|
||||
async def _async_restart_kernel(
|
||||
self, now: bool = False, newports: bool = False, **kw: t.Any
|
||||
) -> None:
|
||||
"""Restarts a kernel with the arguments that were used to launch it.
|
||||
|
||||
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.
|
||||
|
||||
newports : bool, optional
|
||||
If the old kernel was launched with random ports, this flag decides
|
||||
whether the same ports and connection file will be used again.
|
||||
If False, the same ports and connection file are used. This is
|
||||
the default. If True, new random port numbers are chosen and a
|
||||
new connection file is written. It is still possible that the newly
|
||||
chosen random port numbers happen to be the same as the old ones.
|
||||
|
||||
`**kw` : optional
|
||||
Any options specified here will overwrite those used to launch the
|
||||
kernel.
|
||||
"""
|
||||
if self._launch_args is None:
|
||||
msg = "Cannot restart the kernel. No previous call to 'start_kernel'."
|
||||
raise RuntimeError(msg)
|
||||
|
||||
# Stop currently running kernel.
|
||||
await self._async_shutdown_kernel(now=now, restart=True)
|
||||
|
||||
if newports:
|
||||
self.cleanup_random_ports()
|
||||
|
||||
# Start new kernel.
|
||||
self._launch_args.update(kw)
|
||||
await self._async_start_kernel(**self._launch_args)
|
||||
|
||||
restart_kernel = run_sync(_async_restart_kernel)
|
||||
|
||||
@property
|
||||
def owns_kernel(self) -> bool:
|
||||
return self._owns_kernel
|
||||
|
||||
@property
|
||||
def has_kernel(self) -> bool:
|
||||
"""Has a kernel process been started that we are actively managing."""
|
||||
return self.provisioner is not None and self.provisioner.has_process
|
||||
|
||||
async def _async_send_kernel_sigterm(self, restart: bool = False) -> None:
|
||||
"""similar to _kill_kernel, but with sigterm (not sigkill), but do not block"""
|
||||
if self.has_kernel:
|
||||
assert self.provisioner is not None
|
||||
await self.provisioner.terminate(restart=restart)
|
||||
|
||||
_send_kernel_sigterm = run_sync(_async_send_kernel_sigterm)
|
||||
|
||||
async def _async_kill_kernel(self, restart: bool = False) -> None:
|
||||
"""Kill the running kernel.
|
||||
|
||||
This is a private method, callers should use shutdown_kernel(now=True).
|
||||
"""
|
||||
if self.has_kernel:
|
||||
assert self.provisioner is not None
|
||||
await self.provisioner.kill(restart=restart)
|
||||
|
||||
# Wait until the kernel terminates.
|
||||
try:
|
||||
await asyncio.wait_for(self._async_wait(), timeout=5.0)
|
||||
except asyncio.TimeoutError:
|
||||
# Wait timed out, just log warning but continue - not much more we can do.
|
||||
self.log.warning("Wait for final termination of kernel timed out - continuing...")
|
||||
pass
|
||||
else:
|
||||
# Process is no longer alive, wait and clear
|
||||
if self.has_kernel:
|
||||
await self.provisioner.wait()
|
||||
|
||||
_kill_kernel = run_sync(_async_kill_kernel)
|
||||
|
||||
async def _async_interrupt_kernel(self) -> None:
|
||||
"""Interrupts the kernel by sending it a signal.
|
||||
|
||||
Unlike ``signal_kernel``, this operation is well supported on all
|
||||
platforms.
|
||||
"""
|
||||
if not self.has_kernel and self._ready is not None:
|
||||
if isinstance(self._ready, CFuture):
|
||||
ready = asyncio.ensure_future(t.cast(Future[t.Any], self._ready))
|
||||
else:
|
||||
ready = self._ready
|
||||
# Wait for a shutdown if one is in progress.
|
||||
if self.shutting_down:
|
||||
await ready
|
||||
# Wait for a startup.
|
||||
await ready
|
||||
|
||||
if self.has_kernel:
|
||||
assert self.kernel_spec is not None
|
||||
interrupt_mode = self.kernel_spec.interrupt_mode
|
||||
if interrupt_mode == "signal":
|
||||
await self._async_signal_kernel(signal.SIGINT)
|
||||
|
||||
elif interrupt_mode == "message":
|
||||
msg = self.session.msg("interrupt_request", content={})
|
||||
self._connect_control_socket()
|
||||
self.session.send(self._control_socket, msg)
|
||||
else:
|
||||
msg = "Cannot interrupt kernel. No kernel is running!"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
interrupt_kernel = run_sync(_async_interrupt_kernel)
|
||||
|
||||
async def _async_signal_kernel(self, signum: int) -> None:
|
||||
"""Sends a signal to the process group of the kernel (this
|
||||
usually includes the kernel and any subprocesses spawned by
|
||||
the kernel).
|
||||
|
||||
Note that since only SIGTERM is supported on Windows, this function is
|
||||
only useful on Unix systems.
|
||||
"""
|
||||
if self.has_kernel:
|
||||
assert self.provisioner is not None
|
||||
await self.provisioner.send_signal(signum)
|
||||
else:
|
||||
msg = "Cannot signal kernel. No kernel is running!"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
signal_kernel = run_sync(_async_signal_kernel)
|
||||
|
||||
async def _async_is_alive(self) -> bool:
|
||||
"""Is the kernel process still running?"""
|
||||
if not self.owns_kernel:
|
||||
return True
|
||||
|
||||
if self.has_kernel:
|
||||
assert self.provisioner is not None
|
||||
ret = await self.provisioner.poll()
|
||||
if ret is None:
|
||||
return True
|
||||
return False
|
||||
|
||||
is_alive = run_sync(_async_is_alive)
|
||||
|
||||
async def _async_wait(self, pollinterval: float = 0.1) -> None:
|
||||
# Use busy loop at 100ms intervals, polling until the process is
|
||||
# not alive. If we find the process is no longer alive, complete
|
||||
# its cleanup via the blocking wait(). Callers are responsible for
|
||||
# issuing calls to wait() using a timeout (see _kill_kernel()).
|
||||
while await self._async_is_alive():
|
||||
await asyncio.sleep(pollinterval)
|
||||
|
||||
|
||||
class AsyncKernelManager(KernelManager):
|
||||
"""An async kernel manager."""
|
||||
|
||||
# the class to create with our `client` method
|
||||
client_class: DottedObjectName = DottedObjectName(
|
||||
"jupyter_client.asynchronous.AsyncKernelClient"
|
||||
)
|
||||
client_factory: Type = Type(klass="jupyter_client.asynchronous.AsyncKernelClient")
|
||||
|
||||
# The PyZMQ Context to use for communication with the kernel.
|
||||
context: Instance = Instance(zmq.asyncio.Context)
|
||||
|
||||
@default("context")
|
||||
def _context_default(self) -> zmq.asyncio.Context:
|
||||
self._created_context = True
|
||||
return zmq.asyncio.Context()
|
||||
|
||||
def client( # type:ignore[override]
|
||||
self, **kwargs: t.Any
|
||||
) -> AsyncKernelClient:
|
||||
"""Get a client for the manager."""
|
||||
return super().client(**kwargs) # type:ignore[return-value]
|
||||
|
||||
_launch_kernel = KernelManager._async_launch_kernel # type:ignore[assignment]
|
||||
start_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_start_kernel # type:ignore[assignment]
|
||||
pre_start_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_pre_start_kernel # type:ignore[assignment]
|
||||
post_start_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_post_start_kernel # type:ignore[assignment]
|
||||
request_shutdown: t.Callable[..., t.Awaitable] = KernelManager._async_request_shutdown # type:ignore[assignment]
|
||||
finish_shutdown: t.Callable[..., t.Awaitable] = KernelManager._async_finish_shutdown # type:ignore[assignment]
|
||||
cleanup_resources: t.Callable[..., t.Awaitable] = KernelManager._async_cleanup_resources # type:ignore[assignment]
|
||||
shutdown_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_shutdown_kernel # type:ignore[assignment]
|
||||
restart_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_restart_kernel # type:ignore[assignment]
|
||||
_send_kernel_sigterm = KernelManager._async_send_kernel_sigterm # type:ignore[assignment]
|
||||
_kill_kernel = KernelManager._async_kill_kernel # type:ignore[assignment]
|
||||
interrupt_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_interrupt_kernel # type:ignore[assignment]
|
||||
signal_kernel: t.Callable[..., t.Awaitable] = KernelManager._async_signal_kernel # type:ignore[assignment]
|
||||
is_alive: t.Callable[..., t.Awaitable] = KernelManager._async_is_alive # type:ignore[assignment]
|
||||
|
||||
|
||||
KernelManagerABC.register(KernelManager)
|
||||
|
||||
|
||||
def start_new_kernel(
|
||||
startup_timeout: float = 60, kernel_name: str = "python", **kwargs: t.Any
|
||||
) -> t.Tuple[KernelManager, BlockingKernelClient]:
|
||||
"""Start a new kernel, and return its Manager and Client"""
|
||||
km = KernelManager(kernel_name=kernel_name)
|
||||
km.start_kernel(**kwargs)
|
||||
kc = km.client()
|
||||
kc.start_channels()
|
||||
try:
|
||||
kc.wait_for_ready(timeout=startup_timeout)
|
||||
except RuntimeError:
|
||||
kc.stop_channels()
|
||||
km.shutdown_kernel()
|
||||
raise
|
||||
|
||||
return km, kc
|
||||
|
||||
|
||||
async def start_new_async_kernel(
|
||||
startup_timeout: float = 60, kernel_name: str = "python", **kwargs: t.Any
|
||||
) -> t.Tuple[AsyncKernelManager, AsyncKernelClient]:
|
||||
"""Start a new kernel, and return its Manager and Client"""
|
||||
km = AsyncKernelManager(kernel_name=kernel_name)
|
||||
await km.start_kernel(**kwargs)
|
||||
kc = km.client()
|
||||
kc.start_channels()
|
||||
try:
|
||||
await kc.wait_for_ready(timeout=startup_timeout)
|
||||
except RuntimeError:
|
||||
kc.stop_channels()
|
||||
await km.shutdown_kernel()
|
||||
raise
|
||||
|
||||
return (km, kc)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def run_kernel(**kwargs: t.Any) -> t.Iterator[KernelClient]:
|
||||
"""Context manager to create a kernel in a subprocess.
|
||||
|
||||
The kernel is shut down when the context exits.
|
||||
|
||||
Returns
|
||||
-------
|
||||
kernel_client: connected KernelClient instance
|
||||
"""
|
||||
km, kc = start_new_kernel(**kwargs)
|
||||
try:
|
||||
yield kc
|
||||
finally:
|
||||
kc.stop_channels()
|
||||
km.shutdown_kernel(now=True)
|
||||
Reference in New Issue
Block a user