Merge pull request #38144 from BerriAI/litellm_fix_logging_worker_loop_drop

fix(logging_worker): carry queued tasks across event-loop change instead of dropping them

litellm/litellm_core_utils/logging_worker.py

@@ -5,7 +5,7 @@ import asyncio
 import atexit
 import contextvars
 import logging
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Iterator
 from typing import Final
 
 from typing_extensions import TypedDict
@@ -61,6 +61,19 @@ class LoggingWorker:
         # Register cleanup handler to flush remaining events on exit
         atexit.register(self._flush_on_exit)
 
+    @staticmethod
+    def _drain_pending(queue: "asyncio.Queue[LoggingTask]") -> tuple[LoggingTask, ...]:
+        """Pop every task still queued, without awaiting them, so they can be moved to another queue."""
+
+        def _pop_until_empty() -> Iterator[LoggingTask]:
+            while True:
+                try:
+                    yield queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    return
+
+        return tuple(_pop_until_empty())
+
     def _ensure_queue(self) -> None:
         """Initialize the queue if it doesn't exist or if event loop has changed."""
         try:
@@ -69,14 +82,27 @@ class LoggingWorker:
             # No running loop, can't initialize
             return
 
-        # Check if we need to reinitialize due to event loop change
+        # The queue, semaphore and worker task are all bound to the loop that created them. On a
+        # loop change we hand the still-pending tasks to a fresh queue instead of dropping them,
+        # so queued spend-logging coroutines are not silently discarded (and never left un-awaited).
         if self._queue is not None and self._bound_loop is not current_loop:
-            verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
-            # Clear old state - these are bound to the old loop
-            self._queue = None
+            carried_over: Final = self._drain_pending(self._queue)
+            new_queue: Final[asyncio.Queue[LoggingTask]] = asyncio.Queue(maxsize=self.max_queue_size)
+            for carried_task in carried_over:
+                new_queue.put_nowait(carried_task)
+            if carried_over:
+                verbose_logger.warning(
+                    "LoggingWorker: event loop changed; carried %d pending logging task(s) onto the new loop",
+                    len(carried_over),
+                )
+            else:
+                verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
             self._sem = None
             self._worker_task = None
             self._running_tasks.clear()
+            self._queue = new_queue
+            self._bound_loop = current_loop
+            return
 
         if self._queue is None:
             self._queue = asyncio.Queue(maxsize=self.max_queue_size)

tests/test_litellm/litellm_core_utils/test_logging_worker.py

@@ -413,3 +413,42 @@ class TestLoggingWorker:
         assert worker2._bound_loop is not None
 
         await worker2.stop()
+
+    def test_event_loop_change_carries_pending_tasks_over(self):
+        """Regression (LIT-6028): a loop change must not silently drop queued coroutines.
+
+        Before the fix ``_ensure_queue`` nulled ``self._queue`` on a loop change, discarding
+        every pending ``LoggingTask`` (each an un-awaited spend-logging coroutine). The tasks
+        must instead be moved onto the queue bound to the new loop and still execute there.
+        """
+        worker = LoggingWorker(timeout=1.0, max_queue_size=10)
+        executed: list[int] = []
+
+        async def spend_log(index: int) -> None:
+            executed.append(index)
+
+        async def enqueue_on_first_loop() -> None:
+            worker._ensure_queue()
+            for i in range(5):
+                worker.enqueue(spend_log(i))
+            assert worker._queue is not None
+            assert worker._queue.qsize() == 5
+
+        asyncio.run(enqueue_on_first_loop())
+
+        stale_queue = worker._queue
+        assert stale_queue is not None
+
+        async def rebind_on_second_loop() -> None:
+            worker._ensure_queue()
+            assert worker._queue is not None
+            # A fresh queue bound to the new loop, holding every carried-over task (not dropped).
+            assert worker._queue is not stale_queue
+            assert worker._queue.qsize() == 5
+            while not worker._queue.empty():
+                task = worker._queue.get_nowait()
+                await task["context"].run(asyncio.create_task, task["coroutine"])
+
+        asyncio.run(rebind_on_second_loop())
+
+        assert sorted(executed) == [0, 1, 2, 3, 4]