🚸 Terminate early on failed computation (#14)

* 🚸 Terminate early on failed computation

* ✅ Add test for failing imap

* ✅ Add test for failing map

setup.cfg

@@ -11,6 +11,7 @@ ignore = D10,E203,E501,W503
 max-line-length = 88
 inline-quotes = double
 docstring-convention = google
+max-cognitive-complexity = 10
 
 [coverage:run]
 branch = True

src/mapply/parallel.py

@@ -76,6 +76,9 @@ def multiprocessing_imap(
 
     Returns:
         Results in same order as input iterable.
+
+    Raises:
+        Exception: Any error occurred during computation (will terminate the pool early).
     """
     n_chunks: Optional[int] = tqdm(iterable, disable=True).__len__()  # doesn't exhaust
     func = partial(func, *args, **kwargs)
@@ -97,6 +100,12 @@ def multiprocessing_imap(
 
     try:
         return list(stage)
+    except Exception:
+        if pool:
+            logger.debug("Terminating ProcessPool")
+            pool.close()
+            pool.terminate()
+        raise
     finally:
         if pool:
             logger.debug("Closing ProcessPool")

tests/test_parallel.py

@@ -1,7 +1,11 @@
+import pytest
+
 from mapply.parallel import multiprocessing_imap
 
 
 def foo(x, power):
+    if not isinstance(power, float):
+        raise ValueError("To check we reraise errors from subprocesses")
     return pow(x, power)
 
 
@@ -19,3 +23,13 @@ def test_multiprocessing_imap(size=100, power=1.1):
     assert multicore_list1 == multicore_list2
     assert multicore_list1 == multicore_list3
     assert multicore_list1 == [foo(x, power=power) for x in range(size)]
+    with pytest.raises(ValueError, match="reraise"):
+        # hit with ProcessPool
+        multiprocessing_imap(
+            foo, range(size), power=None, progressbar=False, n_workers=2
+        )
+    with pytest.raises(ValueError, match="reraise"):
+        # hit without ProcessPool
+        multiprocessing_imap(
+            foo, range(size), power=None, progressbar=False, n_workers=1
+        )