diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index cb3d8bf4fe5..2f8e7bdccea 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -485,10 +485,8 @@ class LoggingWorker: callback hasn't finished yet — ``queue.empty()`` would return True in that window and cause us to skip the wait. - ``start()`` runs first so that, after an event loop change, the tasks - still on the previous loop's queue move onto this loop and a worker - here drains them; joining the old queue directly would wait on a - counter nothing on this loop ever decrements. + ``start()`` runs first so a queue left behind by a previous event loop + is carried onto this one and drained here instead of joined forever. """ if self._queue is None: return diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 891bd0686b2..336067976fa 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -207,10 +207,29 @@ class TestLoggingWorker: assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape assert fired == [], "precondition: the callback never ran before the first loop closed" - async def flush_on_second_loop(): + async def flush_twice_on_second_loop(): + await asyncio.wait_for(worker.flush(), timeout=5) await asyncio.wait_for(worker.flush(), timeout=5) - asyncio.run(flush_on_second_loop()) + asyncio.run(flush_twice_on_second_loop()) + + assert fired == [True] + + def test_flush_starts_a_worker_when_the_queue_has_none(self): + """``flush()`` must drain a queue that exists on the current loop without a running worker.""" + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + fired = [] + + async def marker(): + fired.append(True) + + async def enqueue_then_flush(): + worker._ensure_queue() + worker.enqueue(marker()) + assert worker._worker_task is None, "precondition: nothing is draining the queue yet" + await asyncio.wait_for(worker.flush(), timeout=3) + + asyncio.run(enqueue_then_flush()) assert fired == [True]