Merge 50d9d5de85 into 72228db5c6

litellm/litellm_core_utils/logging_worker.py

@@ -9,7 +9,7 @@ import logging
 from collections.abc import Awaitable, Callable, Coroutine, Iterator
 from typing import Final
 
-from typing_extensions import TypedDict
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_logger
 from litellm.constants import (
@@ -47,8 +47,38 @@ class LoggingTask(TypedDict):
     the original task's context.
     """
 
-    coroutine: Coroutine
-    context: contextvars.Context
+    coroutine: ReadOnly[Coroutine[object, object, object]]
+    context: ReadOnly[contextvars.Context]
+
+
+class _LoggingTaskCompletion:
+    __slots__ = ("_discard", "_queue", "_sem", "_settle", "_started", "_task")
+
+    def __init__(
+        self,
+        settle: Callable[[LoggingTask, asyncio.Queue[LoggingTask], asyncio.Semaphore | None, bool], None],
+        discard: Callable[[asyncio.Task[None]], None],
+        task: LoggingTask,
+        queue: asyncio.Queue[LoggingTask],
+        sem: asyncio.Semaphore | None,
+    ) -> None:
+        self._settle = settle
+        self._discard = discard
+        self._task = task
+        self._queue = queue
+        self._sem = sem
+        self._started = False
+
+    def start(self) -> None:
+        self._started = True
+
+    def finish(self) -> None:
+        self._settle(self._task, self._queue, self._sem, False)
+
+    def done(self, processor: asyncio.Task[None]) -> None:
+        self._discard(processor)
+        if not self._started:
+            self._settle(self._task, self._queue, self._sem, True)
 
 
 class LoggingWorker:
@@ -80,6 +110,7 @@ class LoggingWorker:
         self._pending_enqueues: dict[int, LoggingTask] = {}  # mutable-ok: retries and shutdown share callback ownership
         self._retry_sleep = retry_sleep
         self._dequeued_tasks: dict[int, LoggingTask] = {}  # mutable-ok: refs so flush can rescue never-started tasks
+        self._interrupted_tasks: dict[int, LoggingTask] = {}  # mutable-ok: pending processor recovery
         self._sem: asyncio.Semaphore | None = None
         self._bound_loop: asyncio.AbstractEventLoop | None = None
         self._last_aggressive_clear_time: float = 0.0
@@ -170,6 +201,7 @@ class LoggingWorker:
                 )
             else:
                 verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
+            self._interrupted_tasks.clear()
             self._sem = None
             self._worker_task = None
             self._running_tasks.clear()
@@ -196,29 +228,53 @@ class LoggingWorker:
             for pending in tuple(self._pending_enqueues.values()):
                 self._schedule_delayed_enqueue_retry(pending)
 
-    async def _process_log_task(
-        self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]"
-    ) -> None:
-        """Runs the logging task and handles cleanup. Releases semaphore when done."""
+    async def _process_log_task(self, task: LoggingTask, completion: _LoggingTaskCompletion) -> None:
+        completion.start()
         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)
-                except Exception as 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)
-                finally:
-                    self._untrack_dequeued(task)
-                    queue.task_done()
+            except Exception as e:
+                verbose_logger.exception("LoggingWorker error: %s", e)
         finally:
-            # Always release semaphore, even if queue is None
+            completion.finish()
+
+    def _finish_log_task(
+        self,
+        task: LoggingTask,
+        queue: asyncio.Queue[LoggingTask],
+        sem: asyncio.Semaphore | None,
+        retry_unstarted: bool,
+    ) -> None:
+        if sem is not None:
             sem.release()
+        if queue is self._queue and id(task) in self._dequeued_tasks:
+            if retry_unstarted and inspect.getcoroutinestate(task["coroutine"]) == inspect.CORO_CREATED:
+                try:
+                    queue.put_nowait(task)
+                except asyncio.QueueFull:
+                    self._interrupted_tasks[id(task)] = task
+                    return
+            self._untrack_dequeued(task)
+        queue.task_done()
+
+    def _start_log_task(
+        self, task: LoggingTask, queue: asyncio.Queue[LoggingTask], sem: asyncio.Semaphore | None
+    ) -> asyncio.Task[None]:
+        completion: Final = _LoggingTaskCompletion(self._finish_log_task, self._running_tasks.discard, task, queue, sem)
+        coroutine: Final = (
+            self._process_single_task(task, completion) if sem is None else self._process_log_task(task, completion)
+        )
+        processor: Final = asyncio.create_task(coroutine)
+        processor.add_done_callback(completion.done)
+        if sem is not None:
+            self._running_tasks.add(processor)
+        return processor
 
     def _record_callback_timeout(self, coroutine: Coroutine) -> None:
         """Count a callback timeout and arm a debounced summary, so a burst of timeouts
@@ -253,23 +309,21 @@ class LoggingWorker:
         """Main worker loop that gets tasks and schedules them to run concurrently."""
         try:
             queue: Final = self._queue
-            if queue is None or self._sem is None:
+            sem: Final = self._sem
+            if queue is None or sem is None:
                 return
 
             while True:
                 # Acquire semaphore before removing task from queue to prevent
                 # unbounded growth of waiting tasks
-                await self._sem.acquire()
+                await sem.acquire()
                 try:
-                    task = await queue.get()
+                    task = self._interrupted_tasks.popitem()[1] if self._interrupted_tasks else await queue.get()
                     self._track_dequeued(task)
                     # Track each spawned coroutine so we can cancel on shutdown.
-                    processing_task = asyncio.create_task(self._process_log_task(task, self._sem, queue))
-                    self._running_tasks.add(processing_task)
-                    processing_task.add_done_callback(self._running_tasks.discard)
-                except Exception:
-                    # If task creation fails, release semaphore to prevent deadlock
-                    self._sem.release()
+                    self._start_log_task(task, queue, sem)
+                except (Exception, asyncio.CancelledError):
+                    sem.release()
                     raise
 
         except asyncio.CancelledError:
@@ -435,33 +489,27 @@ class LoggingWorker:
             if self._queue is None:
                 return
 
+            queue: Final = self._queue
             extracted_tasks: Final = self._extract_tasks_from_queue()
             if new_task is not None:
                 self._enqueue_task(new_task)
             if extracted_tasks:
-                await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks))
+                processors: Final = tuple(self._start_log_task(task, queue, None) for task in extracted_tasks)
+                await asyncio.gather(*processors)
         except Exception as e:
             verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e)
         finally:
             # Always reset the flag even if an error occurs
             self._aggressive_clear_in_progress = False
 
