diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 59f0ee0d7ec..cb0baa82767 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -143,14 +143,18 @@ class LoggingWorker: def _close_queued_coroutines(self) -> None: if self._queue is None: return - while True: - try: - logging_task: Final = self._queue.get_nowait() - except asyncio.QueueEmpty: - return + while (logging_task := self._dequeue_task()) is not None: logging_task["coroutine"].close() self._queue.task_done() + def _dequeue_task(self) -> LoggingTask | None: + if self._queue is None: + return None + try: + return self._queue.get_nowait() + except asyncio.QueueEmpty: + return None + def _release_running_tasks(self) -> None: for task in tuple(self._running_tasks): if not task.done(): @@ -158,7 +162,7 @@ class LoggingWorker: task.cancel() except RuntimeError: pass - setattr(task, "_log_destroy_pending", False) + task._log_destroy_pending = False # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # asyncio private shutdown flag if coroutine := self._task_coroutines.get(task): coroutine.close() self._running_tasks.clear() @@ -280,6 +284,7 @@ class LoggingWorker: try: self._queue.put_nowait(task) + self._start_queued_tasks() except asyncio.QueueFull: # Still full - handle it appropriately (clear or retry again) self._handle_queue_full(task) 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 74bba4b3ec8..af419ea6efe 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -184,9 +184,10 @@ class TestLoggingWorker: @pytest.mark.asyncio async def test_clear_queue_with_time_limit(self, logging_worker): """Test that clear_queue respects the time limit.""" + # Create several mock coroutines that take time to complete - slow_coro = AsyncMock() - slow_coro.return_value = asyncio.sleep(0.5) # Takes 500ms + async def slow_coro() -> None: + await asyncio.sleep(0.5) # Initialize the worker and add items logging_worker._ensure_queue() @@ -388,6 +389,25 @@ class TestLoggingWorker: assert len(processed) >= 4, f"Expected 4+ tasks processed, got {len(processed)}" + @pytest.mark.asyncio + async def test_delayed_enqueue_retry_restarts_workers_after_queue_drains(self): + worker = LoggingWorker(timeout=1.0, max_queue_size=1, concurrency=1) + worker.start() + processed = asyncio.Event() + + async def retried_task() -> None: + processed.set() + + try: + await worker._retry_enqueue_task( + {"coroutine": retried_task(), "context": contextvars.copy_context()}, + 0.0, + ) + + await asyncio.wait_for(processed.wait(), timeout=1.0) + finally: + await worker.stop() + @pytest.mark.asyncio async def test_event_loop_change_handling(self): """Test that LoggingWorker handles event loop changes correctly.