Recover from error raised by the offset committer callback

[ML-10538](https://iguazio.atlassian.net/browse/ML-10538)

Author: gtopper

storey/sources.py

@@ -289,7 +289,9 @@ async def _run_loop(self):
                     or num_offsets_not_committed > 1
                     and time.monotonic() >= last_commit_time + self._max_time_before_commit
                 ):
-                    num_offsets_not_committed = await _commit_handled_events(self._outstanding_offsets, committer)
+                    num_offsets_not_committed = await _commit_handled_events(
+                        self._outstanding_offsets, committer, self.logger
+                    )
                     events_handled_since_commit = 0
                     last_commit_time = time.monotonic()
                 # Due to the last event not being garbage collected, we tolerate a single unhandled event
@@ -300,7 +302,9 @@ async def _run_loop(self):
                         break
                     except queue.Empty:
                         pass
-                    num_offsets_not_committed = await _commit_handled_events(self._outstanding_offsets, committer)
+                    num_offsets_not_committed = await _commit_handled_events(
+                        self._outstanding_offsets, committer, self.logger
+                    )
                     events_handled_since_commit = 0
                     last_commit_time = time.monotonic()
             if event is None:
@@ -318,7 +322,7 @@ async def _run_loop(self):
                 if event is _termination_obj:
                     # We can commit all at this point because termination of
                     # all downstream steps completed successfully.
-                    await _commit_handled_events(self._outstanding_offsets, committer, commit_all=True)
+                    await _commit_handled_events(self._outstanding_offsets, committer, self.logger, commit_all=True)
                     self._termination_future.set_result(termination_result)
             except BaseException as ex:
                 if self.logger:
@@ -510,7 +514,7 @@ async def await_termination(self):
         return await self._loop_task
 
 
-async def _commit_handled_events(outstanding_offsets_by_qualified_shard, committer, commit_all=False):
+async def _commit_handled_events(outstanding_offsets_by_qualified_shard, committer, logger, commit_all=False):
     num_offsets_not_handled = 0
     if not commit_all:
         gc.collect()
@@ -530,7 +534,12 @@ async def _commit_handled_events(outstanding_offsets_by_qualified_shard, committ
                 num_to_clear += 1
         if last_handled_offset is not None:
             path, shard_id = qualified_shard
-            await committer(QualifiedOffset(path, shard_id, last_handled_offset))
+            try:
+                await committer(QualifiedOffset(path, shard_id, last_handled_offset))
+            except BaseException:
+                if logger:
+                    logger.error(f"Failed to commit offsets due to error: {traceback.format_exc()}")
+                return num_offsets_not_handled + num_to_clear
             outstanding_offsets_by_qualified_shard[qualified_shard] = offsets[num_to_clear:]
     return num_offsets_not_handled
 
@@ -601,7 +610,9 @@ async def _run_loop(self):
                     or num_offsets_not_handled > 0
                     and time.monotonic() >= last_commit_time + self._max_time_before_commit
                 ):
-                    num_offsets_not_handled = await _commit_handled_events(self._outstanding_offsets, committer)
+                    num_offsets_not_handled = await _commit_handled_events(
+                        self._outstanding_offsets, committer, self.logger
+                    )
                     events_handled_since_commit = 0
                     last_commit_time = time.monotonic()
                 # In case we can't block because there are outstanding events
@@ -611,7 +622,9 @@ async def _run_loop(self):
                         break
                     except TimeoutError:
                         pass
-                    num_offsets_not_handled = await _commit_handled_events(self._outstanding_offsets, committer)
+                    num_offsets_not_handled = await _commit_handled_events(
+                        self._outstanding_offsets, committer, self.logger
+                    )
                     events_handled_since_commit = 0
                     last_commit_time = time.monotonic()
             if not event:
@@ -629,7 +642,7 @@ async def _run_loop(self):
                 if event is _termination_obj:
                     # We can commit all at this point because termination of
                     # all downstream steps completed successfully.
