Repository: mlrun/storey (language: python)

Issue:
ML-1506: Fix ConcurrentJobExecution max_in_flight.
Also improve error traces.


Discussion / hints:

Fix patch:
diff --git a/storey/flow.py b/storey/flow.py
index 7fd0ef07..3b297fe2 100644
--- a/storey/flow.py
+++ b/storey/flow.py
@@ -617,6 +617,7 @@ def __init__(self, max_in_flight=None, retries=None, backoff_factor=None, **kwar
 
     def _init(self):
         self._q = None
+        self._lazy_init_complete = False
 
     async def _worker(self):
         event = None
@@ -670,8 +671,11 @@ async def _process_event_with_retries(self, event):
                     await asyncio.sleep(backoff_value)
 
     async def _do(self, event):
-        if not self._q and self._queue_size > 0:
+        if not self._lazy_init_complete:
             await self._lazy_init()
+            self._lazy_init_complete = True
+
+        if not self._q and self._queue_size > 0:
             self._q = asyncio.queues.Queue(self._queue_size)
             self._worker_awaitable = asyncio.get_running_loop().create_task(self._worker())
 


Test patch:
diff --git a/tests/test_flow.py b/tests/test_flow.py
index 7db4af84..0ec987ed 100644
--- a/tests/test_flow.py
+++ b/tests/test_flow.py
@@ -2927,6 +2927,10 @@ class _TestConcurrentExecution(_ConcurrentJobExecution):
         def __init__(self, **kwargs):
             super().__init__(**kwargs)
             self._ongoing_processing = 0
+            self.lazy_init_called = 0
+
+        async def _lazy_init(self):
+            self.lazy_init_called += 1
 
         async def _process_event(self, event):
             self._ongoing_processing += 1
@@ -2937,9 +2941,10 @@ async def _process_event(self, event):
         async def _handle_completed(self, event, response):
             pass
 
+    concurrent_step = _TestConcurrentExecution(max_in_flight=max_in_flight)
     controller = build_flow([
         SyncEmitSource(),
-        _TestConcurrentExecution(max_in_flight=max_in_flight),
+        concurrent_step,
     ]).run()
 
     for i in range(max_in_flight + 1):
@@ -2947,6 +2952,8 @@ async def _handle_completed(self, event, response):
     controller.terminate()
     controller.await_termination()
 
+    assert concurrent_step.lazy_init_called == 1
+
 
 class MockLogger:
     def __init__(self):