Handling exceptions

Author: jan-janssen

src/executorlib/standalone/batched.py

@@ -357,3 +357,37 @@ def _update_waiting_task(
     if len(wait_lst) == len(wait_tmp_lst):
         sleep(refresh_rate)
     return wait_tmp_lst
+
+
+def batched_futures(
+    lst: list[Future], nested_skip_lst: list[Future[list]], n: int
+) -> list[list]:
+    """
+    Batch n completed future objects. If the number of completed futures is smaller than n and the end of the batch is
+    not reached yet, then an empty list is returned. If n future objects are done, which are not included in the skip_set
+    then they are returned as batch.
+
+    Args:
+        lst (list): list of all future objects
+        nested_skip_lst (list): nest list of individual results already assigned to previous batches
+        n (int): batch size
+
+    Returns:
+        list: results of the batched futures
+    """
+    skip_set = {id(item) for f in nested_skip_lst for item in f.result()}
+
+    done_lst = []
+    failed_lst = []
+    n_expected = min(n, len(lst) - len(skip_set))
+    for v in lst:
+        if v.done():
+            if check_exception_was_raised(future_obj=v):
+                failed_lst.append(v)
+            elif id(v.result()) not in skip_set:
+                done_lst.append(v.result())
+                if len(done_lst) == n_expected:
+                    return done_lst
+    if len(failed_lst) == n_expected:
+        return failed_lst
+    return []

src/executorlib/task_scheduler/interactive/dependency.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from executorlib.task_scheduler.interactive.blockallocation import BlockAllocationTaskScheduler
+from executorlib.task_scheduler.interactive.dependency import batched_futures
 from executorlib.standalone.interactive.spawner import MpiExecSpawner
 
 
@@ -141,3 +142,42 @@ def finished(self, future):
                 foo.running,
                 msg="After task completion, we expect the callback to modify the class",
             )
+
+
+class TestBatched(unittest.TestCase):
+    def test_batched_futures(self):
+        lst = []
+        for i in range(10):
+            f = Future()
+            f.set_result(i)
+            lst.append(f)
+        batched_lst = [Future(), Future(), Future()]
+        batched_lst[0].set_result([0, 1, 2])
+        batched_lst[1].set_result([3, 4, 5])
+        batched_lst[2].set_result([6, 7, 8])
+        self.assertEqual(batched_futures(lst=lst, n=3, nested_skip_lst=set()), [0, 1, 2])
+        self.assertEqual(batched_futures(lst=lst, nested_skip_lst=batched_lst[:1], n=3), [3, 4, 5])
+        self.assertEqual(batched_futures(lst=lst, nested_skip_lst=batched_lst[:2], n=3), [6, 7, 8])
+        self.assertEqual(batched_futures(lst=lst, nested_skip_lst=batched_lst, n=3), [9])
+
+    def test_batched_futures_duplicated(self):
+        lst = []
+        for i in range(1,4):
+            for _ in range(3):
+                f = Future()
+                f.set_result(i)
+                lst.append(f)
+        batched_lst = [Future(), Future(), Future()]
+        batched_lst[0].set_result([1, 1, 1])
+        batched_lst[1].set_result([2, 2, 2])
+        batched_lst[2].set_result([3, 3, 3])
+        self.assertEqual(batched_futures(lst=lst, n=3, nested_skip_lst=set()), [1, 1, 1])
+        self.assertEqual(batched_futures(lst=lst, nested_skip_lst=batched_lst[:1], n=3), [2, 2, 2])
+        self.assertEqual(batched_futures(lst=lst, nested_skip_lst=batched_lst[:2], n=3), [3, 3, 3])
+
+    def test_batched_futures_not_finished(self):
+        lst = []
+        for _ in list(range(10)):
+            f = Future()
+            lst.append(f)
+        self.assertEqual(batched_futures(lst=lst, n=3, nested_skip_lst=set()), [])