"""A kernel manager for multiple kernels""" # Copyright (c) Jupyter Development Team. # Distributed under the terms of the Modified BSD License. import asyncio import os import socket import typing as t import uuid from functools import wraps import zmq from traitlets import Any, Bool, Dict, DottedObjectName, Instance, Unicode, default, observe from traitlets.config.configurable import LoggingConfigurable from traitlets.utils.importstring import import_item from .kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager from .manager import KernelManager from .utils import ensure_async, run_sync class DuplicateKernelError(Exception): pass def kernel_method(f: t.Callable) -> t.Callable: """decorator for proxying MKM.method(kernel_id) to individual KMs by ID""" @wraps(f) def wrapped( self: t.Any, kernel_id: str, *args: t.Any, **kwargs: t.Any ) -> t.Union[t.Callable, t.Awaitable]: # get the kernel km = self.get_kernel(kernel_id) method = getattr(km, f.__name__) # call the kernel's method r = method(*args, **kwargs) # last thing, call anything defined in the actual class method # such as logging messages f(self, kernel_id, *args, **kwargs) # return the method result return r return wrapped class MultiKernelManager(LoggingConfigurable): """A class for managing multiple kernels.""" default_kernel_name = Unicode( NATIVE_KERNEL_NAME, help="The name of the default kernel to start" ).tag(config=True) kernel_spec_manager = Instance(KernelSpecManager, allow_none=True) kernel_manager_class = DottedObjectName( "jupyter_client.ioloop.IOLoopKernelManager", help="""The kernel manager class. This is configurable to allow subclassing of the KernelManager for customized behavior. """, ).tag(config=True) @observe("kernel_manager_class") def _kernel_manager_class_changed(self, change): self.kernel_manager_factory = self._create_kernel_manager_factory() kernel_manager_factory = Any(help="this is kernel_manager_class after import") @default("kernel_manager_factory") def _kernel_manager_factory_default(self): return self._create_kernel_manager_factory() def _create_kernel_manager_factory(self) -> t.Callable: kernel_manager_ctor = import_item(self.kernel_manager_class) def create_kernel_manager(*args: t.Any, **kwargs: t.Any) -> KernelManager: if self.shared_context: if self.context.closed: # recreate context if closed self.context = self._context_default() kwargs.setdefault("context", self.context) km = kernel_manager_ctor(*args, **kwargs) return km return create_kernel_manager shared_context = Bool( True, help="Share a single zmq.Context to talk to all my kernels", ).tag(config=True) context = Instance("zmq.Context") _created_context = Bool(False) _pending_kernels = Dict() @property def _starting_kernels(self): """A shim for backwards compatibility.""" return self._pending_kernels @default("context") # type:ignore[misc] def _context_default(self) -> zmq.Context: self._created_context = True return zmq.Context() connection_dir = Unicode("") _kernels = Dict() def __del__(self): """Handle garbage collection. Destroy context if applicable.""" if self._created_context and self.context and not self.context.closed: if self.log: self.log.debug("Destroying zmq context for %s", self) self.context.destroy() try: super_del = super().__del__ # type:ignore[misc] except AttributeError: pass else: super_del() def list_kernel_ids(self) -> t.List[str]: """Return a list of the kernel ids of the active kernels.""" # Create a copy so we can iterate over kernels in operations # that delete keys. return list(self._kernels.keys()) def __len__(self) -> int: """Return the number of running kernels.""" return len(self.list_kernel_ids()) def __contains__(self, kernel_id: str) -> bool: return kernel_id in self._kernels def pre_start_kernel( self, kernel_name: t.Optional[str], kwargs: t.Any ) -> t.Tuple[KernelManager, str, str]: # kwargs should be mutable, passing it as a dict argument. kernel_id = kwargs.pop("kernel_id", self.new_kernel_id(**kwargs)) if kernel_id in self: raise DuplicateKernelError("Kernel already exists: %s" % kernel_id) if kernel_name is None: kernel_name = self.default_kernel_name # kernel_manager_factory is the constructor for the KernelManager # subclass we are using. It can be configured as any Configurable, # including things like its transport and ip. constructor_kwargs = {} if self.kernel_spec_manager: constructor_kwargs["kernel_spec_manager"] = self.kernel_spec_manager km = self.kernel_manager_factory( connection_file=os.path.join(self.connection_dir, "kernel-%s.json" % kernel_id), parent=self, log=self.log, kernel_name=kernel_name, **constructor_kwargs, ) return km, kernel_name, kernel_id async def _add_kernel_when_ready( self, kernel_id: str, km: KernelManager, kernel_awaitable: t.Awaitable ) -> None: try: await kernel_awaitable self._kernels[kernel_id] = km self._pending_kernels.pop(kernel_id, None) except Exception as e: self.log.exception(e) async def _remove_kernel_when_ready( self, kernel_id: str, kernel_awaitable: t.Awaitable ) -> None: try: await kernel_awaitable self.remove_kernel(kernel_id) self._pending_kernels.pop(kernel_id, None) except Exception as e: self.log.exception(e) def _using_pending_kernels(self): """Returns a boolean; a clearer method for determining if this multikernelmanager is using pending kernels or not """ return getattr(self, 'use_pending_kernels', False) async def _async_start_kernel( self, *, kernel_name: t.Optional[str] = None, **kwargs: t.Any ) -> str: """Start a new kernel. The caller can pick a kernel_id by passing one in as a keyword arg, otherwise one will be generated using new_kernel_id(). The kernel ID for the newly started kernel is returned. """ km, kernel_name, kernel_id = self.pre_start_kernel(kernel_name, kwargs) if not isinstance(km, KernelManager): self.log.warning( "Kernel manager class ({km_class}) is not an instance of 'KernelManager'!".format( km_class=self.kernel_manager_class.__class__ ) ) kwargs['kernel_id'] = kernel_id # Make kernel_id available to manager and provisioner starter = ensure_async(km.start_kernel(**kwargs)) task = asyncio.create_task(self._add_kernel_when_ready(kernel_id, km, starter)) self._pending_kernels[kernel_id] = task # Handling a Pending Kernel if self._using_pending_kernels(): # If using pending kernels, do not block # on the kernel start. self._kernels[kernel_id] = km else: await task # raise an exception if one occurred during kernel startup. if km.ready.exception(): raise km.ready.exception() # type: ignore return kernel_id start_kernel = run_sync(_async_start_kernel) async def _async_shutdown_kernel( self, kernel_id: str, now: t.Optional[bool] = False, restart: t.Optional[bool] = False, ) -> None: """Shutdown a kernel by its kernel uuid. Parameters ========== kernel_id : uuid The id of the kernel to shutdown. now : bool Should the kernel be shutdown forcibly using a signal. restart : bool Will the kernel be restarted? """ self.log.info("Kernel shutdown: %s", kernel_id) # If the kernel is still starting, wait for it to be ready. if kernel_id in self._pending_kernels: task = self._pending_kernels[kernel_id] try: await task km = self.get_kernel(kernel_id) await t.cast(asyncio.Future, km.ready) except asyncio.CancelledError: pass except Exception: self.remove_kernel(kernel_id) return km = self.get_kernel(kernel_id) # If a pending kernel raised an exception, remove it. if not km.ready.cancelled() and km.ready.exception(): self.remove_kernel(kernel_id) return stopper = ensure_async(km.shutdown_kernel(now, restart)) fut = asyncio.ensure_future(self._remove_kernel_when_ready(kernel_id, stopper)) self._pending_kernels[kernel_id] = fut # Await the kernel if not using pending kernels. if not self._using_pending_kernels(): await fut # raise an exception if one occurred during kernel shutdown. if km.ready.exception(): raise km.ready.exception() # type: ignore shutdown_kernel = run_sync(_async_shutdown_kernel) @kernel_method def request_shutdown(self, kernel_id: str, restart: t.Optional[bool] = False) -> None: """Ask a kernel to shut down by its kernel uuid""" @kernel_method def finish_shutdown( self, kernel_id: str, waittime: t.Optional[float] = None, pollinterval: t.Optional[float] = 0.1, ) -> None: """Wait for a kernel to finish shutting down, and kill it if it doesn't""" self.log.info("Kernel shutdown: %s", kernel_id) @kernel_method def cleanup_resources(self, kernel_id: str, restart: bool = False) -> None: """Clean up a kernel's resources""" def remove_kernel(self, kernel_id: str) -> KernelManager: """remove a kernel from our mapping. Mainly so that a kernel can be removed if it is already dead, without having to call shutdown_kernel. The kernel object is returned, or `None` if not found. """ return self._kernels.pop(kernel_id, None) async def _async_shutdown_all(self, now: bool = False) -> None: """Shutdown all kernels.""" kids = self.list_kernel_ids() kids += list(self._pending_kernels) kms = list(self._kernels.values()) futs = [self._async_shutdown_kernel(kid, now=now) for kid in set(kids)] await asyncio.gather(*futs) # If using pending kernels, the kernels will not have been fully shut down. if self._using_pending_kernels(): for km in kms: try: await km.ready except asyncio.CancelledError: self._pending_kernels[km.kernel_id].cancel() except Exception: # Will have been logged in _add_kernel_when_ready pass shutdown_all = run_sync(_async_shutdown_all) def interrupt_kernel(self, kernel_id: str) -> None: """Interrupt (SIGINT) the kernel by its uuid. Parameters ========== kernel_id : uuid The id of the kernel to interrupt. """ kernel = self.get_kernel(kernel_id) if not kernel.ready.done(): msg = "Kernel is in a pending state. Cannot interrupt." raise RuntimeError(msg) out = kernel.interrupt_kernel() self.log.info("Kernel interrupted: %s", kernel_id) return out @kernel_method def signal_kernel(self, kernel_id: str, signum: int) -> None: """Sends a signal to the kernel by its uuid. Note that since only SIGTERM is supported on Windows, this function is only useful on Unix systems. Parameters ========== kernel_id : uuid The id of the kernel to signal. signum : int Signal number to send kernel. """ self.log.info("Signaled Kernel %s with %s", kernel_id, signum) async def _async_restart_kernel(self, kernel_id: str, now: bool = False) -> None: """Restart a kernel by its uuid, keeping the same ports. Parameters ========== kernel_id : uuid The id of the kernel to interrupt. 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. """ kernel = self.get_kernel(kernel_id) if self._using_pending_kernels() and not kernel.ready.done(): msg = "Kernel is in a pending state. Cannot restart." raise RuntimeError(msg) await ensure_async(kernel.restart_kernel(now=now)) self.log.info("Kernel restarted: %s", kernel_id) restart_kernel = run_sync(_async_restart_kernel) @kernel_method def is_alive(self, kernel_id: str) -> bool: # type:ignore[empty-body] """Is the kernel alive. This calls KernelManager.is_alive() which calls Popen.poll on the actual kernel subprocess. Parameters ========== kernel_id : uuid The id of the kernel. """ def _check_kernel_id(self, kernel_id: str) -> None: """check that a kernel id is valid""" if kernel_id not in self: raise KeyError("Kernel with id not found: %s" % kernel_id) def get_kernel(self, kernel_id: str) -> KernelManager: """Get the single KernelManager object for a kernel by its uuid. Parameters ========== kernel_id : uuid The id of the kernel. """ self._check_kernel_id(kernel_id) return self._kernels[kernel_id] @kernel_method def add_restart_callback( self, kernel_id: str, callback: t.Callable, event: str = "restart" ) -> None: """add a callback for the KernelRestarter""" @kernel_method def remove_restart_callback( self, kernel_id: str, callback: t.Callable, event: str = "restart" ) -> None: """remove a callback for the KernelRestarter""" @kernel_method def get_connection_info(self, kernel_id: str) -> t.Dict[str, t.Any]: # type:ignore[empty-body] """Return a dictionary of connection data for a kernel. Parameters ========== kernel_id : uuid The id of the kernel. Returns ======= connection_dict : dict A dict of the information needed to connect to a kernel. This includes the ip address and the integer port numbers of the different channels (stdin_port, iopub_port, shell_port, hb_port). """ @kernel_method def connect_iopub( # type:ignore[empty-body] self, kernel_id: str, identity: t.Optional[bytes] = None ) -> socket.socket: """Return a zmq Socket connected to the iopub channel. Parameters ========== kernel_id : uuid The id of the kernel identity : bytes (optional) The zmq identity of the socket Returns ======= stream : zmq Socket or ZMQStream """ @kernel_method def connect_shell( # type:ignore[empty-body] self, kernel_id: str, identity: t.Optional[bytes] = None ) -> socket.socket: """Return a zmq Socket connected to the shell channel. Parameters ========== kernel_id : uuid The id of the kernel identity : bytes (optional) The zmq identity of the socket Returns ======= stream : zmq Socket or ZMQStream """ @kernel_method def connect_control( # type:ignore[empty-body] self, kernel_id: str, identity: t.Optional[bytes] = None ) -> socket.socket: """Return a zmq Socket connected to the control channel. Parameters ========== kernel_id : uuid The id of the kernel identity : bytes (optional) The zmq identity of the socket Returns ======= stream : zmq Socket or ZMQStream """ @kernel_method def connect_stdin( # type:ignore[empty-body] self, kernel_id: str, identity: t.Optional[bytes] = None ) -> socket.socket: """Return a zmq Socket connected to the stdin channel. Parameters ========== kernel_id : uuid The id of the kernel identity : bytes (optional) The zmq identity of the socket Returns ======= stream : zmq Socket or ZMQStream """ @kernel_method def connect_hb( # type:ignore[empty-body] self, kernel_id: str, identity: t.Optional[bytes] = None ) -> socket.socket: """Return a zmq Socket connected to the hb channel. Parameters ========== kernel_id : uuid The id of the kernel identity : bytes (optional) The zmq identity of the socket Returns ======= stream : zmq Socket or ZMQStream """ def new_kernel_id(self, **kwargs: t.Any) -> str: """ Returns the id to associate with the kernel for this request. Subclasses may override this method to substitute other sources of kernel ids. :param kwargs: :return: string-ized version 4 uuid """ return str(uuid.uuid4()) class AsyncMultiKernelManager(MultiKernelManager): kernel_manager_class = DottedObjectName( "jupyter_client.ioloop.AsyncIOLoopKernelManager", config=True, help="""The kernel manager class. This is configurable to allow subclassing of the AsyncKernelManager for customized behavior. """, ) use_pending_kernels = Bool( False, help="""Whether to make kernels available before the process has started. The kernel has a `.ready` future which can be awaited before connecting""", ).tag(config=True) context = Instance("zmq.asyncio.Context") @default("context") # type:ignore[misc] def _context_default(self) -> zmq.asyncio.Context: self._created_context = True return zmq.asyncio.Context() start_kernel: t.Callable[ ..., t.Awaitable ] = MultiKernelManager._async_start_kernel # type:ignore[assignment] restart_kernel: t.Callable[ ..., t.Awaitable ] = MultiKernelManager._async_restart_kernel # type:ignore[assignment] shutdown_kernel: t.Callable[ ..., t.Awaitable ] = MultiKernelManager._async_shutdown_kernel # type:ignore[assignment] shutdown_all: t.Callable[ ..., t.Awaitable ] = MultiKernelManager._async_shutdown_all # type:ignore[assignment]