Repository: mlrun/storey (language: python)

Issue:
Fix completion on error in concurrent execution step.



Discussion / hints:

Fix patch:
diff --git a/storey/flow.py b/storey/flow.py
index 53f17744..7458ea86 100644
--- a/storey/flow.py
+++ b/storey/flow.py
@@ -603,9 +603,13 @@ def __init__(self, status, body):
 
 
 class _ConcurrentJobExecution(Flow):
-    def __init__(self, max_in_flight=8, **kwargs):
+    _BACKOFF_MAX = 120
+
+    def __init__(self, max_in_flight=8, retries=0, backoff_factor=1, **kwargs):
         Flow.__init__(self, **kwargs)
         self._max_in_flight = max_in_flight
+        self._retries = retries
+        self._backoff_factor = backoff_factor
 
     def _init(self):
         self._q = None
@@ -643,17 +647,23 @@ async def _cleanup(self):
     async def _lazy_init(self):
         pass
 
-    async def _safe_process_event(self, event):
-        if event._awaitable_result:
+    async def _process_event_with_retries(self, event):
+        times_attempted = 0
+        max_attempts = self._retries + 1
+        while True:
             try:
                 return await self._process_event(event)
-            except BaseException as ex:
-                none_or_coroutine = event._awaitable_result._set_error(ex)
-                if none_or_coroutine:
-                    await none_or_coroutine
-                raise ex
-        else:
-            return await self._process_event(event)
+            except Exception as ex:
+                times_attempted += 1
+                attempts_left = max_attempts - times_attempted
+                if self.logger:
+                    self.logger.warn(f'{self.name} failed to process event ({attempts_left} retries left): {ex}')
+                if attempts_left <= 0:
+                    raise ex
+                backoff_value = self._backoff_factor * (2 ** (times_attempted - 1))
+                backoff_value = min(self._BACKOFF_MAX, backoff_value)
+                if backoff_value >= 0:
+                    await asyncio.sleep(backoff_value)
 
     async def _do(self, event):
         if not self._q:
@@ -670,7 +680,7 @@ async def _do(self, event):
             await self._worker_awaitable
             return await self._do_downstream(_termination_obj)
         else:
-            task = self._safe_process_event(event)
+            task = self._process_event_with_retries(event)
             await self._q.put((event, asyncio.get_running_loop().create_task(task)))
             if self._worker_awaitable.done():
                 await self._worker_awaitable


Test patch:
diff --git a/tests/test_flow.py b/tests/test_flow.py
index b49cd50b..5ae71b68 100644
--- a/tests/test_flow.py
+++ b/tests/test_flow.py
@@ -3,6 +3,7 @@
 import math
 import os
 import queue
+import time
 import traceback
 import uuid
 from datetime import datetime
@@ -2778,15 +2779,14 @@ def my_func(param1, param2):
     assert len(dictionary) == 1
 
 
-class _ErrorInConcurrentExecution(_ConcurrentJobExecution):
-    async def _process_event(self, event):
-        pass
-
-    async def _handle_completed(self, event, response):
-        raise ATestException()
+def test_completion_on_error_in_concurrent_execution_step():
+    class _ErrorInConcurrentExecution(_ConcurrentJobExecution):
+        async def _process_event(self, event):
+            pass
 
+        async def _handle_completed(self, event, response):
+            raise ATestException()
 
-def test_completion_on_error_in_concurrent_execution_step():
     controller = build_flow([
         SyncEmitSource(),
         _ErrorInConcurrentExecution(),
@@ -2799,3 +2799,45 @@ def test_completion_on_error_in_concurrent_execution_step():
             awaitable_result.await_result()
     finally:
         controller.terminate()
+
+
+@pytest.mark.parametrize('backoff_factor', [(0, 2, 0), (1, 2, 3), (0, 1, None)])
+def test_completion_after_retry_in_concurrent_execution_step(backoff_factor):
+    backoff_factor, retries, expected_sleep = backoff_factor
+
+    class _ErrorInConcurrentExecution(_ConcurrentJobExecution):
+        def __init__(self, **kwargs):
+            super().__init__(**kwargs)
+            self._nums_called = 0
+
+        async def _process_event(self, event):
+            self._nums_called += 1
+            if self._nums_called <= 2:  # fail twice
+                raise ATestException()
+
+        async def _handle_completed(self, event, response):
+            return await self._do_downstream(event)
+
+    controller = build_flow([
+        SyncEmitSource(),
+        _ErrorInConcurrentExecution(retries=retries, backoff_factor=backoff_factor),
+        Complete()
+    ]).run()
+
+    awaitable_result = controller.emit(1)
+    try:
+        start = time.time()
+        if expected_sleep is None:
+            with pytest.raises(ATestException):
+                awaitable_result.await_result()
+        else:
+            awaitable_result.await_result()
+        end = time.time()
+    finally:
+        controller.terminate()
+    if expected_sleep is None:
+        with pytest.raises(ATestException):
+            controller.await_termination()
+    else:
+        controller.await_termination()
+        assert end - start > expected_sleep