kernelmanager.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928
  1. """A MultiKernelManager for use in the Jupyter server
  2. - raises HTTPErrors
  3. - creates REST API models
  4. """
  5. # Copyright (c) Jupyter Development Team.
  6. # Distributed under the terms of the Modified BSD License.
  7. from __future__ import annotations
  8. import asyncio
  9. import os
  10. import pathlib # noqa: TCH003
  11. import sys
  12. import typing as t
  13. import warnings
  14. from collections import defaultdict
  15. from datetime import datetime, timedelta
  16. from functools import partial, wraps
  17. from jupyter_client.ioloop.manager import AsyncIOLoopKernelManager
  18. from jupyter_client.multikernelmanager import AsyncMultiKernelManager, MultiKernelManager
  19. from jupyter_client.session import Session
  20. from jupyter_core.paths import exists
  21. from jupyter_core.utils import ensure_async
  22. from jupyter_events import EventLogger
  23. from jupyter_events.schema_registry import SchemaRegistryException
  24. if sys.version_info >= (3, 12):
  25. from typing import override
  26. else:
  27. from overrides import overrides as override
  28. from tornado import web
  29. from tornado.concurrent import Future
  30. from tornado.ioloop import IOLoop, PeriodicCallback
  31. from traitlets import (
  32. Any,
  33. Bool,
  34. Dict,
  35. Float,
  36. Instance,
  37. Integer,
  38. List,
  39. TraitError,
  40. Unicode,
  41. default,
  42. validate,
  43. )
  44. from jupyter_server import DEFAULT_EVENTS_SCHEMA_PATH
  45. from jupyter_server._tz import isoformat, utcnow
  46. from jupyter_server.prometheus.metrics import KERNEL_CURRENTLY_RUNNING_TOTAL
  47. from jupyter_server.utils import ApiPath, import_item, to_os_path
  48. class MappingKernelManager(MultiKernelManager):
  49. """A KernelManager that handles
  50. - File mapping
  51. - HTTP error handling
  52. - Kernel message filtering
  53. """
  54. @default("kernel_manager_class")
  55. def _default_kernel_manager_class(self):
  56. return "jupyter_client.ioloop.IOLoopKernelManager"
  57. kernel_argv = List(Unicode())
  58. root_dir = Unicode(config=True)
  59. _kernel_connections = Dict()
  60. _kernel_ports: dict[str, list[int]] = Dict() # type: ignore[assignment]
  61. _culler_callback = None
  62. _initialized_culler = False
  63. @default("root_dir")
  64. def _default_root_dir(self):
  65. if not self.parent:
  66. return os.getcwd()
  67. return self.parent.root_dir
  68. @validate("root_dir")
  69. def _update_root_dir(self, proposal):
  70. """Do a bit of validation of the root dir."""
  71. value = proposal["value"]
  72. if not os.path.isabs(value):
  73. # If we receive a non-absolute path, make it absolute.
  74. value = os.path.abspath(value)
  75. if not exists(value) or not os.path.isdir(value):
  76. raise TraitError("kernel root dir %r is not a directory" % value)
  77. return value
  78. cull_idle_timeout = Integer(
  79. 0,
  80. config=True,
  81. help="""Timeout (in seconds) after which a kernel is considered idle and ready to be culled.
  82. Values of 0 or lower disable culling. Very short timeouts may result in kernels being culled
  83. for users with poor network connections.""",
  84. )
  85. cull_interval_default = 300 # 5 minutes
  86. cull_interval = Integer(
  87. cull_interval_default,
  88. config=True,
  89. help="""The interval (in seconds) on which to check for idle kernels exceeding the cull timeout value.""",
  90. )
  91. cull_connected = Bool(
  92. False,
  93. config=True,
  94. help="""Whether to consider culling kernels which have one or more connections.
  95. Only effective if cull_idle_timeout > 0.""",
  96. )
  97. cull_busy = Bool(
  98. False,
  99. config=True,
  100. help="""Whether to consider culling kernels which are busy.
  101. Only effective if cull_idle_timeout > 0.""",
  102. )
  103. buffer_offline_messages = Bool(
  104. True,
  105. config=True,
  106. help="""Whether messages from kernels whose frontends have disconnected should be buffered in-memory.
  107. When True (default), messages are buffered and replayed on reconnect,
  108. avoiding lost messages due to interrupted connectivity.
  109. Disable if long-running kernels will produce too much output while
  110. no frontends are connected.
  111. """,
  112. )
  113. kernel_info_timeout = Float(
  114. 60,
  115. config=True,
  116. help="""Timeout for giving up on a kernel (in seconds).
  117. On starting and restarting kernels, we check whether the
  118. kernel is running and responsive by sending kernel_info_requests.
  119. This sets the timeout in seconds for how long the kernel can take
  120. before being presumed dead.
  121. This affects the MappingKernelManager (which handles kernel restarts)
  122. and the ZMQChannelsHandler (which handles the startup).
  123. """,
  124. )
  125. _kernel_buffers = Any()
  126. @default("_kernel_buffers")
  127. def _default_kernel_buffers(self):
  128. return defaultdict(lambda: {"buffer": [], "session_key": "", "channels": {}})
  129. last_kernel_activity = Instance(
  130. datetime,
  131. help="The last activity on any kernel, including shutting down a kernel",
  132. )
  133. def __init__(self, **kwargs):
  134. """Initialize a kernel manager."""
  135. self.pinned_superclass = MultiKernelManager
  136. self._pending_kernel_tasks = {}
  137. self.pinned_superclass.__init__(self, **kwargs)
  138. self.last_kernel_activity = utcnow()
  139. allowed_message_types = List(
  140. trait=Unicode(),
  141. config=True,
  142. help="""White list of allowed kernel message types.
  143. When the list is empty, all message types are allowed.
  144. """,
  145. )
  146. allow_tracebacks = Bool(
  147. True, config=True, help=("Whether to send tracebacks to clients on exceptions.")
  148. )
  149. traceback_replacement_message = Unicode(
  150. "An exception occurred at runtime, which is not shown due to security reasons.",
  151. config=True,
  152. help=("Message to print when allow_tracebacks is False, and an exception occurs"),
  153. )
  154. # -------------------------------------------------------------------------
  155. # Methods for managing kernels and sessions
  156. # -------------------------------------------------------------------------
  157. def _handle_kernel_died(self, kernel_id):
  158. """notice that a kernel died"""
  159. self.log.warning("Kernel %s died, removing from map.", kernel_id)
  160. self.remove_kernel(kernel_id)
  161. def cwd_for_path(self, path, **kwargs):
  162. """Turn API path into absolute OS path."""
  163. os_path = to_os_path(path, self.root_dir)
  164. # in the case of documents and kernels not being on the same filesystem,
  165. # walk up to root_dir if the paths don't exist
  166. while not os.path.isdir(os_path) and os_path != self.root_dir:
  167. os_path = os.path.dirname(os_path)
  168. return os_path
  169. async def _remove_kernel_when_ready(self, kernel_id, kernel_awaitable):
  170. """Remove a kernel when it is ready."""
  171. await super()._remove_kernel_when_ready(kernel_id, kernel_awaitable)
  172. self._kernel_connections.pop(kernel_id, None)
  173. self._kernel_ports.pop(kernel_id, None)
  174. # TODO: DEC 2022: Revise the type-ignore once the signatures have been changed upstream
  175. # https://github.com/jupyter/jupyter_client/pull/905
  176. async def _async_start_kernel( # type:ignore[override]
  177. self, *, kernel_id: str | None = None, path: ApiPath | None = None, **kwargs: str
  178. ) -> str:
  179. """Start a kernel for a session and return its kernel_id.
  180. Parameters
  181. ----------
  182. kernel_id : uuid (str)
  183. The uuid to associate the new kernel with. If this
  184. is not None, this kernel will be persistent whenever it is
  185. requested.
  186. path : API path
  187. The API path (unicode, '/' delimited) for the cwd.
  188. Will be transformed to an OS path relative to root_dir.
  189. kernel_name : str
  190. The name identifying which kernel spec to launch. This is ignored if
  191. an existing kernel is returned, but it may be checked in the future.
  192. """
  193. if kernel_id is None or kernel_id not in self:
  194. if path is not None:
  195. kwargs["cwd"] = self.cwd_for_path(path, env=kwargs.get("env", {}))
  196. if kernel_id is not None:
  197. assert kernel_id is not None, "Never Fail, but necessary for mypy "
  198. kwargs["kernel_id"] = kernel_id
  199. kernel_id = await self.pinned_superclass._async_start_kernel(self, **kwargs)
  200. self._kernel_connections[kernel_id] = 0
  201. # add busy/activity markers:
  202. kernel = self.get_kernel(kernel_id)
  203. kernel.execution_state = "starting" # type:ignore[attr-defined]
  204. kernel.reason = "" # type:ignore[attr-defined]
  205. kernel.last_activity = utcnow() # type:ignore[attr-defined]
  206. self.log.info("Kernel started: %s", kernel_id)
  207. self.log.debug(
  208. "Kernel args (excluding env): %r", {k: v for k, v in kwargs.items() if k != "env"}
  209. )
  210. env = kwargs.get("env", None)
  211. if env and isinstance(env, dict): # type:ignore[unreachable]
  212. self.log.debug("Kernel argument 'env' passed with: %r", list(env.keys())) # type:ignore[unreachable]
  213. task = asyncio.create_task(self._finish_kernel_start(kernel_id))
  214. if not getattr(self, "use_pending_kernels", None):
  215. await task
  216. else:
  217. self._pending_kernel_tasks[kernel_id] = task
  218. # Increase the metric of number of kernels running
  219. # for the relevant kernel type by 1
  220. KERNEL_CURRENTLY_RUNNING_TOTAL.labels(type=self._kernels[kernel_id].kernel_name).inc()
  221. else:
  222. self.log.info("Using existing kernel: %s", kernel_id)
  223. # Initialize culling if not already
  224. if not self._initialized_culler:
  225. self.initialize_culler()
  226. assert kernel_id is not None
  227. return kernel_id
  228. # see https://github.com/jupyter-server/jupyter_server/issues/1165
  229. # this assignment is technically incorrect, but might need a change of API
  230. # in jupyter_client.
  231. start_kernel = _async_start_kernel # type:ignore[assignment]
  232. async def _finish_kernel_start(self, kernel_id):
  233. """Handle a kernel that finishes starting."""
  234. km = self.get_kernel(kernel_id)
  235. if hasattr(km, "ready"):
  236. ready = km.ready
  237. if not isinstance(ready, asyncio.Future):
  238. ready = asyncio.wrap_future(ready)
  239. try:
  240. await ready
  241. except Exception:
  242. self.log.exception("Error waiting for kernel manager ready")
  243. return
  244. self._kernel_ports[kernel_id] = km.ports
  245. self.start_watching_activity(kernel_id)
  246. # register callback for failed auto-restart
  247. self.add_restart_callback(
  248. kernel_id,
  249. lambda: self._handle_kernel_died(kernel_id),
  250. "dead",
  251. )
  252. def ports_changed(self, kernel_id):
  253. """Used by ZMQChannelsHandler to determine how to coordinate nudge and replays.
  254. Ports are captured when starting a kernel (via MappingKernelManager). Ports
  255. are considered changed (following restarts) if the referenced KernelManager
  256. is using a set of ports different from those captured at startup. If changes
  257. are detected, the captured set is updated and a value of True is returned.
  258. NOTE: Use is exclusive to ZMQChannelsHandler because this object is a singleton
  259. instance while ZMQChannelsHandler instances are per WebSocket connection that
  260. can vary per kernel lifetime.
  261. """
  262. changed_ports = self._get_changed_ports(kernel_id)
  263. if changed_ports:
  264. # If changed, update captured ports and return True, else return False.
  265. self.log.debug("Port change detected for kernel: %s", kernel_id)
  266. self._kernel_ports[kernel_id] = changed_ports
  267. return True
  268. return False
  269. def _get_changed_ports(self, kernel_id):
  270. """Internal method to test if a kernel's ports have changed and, if so, return their values.
  271. This method does NOT update the captured ports for the kernel as that can only be done
  272. by ZMQChannelsHandler, but instead returns the new list of ports if they are different
  273. than those captured at startup. This enables the ability to conditionally restart
  274. activity monitoring immediately following a kernel's restart (if ports have changed).
  275. """
  276. # Get current ports and return comparison with ports captured at startup.
  277. km = self.get_kernel(kernel_id)
  278. assert isinstance(km.ports, list)
  279. assert isinstance(self._kernel_ports[kernel_id], list)
  280. if km.ports != self._kernel_ports[kernel_id]:
  281. return km.ports
  282. return None
  283. def start_buffering(self, kernel_id, session_key, channels):
  284. """Start buffering messages for a kernel
  285. Parameters
  286. ----------
  287. kernel_id : str
  288. The id of the kernel to stop buffering.
  289. session_key : str
  290. The session_key, if any, that should get the buffer.
  291. If the session_key matches the current buffered session_key,
  292. the buffer will be returned.
  293. channels : dict({'channel': ZMQStream})
  294. The zmq channels whose messages should be buffered.
  295. """
  296. if not self.buffer_offline_messages:
  297. for stream in channels.values():
  298. stream.close()
  299. return
  300. self.log.info("Starting buffering for %s", session_key)
  301. self._check_kernel_id(kernel_id)
  302. # clear previous buffering state
  303. self.stop_buffering(kernel_id)
  304. buffer_info = self._kernel_buffers[kernel_id]
  305. # record the session key because only one session can buffer
  306. buffer_info["session_key"] = session_key
  307. # TODO: the buffer should likely be a memory bounded queue, we're starting with a list to keep it simple
  308. buffer_info["buffer"] = []
  309. buffer_info["channels"] = channels
  310. # forward any future messages to the internal buffer
  311. def buffer_msg(channel, msg_parts):
  312. self.log.debug("Buffering msg on %s:%s", kernel_id, channel)
  313. buffer_info["buffer"].append((channel, msg_parts))
  314. for channel, stream in channels.items():
  315. stream.on_recv(partial(buffer_msg, channel))
  316. def get_buffer(self, kernel_id, session_key):
  317. """Get the buffer for a given kernel
  318. Parameters
  319. ----------
  320. kernel_id : str
  321. The id of the kernel to stop buffering.
  322. session_key : str, optional
  323. The session_key, if any, that should get the buffer.
  324. If the session_key matches the current buffered session_key,
  325. the buffer will be returned.
  326. """
  327. self.log.debug("Getting buffer for %s", kernel_id)
  328. if kernel_id not in self._kernel_buffers:
  329. return None
  330. buffer_info = self._kernel_buffers[kernel_id]
  331. if buffer_info["session_key"] == session_key:
  332. # remove buffer
  333. self._kernel_buffers.pop(kernel_id)
  334. # only return buffer_info if it's a match
  335. return buffer_info
  336. else:
  337. self.stop_buffering(kernel_id)
  338. def stop_buffering(self, kernel_id):
  339. """Stop buffering kernel messages
  340. Parameters
  341. ----------
  342. kernel_id : str
  343. The id of the kernel to stop buffering.
  344. """
  345. self.log.debug("Clearing buffer for %s", kernel_id)
  346. self._check_kernel_id(kernel_id)
  347. if kernel_id not in self._kernel_buffers:
  348. return
  349. buffer_info = self._kernel_buffers.pop(kernel_id)
  350. # close buffering streams
  351. for stream in buffer_info["channels"].values():
  352. if not stream.socket.closed:
  353. stream.on_recv(None)
  354. stream.close()
  355. msg_buffer = buffer_info["buffer"]
  356. if msg_buffer:
  357. self.log.info(
  358. "Discarding %s buffered messages for %s",
  359. len(msg_buffer),
  360. buffer_info["session_key"],
  361. )
  362. async def _async_shutdown_kernel(self, kernel_id, now=False, restart=False):
  363. """Shutdown a kernel by kernel_id"""
  364. self._check_kernel_id(kernel_id)
  365. # Decrease the metric of number of kernels
  366. # running for the relevant kernel type by 1
  367. KERNEL_CURRENTLY_RUNNING_TOTAL.labels(type=self._kernels[kernel_id].kernel_name).dec()
  368. if kernel_id in self._pending_kernel_tasks:
  369. task = self._pending_kernel_tasks.pop(kernel_id)
  370. task.cancel()
  371. self.stop_watching_activity(kernel_id)
  372. self.stop_buffering(kernel_id)
  373. return await self.pinned_superclass._async_shutdown_kernel(
  374. self, kernel_id, now=now, restart=restart
  375. )
  376. shutdown_kernel = _async_shutdown_kernel
  377. async def _async_restart_kernel(self, kernel_id, now=False):
  378. """Restart a kernel by kernel_id"""
  379. self._check_kernel_id(kernel_id)
  380. await self.pinned_superclass._async_restart_kernel(self, kernel_id, now=now)
  381. kernel = self.get_kernel(kernel_id)
  382. # return a Future that will resolve when the kernel has successfully restarted
  383. channel = kernel.connect_shell()
  384. future: Future[Any] = Future()
  385. def finish():
  386. """Common cleanup when restart finishes/fails for any reason."""
  387. if not channel.closed(): # type:ignore[operator]
  388. channel.close()
  389. loop.remove_timeout(timeout)
  390. kernel.remove_restart_callback(on_restart_failed, "dead")
  391. kernel._pending_restart_cleanup = None # type:ignore[attr-defined]
  392. def on_reply(msg):
  393. self.log.debug("Kernel info reply received: %s", kernel_id)
  394. finish()
  395. if not future.done():
  396. future.set_result(msg)
  397. def on_timeout():
  398. self.log.warning("Timeout waiting for kernel_info_reply: %s", kernel_id)
  399. finish()
  400. if not future.done():
  401. future.set_exception(TimeoutError("Timeout waiting for restart"))
  402. def on_restart_failed():
  403. self.log.warning("Restarting kernel failed: %s", kernel_id)
  404. finish()
  405. if not future.done():
  406. future.set_exception(RuntimeError("Restart failed"))
  407. kernel.add_restart_callback(on_restart_failed, "dead")
  408. kernel._pending_restart_cleanup = finish # type:ignore[attr-defined]
  409. kernel.session.send(channel, "kernel_info_request")
  410. channel.on_recv(on_reply) # type:ignore[operator]
  411. loop = IOLoop.current()
  412. timeout = loop.add_timeout(loop.time() + self.kernel_info_timeout, on_timeout)
  413. # Re-establish activity watching if ports have changed...
  414. if self._get_changed_ports(kernel_id) is not None:
  415. self.stop_watching_activity(kernel_id)
  416. self.start_watching_activity(kernel_id)
  417. return future
  418. restart_kernel = _async_restart_kernel
  419. def notify_connect(self, kernel_id):
  420. """Notice a new connection to a kernel"""
  421. if kernel_id in self._kernel_connections:
  422. self._kernel_connections[kernel_id] += 1
  423. def notify_disconnect(self, kernel_id):
  424. """Notice a disconnection from a kernel"""
  425. if kernel_id in self._kernel_connections:
  426. self._kernel_connections[kernel_id] -= 1
  427. def kernel_model(self, kernel_id):
  428. """Return a JSON-safe dict representing a kernel
  429. For use in representing kernels in the JSON APIs.
  430. """
  431. self._check_kernel_id(kernel_id)
  432. kernel = self._kernels[kernel_id]
  433. model = {
  434. "id": kernel_id,
  435. "name": kernel.kernel_name,
  436. "last_activity": isoformat(kernel.last_activity),
  437. "execution_state": kernel.execution_state,
  438. "connections": self._kernel_connections.get(kernel_id, 0),
  439. }
  440. if getattr(kernel, "reason", None):
  441. model["reason"] = kernel.reason
  442. return model
  443. def list_kernels(self):
  444. """Returns a list of kernel_id's of kernels running."""
  445. kernels = []
  446. kernel_ids = self.pinned_superclass.list_kernel_ids(self)
  447. for kernel_id in kernel_ids:
  448. try:
  449. model = self.kernel_model(kernel_id)
  450. kernels.append(model)
  451. except (web.HTTPError, KeyError):
  452. # Probably due to a (now) non-existent kernel, continue building the list
  453. pass
  454. return kernels
  455. # override _check_kernel_id to raise 404 instead of KeyError
  456. def _check_kernel_id(self, kernel_id):
  457. """Check a that a kernel_id exists and raise 404 if not."""
  458. if kernel_id not in self:
  459. raise web.HTTPError(404, "Kernel does not exist: %s" % kernel_id)
  460. # monitoring activity:
  461. untracked_message_types = List(
  462. trait=Unicode(),
  463. config=True,
  464. default_value=[
  465. "comm_info_request",
  466. "comm_info_reply",
  467. "kernel_info_request",
  468. "kernel_info_reply",
  469. "shutdown_request",
  470. "shutdown_reply",
  471. "interrupt_request",
  472. "interrupt_reply",
  473. "debug_request",
  474. "debug_reply",
  475. "stream",
  476. "display_data",
  477. "update_display_data",
  478. "execute_input",
  479. "execute_result",
  480. "error",
  481. "status",
  482. "clear_output",
  483. "debug_event",
  484. "input_request",
  485. "input_reply",
  486. ],
  487. help="""List of kernel message types excluded from user activity tracking.
  488. This should be a superset of the message types sent on any channel other
  489. than the shell channel.""",
  490. )
  491. def track_message_type(self, message_type):
  492. return message_type not in self.untracked_message_types
  493. def start_watching_activity(self, kernel_id):
  494. """Start watching IOPub messages on a kernel for activity.
  495. - update last_activity on every message
  496. - record execution_state from status messages
  497. """
  498. kernel = self._kernels[kernel_id]
  499. # add busy/activity markers:
  500. kernel.execution_state = "starting"
  501. kernel.reason = ""
  502. kernel.last_activity = utcnow()
  503. kernel._activity_stream = kernel.connect_iopub()
  504. session = Session(
  505. config=kernel.session.config,
  506. key=kernel.session.key,
  507. )
  508. def record_activity(msg_list):
  509. """Record an IOPub message arriving from a kernel"""
  510. idents, fed_msg_list = session.feed_identities(msg_list)
  511. msg = session.deserialize(fed_msg_list, content=False)
  512. msg_type = msg["header"]["msg_type"]
  513. parent_header = msg.get("parent_header")
  514. parent_msg_type = None if parent_header is None else parent_header.get("msg_type")
  515. if (
  516. self.track_message_type(msg_type)
  517. or self.track_message_type(parent_msg_type)
  518. or kernel.execution_state == "busy"
  519. ):
  520. self.last_kernel_activity = kernel.last_activity = utcnow()
  521. if msg_type == "status":
  522. msg = session.deserialize(fed_msg_list)
  523. execution_state = msg["content"]["execution_state"]
  524. if self.track_message_type(parent_msg_type):
  525. kernel.execution_state = execution_state
  526. elif kernel.execution_state == "starting" and execution_state != "starting":
  527. # We always normalize post-starting execution state to "idle"
  528. # unless we know that the status is in response to one of our
  529. # tracked message types.
  530. kernel.execution_state = "idle"
  531. self.log.debug(
  532. "activity on %s: %s (%s)",
  533. kernel_id,
  534. msg_type,
  535. kernel.execution_state,
  536. )
  537. else:
  538. self.log.debug("activity on %s: %s", kernel_id, msg_type)
  539. kernel._activity_stream.on_recv(record_activity)
  540. def stop_watching_activity(self, kernel_id):
  541. """Stop watching IOPub messages on a kernel for activity."""
  542. kernel = self._kernels[kernel_id]
  543. if getattr(kernel, "_activity_stream", None):
  544. if not kernel._activity_stream.socket.closed:
  545. kernel._activity_stream.close()
  546. kernel._activity_stream = None
  547. if getattr(kernel, "_pending_restart_cleanup", None):
  548. kernel._pending_restart_cleanup()
  549. def initialize_culler(self):
  550. """Start idle culler if 'cull_idle_timeout' is greater than zero.
  551. Regardless of that value, set flag that we've been here.
  552. """
  553. if (
  554. not self._initialized_culler
  555. and self.cull_idle_timeout > 0
  556. and self._culler_callback is None
  557. ):
  558. _ = IOLoop.current()
  559. if self.cull_interval <= 0: # handle case where user set invalid value
  560. self.log.warning(
  561. "Invalid value for 'cull_interval' detected (%s) - using default value (%s).",
  562. self.cull_interval,
  563. self.cull_interval_default,
  564. )
  565. self.cull_interval = self.cull_interval_default
  566. self._culler_callback = PeriodicCallback(self.cull_kernels, 1000 * self.cull_interval)
  567. self.log.info(
  568. "Culling kernels with idle durations > %s seconds at %s second intervals ...",
  569. self.cull_idle_timeout,
  570. self.cull_interval,
  571. )
  572. if self.cull_busy:
  573. self.log.info("Culling kernels even if busy")
  574. if self.cull_connected:
  575. self.log.info("Culling kernels even with connected clients")
  576. self._culler_callback.start()
  577. self._initialized_culler = True
  578. async def cull_kernels(self):
  579. """Handle culling kernels."""
  580. self.log.debug(
  581. "Polling every %s seconds for kernels idle > %s seconds...",
  582. self.cull_interval,
  583. self.cull_idle_timeout,
  584. )
  585. """Create a separate list of kernels to avoid conflicting updates while iterating"""
  586. for kernel_id in list(self._kernels):
  587. try:
  588. await self.cull_kernel_if_idle(kernel_id)
  589. except Exception as e:
  590. self.log.exception(
  591. "The following exception was encountered while checking the idle duration of kernel %s: %s",
  592. kernel_id,
  593. e,
  594. )
  595. async def cull_kernel_if_idle(self, kernel_id):
  596. """Cull a kernel if it is idle."""
  597. kernel = self._kernels[kernel_id]
  598. if getattr(kernel, "execution_state", None) == "dead":
  599. self.log.warning(
  600. "Culling '%s' dead kernel '%s' (%s).",
  601. kernel.execution_state,
  602. kernel.kernel_name,
  603. kernel_id,
  604. )
  605. await ensure_async(self.shutdown_kernel(kernel_id))
  606. return
  607. kernel_spec_metadata = kernel.kernel_spec.metadata
  608. cull_idle_timeout = kernel_spec_metadata.get("cull_idle_timeout", self.cull_idle_timeout)
  609. if hasattr(
  610. kernel, "last_activity"
  611. ): # last_activity is monkey-patched, so ensure that has occurred
  612. self.log.debug(
  613. "kernel_id=%s, kernel_name=%s, last_activity=%s",
  614. kernel_id,
  615. kernel.kernel_name,
  616. kernel.last_activity,
  617. )
  618. dt_now = utcnow()
  619. dt_idle = dt_now - kernel.last_activity
  620. # Compute idle properties
  621. is_idle_time = dt_idle > timedelta(seconds=cull_idle_timeout)
  622. is_idle_execute = self.cull_busy or (kernel.execution_state != "busy")
  623. connections = self._kernel_connections.get(kernel_id, 0)
  624. is_idle_connected = self.cull_connected or not connections
  625. # Cull the kernel if all three criteria are met
  626. if is_idle_time and is_idle_execute and is_idle_connected:
  627. idle_duration = int(dt_idle.total_seconds())
  628. self.log.warning(
  629. "Culling '%s' kernel '%s' (%s) with %d connections due to %s seconds of inactivity.",
  630. kernel.execution_state,
  631. kernel.kernel_name,
  632. kernel_id,
  633. connections,
  634. idle_duration,
  635. )
  636. await ensure_async(self.shutdown_kernel(kernel_id))
  637. # AsyncMappingKernelManager inherits as much as possible from MappingKernelManager,
  638. # overriding only what is different.
  639. class AsyncMappingKernelManager(MappingKernelManager, AsyncMultiKernelManager): # type:ignore[misc]
  640. """An asynchronous mapping kernel manager."""
  641. @default("kernel_manager_class")
  642. def _default_kernel_manager_class(self):
  643. return "jupyter_server.services.kernels.kernelmanager.ServerKernelManager"
  644. @validate("kernel_manager_class")
  645. def _validate_kernel_manager_class(self, proposal):
  646. """A validator for the kernel manager class."""
  647. km_class_value = proposal.value
  648. km_class = import_item(km_class_value)
  649. if not issubclass(km_class, ServerKernelManager):
  650. warnings.warn(
  651. f"KernelManager class '{km_class}' is not a subclass of 'ServerKernelManager'. Custom "
  652. "KernelManager classes should derive from 'ServerKernelManager' beginning with jupyter-server 2.0 "
  653. "or risk missing functionality. Continuing...",
  654. FutureWarning,
  655. stacklevel=3,
  656. )
  657. return km_class_value
  658. def __init__(self, **kwargs):
  659. """Initialize an async mapping kernel manager."""
  660. self.pinned_superclass = MultiKernelManager
  661. self._pending_kernel_tasks = {}
  662. self.pinned_superclass.__init__(self, **kwargs)
  663. self.last_kernel_activity = utcnow()
  664. def emit_kernel_action_event(success_msg: str = "") -> t.Callable[..., t.Any]:
  665. """Decorate kernel action methods to
  666. begin emitting jupyter kernel action events.
  667. Parameters
  668. ----------
  669. success_msg: str
  670. A formattable string that's passed to the message field of
  671. the emitted event when the action succeeds. You can include
  672. the kernel_id, kernel_name, or action in the message using
  673. a formatted string argument,
  674. e.g. "{kernel_id} succeeded to {action}."
  675. error_msg: str
  676. A formattable string that's passed to the message field of
  677. the emitted event when the action fails. You can include
  678. the kernel_id, kernel_name, or action in the message using
  679. a formatted string argument,
  680. e.g. "{kernel_id} failed to {action}."
  681. """
  682. def wrap_method(method):
  683. @wraps(method)
  684. async def wrapped_method(self, *args, **kwargs):
  685. """"""
  686. # Get the method name from the
  687. action = method.__name__.replace("_kernel", "")
  688. # If the method succeeds, emit a success event.
  689. try:
  690. out = await method(self, *args, **kwargs)
  691. data = {
  692. "kernel_name": self.kernel_name,
  693. "action": action,
  694. "status": "success",
  695. "msg": success_msg.format(
  696. kernel_id=self.kernel_id, kernel_name=self.kernel_name, action=action
  697. ),
  698. }
  699. if self.kernel_id:
  700. data["kernel_id"] = self.kernel_id
  701. self.emit(
  702. schema_id="https://events.jupyter.org/jupyter_server/kernel_actions/v1",
  703. data=data,
  704. )
  705. return out
  706. # If the method fails, emit a failed event.
  707. except Exception as err:
  708. data = {
  709. "kernel_name": self.kernel_name,
  710. "action": action,
  711. "status": "error",
  712. "msg": str(err),
  713. }
  714. if self.kernel_id:
  715. data["kernel_id"] = self.kernel_id
  716. # If the exception is an HTTPError (usually via a gateway request)
  717. # log the status_code and HTTPError log_message.
  718. if isinstance(err, web.HTTPError):
  719. msg = err.log_message or ""
  720. data["status_code"] = err.status_code
  721. data["msg"] = msg
  722. self.emit(
  723. schema_id="https://events.jupyter.org/jupyter_server/kernel_actions/v1",
  724. data=data,
  725. )
  726. raise err
  727. return wrapped_method
  728. return wrap_method
  729. class ServerKernelManager(AsyncIOLoopKernelManager):
  730. """A server-specific kernel manager."""
  731. # Define activity-related attributes:
  732. execution_state = Unicode(
  733. None, allow_none=True, help="The current execution state of the kernel"
  734. )
  735. reason = Unicode("", help="The reason for the last failure against the kernel")
  736. last_activity = Instance(datetime, help="The last activity on the kernel")
  737. # A list of pathlib objects, each pointing at an event
  738. # schema to register with this kernel manager's eventlogger.
  739. # This trait should not be overridden.
  740. @property
  741. def core_event_schema_paths(self) -> list[pathlib.Path]:
  742. return [DEFAULT_EVENTS_SCHEMA_PATH / "kernel_actions" / "v1.yaml"]
  743. # This trait is intended for subclasses to override and define
  744. # custom event schemas.
  745. extra_event_schema_paths: List[str] = List(
  746. default_value=[],
  747. help="""
  748. A list of pathlib.Path objects pointing at to register with
  749. the kernel manager's eventlogger.
  750. """,
  751. ).tag(config=True)
  752. event_logger = Instance(EventLogger)
  753. @default("event_logger")
  754. def _default_event_logger(self):
  755. """Initialize the logger and ensure all required events are present."""
  756. if (
  757. self.parent is not None
  758. and self.parent.parent is not None
  759. and hasattr(self.parent.parent, "event_logger")
  760. ):
  761. logger = self.parent.parent.event_logger
  762. else:
  763. # If parent does not have an event logger, create one.
  764. logger = EventLogger()
  765. # Ensure that all the expected schemas are registered. If not, register them.
  766. schemas = self.core_event_schema_paths + self.extra_event_schema_paths
  767. for schema_path in schemas:
  768. # Try registering the event.
  769. try:
  770. logger.register_event_schema(schema_path)
  771. # Pass if it already exists.
  772. except SchemaRegistryException:
  773. pass
  774. return logger
  775. def emit(self, schema_id, data):
  776. """Emit an event from the kernel manager."""
  777. self.event_logger.emit(schema_id=schema_id, data=data)
  778. @override
  779. @emit_kernel_action_event(
  780. success_msg="Kernel {kernel_id} was started.",
  781. )
  782. async def start_kernel(self, *args, **kwargs):
  783. return await super().start_kernel(*args, **kwargs)
  784. @override
  785. @emit_kernel_action_event(
  786. success_msg="Kernel {kernel_id} was shutdown.",
  787. )
  788. async def shutdown_kernel(self, *args, **kwargs):
  789. return await super().shutdown_kernel(*args, **kwargs)
  790. @override
  791. @emit_kernel_action_event(
  792. success_msg="Kernel {kernel_id} was restarted.",
  793. )
  794. async def restart_kernel(self, *args, **kwargs):
  795. return await super().restart_kernel(*args, **kwargs)
  796. @override
  797. @emit_kernel_action_event(
  798. success_msg="Kernel {kernel_id} was interrupted.",
  799. )
  800. async def interrupt_kernel(self, *args, **kwargs):
  801. return await super().interrupt_kernel(*args, **kwargs)