multikernelmanager.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374
  1. """A kernel manager for multiple kernels"""
  2. # Copyright (c) Jupyter Development Team.
  3. # Distributed under the terms of the Modified BSD License.
  4. from __future__ import absolute_import
  5. import os
  6. import uuid
  7. import zmq
  8. from traitlets.config.configurable import LoggingConfigurable
  9. from ipython_genutils.importstring import import_item
  10. from traitlets import (
  11. Any, Bool, Dict, DottedObjectName, Instance, Unicode, default, observe
  12. )
  13. from ipython_genutils.py3compat import unicode_type
  14. from .kernelspec import NATIVE_KERNEL_NAME, KernelSpecManager
  15. class DuplicateKernelError(Exception):
  16. pass
  17. def kernel_method(f):
  18. """decorator for proxying MKM.method(kernel_id) to individual KMs by ID"""
  19. def wrapped(self, kernel_id, *args, **kwargs):
  20. # get the kernel
  21. km = self.get_kernel(kernel_id)
  22. method = getattr(km, f.__name__)
  23. # call the kernel's method
  24. r = method(*args, **kwargs)
  25. # last thing, call anything defined in the actual class method
  26. # such as logging messages
  27. f(self, kernel_id, *args, **kwargs)
  28. # return the method result
  29. return r
  30. return wrapped
  31. class MultiKernelManager(LoggingConfigurable):
  32. """A class for managing multiple kernels."""
  33. default_kernel_name = Unicode(NATIVE_KERNEL_NAME, config=True,
  34. help="The name of the default kernel to start"
  35. )
  36. kernel_spec_manager = Instance(KernelSpecManager, allow_none=True)
  37. kernel_manager_class = DottedObjectName(
  38. "jupyter_client.ioloop.IOLoopKernelManager", config=True,
  39. help="""The kernel manager class. This is configurable to allow
  40. subclassing of the KernelManager for customized behavior.
  41. """
  42. )
  43. @observe('kernel_manager_class')
  44. def _kernel_manager_class_changed(self, name, old, new):
  45. self.kernel_manager_factory = self._create_kernel_manager_factory()
  46. kernel_manager_factory = Any(help="this is kernel_manager_class after import")
  47. @default('kernel_manager_factory')
  48. def _kernel_manager_factory_default(self):
  49. return self._create_kernel_manager_factory()
  50. def _create_kernel_manager_factory(self):
  51. kernel_manager_ctor = import_item(self.kernel_manager_class)
  52. def create_kernel_manager(*args, **kwargs):
  53. if self.shared_context:
  54. if self.context.closed:
  55. # recreate context if closed
  56. self.context = self._context_default()
  57. kwargs.setdefault("context", self.context)
  58. km = kernel_manager_ctor(*args, **kwargs)
  59. return km
  60. return create_kernel_manager
  61. shared_context = Bool(
  62. True,
  63. config=True,
  64. help="Share a single zmq.Context to talk to all my kernels",
  65. )
  66. _created_context = Bool(False)
  67. context = Instance('zmq.Context')
  68. @default("context")
  69. def _context_default(self):
  70. self._created_context = True
  71. return zmq.Context()
  72. def __del__(self):
  73. if self._created_context and self.context and not self.context.closed:
  74. if self.log:
  75. self.log.debug("Destroying zmq context for %s", self)
  76. self.context.destroy()
  77. try:
  78. super_del = super().__del__
  79. except AttributeError:
  80. pass
  81. else:
  82. super_del()
  83. connection_dir = Unicode('')
  84. _kernels = Dict()
  85. def list_kernel_ids(self):
  86. """Return a list of the kernel ids of the active kernels."""
  87. # Create a copy so we can iterate over kernels in operations
  88. # that delete keys.
  89. return list(self._kernels.keys())
  90. def __len__(self):
  91. """Return the number of running kernels."""
  92. return len(self.list_kernel_ids())
  93. def __contains__(self, kernel_id):
  94. return kernel_id in self._kernels
  95. def start_kernel(self, kernel_name=None, **kwargs):
  96. """Start a new kernel.
  97. The caller can pick a kernel_id by passing one in as a keyword arg,
  98. otherwise one will be picked using a uuid.
  99. The kernel ID for the newly started kernel is returned.
  100. """
  101. kernel_id = kwargs.pop('kernel_id', unicode_type(uuid.uuid4()))
  102. if kernel_id in self:
  103. raise DuplicateKernelError('Kernel already exists: %s' % kernel_id)
  104. if kernel_name is None:
  105. kernel_name = self.default_kernel_name
  106. # kernel_manager_factory is the constructor for the KernelManager
  107. # subclass we are using. It can be configured as any Configurable,
  108. # including things like its transport and ip.
  109. constructor_kwargs = {}
  110. if self.kernel_spec_manager:
  111. constructor_kwargs['kernel_spec_manager'] = self.kernel_spec_manager
  112. km = self.kernel_manager_factory(connection_file=os.path.join(
  113. self.connection_dir, "kernel-%s.json" % kernel_id),
  114. parent=self, log=self.log, kernel_name=kernel_name,
  115. **constructor_kwargs
  116. )
  117. km.start_kernel(**kwargs)
  118. self._kernels[kernel_id] = km
  119. return kernel_id
  120. @kernel_method
  121. def shutdown_kernel(self, kernel_id, now=False, restart=False):
  122. """Shutdown a kernel by its kernel uuid.
  123. Parameters
  124. ==========
  125. kernel_id : uuid
  126. The id of the kernel to shutdown.
  127. now : bool
  128. Should the kernel be shutdown forcibly using a signal.
  129. restart : bool
  130. Will the kernel be restarted?
  131. """
  132. self.log.info("Kernel shutdown: %s" % kernel_id)
  133. self.remove_kernel(kernel_id)
  134. @kernel_method
  135. def request_shutdown(self, kernel_id, restart=False):
  136. """Ask a kernel to shut down by its kernel uuid"""
  137. @kernel_method
  138. def finish_shutdown(self, kernel_id, waittime=None, pollinterval=0.1):
  139. """Wait for a kernel to finish shutting down, and kill it if it doesn't
  140. """
  141. self.log.info("Kernel shutdown: %s" % kernel_id)
  142. @kernel_method
  143. def cleanup(self, kernel_id, connection_file=True):
  144. """Clean up a kernel's resources"""
  145. def remove_kernel(self, kernel_id):
  146. """remove a kernel from our mapping.
  147. Mainly so that a kernel can be removed if it is already dead,
  148. without having to call shutdown_kernel.
  149. The kernel object is returned.
  150. """
  151. return self._kernels.pop(kernel_id)
  152. def shutdown_all(self, now=False):
  153. """Shutdown all kernels."""
  154. kids = self.list_kernel_ids()
  155. for kid in kids:
  156. self.request_shutdown(kid)
  157. for kid in kids:
  158. self.finish_shutdown(kid)
  159. self.cleanup(kid)
  160. self.remove_kernel(kid)
  161. @kernel_method
  162. def interrupt_kernel(self, kernel_id):
  163. """Interrupt (SIGINT) the kernel by its uuid.
  164. Parameters
  165. ==========
  166. kernel_id : uuid
  167. The id of the kernel to interrupt.
  168. """
  169. self.log.info("Kernel interrupted: %s" % kernel_id)
  170. @kernel_method
  171. def signal_kernel(self, kernel_id, signum):
  172. """Sends a signal to the kernel by its uuid.
  173. Note that since only SIGTERM is supported on Windows, this function
  174. is only useful on Unix systems.
  175. Parameters
  176. ==========
  177. kernel_id : uuid
  178. The id of the kernel to signal.
  179. """
  180. self.log.info("Signaled Kernel %s with %s" % (kernel_id, signum))
  181. @kernel_method
  182. def restart_kernel(self, kernel_id, now=False):
  183. """Restart a kernel by its uuid, keeping the same ports.
  184. Parameters
  185. ==========
  186. kernel_id : uuid
  187. The id of the kernel to interrupt.
  188. """
  189. self.log.info("Kernel restarted: %s" % kernel_id)
  190. @kernel_method
  191. def is_alive(self, kernel_id):
  192. """Is the kernel alive.
  193. This calls KernelManager.is_alive() which calls Popen.poll on the
  194. actual kernel subprocess.
  195. Parameters
  196. ==========
  197. kernel_id : uuid
  198. The id of the kernel.
  199. """
  200. def _check_kernel_id(self, kernel_id):
  201. """check that a kernel id is valid"""
  202. if kernel_id not in self:
  203. raise KeyError("Kernel with id not found: %s" % kernel_id)
  204. def get_kernel(self, kernel_id):
  205. """Get the single KernelManager object for a kernel by its uuid.
  206. Parameters
  207. ==========
  208. kernel_id : uuid
  209. The id of the kernel.
  210. """
  211. self._check_kernel_id(kernel_id)
  212. return self._kernels[kernel_id]
  213. @kernel_method
  214. def add_restart_callback(self, kernel_id, callback, event='restart'):
  215. """add a callback for the KernelRestarter"""
  216. @kernel_method
  217. def remove_restart_callback(self, kernel_id, callback, event='restart'):
  218. """remove a callback for the KernelRestarter"""
  219. @kernel_method
  220. def get_connection_info(self, kernel_id):
  221. """Return a dictionary of connection data for a kernel.
  222. Parameters
  223. ==========
  224. kernel_id : uuid
  225. The id of the kernel.
  226. Returns
  227. =======
  228. connection_dict : dict
  229. A dict of the information needed to connect to a kernel.
  230. This includes the ip address and the integer port
  231. numbers of the different channels (stdin_port, iopub_port,
  232. shell_port, hb_port).
  233. """
  234. @kernel_method
  235. def connect_iopub(self, kernel_id, identity=None):
  236. """Return a zmq Socket connected to the iopub channel.
  237. Parameters
  238. ==========
  239. kernel_id : uuid
  240. The id of the kernel
  241. identity : bytes (optional)
  242. The zmq identity of the socket
  243. Returns
  244. =======
  245. stream : zmq Socket or ZMQStream
  246. """
  247. @kernel_method
  248. def connect_shell(self, kernel_id, identity=None):
  249. """Return a zmq Socket connected to the shell channel.
  250. Parameters
  251. ==========
  252. kernel_id : uuid
  253. The id of the kernel
  254. identity : bytes (optional)
  255. The zmq identity of the socket
  256. Returns
  257. =======
  258. stream : zmq Socket or ZMQStream
  259. """
  260. @kernel_method
  261. def connect_control(self, kernel_id, identity=None):
  262. """Return a zmq Socket connected to the control channel.
  263. Parameters
  264. ==========
  265. kernel_id : uuid
  266. The id of the kernel
  267. identity : bytes (optional)
  268. The zmq identity of the socket
  269. Returns
  270. =======
  271. stream : zmq Socket or ZMQStream
  272. """
  273. @kernel_method
  274. def connect_stdin(self, kernel_id, identity=None):
  275. """Return a zmq Socket connected to the stdin channel.
  276. Parameters
  277. ==========
  278. kernel_id : uuid
  279. The id of the kernel
  280. identity : bytes (optional)
  281. The zmq identity of the socket
  282. Returns
  283. =======
  284. stream : zmq Socket or ZMQStream
  285. """
  286. @kernel_method
  287. def connect_hb(self, kernel_id, identity=None):
  288. """Return a zmq Socket connected to the hb channel.
  289. Parameters
  290. ==========
  291. kernel_id : uuid
  292. The id of the kernel
  293. identity : bytes (optional)
  294. The zmq identity of the socket
  295. Returns
  296. =======
  297. stream : zmq Socket or ZMQStream
  298. """