You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

332 lines
14 KiB

4 years ago
  1. #############################################################################
  2. # Copyright (c) 2018, Voila Contributors #
  3. # Copyright (c) 2018, QuantStack #
  4. # #
  5. # Distributed under the terms of the BSD 3-Clause License. #
  6. # #
  7. # The full license is in the file LICENSE, distributed with this software. #
  8. #############################################################################
  9. import collections
  10. import logging
  11. try:
  12. from time import monotonic # Py 3
  13. except ImportError:
  14. from time import time as monotonic # Py 2
  15. from nbconvert.preprocessors import ClearOutputPreprocessor
  16. from nbconvert.preprocessors.execute import CellExecutionError, ExecutePreprocessor
  17. from nbformat.v4 import output_from_msg
  18. import zmq
  19. from traitlets import Unicode
  20. from ipykernel.jsonutil import json_clean
  21. try:
  22. TimeoutError # Py 3
  23. except NameError:
  24. TimeoutError = RuntimeError # Py 2
  25. def strip_code_cell_warnings(cell):
  26. """Strip any warning outputs and traceback from a code cell."""
  27. if cell['cell_type'] != 'code':
  28. return cell
  29. outputs = cell['outputs']
  30. cell['outputs'] = [
  31. output for output in outputs
  32. if output['output_type'] != 'stream' or output['name'] != 'stderr'
  33. ]
  34. return cell
  35. def should_strip_error(config):
  36. """Return True if errors should be stripped from the Notebook, False otherwise, depending on the current config."""
  37. return 'Voila' not in config or 'log_level' not in config['Voila'] or config['Voila']['log_level'] != logging.DEBUG
  38. class OutputWidget:
  39. """This class mimics a front end output widget"""
  40. def __init__(self, comm_id, state, kernel_client, executor):
  41. self.comm_id = comm_id
  42. self.state = state
  43. self.kernel_client = kernel_client
  44. self.executor = executor
  45. self.topic = ('comm-%s' % self.comm_id).encode('ascii')
  46. self.outputs = self.state['outputs']
  47. self.clear_before_next_output = False
  48. def clear_output(self, outs, msg, cell_index):
  49. self.parent_header = msg['parent_header']
  50. content = msg['content']
  51. if content.get('wait'):
  52. self.clear_before_next_output = True
  53. else:
  54. self.outputs = []
  55. # sync back the state to the kernel
  56. self.sync_state()
  57. if hasattr(self.executor, 'widget_state'):
  58. # sync the state to the nbconvert state as well, since that is used for testing
  59. self.executor.widget_state[self.comm_id]['outputs'] = self.outputs
  60. def sync_state(self):
  61. state = {'outputs': self.outputs}
  62. msg = {'method': 'update', 'state': state, 'buffer_paths': []}
  63. self.send(msg)
  64. def _publish_msg(self, msg_type, data=None, metadata=None, buffers=None, **keys):
  65. """Helper for sending a comm message on IOPub"""
  66. data = {} if data is None else data
  67. metadata = {} if metadata is None else metadata
  68. content = json_clean(dict(data=data, comm_id=self.comm_id, **keys))
  69. msg = self.kernel_client.session.msg(msg_type, content=content, parent=self.parent_header, metadata=metadata)
  70. self.kernel_client.shell_channel.send(msg)
  71. def send(self, data=None, metadata=None, buffers=None):
  72. self._publish_msg('comm_msg', data=data, metadata=metadata, buffers=buffers)
  73. def output(self, outs, msg, display_id, cell_index):
  74. if self.clear_before_next_output:
  75. self.outputs = []
  76. self.clear_before_next_output = False
  77. self.parent_header = msg['parent_header']
  78. output = output_from_msg(msg)
  79. if self.outputs:
  80. # try to coalesce/merge output text
  81. last_output = self.outputs[-1]
  82. if (last_output['output_type'] == 'stream' and
  83. output['output_type'] == 'stream' and
  84. last_output['name'] == output['name']):
  85. last_output['text'] += output['text']
  86. else:
  87. self.outputs.append(output)
  88. else:
  89. self.outputs.append(output)
  90. self.sync_state()
  91. if hasattr(self.executor, 'widget_state'):
  92. # sync the state to the nbconvert state as well, since that is used for testing
  93. self.executor.widget_state[self.comm_id]['outputs'] = self.outputs
  94. def set_state(self, state):
  95. if 'msg_id' in state:
  96. msg_id = state.get('msg_id')
  97. if msg_id:
  98. self.executor.register_output_hook(msg_id, self)
  99. self.msg_id = msg_id
  100. else:
  101. self.executor.remove_output_hook(self.msg_id, self)
  102. self.msg_id = msg_id
  103. class VoilaExecutePreprocessor(ExecutePreprocessor):
  104. """Execute, but respect the output widget behaviour"""
  105. cell_error_instruction = Unicode(
  106. 'Please run Voila with --debug to see the error message.',
  107. config=True,
  108. help=(
  109. 'instruction given to user to debug cell errors'
  110. )
  111. )
  112. def __init__(self, **kwargs):
  113. super(VoilaExecutePreprocessor, self).__init__(**kwargs)
  114. self.output_hook_stack = collections.defaultdict(list) # maps to list of hooks, where the last is used
  115. self.output_objects = {}
  116. def preprocess(self, nb, resources, km=None):
  117. try:
  118. result = super(VoilaExecutePreprocessor, self).preprocess(nb, resources=resources, km=km)
  119. except CellExecutionError as e:
  120. self.log.error(e)
  121. result = (nb, resources)
  122. # Strip errors and traceback if not in debug mode
  123. if should_strip_error(self.config):
  124. self.strip_notebook_errors(nb)
  125. return result
  126. def preprocess_cell(self, cell, resources, cell_index, store_history=True):
  127. try:
  128. # TODO: pass store_history as a 5th argument when we can require nbconver >=5.6.1
  129. # result = super(VoilaExecutePreprocessor, self).preprocess_cell(cell, resources, cell_index, store_history)
  130. result = super(VoilaExecutePreprocessor, self).preprocess_cell(cell, resources, cell_index)
  131. except CellExecutionError as e:
  132. self.log.error(e)
  133. result = (cell, resources)
  134. # Strip errors and traceback if not in debug mode
  135. if should_strip_error(self.config):
  136. strip_code_cell_warnings(cell)
  137. self.strip_code_cell_errors(cell)
  138. return result
  139. def register_output_hook(self, msg_id, hook):
  140. # mimics
  141. # https://jupyterlab.github.io/jupyterlab/services/interfaces/kernel.ikernelconnection.html#registermessagehook
  142. self.output_hook_stack[msg_id].append(hook)
  143. def remove_output_hook(self, msg_id, hook):
  144. # mimics
  145. # https://jupyterlab.github.io/jupyterlab/services/interfaces/kernel.ikernelconnection.html#removemessagehook
  146. removed_hook = self.output_hook_stack[msg_id].pop()
  147. assert removed_hook == hook
  148. def output(self, outs, msg, display_id, cell_index):
  149. parent_msg_id = msg['parent_header'].get('msg_id')
  150. if self.output_hook_stack[parent_msg_id]:
  151. hook = self.output_hook_stack[parent_msg_id][-1]
  152. hook.output(outs, msg, display_id, cell_index)
  153. return
  154. super(VoilaExecutePreprocessor, self).output(outs, msg, display_id, cell_index)
  155. def handle_comm_msg(self, outs, msg, cell_index):
  156. super(VoilaExecutePreprocessor, self).handle_comm_msg(outs, msg, cell_index)
  157. self.log.debug('comm msg: %r', msg)
  158. if msg['msg_type'] == 'comm_open' and msg['content'].get('target_name') == 'jupyter.widget':
  159. content = msg['content']
  160. data = content['data']
  161. state = data['state']
  162. comm_id = msg['content']['comm_id']
  163. if state['_model_module'] == '@jupyter-widgets/output' and state['_model_name'] == 'OutputModel':
  164. self.output_objects[comm_id] = OutputWidget(comm_id, state, self.kc, self)
  165. elif msg['msg_type'] == 'comm_msg':
  166. content = msg['content']
  167. data = content['data']
  168. if 'state' in data:
  169. state = data['state']
  170. comm_id = msg['content']['comm_id']
  171. if comm_id in self.output_objects:
  172. self.output_objects[comm_id].set_state(state)
  173. def clear_output(self, outs, msg, cell_index):
  174. parent_msg_id = msg['parent_header'].get('msg_id')
  175. if self.output_hook_stack[parent_msg_id]:
  176. hook = self.output_hook_stack[parent_msg_id][-1]
  177. hook.clear_output(outs, msg, cell_index)
  178. return
  179. super(VoilaExecutePreprocessor, self).clear_output(outs, msg, cell_index)
  180. def strip_notebook_errors(self, nb):
  181. """Strip error messages and traceback from a Notebook."""
  182. cells = nb['cells']
  183. code_cells = [cell for cell in cells if cell['cell_type'] == 'code']
  184. for cell in code_cells:
  185. strip_code_cell_warnings(cell)
  186. self.strip_code_cell_errors(cell)
  187. return nb
  188. def strip_code_cell_errors(self, cell):
  189. """Strip any error outputs and traceback from a code cell."""
  190. # There is no 'outputs' key for markdown cells
  191. if cell['cell_type'] != 'code':
  192. return cell
  193. outputs = cell['outputs']
  194. error_outputs = [output for output in outputs if output['output_type'] == 'error']
  195. error_message = 'There was an error when executing cell [{}]. {}'.format(cell['execution_count'], self.cell_error_instruction)
  196. for output in error_outputs:
  197. output['ename'] = 'ExecutionError'
  198. output['evalue'] = 'Execution error'
  199. output['traceback'] = [error_message]
  200. return cell
  201. # make it nbconvert 5.5 compatible
  202. def _get_timeout(self, cell):
  203. if self.timeout_func is not None and cell is not None:
  204. timeout = self.timeout_func(cell)
  205. else:
  206. timeout = self.timeout
  207. if not timeout or timeout < 0:
  208. timeout = None
  209. return timeout
  210. # make it nbconvert 5.5 compatible
  211. def _handle_timeout(self, timeout):
  212. self.log.error(
  213. "Timeout waiting for execute reply (%is)." % timeout)
  214. if self.interrupt_on_timeout:
  215. self.log.error("Interrupting kernel")
  216. self.km.interrupt_kernel()
  217. else:
  218. raise TimeoutError("Cell execution timed out")
  219. def run_cell(self, cell, cell_index=0, store_history=False):
  220. parent_msg_id = self.kc.execute(cell.source, store_history=store_history, stop_on_error=not self.allow_errors)
  221. self.log.debug("Executing cell:\n%s", cell.source)
  222. exec_timeout = self._get_timeout(cell)
  223. deadline = None
  224. if exec_timeout is not None:
  225. deadline = monotonic() + exec_timeout
  226. cell.outputs = []
  227. self.clear_before_next_output = False
  228. # we need to have a reply, and return to idle before we can consider the cell executed
  229. idle = False
  230. execute_reply = None
  231. deadline_passed = 0
  232. deadline_passed_max = 5
  233. while not idle or execute_reply is None:
  234. # we want to timeout regularly, to see if the kernel is still alive
  235. # this is tested in preprocessors/test/test_execute.py#test_kernel_death
  236. # this actually fakes the kernel death, and we might be able to use the xlist
  237. # to detect a disconnected kernel
  238. timeout = min(1, deadline - monotonic())
  239. # if we interrupt on timeout, we allow 1 seconds to pass till deadline
  240. # to make sure we get the interrupt message
  241. if timeout >= 0.0:
  242. # we include 0, which simply is a poll to see if we have messages left
  243. rlist = zmq.select([self.kc.iopub_channel.socket, self.kc.shell_channel.socket], [], [], timeout)[0]
  244. if not rlist:
  245. self._check_alive()
  246. if monotonic() > deadline:
  247. self._handle_timeout(exec_timeout)
  248. deadline_passed += 1
  249. assert self.interrupt_on_timeout
  250. if deadline_passed <= deadline_passed_max:
  251. # when we interrupt, we want to do this multiple times, so we give some
  252. # extra time to handle the interrupt message
  253. deadline += self.iopub_timeout
  254. if self.kc.shell_channel.socket in rlist:
  255. msg = self.kc.shell_channel.get_msg(block=False)
  256. if msg['parent_header'].get('msg_id') == parent_msg_id:
  257. execute_reply = msg
  258. if self.kc.iopub_channel.socket in rlist:
  259. msg = self.kc.iopub_channel.get_msg(block=False)
  260. if msg['parent_header'].get('msg_id') == parent_msg_id:
  261. if msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
  262. idle = True
  263. else:
  264. self.process_message(msg, cell, cell_index)
  265. else:
  266. self.log.debug("Received message for which we were not the parent: %s", msg)
  267. else:
  268. self._handle_timeout(exec_timeout)
  269. break
  270. return execute_reply, cell.outputs
  271. def executenb(nb, cwd=None, km=None, **kwargs):
  272. resources = {}
  273. if cwd is not None:
  274. resources['metadata'] = {'path': cwd} # pragma: no cover
  275. # Clear any stale output, in case of exception
  276. nb, resources = ClearOutputPreprocessor().preprocess(nb, resources)
  277. ep = VoilaExecutePreprocessor(**kwargs)
  278. return ep.preprocess(nb, resources, km=km)[0]