Fixing #659 (output cutoff and iopub timeout)

nbconvert/preprocessors/execute.py

@@ -6,6 +6,10 @@
 import base64
 from textwrap import dedent
 from contextlib import contextmanager
+try:
+    from time import monotonic # Py 3
+except ImportError:
+    from time import time as monotonic # Py 2
 
 try:
     from queue import Empty  # Py 3
@@ -449,6 +453,36 @@ def _update_display_id(self, display_id, msg):
                 outputs[output_idx]['data'] = out['data']
                 outputs[output_idx]['metadata'] = out['metadata']
 
+    def _poll_for_reply(self, msg_id, cell=None, timeout=None):
+        try:
+            # check with timeout if kernel is still alive
+            msg = self.kc.shell_channel.get_msg(timeout=timeout)
+            if msg['parent_header'].get('msg_id') == msg_id:
+                return msg
+        except Empty:
+            # received no message, check if kernel is still alive
+            self._check_alive()
+            # kernel still alive, wait for a message
+
+    def _get_timeout(self, cell):
+        if self.timeout_func is not None and cell is not None:
+            timeout = self.timeout_func(cell)
+        else:
+            timeout = self.timeout
+
+        if not timeout or timeout < 0:
+            timeout = None
+
+        return timeout
+
+    def _handle_timeout(self):
+        self.log.error(
+            "Timeout waiting for execute reply (%is)." % self.timeout)
+        if self.interrupt_on_timeout:
+            self.log.error("Interrupting kernel")
+            self.km.interrupt_kernel()
+        else:
+            raise TimeoutError("Cell execution timed out")
 
     def _check_alive(self):
         if not self.kc.is_alive():
@@ -458,13 +492,7 @@ def _check_alive(self):
 
     def _wait_for_reply(self, msg_id, cell=None):
         # wait for finish, with timeout
-        if self.timeout_func is not None and cell is not None:
-            timeout = self.timeout_func(cell)
-        else:
-            timeout = self.timeout
-
-        if not timeout or timeout < 0:
-            timeout = None
+        timeout = self._get_timeout(cell)
         cummulative_time = 0
         timeout_interval = 5
         while True:
@@ -474,49 +502,77 @@ def _wait_for_reply(self, msg_id, cell=None):
                 self._check_alive()
                 cummulative_time += timeout_interval
                 if timeout and cummulative_time > timeout:
-                    self.log.error(
-                        "Timeout waiting for execute reply (%is)." % self.timeout)
-                    if self.interrupt_on_timeout:
-                        self.log.error("Interrupting kernel")
-                        self.km.interrupt_kernel()
-                        break
-                    else:
-                        raise TimeoutError("Cell execution timed out")
+                    self._handle_timeout()
+                    break
             else:
                 if msg['parent_header'].get('msg_id') == msg_id:
                     return msg
 
+    def _timeout_with_deadline(self, timeout, deadline):
+        if deadline is not None and deadline - monotonic() < timeout:
+            timeout = deadline - monotonic()
+
+        if timeout < 0:
+            timeout = 0
+
+        return timeout
+
+    def _passed_deadline(self, deadline):
+        if deadline is not None and deadline - monotonic() <= 0:
+            self._handle_timeout()
+            return True
+        return False
+
     def run_cell(self, cell, cell_index=0):
         parent_msg_id = self.kc.execute(cell.source)
         self.log.debug("Executing cell:\n%s", cell.source)
-        exec_reply = self._wait_for_reply(parent_msg_id, cell)
+        exec_timeout = self._get_timeout(cell)
+        deadline = None
+        if exec_timeout is not None:
+            deadline = monotonic() + exec_timeout
 
         cell.outputs = []
         self.clear_before_next_output = False
 
-        while True:
-            try:
-                # We've already waited for execute_reply, so all output
-                # should already be waiting. However, on slow networks, like
-                # in certain CI systems, waiting < 1 second might miss messages.
-                # So long as the kernel sends a status:idle message when it
-                # finishes, we won't actually have to wait this long, anyway.
-                msg = self.kc.iopub_channel.get_msg(timeout=self.iopub_timeout)
-            except Empty:
-                self.log.warning("Timeout waiting for IOPub output")
-                if self.raise_on_iopub_timeout:
-                    raise RuntimeError("Timeout waiting for IOPub output")
-                else:
-                    break
+        more_output = True
+        polling_exec_reply = True
+
+        while more_output or polling_exec_reply:
+            if polling_exec_reply:
+                if self._passed_deadline(deadline):
+                    polling_exec_reply = False
+                    continue
+
+                timeout = self._timeout_with_deadline(1, deadline)
+                exec_reply = self._poll_for_reply(parent_msg_id, cell, timeout)
+                if exec_reply is not None:
+                    polling_exec_reply = False
+
+            if more_output:
+                try:
+                    timeout = self.iopub_timeout
+                    if polling_exec_reply:
+                        timeout = self._timeout_with_deadline(timeout, deadline)
+                    msg = self.kc.iopub_channel.get_msg(timeout=timeout)
+                except Empty:
+                    if polling_exec_reply:
+                        continue
+
+                    if self.raise_on_iopub_timeout:
+                        raise TimeoutError("Timeout waiting for IOPub output")
+                    else:
+                        self.log.warning("Timeout waiting for IOPub output")
+                        more_output = False
+                        continue
             if msg['parent_header'].get('msg_id') != parent_msg_id:
                 # not an output from our execution
                 continue
 
