diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index e7a4732ee37..8b25e9d06fe 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -423,13 +423,11 @@ class LoggingWorker: return extracted_tasks: Final = self._extract_tasks_from_queue() - runs: Final = [self._process_single_task(task) for task in extracted_tasks] - if new_task is not None: - try: - self._queue.put_nowait(new_task) - except asyncio.QueueFull: - self._track_dequeued(new_task) - runs.append(self._process_single_task(new_task, counted=False)) + uncounted: Final = None if new_task is None or self._requeue(new_task) else new_task + runs: Final = ( + *(self._process_single_task(task) for task in extracted_tasks), + *(() if uncounted is None else (self._process_single_task(uncounted, counted=False),)), + ) if runs: await asyncio.gather(*runs) except Exception as e: @@ -438,6 +436,17 @@ class LoggingWorker: # Always reset the flag even if an error occurs self._aggressive_clear_in_progress = False + def _requeue(self, task: LoggingTask) -> bool: + """Put a rejected task back on the queue once the clear made room; track it as dequeued otherwise.""" + if self._queue is None: + return False + try: + self._queue.put_nowait(task) + except asyncio.QueueFull: + self._track_dequeued(task) + return False + return True + async def _process_single_task(self, task: LoggingTask, *, counted: bool = True) -> None: """Process a single task; ``counted`` is False for a task the full queue never accepted.""" if self._queue is None: diff --git a/tests/unit/litellm_core_utils/test_logging_worker.py b/tests/unit/litellm_core_utils/test_logging_worker.py index 2840f24934e..30ec851d54f 100644 --- a/tests/unit/litellm_core_utils/test_logging_worker.py +++ b/tests/unit/litellm_core_utils/test_logging_worker.py @@ -167,9 +167,9 @@ async def test_overflow_callback_does_not_undercount_the_unfinished_queue() -> N queue.task_done() it never contributed, otherwise flush() returns while a real callback is still running and the shutdown drain cancels it with budget to spare. """ - worker = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1) - gates = {name: asyncio.Event() for name in ("first", "second", "third")} - completed: list[str] = [] + worker: Final = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1) + gates: Final = {name: asyncio.Event() for name in ("first", "second", "third")} + completed: Final[list[str]] = [] async def gated(name: str) -> None: await gates[name].wait()