Issue #21817: When an exception is raised in a task submitted to a ProcessPoolExecutor, the remote traceback is now displayed in the parent process.

Patch by Claudiu Popa.

2 files changed, 50 insertions, 2 deletions

diff --git a/Lib/concurrent/futures/process.py b/Lib/concurrent/futures/process.py
index fc64dbe..3dd6da1 100644
--- a/Lib/concurrent/futures/process.py
+++ b/Lib/concurrent/futures/process.py
@@ -57,6 +57,7 @@ import threading
 import weakref
 from functools import partial
 import itertools
+import traceback
 
 # Workers are created as daemon threads and processes. This is done to allow the
 # interpreter to exit when there are still idle processes in a
@@ -90,6 +91,27 @@ def _python_exit():
 # (Futures in the call queue cannot be cancelled).
 EXTRA_QUEUED_CALLS = 1
 
+# Hack to embed stringification of remote traceback in local traceback
+
+class _RemoteTraceback(Exception):
+    def __init__(self, tb):
+        self.tb = tb
+    def __str__(self):
+        return self.tb
+
+class _ExceptionWithTraceback:
+    def __init__(self, exc, tb):
+        tb = traceback.format_exception(type(exc), exc, tb)
+        tb = ''.join(tb)
+        self.exc = exc
+        self.tb = '\n"""\n%s"""' % tb
+    def __reduce__(self):
+        return _rebuild_exc, (self.exc, self.tb)
+
+def _rebuild_exc(exc, tb):
+    exc.__cause__ = _RemoteTraceback(tb)
+    return exc
+
 class _WorkItem(object):
     def __init__(self, future, fn, args, kwargs):
         self.future = future
@@ -152,8 +174,8 @@ def _process_worker(call_queue, result_queue):
         try:
             r = call_item.fn(*call_item.args, **call_item.kwargs)
         except BaseException as e:
-            result_queue.put(_ResultItem(call_item.work_id,
-                                         exception=e))
+            exc = _ExceptionWithTraceback(e, e.__traceback__)
+            result_queue.put(_ResultItem(call_item.work_id, exception=exc))
         else:
             result_queue.put(_ResultItem(call_item.work_id,
                                          result=r))
diff --git a/Lib/test/test_concurrent_futures.py b/Lib/test/test_concurrent_futures.py
index 7f92618..86802c2 100644
--- a/Lib/test/test_concurrent_futures.py
+++ b/Lib/test/test_concurrent_futures.py
@@ -480,6 +480,32 @@ class ProcessPoolExecutorTest(ProcessPoolMixin, ExecutorTest, unittest.TestCase)
             ref)
         self.assertRaises(ValueError, bad_map)
 
+    @classmethod
+    def _test_traceback(cls):
+        raise RuntimeError(123) # some comment
+
+    def test_traceback(self):
+        # We want ensure that the traceback from the child process is
+        # contained in the traceback raised in the main process.
+        future = self.executor.submit(self._test_traceback)
+        with self.assertRaises(Exception) as cm:
+            future.result()
+
+        exc = cm.exception
+        self.assertIs(type(exc), RuntimeError)
+        self.assertEqual(exc.args, (123,))
+        cause = exc.__cause__
+        self.assertIs(type(cause), futures.process._RemoteTraceback)
+        self.assertIn('raise RuntimeError(123) # some comment', cause.tb)
+
+        with test.support.captured_stderr() as f1:
+            try:
+                raise exc
+            except RuntimeError:
+                sys.excepthook(*sys.exc_info())
+        self.assertIn('raise RuntimeError(123) # some comment',
+                      f1.getvalue())
+
 
 class FutureTests(unittest.TestCase):
     def test_done_callback_with_result(self):