Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""A kernel manager for multiple kernels"""
2# Copyright (c) Jupyter Development Team.
3# Distributed under the terms of the Modified BSD License.
4from __future__ import annotations
6import asyncio
7import json
8import os
9import socket
10import typing as t
11import uuid
12from functools import wraps
13from pathlib import Path
15import zmq
16from traitlets import Any, Bool, Dict, DottedObjectName, Instance, Unicode, default, observe
17from traitlets.config.configurable import LoggingConfigurable
18from traitlets.utils.importstring import import_item
20from .connect import KernelConnectionInfo
21from .kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager
22from .manager import KernelManager
23from .utils import ensure_async, run_sync, utcnow
26class DuplicateKernelError(Exception):
27 pass
30def kernel_method(f: t.Callable) -> t.Callable:
31 """decorator for proxying MKM.method(kernel_id) to individual KMs by ID"""
33 @wraps(f)
34 def wrapped(
35 self: t.Any, kernel_id: str, *args: t.Any, **kwargs: t.Any
36 ) -> t.Callable | t.Awaitable:
37 # get the kernel
38 km = self.get_kernel(kernel_id)
39 method = getattr(km, f.__name__)
40 # call the kernel's method
41 r = method(*args, **kwargs)
42 # last thing, call anything defined in the actual class method
43 # such as logging messages
44 f(self, kernel_id, *args, **kwargs)
45 # return the method result
46 return r
48 return wrapped
51class MultiKernelManager(LoggingConfigurable):
52 """A class for managing multiple kernels."""
54 default_kernel_name = Unicode(
55 NATIVE_KERNEL_NAME, help="The name of the default kernel to start"
56 ).tag(config=True)
58 kernel_spec_manager = Instance(KernelSpecManager, allow_none=True)
60 kernel_manager_class = DottedObjectName(
61 "jupyter_client.ioloop.IOLoopKernelManager",
62 help="""The kernel manager class. This is configurable to allow
63 subclassing of the KernelManager for customized behavior.
64 """,
65 ).tag(config=True)
67 @observe("kernel_manager_class")
68 def _kernel_manager_class_changed(self, change: t.Any) -> None:
69 self.kernel_manager_factory = self._create_kernel_manager_factory()
71 kernel_manager_factory = Any(help="this is kernel_manager_class after import")
73 @default("kernel_manager_factory")
74 def _kernel_manager_factory_default(self) -> t.Callable:
75 return self._create_kernel_manager_factory()
77 def _create_kernel_manager_factory(self) -> t.Callable:
78 kernel_manager_ctor = import_item(self.kernel_manager_class)
80 def create_kernel_manager(*args: t.Any, **kwargs: t.Any) -> KernelManager:
81 if self.shared_context:
82 if self.context.closed:
83 # recreate context if closed
84 self.context = self._context_default()
85 kwargs.setdefault("context", self.context)
86 km = kernel_manager_ctor(*args, **kwargs)
87 return km
89 return create_kernel_manager
91 shared_context = Bool(
92 True,
93 help="Share a single zmq.Context to talk to all my kernels",
94 ).tag(config=True)
96 context = Instance("zmq.Context")
98 _created_context = Bool(False)
100 _pending_kernels = Dict()
102 @property
103 def _starting_kernels(self) -> dict:
104 """A shim for backwards compatibility."""
105 return self._pending_kernels
107 @default("context")
108 def _context_default(self) -> zmq.Context:
109 self._created_context = True
110 return zmq.Context()
112 connection_dir = Unicode("")
113 external_connection_dir = Unicode(None, allow_none=True)
115 _kernels = Dict()
117 def __init__(self, *args: t.Any, **kwargs: t.Any) -> None:
118 super().__init__(*args, **kwargs)
119 self.kernel_id_to_connection_file: dict[str, Path] = {}
121 def __del__(self) -> None:
122 """Handle garbage collection. Destroy context if applicable."""
123 if self._created_context and self.context and not self.context.closed:
124 if self.log:
125 self.log.debug("Destroying zmq context for %s", self)
126 self.context.destroy()
127 try:
128 super_del = super().__del__ # type:ignore[misc]
129 except AttributeError:
130 pass
131 else:
132 super_del()
134 def list_kernel_ids(self) -> list[str]:
135 """Return a list of the kernel ids of the active kernels."""
136 if self.external_connection_dir is not None:
137 external_connection_dir = Path(self.external_connection_dir)
138 if external_connection_dir.is_dir():
139 connection_files = [p for p in external_connection_dir.iterdir() if p.is_file()]
141 # remove kernels (whose connection file has disappeared) from our list
142 k = list(self.kernel_id_to_connection_file.keys())
143 v = list(self.kernel_id_to_connection_file.values())
144 for connection_file in list(self.kernel_id_to_connection_file.values()):
145 if connection_file not in connection_files:
146 kernel_id = k[v.index(connection_file)]
147 del self.kernel_id_to_connection_file[kernel_id]
148 del self._kernels[kernel_id]
150 # add kernels (whose connection file appeared) to our list
151 for connection_file in connection_files:
152 if connection_file in self.kernel_id_to_connection_file.values():
153 continue
154 try:
155 connection_info: KernelConnectionInfo = json.loads(
156 connection_file.read_text()
157 )
158 except Exception: # noqa: S112
159 continue
160 self.log.debug("Loading connection file %s", connection_file)
161 if not ("kernel_name" in connection_info and "key" in connection_info):
162 continue
163 # it looks like a connection file
164 kernel_id = self.new_kernel_id()
165 self.kernel_id_to_connection_file[kernel_id] = connection_file
166 km = self.kernel_manager_factory(
167 parent=self,
168 log=self.log,
169 owns_kernel=False,
170 )
171 km.load_connection_info(connection_info)
172 km.last_activity = utcnow()
173 km.execution_state = "idle"
174 km.connections = 1
175 km.kernel_id = kernel_id
176 km.kernel_name = connection_info["kernel_name"]
177 km.ready.set_result(None)
179 self._kernels[kernel_id] = km
181 # Create a copy so we can iterate over kernels in operations
182 # that delete keys.
183 return list(self._kernels.keys())
185 def __len__(self) -> int:
186 """Return the number of running kernels."""
187 return len(self.list_kernel_ids())
189 def __contains__(self, kernel_id: str) -> bool:
190 return kernel_id in self._kernels
192 def pre_start_kernel(
193 self, kernel_name: str | None, kwargs: t.Any
194 ) -> tuple[KernelManager, str, str]:
195 # kwargs should be mutable, passing it as a dict argument.
196 kernel_id = kwargs.pop("kernel_id", self.new_kernel_id(**kwargs))
197 if kernel_id in self:
198 raise DuplicateKernelError("Kernel already exists: %s" % kernel_id)
200 if kernel_name is None:
201 kernel_name = self.default_kernel_name
202 # kernel_manager_factory is the constructor for the KernelManager
203 # subclass we are using. It can be configured as any Configurable,
204 # including things like its transport and ip.
205 constructor_kwargs = {}
206 if self.kernel_spec_manager:
207 constructor_kwargs["kernel_spec_manager"] = self.kernel_spec_manager
208 km = self.kernel_manager_factory(
209 connection_file=os.path.join(self.connection_dir, "kernel-%s.json" % kernel_id),
210 parent=self,
211 log=self.log,
212 kernel_name=kernel_name,
213 **constructor_kwargs,
214 )
215 return km, kernel_name, kernel_id
217 def update_env(self, *, kernel_id: str, env: t.Dict[str, str]) -> None:
218 """
219 Allow to update the environment of the given kernel.
221 Forward the update env request to the corresponding kernel.
223 .. version-added: 8.5
224 """
225 if kernel_id in self:
226 self._kernels[kernel_id].update_env(env=env)
228 async def _add_kernel_when_ready(
229 self, kernel_id: str, km: KernelManager, kernel_awaitable: t.Awaitable
230 ) -> None:
231 try:
232 await kernel_awaitable
233 self._kernels[kernel_id] = km
234 self._pending_kernels.pop(kernel_id, None)
235 except Exception as e:
236 self.log.exception(e)
238 async def _remove_kernel_when_ready(
239 self, kernel_id: str, kernel_awaitable: t.Awaitable
240 ) -> None:
241 try:
242 await kernel_awaitable
243 self.remove_kernel(kernel_id)
244 self._pending_kernels.pop(kernel_id, None)
245 except Exception as e:
246 self.log.exception(e)
248 def _using_pending_kernels(self) -> bool:
249 """Returns a boolean; a clearer method for determining if
250 this multikernelmanager is using pending kernels or not
251 """
252 return getattr(self, "use_pending_kernels", False)
254 async def _async_start_kernel(self, *, kernel_name: str | None = None, **kwargs: t.Any) -> str:
255 """Start a new kernel.
257 The caller can pick a kernel_id by passing one in as a keyword arg,
258 otherwise one will be generated using new_kernel_id().
260 The kernel ID for the newly started kernel is returned.
261 """
262 km, kernel_name, kernel_id = self.pre_start_kernel(kernel_name, kwargs)
263 if not isinstance(km, KernelManager):
264 self.log.warning( # type:ignore[unreachable]
265 "Kernel manager class ({km_class}) is not an instance of 'KernelManager'!".format(
266 km_class=self.kernel_manager_class.__class__
267 )
268 )
269 kwargs["kernel_id"] = kernel_id # Make kernel_id available to manager and provisioner
271 starter = ensure_async(km.start_kernel(**kwargs))
272 task = asyncio.create_task(self._add_kernel_when_ready(kernel_id, km, starter))
273 self._pending_kernels[kernel_id] = task
274 # Handling a Pending Kernel
275 if self._using_pending_kernels():
276 # If using pending kernels, do not block
277 # on the kernel start.
278 self._kernels[kernel_id] = km
279 else:
280 await task
281 # raise an exception if one occurred during kernel startup.
282 if km.ready.exception():
283 raise km.ready.exception() # type: ignore[misc]
285 return kernel_id
287 start_kernel = run_sync(_async_start_kernel)
289 async def _async_shutdown_kernel(
290 self,
291 kernel_id: str,
292 now: bool | None = False,
293 restart: bool | None = False,
294 ) -> None:
295 """Shutdown a kernel by its kernel uuid.
297 Parameters
298 ==========
299 kernel_id : uuid
300 The id of the kernel to shutdown.
301 now : bool
302 Should the kernel be shutdown forcibly using a signal.
303 restart : bool
304 Will the kernel be restarted?
305 """
306 self.log.info("Kernel shutdown: %s", kernel_id)
307 # If the kernel is still starting, wait for it to be ready.
308 if kernel_id in self._pending_kernels:
309 task = self._pending_kernels[kernel_id]
310 try:
311 await task
312 km = self.get_kernel(kernel_id)
313 await t.cast(asyncio.Future, km.ready)
314 except asyncio.CancelledError:
315 pass
316 except Exception:
317 self.remove_kernel(kernel_id)
318 return
319 km = self.get_kernel(kernel_id)
320 # If a pending kernel raised an exception, remove it.
321 if not km.ready.cancelled() and km.ready.exception():
322 self.remove_kernel(kernel_id)
323 return
324 stopper = ensure_async(km.shutdown_kernel(now, restart))
325 fut = asyncio.ensure_future(self._remove_kernel_when_ready(kernel_id, stopper))
326 self._pending_kernels[kernel_id] = fut
327 # Await the kernel if not using pending kernels.
328 if not self._using_pending_kernels():
329 await fut
330 # raise an exception if one occurred during kernel shutdown.
331 if km.ready.exception():
332 raise km.ready.exception() # type: ignore[misc]
334 shutdown_kernel = run_sync(_async_shutdown_kernel)
336 @kernel_method
337 def request_shutdown(self, kernel_id: str, restart: bool | None = False) -> None:
338 """Ask a kernel to shut down by its kernel uuid"""
340 @kernel_method
341 def finish_shutdown(
342 self,
343 kernel_id: str,
344 waittime: float | None = None,
345 pollinterval: float | None = 0.1,
346 ) -> None:
347 """Wait for a kernel to finish shutting down, and kill it if it doesn't"""
348 self.log.info("Kernel shutdown: %s", kernel_id)
350 @kernel_method
351 def cleanup_resources(self, kernel_id: str, restart: bool = False) -> None:
352 """Clean up a kernel's resources"""
354 def remove_kernel(self, kernel_id: str) -> KernelManager:
355 """remove a kernel from our mapping.
357 Mainly so that a kernel can be removed if it is already dead,
358 without having to call shutdown_kernel.
360 The kernel object is returned, or `None` if not found.
361 """
362 return self._kernels.pop(kernel_id, None)
364 async def _async_shutdown_all(self, now: bool = False) -> None:
365 """Shutdown all kernels."""
366 kids = self.list_kernel_ids()
367 kids += list(self._pending_kernels)
368 kms = list(self._kernels.values())
369 futs = [self._async_shutdown_kernel(kid, now=now) for kid in set(kids)]
370 await asyncio.gather(*futs)
371 # If using pending kernels, the kernels will not have been fully shut down.
372 if self._using_pending_kernels():
373 for km in kms:
374 try:
375 await km.ready
376 except asyncio.CancelledError:
377 self._pending_kernels[km.kernel_id].cancel()
378 except Exception:
379 # Will have been logged in _add_kernel_when_ready
380 pass
382 shutdown_all = run_sync(_async_shutdown_all)
384 def interrupt_kernel(self, kernel_id: str) -> None:
385 """Interrupt (SIGINT) the kernel by its uuid.
387 Parameters
388 ==========
389 kernel_id : uuid
390 The id of the kernel to interrupt.
391 """
392 kernel = self.get_kernel(kernel_id)
393 if not kernel.ready.done():
394 msg = "Kernel is in a pending state. Cannot interrupt."
395 raise RuntimeError(msg)
396 out = kernel.interrupt_kernel()
397 self.log.info("Kernel interrupted: %s", kernel_id)
398 return out
400 @kernel_method
401 def signal_kernel(self, kernel_id: str, signum: int) -> None:
402 """Sends a signal to the kernel by its uuid.
404 Note that since only SIGTERM is supported on Windows, this function
405 is only useful on Unix systems.
407 Parameters
408 ==========
409 kernel_id : uuid
410 The id of the kernel to signal.
411 signum : int
412 Signal number to send kernel.
413 """
414 self.log.info("Signaled Kernel %s with %s", kernel_id, signum)
416 async def _async_restart_kernel(self, kernel_id: str, now: bool = False) -> None:
417 """Restart a kernel by its uuid, keeping the same ports.
419 Parameters
420 ==========
421 kernel_id : uuid
422 The id of the kernel to interrupt.
423 now : bool, optional
424 If True, the kernel is forcefully restarted *immediately*, without
425 having a chance to do any cleanup action. Otherwise the kernel is
426 given 1s to clean up before a forceful restart is issued.
428 In all cases the kernel is restarted, the only difference is whether
429 it is given a chance to perform a clean shutdown or not.
430 """
431 kernel = self.get_kernel(kernel_id)
432 if self._using_pending_kernels() and not kernel.ready.done():
433 msg = "Kernel is in a pending state. Cannot restart."
434 raise RuntimeError(msg)
435 await ensure_async(kernel.restart_kernel(now=now))
436 self.log.info("Kernel restarted: %s", kernel_id)
438 restart_kernel = run_sync(_async_restart_kernel)
440 @kernel_method
441 def is_alive(self, kernel_id: str) -> bool: # type:ignore[empty-body]
442 """Is the kernel alive.
444 This calls KernelManager.is_alive() which calls Popen.poll on the
445 actual kernel subprocess.
447 Parameters
448 ==========
449 kernel_id : uuid
450 The id of the kernel.
451 """
453 def _check_kernel_id(self, kernel_id: str) -> None:
454 """check that a kernel id is valid"""
455 if kernel_id not in self:
456 raise KeyError("Kernel with id not found: %s" % kernel_id)
458 def get_kernel(self, kernel_id: str) -> KernelManager:
459 """Get the single KernelManager object for a kernel by its uuid.
461 Parameters
462 ==========
463 kernel_id : uuid
464 The id of the kernel.
465 """
466 self._check_kernel_id(kernel_id)
467 return self._kernels[kernel_id]
469 @kernel_method
470 def add_restart_callback(
471 self, kernel_id: str, callback: t.Callable, event: str = "restart"
472 ) -> None:
473 """add a callback for the KernelRestarter"""
475 @kernel_method
476 def remove_restart_callback(
477 self, kernel_id: str, callback: t.Callable, event: str = "restart"
478 ) -> None:
479 """remove a callback for the KernelRestarter"""
481 @kernel_method
482 def get_connection_info(self, kernel_id: str) -> dict[str, t.Any]: # type:ignore[empty-body]
483 """Return a dictionary of connection data for a kernel.
485 Parameters
486 ==========
487 kernel_id : uuid
488 The id of the kernel.
490 Returns
491 =======
492 connection_dict : dict
493 A dict of the information needed to connect to a kernel.
494 This includes the ip address and the integer port
495 numbers of the different channels (stdin_port, iopub_port,
496 shell_port, hb_port).
497 """
499 @kernel_method
500 def connect_iopub( # type:ignore[empty-body]
501 self, kernel_id: str, identity: bytes | None = None
502 ) -> socket.socket:
503 """Return a zmq Socket connected to the iopub channel.
505 Parameters
506 ==========
507 kernel_id : uuid
508 The id of the kernel
509 identity : bytes (optional)
510 The zmq identity of the socket
512 Returns
513 =======
514 stream : zmq Socket or ZMQStream
515 """
517 @kernel_method
518 def connect_shell( # type:ignore[empty-body]
519 self, kernel_id: str, identity: bytes | None = None
520 ) -> socket.socket:
521 """Return a zmq Socket connected to the shell channel.
523 Parameters
524 ==========
525 kernel_id : uuid
526 The id of the kernel
527 identity : bytes (optional)
528 The zmq identity of the socket
530 Returns
531 =======
532 stream : zmq Socket or ZMQStream
533 """
535 @kernel_method
536 def connect_control( # type:ignore[empty-body]
537 self, kernel_id: str, identity: bytes | None = None
538 ) -> socket.socket:
539 """Return a zmq Socket connected to the control channel.
541 Parameters
542 ==========
543 kernel_id : uuid
544 The id of the kernel
545 identity : bytes (optional)
546 The zmq identity of the socket
548 Returns
549 =======
550 stream : zmq Socket or ZMQStream
551 """
553 @kernel_method
554 def connect_stdin( # type:ignore[empty-body]
555 self, kernel_id: str, identity: bytes | None = None
556 ) -> socket.socket:
557 """Return a zmq Socket connected to the stdin channel.
559 Parameters
560 ==========
561 kernel_id : uuid
562 The id of the kernel
563 identity : bytes (optional)
564 The zmq identity of the socket
566 Returns
567 =======
568 stream : zmq Socket or ZMQStream
569 """
571 @kernel_method
572 def connect_hb( # type:ignore[empty-body]
573 self, kernel_id: str, identity: bytes | None = None
574 ) -> socket.socket:
575 """Return a zmq Socket connected to the hb channel.
577 Parameters
578 ==========
579 kernel_id : uuid
580 The id of the kernel
581 identity : bytes (optional)
582 The zmq identity of the socket
584 Returns
585 =======
586 stream : zmq Socket or ZMQStream
587 """
589 def new_kernel_id(self, **kwargs: t.Any) -> str:
590 """
591 Returns the id to associate with the kernel for this request. Subclasses may override
592 this method to substitute other sources of kernel ids.
593 :param kwargs:
594 :return: string-ized version 4 uuid
595 """
596 return str(uuid.uuid4())
599class AsyncMultiKernelManager(MultiKernelManager):
600 kernel_manager_class = DottedObjectName(
601 "jupyter_client.ioloop.AsyncIOLoopKernelManager",
602 config=True,
603 help="""The kernel manager class. This is configurable to allow
604 subclassing of the AsyncKernelManager for customized behavior.
605 """,
606 )
608 use_pending_kernels = Bool(
609 False,
610 help="""Whether to make kernels available before the process has started. The
611 kernel has a `.ready` future which can be awaited before connecting""",
612 ).tag(config=True)
614 context = Instance("zmq.asyncio.Context")
616 @default("context")
617 def _context_default(self) -> zmq.asyncio.Context:
618 self._created_context = True
619 return zmq.asyncio.Context()
621 start_kernel: t.Callable[..., t.Awaitable] = MultiKernelManager._async_start_kernel # type:ignore[assignment]
622 restart_kernel: t.Callable[..., t.Awaitable] = MultiKernelManager._async_restart_kernel # type:ignore[assignment]
623 shutdown_kernel: t.Callable[..., t.Awaitable] = MultiKernelManager._async_shutdown_kernel # type:ignore[assignment]
624 shutdown_all: t.Callable[..., t.Awaitable] = MultiKernelManager._async_shutdown_all # type:ignore[assignment]