Issue #13813: Embed stringification of remote traceback in local

traceback raised when pool task raises an exception.

2 files changed, 54 insertions, 0 deletions

diff --git a/Lib/multiprocessing/pool.py b/Lib/multiprocessing/pool.py
index c0aa717..82a2923 100644
--- a/Lib/multiprocessing/pool.py
+++ b/Lib/multiprocessing/pool.py
@@ -18,6 +18,7 @@ import queue
 import itertools
 import collections
 import time
+import traceback
 
 from multiprocessing import Process, cpu_count, TimeoutError
 from multiprocessing.util import Finalize, debug
@@ -43,6 +44,29 @@ def starmapstar(args):
     return list(itertools.starmap(args[0], args[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
+
+#
 # Code run by worker processes
 #
 
@@ -90,6 +114,7 @@ def worker(inqueue, outqueue, initializer=None, initargs=(), maxtasks=None):
         try:
             result = (True, func(*args, **kwds))
         except Exception as e:
+            e = ExceptionWithTraceback(e, e.__traceback__)
             result = (False, e)
         try:
             put((job, i, result))
diff --git a/Lib/test/test_multiprocessing.py b/Lib/test/test_multiprocessing.py
index 094bbd6..3be72c4 100644
--- a/Lib/test/test_multiprocessing.py
+++ b/Lib/test/test_multiprocessing.py
@@ -1757,6 +1757,35 @@ class _TestPool(BaseTestCase):
                 self.assertEqual(r.get(), expected)
             self.assertRaises(ValueError, p.map_async, sqr, L)
 
+    @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.
+        if self.TYPE == 'processes':
+            with self.Pool(1) as p:
+                try:
+                    p.apply(self._test_traceback)
+                except Exception as e:
+                    exc = e
+                else:
+                    raise AssertionError('expected RuntimeError')
+            self.assertIs(type(exc), RuntimeError)
+            self.assertEqual(exc.args, (123,))
+            cause = exc.__cause__
+            self.assertIs(type(cause), multiprocessing.pool.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())
+
 def raising():
     raise KeyError("key")