-            # Will raise CellExecutionComplete when completed
             try:
+                # Will raise CellExecutionComplete when completed
                 self.process_message(msg, cell, cell_index)
             except CellExecutionComplete:
-                break
+                more_output = False
 
         # Return cell.outputs still for backwards compatability
         return exec_reply, cell.outputs

nbconvert/preprocessors/tests/test_execute.py

@@ -32,6 +32,10 @@
 from testpath import modified_env
 from ipython_genutils.py3compat import string_types
 
+try:
+    from queue import Empty  # Py 3
+except ImportError:
+    from Queue import Empty  # Py 2
 try:
     TimeoutError  # Py 3
 except NameError:
@@ -167,6 +171,7 @@ def test_mock_wrapper(self):
                 shell_channel=MagicMock(get_msg=shell_channel_message_mock()),
                 execute=MagicMock(return_value=parent_id)
             )
+            preprocessor.parent_id = parent_id
             return func(self, preprocessor, cell_mock, message_mock)
         return test_mock_wrapper
     return prepared_wrapper
@@ -510,6 +515,65 @@ def test_busy_message(self, preprocessor, cell_mock, message_mock):
         # Ensure no outputs were generated
         assert cell_mock.outputs == []
 
+    @prepare_cell_mocks({
+        'msg_type': 'stream',
+        'header': {'msg_type': 'stream'},
+        'content': {'name': 'stdout', 'text': 'foo'},
+    }, {
+        'msg_type': 'stream',
+        'header': {'msg_type': 'stream'},
+        'content': {'name': 'stderr', 'text': 'bar'}
+    })
+    def test_deadline_exec_reply(self, preprocessor, cell_mock, message_mock):
+        # exec_reply is never received, so we expect to hit the timeout.
+        preprocessor.kc.shell_channel.get_msg = MagicMock(side_effect=Empty())
+        preprocessor.timeout = 1
+
+        with pytest.raises(TimeoutError):
+            preprocessor.run_cell(cell_mock)
+
+        assert message_mock.call_count == 3
+        # Ensure the output was captured
+        self.assertListEqual(cell_mock.outputs, [
+            {'output_type': 'stream', 'name': 'stdout', 'text': 'foo'},
+            {'output_type': 'stream', 'name': 'stderr', 'text': 'bar'}
+        ])
+
+    @prepare_cell_mocks()
+    def test_deadline_iopub(self, preprocessor, cell_mock, message_mock):
+        # The shell_channel will complete, so we expect only to hit the iopub timeout.
+        message_mock.side_effect = Empty()
+        preprocessor.raise_on_iopub_timeout = True
+
+        with pytest.raises(TimeoutError):
+            preprocessor.run_cell(cell_mock)
+
+    @prepare_cell_mocks({
+        'msg_type': 'stream',
+        'header': {'msg_type': 'stream'},
+        'content': {'name': 'stdout', 'text': 'foo'},
+    }, {
+        'msg_type': 'stream',
+        'header': {'msg_type': 'stream'},
+        'content': {'name': 'stderr', 'text': 'bar'}
+    })
+    def test_eventual_deadline_iopub(self, preprocessor, cell_mock, message_mock):
+        # Process a few messages before raising a timeout from iopub
+        message_mock.side_effect = list(message_mock.side_effect)[:-1] + [Empty()]
+        preprocessor.kc.shell_channel.get_msg = MagicMock(
+            return_value={'parent_header': {'msg_id': preprocessor.parent_id}})
+        preprocessor.raise_on_iopub_timeout = True
+
+        with pytest.raises(TimeoutError):
+            preprocessor.run_cell(cell_mock)
+
+        assert message_mock.call_count == 3
+        # Ensure the output was captured
+        self.assertListEqual(cell_mock.outputs, [
+            {'output_type': 'stream', 'name': 'stdout', 'text': 'foo'},
+            {'output_type': 'stream', 'name': 'stderr', 'text': 'bar'}
+        ])
+
     @prepare_cell_mocks({
         'msg_type': 'execute_input',
         'header': {'msg_type': 'execute_input'},