File size: 4,162 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
"""A kernel manager with a tornado IOLoop"""
# Copyright (c) Jupyter Development Team.
# Distributed under the terms of the Modified BSD License.
import typing as t
import zmq
from tornado import ioloop
from traitlets import Instance, Type
from zmq.eventloop.zmqstream import ZMQStream
from ..manager import AsyncKernelManager, KernelManager
from .restarter import AsyncIOLoopKernelRestarter, IOLoopKernelRestarter
def as_zmqstream(f: t.Any) -> t.Callable:
"""Convert a socket to a zmq stream."""
def wrapped(self: t.Any, *args: t.Any, **kwargs: t.Any) -> t.Any:
save_socket_class = None
# zmqstreams only support sync sockets
if self.context._socket_class is not zmq.Socket:
save_socket_class = self.context._socket_class
self.context._socket_class = zmq.Socket
try:
socket = f(self, *args, **kwargs)
finally:
if save_socket_class:
# restore default socket class
self.context._socket_class = save_socket_class
return ZMQStream(socket, self.loop)
return wrapped
class IOLoopKernelManager(KernelManager):
"""An io loop kernel manager."""
loop = Instance("tornado.ioloop.IOLoop")
def _loop_default(self) -> ioloop.IOLoop:
return ioloop.IOLoop.current()
restarter_class = Type(
default_value=IOLoopKernelRestarter,
klass=IOLoopKernelRestarter,
help=(
"Type of KernelRestarter to use. "
"Must be a subclass of IOLoopKernelRestarter.\n"
"Override this to customize how kernel restarts are managed."
),
config=True,
)
_restarter: t.Any = Instance("jupyter_client.ioloop.IOLoopKernelRestarter", allow_none=True)
def start_restarter(self) -> None:
"""Start the restarter."""
if self.autorestart and self.has_kernel:
if self._restarter is None:
self._restarter = self.restarter_class(
kernel_manager=self, loop=self.loop, parent=self, log=self.log
)
self._restarter.start()
def stop_restarter(self) -> None:
"""Stop the restarter."""
if self.autorestart and self._restarter is not None:
self._restarter.stop()
connect_shell = as_zmqstream(KernelManager.connect_shell)
connect_control = as_zmqstream(KernelManager.connect_control)
connect_iopub = as_zmqstream(KernelManager.connect_iopub)
connect_stdin = as_zmqstream(KernelManager.connect_stdin)
connect_hb = as_zmqstream(KernelManager.connect_hb)
class AsyncIOLoopKernelManager(AsyncKernelManager):
"""An async ioloop kernel manager."""
loop = Instance("tornado.ioloop.IOLoop")
def _loop_default(self) -> ioloop.IOLoop:
return ioloop.IOLoop.current()
restarter_class = Type(
default_value=AsyncIOLoopKernelRestarter,
klass=AsyncIOLoopKernelRestarter,
help=(
"Type of KernelRestarter to use. "
"Must be a subclass of AsyncIOLoopKernelManager.\n"
"Override this to customize how kernel restarts are managed."
),
config=True,
)
_restarter: t.Any = Instance(
"jupyter_client.ioloop.AsyncIOLoopKernelRestarter", allow_none=True
)
def start_restarter(self) -> None:
"""Start the restarter."""
if self.autorestart and self.has_kernel:
if self._restarter is None:
self._restarter = self.restarter_class(
kernel_manager=self, loop=self.loop, parent=self, log=self.log
)
self._restarter.start()
def stop_restarter(self) -> None:
"""Stop the restarter."""
if self.autorestart and self._restarter is not None:
self._restarter.stop()
connect_shell = as_zmqstream(AsyncKernelManager.connect_shell)
connect_control = as_zmqstream(AsyncKernelManager.connect_control)
connect_iopub = as_zmqstream(AsyncKernelManager.connect_iopub)
connect_stdin = as_zmqstream(AsyncKernelManager.connect_stdin)
connect_hb = as_zmqstream(AsyncKernelManager.connect_hb)
|