Merge pull request #663 from minrk/fix-engine-id

Fix engine id, show %px output on error

Author: minrk

ipyparallel/client/asyncresult.py

@@ -990,7 +990,7 @@ def display_outputs(self, groupby="type", result_only=False):
         stderrs = self.stderr
         execute_results = self.execute_result
         output_lists = self.outputs
-        results = self.get()
+        results = self.get(return_exceptions=True)
 
         targets = self.engine_id
 

ipyparallel/client/magics.py

@@ -330,6 +330,10 @@ def result(self, line=''):
         if self.last_result is None:
             raise UsageError(NO_LAST_RESULT)
 
+        if args.save_name:
+            self.shell.user_ns[args.save_name] = self.last_result
+            return
+
         self.last_result.get()
         self.last_result.display_outputs(groupby=args.groupby)
 
@@ -402,37 +406,26 @@ def parallel_execute(
                         # wait for 'quick' results before showing progress
                         tic = time.perf_counter()
                         deadline = tic + progress_after
-                        try:
-                            result.get(timeout=progress_after)
-                            remaining = max(deadline - time.perf_counter(), 0)
-                            result.wait_for_output(timeout=remaining)
-                        except TimeoutError:
-                            pass
-                        except error.CompositeError as e:
-                            if stream_output:
-                                # already streamed, show an abbreviated result
-                                raise error.AlreadyDisplayedError(e) from None
-                            else:
-                                raise
-                        else:
-                            finished_waiting = True
+                        result.wait(timeout=progress_after)
+                        remaining = max(deadline - time.perf_counter(), 0)
+                        result.wait_for_output(timeout=remaining)
+                        finished_waiting = result.done()
 
                     if not finished_waiting:
                         if progress_after >= 0:
                             # not an immediate result, start interactive progress
                             result.wait_interactive()
-                        result.wait_for_output()
-                        try:
-                            result.get()
-                        except error.CompositeError as e:
-                            if stream_output:
-                                # already streamed, show an abbreviated result
-                                raise error.AlreadyDisplayedError(e) from None
-                            else:
-                                raise
+                            result.wait_for_output(1)
+
+                    try:
+                        result.get()
+                    except error.CompositeError as e:
+                        if stream_output:
+                            # already streamed, show an abbreviated result
+                            raise error.AlreadyDisplayedError(e) from None
+                        else:
+                            raise
                 # Skip redisplay if streaming output
-                if not stream_output:
-                    result.display_outputs(groupby)
             except KeyboardInterrupt:
                 if signal_on_interrupt is not None:
                     print(
@@ -444,6 +437,14 @@ def parallel_execute(
                     )
                 else:
                     raise
+            finally:
+                # always redisplay outputs if not streaming,
+                # on both success and error
+
+                if not stream_output:
+                    # wait for at most 1 second for output to be complete
+                    result.wait_for_output(1)
+                    result.display_outputs(groupby)
         else:
             # return AsyncResult only on non-blocking submission
             return result

ipyparallel/engine/kernel.py

@@ -1,6 +1,10 @@
 """IPython kernel for parallel computing"""
+import asyncio
+import inspect
 import sys
+from functools import partial
 
+import ipykernel
 from ipykernel.ipkernel import IPythonKernel
 from traitlets import Integer
 from traitlets import Type
@@ -15,6 +19,11 @@ class IPythonParallelKernel(IPythonKernel):
     """Extend IPython kernel for parallel computing"""
 
     engine_id = Integer(-1)
+
+    @property
+    def int_id(self):
+        return self.engine_id
+
     msg_types = getattr(IPythonKernel, 'msg_types', []) + ['apply_request']
     control_msg_types = getattr(IPythonKernel, 'control_msg_types', []) + [
         'abort_request',
@@ -40,6 +49,47 @@ def __init__(self, **kwargs):
         data_pub.pub_socket = self.iopub_socket
         self.aborted = set()
 
+    def _abort_queues(self):
+        # forward-port ipython/ipykernel#853
+        # may remove after requiring ipykernel 6.9
+
+        # while this flag is true,
+        # execute requests will be aborted
+        self._aborting = True
+        self.log.info("Aborting queue")
+
+        # Callback to signal that we are done aborting
+        def stop_aborting():
+            self.log.info("Finishing abort")
+            self._aborting = False
+
+        # put stop_aborting on the message queue
+        # so that it's handled after processing of already-pending messages
+        if ipykernel.version_info < (6,):
+            # 10 is SHELL priority in ipykernel 5.x
+            streams = self.shell_streams
+            schedule_stop_aborting = partial(self.schedule_dispatch, 10, stop_aborting)
+        else:
+            streams = [self.shell_stream]
+            schedule_stop_aborting = partial(self.schedule_dispatch, stop_aborting)
+
+        # flush streams, so all currently waiting messages
+        # are added to the queue
+        for stream in streams:
+            stream.flush()
+
+        # if we have a delay, give messages this long to arrive on the queue
+        # before we start accepting requests
+        asyncio.get_event_loop().call_later(
+            self.stop_on_error_timeout, schedule_stop_aborting
+        )
+
+        # for compatibility, return a completed Future
+        # so this is still awaitable
+        f = asyncio.Future()
+        f.set_result(None)
+        return f
+
     def should_handle(self, stream, msg, idents):
         """Check whether a shell-channel message should be handled
 
@@ -78,13 +128,20 @@ def finish_metadata(self, parent, metadata, reply_content):
         if reply_content['status'] == 'error':
             if reply_content['ename'] == 'UnmetDependency':
                 metadata['dependencies_met'] = False
-            metadata['engine_info'] = dict(
-                engine_uuid=self.ident,
-                engine_id=self.engine_id,
-            )
+            metadata['engine_info'] = self.get_engine_info()
 
         return metadata
 
+    def get_engine_info(self, method=None):
+        """Return engine_info dict"""
+        engine_info = dict(
+            engine_uuid=self.ident,
+            engine_id=self.engine_id,
+        )
+        if method:
+            engine_info["method"] = method
+        return engine_info
+
     def apply_request(self, stream, ident, parent):
         try:
             content = parent['content']
@@ -168,8 +225,7 @@ def do_apply(self, content, bufs, msg_id, reply_metadata):
             else:
                 self.log.warning("Didn't find a traceback where I expected to")
             shell._last_traceback = None
-            e_info = dict(engine_uuid=self.ident, engine_id=self.int_id, method='apply')
-            reply_content['engine_info'] = e_info
+            reply_content["engine_info"] = self.get_engine_info(method="apply")
 
             self.log.info(
                 "Exception in apply request:\n%s", '\n'.join(reply_content['traceback'])
@@ -181,6 +237,26 @@ def do_apply(self, content, bufs, msg_id, reply_metadata):
 
         return reply_content, result_buf
 
+    async def _do_execute_async(self, *args, **kwargs):
+        super_execute = super().do_execute(*args, **kwargs)
+        if inspect.isawaitable(super_execute):
+            reply_content = await super_execute
+        else:
+            reply_content = super_execute
+        # add engine info
+        if reply_content['status'] == 'error':
+            reply_content["engine_info"] = self.get_engine_info(method="execute")
+        return reply_content
+
+    def do_execute(self, *args, **kwargs):
+        coro = self._do_execute_async(*args, **kwargs)
+        if ipykernel.version_info < (6,):
+            # ipykernel 5 uses gen.maybe_future which doesn't accept async def coroutines,
+            # but it does accept asyncio.Futures
+            return asyncio.ensure_future(coro)
+        else:
+            return coro
+
     # Control messages for msgspec extensions:
 
     def abort_request(self, stream, ident, parent):
@@ -189,7 +265,9 @@ def abort_request(self, stream, ident, parent):
         if isinstance(msg_ids, str):
             msg_ids = [msg_ids]
         if not msg_ids:
-            self._abort_queues()
+            f = self._abort_queues()
+            if inspect.isawaitable(f):
+                asyncio.ensure_future(f)
         for mid in msg_ids:
             self.aborted.add(str(mid))
 

ipyparallel/tests/__init__.py

@@ -124,7 +124,7 @@ def teardown():
             try:
                 f = p.stop()
                 if f:
-                    asyncio.run(f)
+                    asyncio.get_event_loop().run_until_complete(f)
             except Exception as e:
                 print(e)
                 pass

ipyparallel/tests/test_asyncresult.py

@@ -5,6 +5,7 @@
 import time
 from datetime import datetime
 
+import pytest
 from IPython.utils.io import capture_output
 
 import ipyparallel as ipp
@@ -287,6 +288,20 @@ def test_display_empty_streams_engine(self):
         self.assertEqual(io.stderr, '')
         self.assertEqual(io.stdout, '')
 
+    def test_display_output_error(self):
+        """display_outputs shows output on error"""
+        self.minimum_engines(1)
+
+        v = self.client[-1]
+        ar = v.execute("print (5555)\n1/0")
+        ar.get(5, return_exceptions=True)
+        ar.wait_for_output(5)
+        with capture_output() as io:
+            ar.display_outputs()
+        self.assertEqual(io.stderr, '')
+        self.assertEqual('5555\n', io.stdout)
+        assert 'ZeroDivisionError' not in io.stdout
+
     def test_await_data(self):
         """asking for ar.data flushes outputs"""
         self.minimum_engines(1)
@@ -488,3 +503,47 @@ def test_progress(self):
         assert amr.progress == 0
         amr.wait_interactive()
         assert amr.progress == len(amr)
+
+    def test_error_engine_info_apply(self):
+        dv = self.client[:]
+        targets = self.client.ids
+        ar = dv.apply_async(lambda: 1 / 0)
+        try:
+            ar.get()
+        except Exception as e:
+            exc = e
+        else:
+            pytest.fail("Should have raised remote ZeroDivisionError")
+        assert isinstance(exc, ipp.error.CompositeError)
+        expected_engine_info = [
+            {
+                "engine_id": engine_id,
+                "engine_uuid": self.client._engines[engine_id],
+                "method": "apply",
+            }
+            for engine_id in self.client.ids
+        ]
+        engine_infos = [e[-1] for e in exc.elist]
+        assert engine_infos == expected_engine_info
+
+    def test_error_engine_info_execute(self):
+        dv = self.client[:]
+        targets = self.client.ids
+        ar = dv.execute("1 / 0", block=False)
+        try:
+            ar.get()
+        except Exception as e:
+            exc = e
+        else:
+            pytest.fail("Should have raised remote ZeroDivisionError")
+        assert isinstance(exc, ipp.error.CompositeError)
+        expected_engine_info = [
+            {
+                "engine_id": engine_id,
+                "engine_uuid": self.client._engines[engine_id],
+                "method": "execute",
+            }
+            for engine_id in self.client.ids
+        ]
+        engine_infos = [e[-1] for e in exc.elist]
+        assert engine_infos == expected_engine_info