diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 7f00c47c1ff..2564b553a34 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -370,11 +370,10 @@ class LoggingWorker: self._running_tasks.clear() async def flush(self) -> None: - """Flush the logging queue.""" + """Flush queued and in-flight logging work.""" if self._queue is None: return - while not self._queue.empty(): - await self._queue.join() + await self._queue.join() async def clear_queue(self): """ 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 a44b821db87..abfc1a83e57 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -360,3 +360,34 @@ class TestLoggingWorker: assert worker2._bound_loop is not None await worker2.stop() + + @pytest.mark.asyncio + async def test_flush_waits_for_in_flight_task_after_queue_is_empty(self): + """Flush must wait for dequeued work that has not called task_done() yet.""" + worker = LoggingWorker(timeout=1.0, max_queue_size=10, concurrency=1) + worker.start() + + started = asyncio.Event() + release = asyncio.Event() + + async def slow_task(): + started.set() + await release.wait() + + worker.enqueue(slow_task()) + + await asyncio.wait_for(started.wait(), timeout=2.0) + + # The worker has already dequeued the item, so qsize() can be zero + # even though the logging coroutine is still running. + assert worker._queue is not None + assert worker._queue.empty() is True + + flush_task = asyncio.create_task(worker.flush()) + await asyncio.sleep(0.05) + + assert flush_task.done() is False + + release.set() + await asyncio.wait_for(flush_task, timeout=2.0) + await worker.stop()