-    async def _process_single_task(self, task: LoggingTask) -> None:
-        """Process a task extracted from the queue and mark it done."""
-        if self._queue is None:
-            return
-
+    async def _process_single_task(self, task: LoggingTask, completion: _LoggingTaskCompletion) -> None:
+        completion.start()
         try:
-            await asyncio.wait_for(
-                self._create_callback_task(task),
-                timeout=self.timeout,
-            )
+            await asyncio.wait_for(self._create_callback_task(task), timeout=self.timeout)
         except Exception:
-            # Suppress errors during processing to ensure we keep going
             pass
         finally:
-            self._untrack_dequeued(task)
-            self._queue.task_done()
+            completion.finish()
 
     def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None:
         """
@@ -493,6 +541,13 @@ class LoggingWorker:
         self._emit_timeout_summary()
 
         if discard_pending:
+            for task in self._unstarted_dequeued_tasks():
+                task["coroutine"].close()
+            for task in tuple(self._interrupted_tasks.values()):
+                self._untrack_dequeued(task)
+                if self._queue is not None:
+                    self._queue.task_done()
+            self._interrupted_tasks.clear()
             rejected: Final = tuple(self._pending_enqueues.values())
             self._pending_enqueues.clear()
             for task in rejected:
@@ -518,11 +573,6 @@ class LoggingWorker:
         if pending:
             verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending))
 
-        if discard_pending and self._queue is not None:
-            for task in self._unstarted_dequeued_tasks():
-                task["coroutine"].close()
-                self._untrack_dequeued(task)
-                self._queue.task_done()
         if self._worker_task is not None and self._worker_task.done():
             self._worker_task = None
         # Drop references to completed tasks so we can restart cleanly.

tests/unit/litellm_core_utils/test_logging_worker.py

@@ -38,7 +38,7 @@ async def test_optional_work_budget_preserves_callback_context_and_reserves_logg
             await worker.flush()
         else:
             assert worker._queue is not None
-            await worker._process_single_task(worker._queue.get_nowait())
+            await worker._aggressively_clear_queue_async()
         restored_identity, budget = await asyncio.wait_for(results.get(), timeout=2)
         assert restored_identity == "request"
         assert 0 < budget <= worker.timeout / 4
@@ -132,6 +132,102 @@ async def test_stop_does_not_replay_a_callback_dequeued_before_its_processor_sta
     await asyncio.wait_for(worker._queue.join(), timeout=1)
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("concurrency", (1, 100))
+@pytest.mark.parametrize("fill_queue", (False, True))
+async def test_restart_delivers_callbacks_cancelled_before_execution_once(concurrency: int, fill_queue: bool) -> None:
+    worker: Final = LoggingWorker(concurrency=concurrency, max_queue_size=1)
+    completed: Final[asyncio.Queue[int]] = asyncio.Queue()
+
+    async def callback(identity: int) -> None:
+        completed.put_nowait(identity)
+
+    try:
+        for cycle in range(3):
+            worker.ensure_initialized_and_enqueue(callback(cycle * 2))
+            await asyncio.sleep(0)
+            await worker.stop()
+            if fill_queue:
+                worker.enqueue(callback(cycle * 2 + 1))
+            await asyncio.wait_for(worker.flush(), timeout=2)
+            expected: Final = (cycle * 2, cycle * 2 + 1) if fill_queue else (cycle * 2,)
+            assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == list(expected)
+            assert worker.pending_count == 0
+    finally:
+        await worker.stop(discard_pending=True, timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_idle_stop_preserves_capacity_for_later_callbacks() -> None:
+    worker: Final = LoggingWorker(concurrency=1)
+    completed: Final[asyncio.Queue[int]] = asyncio.Queue()
+
+    async def callback(identity: int) -> None:
+        completed.put_nowait(identity)
+
+    try:
+        for cycle in range(3):
+            worker.start()
+            await asyncio.sleep(0)
+            await worker.stop()
+            worker.ensure_initialized_and_enqueue(callback(cycle))
+            await asyncio.wait_for(worker.flush(), timeout=2)
+            assert completed.get_nowait() == cycle
+    finally:
+        await worker.stop(discard_pending=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("discard", (False, True))
+async def test_overflow_stop_before_processors_start_preserves_delivery_policy(discard: bool) -> None:
+    worker: Final = LoggingWorker(max_queue_size=2, concurrency=1)
+    completed: Final[asyncio.Queue[int]] = asyncio.Queue()
+
+    async def callback(identity: int) -> None:
+        completed.put_nowait(identity)
+
+    try:
+        worker.start()
+        for identity in range(3):
+            worker.enqueue(callback(identity))
+        await asyncio.sleep(0)
+        await worker.stop(discard_pending=discard)
+        await asyncio.wait_for(worker.flush(), timeout=2)
+        assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == ([] if discard else [0, 1, 2])
+        worker.ensure_initialized_and_enqueue(callback(3))
+        await asyncio.wait_for(worker.flush(), timeout=2)
+        assert completed.get_nowait() == 3
+        assert worker.pending_count == 0
+    finally:
+        await worker.stop(discard_pending=True)
+
+
+@pytest.mark.asyncio
+async def test_restart_does_not_replay_a_callback_cancelled_after_it_started() -> None:
+    worker: Final = LoggingWorker(concurrency=1)
+    started: Final = asyncio.Event()
+    completed: Final[asyncio.Queue[str]] = asyncio.Queue()
+
+    async def callback() -> None:
+        completed.put_nowait("started")
+        started.set()
+        await asyncio.Event().wait()
+
+    async def next_callback() -> None:
+        completed.put_nowait("next")
+
+    try:
+        worker.ensure_initialized_and_enqueue(callback())
+        await started.wait()
+        await worker.stop()
+        worker.ensure_initialized_and_enqueue(next_callback())
+        await asyncio.wait_for(worker.flush(), timeout=2)
+        assert tuple(completed.get_nowait() for _ in range(completed.qsize())) == ("started", "next")
+        assert worker.pending_count == 0
+    finally:
+        await worker.stop(discard_pending=True)
+
+
 @pytest.mark.asyncio
 async def test_stop_timeout_keeps_callbacks_observable_until_queue_cleanup_finishes() -> None:
     worker: Final = LoggingWorker(concurrency=1)
@@ -428,7 +524,7 @@ class TestLoggingWorker:
         asyncio.run(short_lived_script())
 
         assert worker._queue is not None
-        assert worker._queue.qsize() == 0, "precondition: the worker loop dequeued the task before loop close"
+        assert worker.pending_count == 1
         assert fired == [], "precondition: the callback never ran before loop close"
 
         worker._flush_on_exit()
@@ -537,8 +633,7 @@ class TestLoggingWorker:
 
         asyncio.run(enqueue_on_first_loop())
         assert worker._queue is not None
-        expected_shape = (1, 0) if stranded == "still_queued" else (0, 1)
-        assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape
+        assert worker.pending_count == 1
         assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed"
 
         async def flush_twice_on_second_loop():