| | """A kernel manager with a tornado IOLoop""" |
| | |
| | |
| | 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 |
| | |
| | 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: |
| | |
| | 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) |
| |
|
