diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 8f217b1ea5e..470889ca43d 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -202,16 +202,17 @@ class LoggingWorker: """Runs the logging task and handles cleanup. Releases semaphore when done.""" try: if self._queue is not None: - # Run the coroutine in its original context - callback_task: Final = self._create_callback_task(task) try: - await asyncio.wait_for(callback_task, timeout=self.timeout) - except asyncio.TimeoutError as e: - if callback_task.cancelled(): - self._record_callback_timeout(task["coroutine"]) - else: - verbose_logger.exception("LoggingWorker error: %s", e) + callback_task: Final = self._create_callback_task(task) + try: + await asyncio.wait_for(callback_task, timeout=self.timeout) + except asyncio.TimeoutError as e: + if callback_task.cancelled(): + self._record_callback_timeout(task["coroutine"]) + else: + verbose_logger.exception("LoggingWorker error: %s", e) except Exception as e: + task["coroutine"].close() verbose_logger.exception("LoggingWorker error: %s", e) finally: self._untrack_dequeued(task) diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index fccc11617bb..67a20bd53b8 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -4,6 +4,7 @@ Tests for the LoggingWorker class to ensure graceful shutdown handling. import asyncio import contextvars +import inspect import io import logging from typing import Final @@ -564,6 +565,47 @@ class TestLoggingWorker: assert callback.await_count == 1 + def test_flush_returns_after_a_task_the_worker_cannot_start(self): + """ + Regression: when the event loop refuses to start a callback, the queued item was never + marked done, so ``flush()`` waited on ``queue.join()`` until its caller gave up. The + callback that never ran is closed where it is dropped, instead of waiting for the shutdown + cleanup to find it + """ + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + callback = AsyncMock() + + async def never_started(): + pass + + refused_coroutine = never_started() + + async def enqueue_then_flush_and_stop(): + loop = asyncio.get_running_loop() + previous_factory = loop.get_task_factory() + + def refuse_that_one_callback(running_loop, coro, **kwargs): + if coro is refused_coroutine: + raise RuntimeError("the loop refused to start this callback") + if previous_factory is not None: + return previous_factory(running_loop, coro, **kwargs) + return asyncio.Task(coro, loop=running_loop, **kwargs) + + loop.set_task_factory(refuse_that_one_callback) + try: + worker.ensure_initialized_and_enqueue(refused_coroutine) + worker.ensure_initialized_and_enqueue(callback()) + await asyncio.wait_for(worker.flush(), timeout=3) + await worker.stop(discard_pending=True) + finally: + loop.set_task_factory(previous_factory) + + asyncio.run(enqueue_then_flush_and_stop()) + + assert callback.await_count == 1 + assert worker.pending_count == 0 + assert inspect.getcoroutinestate(refused_coroutine) == inspect.CORO_CLOSED + def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self): """A callback raising CancelledError must not abort the atexit flush of later events.""" worker = LoggingWorker(timeout=1.0, max_queue_size=10)