From 5223c8f7cc3a2670e44e2753518a3b2f039f673d Mon Sep 17 00:00:00 2001 From: Sami Jawhar Date: Fri, 30 Oct 2026 06:49:36 +0000 Subject: [PATCH] fix(logging): mark a logging task done when the worker cannot start it LoggingWorker._process_log_task created the callback task outside the try whose finally untracks the item and calls queue.task_done(). When the event loop refused to start that task, the queue's unfinished counter never came back down and the item stayed tracked as dequeued, so every later flush() waited on queue.join() until its caller gave up. On a proxy shutdown that spends the whole success-callback budget and logs that accounting may be incomplete _process_single_task and clear_queue already create the task inside their guard, and now the main path does too, so the failure is logged, the callback that never ran is closed there rather than at shutdown, and the item is counted done Omp-Session: 01a14c3d-9d7f-7079-97aa-f810a884c050 --- litellm/litellm_core_utils/logging_worker.py | 17 ++++---- .../litellm_core_utils/test_logging_worker.py | 42 +++++++++++++++++++ 2 files changed, 51 insertions(+), 8 deletions(-) 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)