-                    await _commit_handled_events(self._outstanding_offsets, committer, commit_all=True)
+                    await _commit_handled_events(self._outstanding_offsets, committer, self.logger, commit_all=True)
                     return termination_result
             except BaseException as ex:
                 if self.logger:

tests/test_flow.py

@@ -146,9 +146,19 @@ async def explicit_ack(self, qualified_offset):
         self.offsets[qualified_shard] = offset
 
 
+class BadCommitter:
+    def __init__(self):
+        self.offsets = {}
+
+    async def explicit_ack(self, qualified_offset):
+        raise RuntimeError("Something went wrong")
+
+
 class CommitterContext:
-    def __init__(self, platform):
+    def __init__(self, platform, logger=None, verbose=None):
         self.platform = platform
+        self.logger = logger
+        self.verbose = verbose
 
 
 class EventHoarder(storey.Flow):
@@ -441,6 +451,97 @@ def test_offset_not_committed_prematurely():
     assert offsets == {("/", i): num_records_per_shard for i in range(num_shards)}
 
 
+async def async_offset_commit_error():
+    platform = BadCommitter()
+    logger = MockLogger()
+    context = CommitterContext(platform, logger=logger)
+
+    controller = build_flow(
+        [
+            AsyncEmitSource(context=context, explicit_ack=True, max_wait_before_commit=1),
+            Map(lambda x: x + 1),
+            Filter(lambda x: x < 3),
+            FlatMap(lambda x: [x, x * 10]),
+            Reduce(0, lambda acc, x: acc + x),
+        ]
+    ).run()
+
+    num_shards = 10
+    num_records_per_shard = 10
+
+    for offset in range(1, num_records_per_shard + 1):
+        for shard in range(num_shards):
+            event = Event(shard)
+            event.shard_id = shard
+            event.offset = offset
+            await controller.emit(event)
+    del event
+
+    # Make sure that offsets are committed even before termination
+    await asyncio.sleep(2)
+    offsets = copy.copy(platform.offsets)
+
+    try:
+        assert offsets == {}
+    finally:
+        termination_result = await controller.terminate(wait=True)
+
+    assert termination_result == 330
+
+    log_level, (log_message,), _ = logger.logs[0]
+    assert log_level == "error"
+    assert "Failed to commit offsets due to error" in log_message
+    assert "RuntimeError: Something went wrong" in log_message
+
+
+# ML-10538
+def test_async_offset_commit_error():
+    asyncio.run(async_offset_commit_error())
+
+
+def test_offset_commit_error():
+    platform = BadCommitter()
+    logger = MockLogger()
+    context = CommitterContext(platform, logger=logger)
+
+    controller = build_flow(
+        [
+            SyncEmitSource(context=context, explicit_ack=True, max_wait_before_commit=1),
+            Map(lambda x: x + 1),
+            Filter(lambda x: x < 3),
+            FlatMap(lambda x: [x, x * 10]),
+            Reduce(0, lambda acc, x: acc + x),
+        ]
+    ).run()
+
+    num_shards = 10
+    num_records_per_shard = 10
+
+    for offset in range(1, num_records_per_shard + 1):
+        for shard in range(num_shards):
+            event = Event(shard)
+            event.shard_id = shard
+            event.offset = offset
+            controller.emit(event)
+    del event
+
+    # Make sure that offsets are committed even before termination
+    time.sleep(2)
+    offsets = copy.copy(platform.offsets)
+
+    try:
+        assert offsets == {}
+    finally:
+        termination_result = controller.terminate(wait=True)
+
+    assert termination_result == 330
+
+    log_level, (log_message,), _ = logger.logs[0]
+    assert log_level == "error"
+    assert "Failed to commit offsets due to error" in log_message
+    assert "RuntimeError: Something went wrong" in log_message
+
+
 def test_multiple_upstreams():
     source = SyncEmitSource()
     map1 = Map(lambda x: x + 1)