Always commit offsets on termination

Even when an error occurs.

Added a regression test.

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

Author: gtopper

storey/sources.py

@@ -662,7 +662,6 @@ 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, self.logger, commit_all=True)
                     return termination_result
             except BaseException as ex:
                 if self.logger:
@@ -680,6 +679,8 @@ async def _run_loop(self):
                     await self._q.get()
                 self._raise_on_error()
             finally:
+                # Commit on termination regardless of errors (ML-11919)
+                await _commit_handled_events(self._outstanding_offsets, committer, self.logger, commit_all=True)
                 if event is _termination_obj or self._ex:
                     for closeable in self._closeables:
                         try:

tests/test_flow.py

@@ -176,6 +176,13 @@ async def _do(self, event):
         return await self._do_downstream(event)
 
 
+class ErrorOnTermination(storey.Flow):
+    async def _do(self, event):
+        if event is storey.dtypes._termination_obj:
+            raise ATestException("We raise this error on termination on purpose")
+        return await self._do_downstream(event)
+
+
 def test_offset_commit():
     platform = Committer()
     context = CommitterContext(platform)
@@ -544,6 +551,40 @@ def test_offset_commit_error():
     assert "RuntimeError: Something went wrong" in log_message
 
 
+async def async_offset_commit_error_on_termination():
+    platform = Committer()
+    logger = MockLogger()
+    context = CommitterContext(platform, logger=logger)
+
+    controller = build_flow(
+        [
+            AsyncEmitSource(context=context, explicit_ack=True, max_wait_before_commit=1),
+            ErrorOnTermination(),
+        ]
+    ).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
+
+    with pytest.raises(ATestException):
+        await controller.terminate(wait=True)
+
+    assert platform.offsets == {("/", shard): 10 for shard in range(num_shards)}
+
+
+# ML-11919
+def test_async_offset_commit_error_on_termination():
+    asyncio.run(async_offset_commit_error_on_termination())
+
+
 def test_multiple_upstreams():
     source = SyncEmitSource()
     map1 = Map(lambda x: x + 1)