Vine: Improve how errors are propagated with FuturesExecutor (#3922)

* raise task exceptions in python task futures

* wrap FutureFunctionCall tracebacks in FunctionCallNoResult

* Raise TimeoutError in Future when task result is pending

* yield futures and raise TimeoutError in as_completed

* Improve error handling in wait

* Import TimeoutError for Python <3.10 compatibility

In Python 3.9 and older, concurrent.futures used its own TimeoutError
class rather than the builtin. TimeoutError is aliased to the builtin in
layer Python versions, so for backwards compatibility we import it in
the test file.

See python/cpython#30197

taskvine/src/bindings/python3/ndcctools/taskvine/futures.py

@@ -1,6 +1,7 @@
 
 from . import cvine
 import hashlib
+from collections import deque
 from concurrent.futures import Executor
 from concurrent.futures import Future
 from concurrent.futures import FIRST_COMPLETED
@@ -14,7 +15,6 @@
 from .task import (
     PythonTask,
     FunctionCall,
-    PythonTaskNoResult,
     FunctionCallNoResult,
 )
 from .manager import (
@@ -53,44 +53,44 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED):
         if not f._is_submitted:
             f.module_manager.submit(f._task)
 
-    time_init = time.time()
-    if timeout is None:
-        time_check = float('inf')
-    else:
-        time_check = timeout
+    start = time.perf_counter()
+    time_check = float('inf') if timeout is None else timeout
+    result_timeout = min(timeout, 5) if timeout is not None else 5
+
     done = False
-    while time.time() - time_init < time_check and not done:
-        done = True
+    while time.perf_counter() - start < time_check and not done:
         for f in fs:
-
             # skip if future is complete
             if f in results.done:
                 continue
 
-            # check for completion
-            result = f.result(timeout=5)
-
-            # add to set of finished tasks and break when needed.
-            if result != RESULT_PENDING:
+            try:
+                # check for completion
+                f.result(timeout=result_timeout)
+            except TimeoutError:
+                # TimeoutError's are expected since we are polling the
+                # future's status. If a timeout happens, we do nothing.
+                pass
+            except Exception:
+                # Future.result() raises the task's exception but that is
+                # not relevant here---just if the task has finished.
+                results.done.add(f)
+                if return_when in (FIRST_EXCEPTION, FIRST_COMPLETED):
+                    done = True
+                    break
+            else:
                 results.done.add(f)
-
-                # if this is the first completed task, break.
                 if return_when == FIRST_COMPLETED:
                     done = True
                     break
 
-            if isinstance(result, Exception) and return_when == FIRST_EXCEPTION:
+            if len(results.done) == len(fs):
                 done = True
                 break
 
-            # set done to false to finish loop.
-            else:
-                done = False
-
             # check form timeout
-            if timeout is not None:
-                if time.time() - time_init > timeout:
-                    break
+            if time.perf_counter() - start > time_check:
+                break
 
     # add incomplete futures to set
     for f in fs:
@@ -101,48 +101,45 @@ def wait(fs, timeout=None, return_when=ALL_COMPLETED):
 
 
 def as_completed(fs, timeout=None):
+    fs = deque(fs)
 
-    results = set()
-
-    # submit tasks if they have not been subitted
+    # submit tasks if they have not been submitted
     for f in fs:
         if not f._is_submitted:
             f.module_manager.submit(f._task)
 
-    time_init = time.time()
-    if timeout is None:
-        time_check = float('inf')
-    else:
-        time_check = timeout
-
-    done = False
-    while time.time() - time_init < time_check and not done:
-        for f in fs:
-            done = True
-            # skip if future is complete
-            if f in results:
-                continue
-
-            # check for completion
-            result = f.result(timeout=5)
-
-            # add to set of finished tasks
-            if result != RESULT_PENDING:
-                results.add(f)
-
-            # set done to false to finish loop.
+    start = time.perf_counter()
+    result_timeout = min(timeout, 5) if timeout is not None else 5
+
+    def _iterator():
+        # iterate of queue of futures, yeilding completed futures and
+        # requeuing non-completed futures until all futures are yielded or
+        # the timeout is reached.
+        while fs:
+            f = fs.popleft()
+
+            try:
+                result = f.result(timeout=result_timeout)
+            except TimeoutError:
+                # TimeoutError's are expected since we are polling the
+                # future's status. If a timeout happens, add the future to
+                # the back of the queue.
+                fs.append(f)
+            except Exception:
+                # Future.result() raises the task's exception but that is
+                # not relevant here---just if the task has finished.
+                yield f
             else:
-                done = False
+                assert result != RESULT_PENDING
+                yield f
 
-            # check form timeout
-            if timeout is not None:
-                if time.time() - time_init > timeout:
-                    break
-    for f in fs:
-        if f not in results:
-            results.add(TimeoutError)
+            if (
+                fs and timeout is not None
+                and time.perf_counter() - start > timeout
+            ):
+                raise TimeoutError()
 
-    return iter(results)
+    return _iterator()
 
 
 ##
@@ -233,6 +230,7 @@ def __init__(self, task):
         self._task = task
         self._callback_fns = []
         self._result = None
+        self._exception = None
         self._is_submitted = False
         self._ran_functions = False
 
@@ -266,15 +264,30 @@ def result(self, timeout="wait_forever"):
             timeout = "wait_forever"
         result = self._task.output(timeout=timeout)
         if result == RESULT_PENDING:
-            return RESULT_PENDING
+            raise TimeoutError()
+
+        if isinstance(result, Exception):
+            self._exception = result
         else:
             self._result = result
-            self._state = FINISHED
-            if self._callback_fns and not self._ran_functions:
-                self._ran_functions = True
-                for fn in self._callback_fns:
-                    fn(self)
-            return result
+
+        self._state = FINISHED
+        if self._callback_fns and not self._ran_functions:
+            self._ran_functions = True
+            for fn in self._callback_fns:
+                fn(self)
+
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def exception(self, timeout="wait_forever"):
+        try:
+            self.result()
+        except Exception as e:
+            return e
+        else:
+            return None
 
     def add_done_callback(self, fn):
         self._callback_fns.append(fn)
@@ -318,11 +331,10 @@ def output(self, timeout="wait_forever"):
                     if output['Success']:
                         self._saved_output = output['Result']
                     else:
-                        self._saved_output = output['Reason']
+                        self._saved_output = FunctionCallNoResult(output['Reason'])
 
                 except Exception as e:
                     self._saved_output = e
-                    raise e
             else:
                 self._saved_output = FunctionCallNoResult()
             self._output_loaded = True
@@ -388,15 +400,14 @@ def output(self, timeout="wait_forever"):
 
         # fetch output file and load output
         if not self._output_loaded and self._has_retrieved:
-            if self.successful():
-                try:
-                    self._module_manager.fetch_file(self._output_file)
-                    self._output = cloudpickle.loads(self._output_file.contents())
-                except Exception as e:
-                    self._output = e
-                    raise e
-            else:
-                self._output = PythonTaskNoResult()
+            try:
+                self._module_manager.fetch_file(self._output_file)
+                # _output can be either the return value of a successful
+                # task or the exception object of a failed task.
+                self._output = cloudpickle.loads(self._output_file.contents())
+            except Exception as e:
+                # handle output file fetch/deserialization failures
+                self._output = e
             self._output_loaded = True
 
         return self._output

taskvine/test/vine_python_future_funcall.py

@@ -22,6 +22,10 @@ def my_sum(x, y, negate=False):
     return s
 
 
+def failure():
+    raise Exception('Expected failure.')
+
+
 def main():
     executor = vine.FuturesExecutor(
         port=[9123, 9129], manager_name="test-executor", factory=False
@@ -33,7 +37,7 @@ def main():
     # Create library task
     print("creating library from functions...")
     libtask = executor.create_library_from_functions(
-        "test-library", my_sum, hoisting_modules=None, add_env=False
+        "test-library", my_sum, failure, hoisting_modules=None, add_env=False
     )
 
     # Install library on executor.manager
@@ -59,6 +63,21 @@ def main():
     c = a + b
     assert res == c
 
+    # Test failure handling
+    f1 = executor.future_funcall("test-library", "failure")
+    future = executor.submit(f1)
+
+    try:
+        future.result()
+    except Exception as e:
+        # FutureFunctionCall wraps the exception string in a
+        # FunctionCallNoResult so we check the original error message is
+        # present in the raised error.
+        assert "Expected failure." in str(e)
+        assert future.exception() is not None
+    else:
+        raise RuntimeError("Future did not raise exception.")
+
 
 if __name__ == "__main__":
     main()

taskvine/test/vine_python_future_module.py

@@ -2,6 +2,10 @@
 
 import sys
 import ndcctools.taskvine as vine
+from concurrent.futures import FIRST_COMPLETED
+from concurrent.futures import FIRST_EXCEPTION
+from concurrent.futures import ALL_COMPLETED
+from concurrent.futures import TimeoutError
 
 port_file = None
 try:
@@ -22,6 +26,15 @@ def my_sum(x, y, negate=False):
     return s
 
 
+def my_exception():
+    raise Exception("Expected failure.")
+
+
+def my_timeout():
+    import time
+    time.sleep(60)
+
+
 def main():
     executor = vine.FuturesExecutor(
         port=[9123, 9129], manager_name="vine_matrtix_build_test", factory=False
@@ -67,9 +80,77 @@ def main():
     c = executor.submit(t3)
 
     print("waiting for result...")
-    results = vine.futures.as_completed([a, b, c])
+    results = list(vine.futures.as_completed([a, b, c]))
     print(f"results = {results}")
 
+    # Test timeouts
+    t1 = executor.future_task(my_timeout)
+    t1.set_cores(1)
+    future = executor.submit(t1)
+
+    try:
+        future.result(timeout=1)
+    except TimeoutError:
+        future.cancel()
+        print("timeout raised correctly")
+    else:
+        raise RuntimeError("TimeoutError was not raised correctly.")
+
+    # # Test error handling with wait
+    t1 = executor.future_task(my_sum, 7, 4)
+    t1.set_cores(1)
+    a = executor.submit(t1)
+
+    t2 = executor.future_task(my_timeout)
+    t2.set_cores(1)
+    b = executor.submit(t2)
+
+    results = vine.futures.wait([a, b], return_when=FIRST_COMPLETED)
+    assert len(results.done) == 1
+    assert len(results.not_done) == 1
+    assert results.done.pop().result() == 11
+
+    results = vine.futures.wait([a, b], timeout=2, return_when=ALL_COMPLETED)
+    assert len(results.done) == 1
+    assert len(results.not_done) == 1
+    assert results.done.pop().result() == 11
+
+    t3 = executor.future_task(my_exception)
+    t3.set_cores(1)
+    c = executor.submit(t3)
+
+    results = vine.futures.wait([b, c], return_when=FIRST_EXCEPTION)
+    assert len(results.done) == 1
+    assert len(results.not_done) == 1
+    assert results.done.pop().exception() is not None
+
+    # Cancel the task that is still sleeping
+    b.cancel()
+
+    # Test timeouts with as_completed
+    t1 = executor.future_task(my_sum, 7, 4)
+    t1.set_cores(1)
+    a = executor.submit(t1)
+
+    t2 = executor.future_task(my_timeout)
+    t2.set_cores(1)
+    b = executor.submit(t2)
+
+    iterator = vine.futures.as_completed([a, b], timeout=5)
+
+    # task 1 should complete correctly within the timeout and be yielded first
+    a_future = next(iterator)
+    assert a_future.result() == 11
+
+    try:
+        next(iterator)
+    except TimeoutError:
+        b.cancel()
+        print("as_completed raised timeout correctly")
+    else:
+        raise RuntimeError("TimeoutError was not raised correctly.")
+
+
 if __name__ == "__main__":
     main()