Raise new group exceptions on callback failure instead of just processing them

with the default exception hook.

This is a breaking change to the API, but only in cases where the existing code
would be completely broken anyway, so hopefully it isn't a problem.

TODO: docs need updating

Author: duaneg

Lib/multiprocessing/managers.py

@@ -1232,7 +1232,7 @@ def __isub__(self, value):
 
 BasePoolProxy = MakeProxyType('PoolProxy', (
     'apply', 'apply_async', 'close', 'imap', 'imap_unordered', 'join',
-    'map', 'map_async', 'starmap', 'starmap_async', 'terminate',
+    'map', 'map_async', 'starmap', 'starmap_async', 'terminate', '_check_error'
     ))
 BasePoolProxy._method_to_typeid_ = {
     'apply_async': 'AsyncResult',
@@ -1246,6 +1246,7 @@ def __enter__(self):
         return self
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.terminate()
+        self._check_error(exc_val)
 
 #
 # Definition of SyncManager

Lib/multiprocessing/pool.py

@@ -7,7 +7,7 @@
 # Licensed to PSF under a Contributor Agreement.
 #
 
-__all__ = ['Pool', 'ThreadPool']
+__all__ = ['BrokenPoolError', 'CallbackError', 'Pool', 'ThreadPool']
 
 #
 # Imports
@@ -69,6 +69,14 @@ def __init__(self, exc, tb):
     def __reduce__(self):
         return rebuild_exc, (self.exc, self.tb)
 
+class BrokenPoolError(ExceptionGroup):
+    def __init__(self, msg, exc):
+        super().__init__(msg, exc)
+
+class CallbackError(ExceptionGroup):
+    def __init__(self, msg, exc):
+        super().__init__(msg, exc)
+
 def rebuild_exc(exc, tb):
     exc.__cause__ = RemoteTraceback(tb)
     return exc
@@ -198,6 +206,7 @@ def __init__(self, processes=None, initializer=None, initargs=(),
         self._maxtasksperchild = maxtasksperchild
         self._initializer = initializer
         self._initargs = initargs
+        self._errors = []
 
         if processes is None:
             processes = os.process_cpu_count() or 1
@@ -349,9 +358,17 @@ def _setup_queues(self):
         self._quick_get = self._outqueue._reader.recv
 
     def _check_running(self):
+        self._check_error()
         if self._state != RUN:
             raise ValueError("Pool not running")
 
+    def _check_error(self, exc=None):
+        if self._errors:
+            errs = list(self._errors)
+            if exc is not None and not isinstance(exc, CallbackError):
+                errs.append(exc)
+            raise BrokenPoolError("Callback(s) failed", errs) from None
+
     def apply(self, func, args=(), kwds={}):
         '''
         Equivalent of `func(*args, **kwds)`.
@@ -737,6 +754,11 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.terminate()
+        self._check_error(exc_val)
+
+    def _error(self, error):
+        util.debug('callback error', exc_info=error)
+        self._errors.append(error)
 
 #
 # Class whose instances are returned by `Pool.apply_async()`
@@ -751,6 +773,7 @@ def __init__(self, pool, callback, error_callback):
         self._cache = pool._cache
         self._callback = callback
         self._error_callback = error_callback
+        self._cb_error = None
         self._cache[self._job] = self
 
     def ready(self):
@@ -768,30 +791,31 @@ def get(self, timeout=None):
         self.wait(timeout)
         if not self.ready():
             raise TimeoutError
-        if self._success:
+        if self._cb_error:
+            raise self._cb_error
+        elif self._success:
             return self._value
         else:
             raise self._value
 
     def _set(self, i, obj):
         self._success, self._value = obj
         if self._callback and self._success:
-            self._handle_exceptions(self._callback, self._value)
+            try:
+                self._callback(self._value)
+            except Exception as e:
+                self._cb_error = CallbackError("apply callback", [e])
+                self._pool._error(self._cb_error)
         if self._error_callback and not self._success:
-            self._handle_exceptions(self._error_callback, self._value)
+            try:
+                self._error_callback(self._value)
+            except Exception as e:
+                self._cb_error = CallbackError("apply error callback", [e, self._value])
+                self._pool._error(self._cb_error)
         self._event.set()
         del self._cache[self._job]
         self._pool = None
 
-    @staticmethod
-    def _handle_exceptions(callback, args):
-        try:
-            return callback(args)
-        except Exception as e:
-            args = threading.ExceptHookArgs([type(e), e, e.__traceback__, None])
-            threading.excepthook(args)
-            del args
-
     __class_getitem__ = classmethod(types.GenericAlias)
 
 AsyncResult = ApplyResult       # create alias -- see #17805
@@ -822,7 +846,11 @@ def _set(self, i, success_result):
             self._value[i*self._chunksize:(i+1)*self._chunksize] = result
             if self._number_left == 0:
                 if self._callback:
-                    self._handle_exceptions(self._callback, self._value)
+                    try:
+                        self._callback(self._value)
+                    except Exception as e:
+                        self._cb_error = CallbackError("map callback", [e])
+                        self._pool._error(self._cb_error)
                 del self._cache[self._job]
                 self._event.set()
                 self._pool = None
@@ -834,7 +862,12 @@ def _set(self, i, success_result):
             if self._number_left == 0:
                 # only consider the result ready once all jobs are done
                 if self._error_callback:
-                    self._handle_exceptions(self._error_callback, self._value)
+                    try:
+                        self._error_callback(self._value)
+                    except Exception as e:
+                        self._cb_error = CallbackError("map error callback",
+                                                       [e, self._value])
+                        self._pool._error(self._cb_error)
                 del self._cache[self._job]
                 self._event.set()
                 self._pool = None

Lib/test/_test_multiprocessing.py

@@ -54,6 +54,7 @@
 import multiprocessing.pool
 import multiprocessing.queues
 from multiprocessing.connection import wait
+from multiprocessing.pool import BrokenPoolError, CallbackError
 
 from multiprocessing import util
 
@@ -3083,9 +3084,89 @@ def test_resource_warning(self):
             pool = None
             support.gc_collect()
 
-def raising():
+    def test_callback_errors(self):
+        def _apply(pool, target, **kwargs):
+            return pool.apply_async(target, **kwargs)
+
+        def _map(pool, target, **kwargs):
+            return pool.map_async(target, range(5), **kwargs)
+
+        for func in [_apply, _map]:
+            with self.subTest(func=func):
+
+                # Fail upon trying to reuse a broken pool after callback failure:
+                # - BrokenPoolError containing:
+                #   - CallbackError containing:
+                #     - Error thrown from the callback
+                with self.assertRaises(BrokenPoolError) as pool_ctx:
+                    with self.Pool(1) as pool:
+                        res = func(pool, noop, callback=raising)
+                        with self.assertRaises(CallbackError) as res_ctx:
+                            res.get()
+                    self._check_subexceptions(res_ctx.exception, [KeyError])
+                    pool.apply_async(noop)
+                self._check_subexceptions(pool_ctx.exception, [CallbackError])
+                self._check_subexceptions(pool_ctx.exception.exceptions[0], [KeyError])
+
+                # Fail upon trying to reuse a broken pool after error callback failures:
+                # - BrokenPoolError containing:
+                #   - 3x CallbackError each containing:
+                #     - Error thrown from the callback
+                #     - Original error
+                with self.assertRaises(BrokenPoolError) as pool_ctx:
+                    with self.Pool(3) as pool:
+                        res = [func(pool, raising2, error_callback=raising)
+                               for _ in range(3)]
+                        for r in res:
+                            with self.assertRaises(CallbackError) as res_ctx:
+                                r.get()
+                            self._check_subexceptions(res_ctx.exception,
+                                                      [KeyError, IndexError])
+                    pool.apply_async(noop)
+                self._check_subexceptions(pool_ctx.exception, [CallbackError] * 3)
+                for se in pool_ctx.exception.exceptions:
+                    self._check_subexceptions(se, [KeyError, IndexError])
+
+                # Exiting the context manager with a "normal" error and a failed callback
+                # - BrokenPoolError containing:
+                #   - CallbackError containing:
+                #     - Error thrown from the callback
+                #   - Exception that caused the context manager to exit
+                with self.assertRaises(BrokenPoolError) as pool_ctx:
+                    with self.Pool(1) as pool:
+                        res = func(pool, noop, callback=raising)
+                        with self.assertRaises(CallbackError) as res_ctx:
+                            res.get()
+                        raise IndexError()
+                self._check_subexceptions(pool_ctx.exception,
+                                          [CallbackError, IndexError])
+
+                # Exiting the context manager directly with a callback failure error
+                # - BrokenPoolError containing:
+                #   - CallbackError instance containing:
+                #     - Error thrown from the callback
+                # Note that only one instance of the error is present: it was
+                # *not* added again as it was above, since it is a CallbackError
+                with self.assertRaises(BrokenPoolError) as pool_ctx:
+                    with self.Pool(1) as pool:
+                        func(pool, noop, callback=raising).get()
+                self._check_subexceptions(pool_ctx.exception, [CallbackError])
+                self._check_subexceptions(pool_ctx.exception.exceptions[0], [KeyError])
+
+    def _check_subexceptions(self, group, sub_types):
+        self.assertEqual(len(group.exceptions), len(sub_types))
+        for sub_exc, sub_type in zip(group.exceptions, sub_types):
+            self.assertIsInstance(sub_exc, sub_type)
+
+def noop(*args):
+    pass
+
+def raising(*args):
     raise KeyError("key")
 
+def raising2(*args):
+    raise IndexError()
+
 def unpickleable_result():
     return lambda: 42