From d7ff8deec1c69ae1497b27225b2ee383dafd2f92 Mon Sep 17 00:00:00 2001 From: King Star Date: Fri, 14 Aug 2026 03:43:23 +0800 Subject: [PATCH] fix(logging): release stale loop callbacks --- litellm/litellm_core_utils/logging_worker.py | 49 ++++++++++++++-- .../litellm_core_utils/test_logging_worker.py | 58 +++++++++++++++++-- 2 files changed, 96 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index f6bca6e4006..59f0ee0d7ec 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -56,6 +56,10 @@ class LoggingWorker: self._queue: asyncio.Queue[LoggingTask] | None = None self._worker_task: asyncio.Task[None] | None = None self._running_tasks: set[asyncio.Task[None]] = set() # mutable-ok: tracks live tasks for cancellation + self._task_coroutines: dict[ # mutable-ok: closes queued callbacks when wrapper tasks are cancelled + asyncio.Task[None], + Coroutine[object, object, object], + ] = {} self._bound_loop: asyncio.AbstractEventLoop | None = None self._stopping: bool = False self._last_aggressive_clear_time: float = 0.0 @@ -76,9 +80,9 @@ class LoggingWorker: 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._close_queued_coroutines() self._queue = None - self._worker_task = None - self._running_tasks.clear() + self._release_running_tasks() if self._queue is None: self._queue = asyncio.Queue(maxsize=self.max_queue_size) @@ -123,16 +127,44 @@ class LoggingWorker: self._process_log_task(logging_task), ) self._running_tasks.add(processing_task) + self._task_coroutines[processing_task] = logging_task["coroutine"] self._worker_task = processing_task processing_task.add_done_callback(self._processing_task_done) return True def _processing_task_done(self, task: asyncio.Task[None]) -> None: self._running_tasks.discard(task) + if coroutine := self._task_coroutines.pop(task, None): + coroutine.close() if self._worker_task is task: self._worker_task = next(iter(self._running_tasks), None) self._start_queued_tasks() + def _close_queued_coroutines(self) -> None: + if self._queue is None: + return + while True: + try: + logging_task: Final = self._queue.get_nowait() + except asyncio.QueueEmpty: + return + logging_task["coroutine"].close() + self._queue.task_done() + + def _release_running_tasks(self) -> None: + for task in tuple(self._running_tasks): + if not task.done(): + try: + task.cancel() + except RuntimeError: + pass + setattr(task, "_log_destroy_pending", False) + if coroutine := self._task_coroutines.get(task): + coroutine.close() + self._running_tasks.clear() + self._task_coroutines.clear() + self._worker_task = None + def enqueue(self, coroutine: Coroutine[object, object, object]) -> None: """ Add a coroutine to the logging queue. @@ -353,9 +385,8 @@ class LoggingWorker: # Wait for cancellation to settle; ignore errors raised during shutdown. await asyncio.gather(*tasks_to_cancel, return_exceptions=True) - self._worker_task = None # Drop references to completed tasks so we can restart cleanly. - self._running_tasks.clear() + self._release_running_tasks() await self.clear_queue() self._stopping = False @@ -522,6 +553,14 @@ class LoggingWorker: def flush_on_exit(self) -> None: self._flush_on_exit() + def flush_stale_loop(self) -> None: + self._flush_on_exit() + self._close_queued_coroutines() + self._release_running_tasks() + self._queue = None + self._bound_loop = None + self._stopping = False + class LoopAwareLoggingWorker: def __init__( @@ -567,7 +606,7 @@ class LoopAwareLoggingWorker: self._workers[id(loop)] = (weakref.ref(loop), worker) for stale_worker in stale_workers: - stale_worker.flush_on_exit() + stale_worker.flush_stale_loop() return worker def start(self) -> None: diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index d2bcb486a84..74bba4b3ec8 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -316,17 +316,17 @@ class TestLoggingWorker: task1_result = next((r for r in task_results if r["task_id"] == "task_1"), None) assert task1_result is not None, "Task 1 result not found" assert task1_result["context_accessible"] is True, "Task 1 should have access to context variable" - assert ( - task1_result["context_value"] == "context_1" - ), f"Task 1 should see 'context_1', got: {task1_result['context_value']}" + assert task1_result["context_value"] == "context_1", ( + f"Task 1 should see 'context_1', got: {task1_result['context_value']}" + ) # Task 2 should see "context_2" task2_result = next((r for r in task_results if r["task_id"] == "task_2"), None) assert task2_result is not None, "Task 2 result not found" assert task2_result["context_accessible"] is True, "Task 2 should have access to context variable" - assert ( - task2_result["context_value"] == "context_2" - ), f"Task 2 should see 'context_2', got: {task2_result['context_value']}" + assert task2_result["context_value"] == "context_2", ( + f"Task 2 should see 'context_2', got: {task2_result['context_value']}" + ) # Task 3 should not have access to the context variable task3_result = next((r for r in task_results if r["task_id"] == "task_3"), None) @@ -490,3 +490,49 @@ def test_loop_aware_worker_does_not_strand_tasks_when_event_loops_close(): loop.close() assert errors.empty() + + +def test_loop_aware_worker_releases_active_tasks_from_closed_loops(): + worker = LoopAwareLoggingWorker(timeout=30.0, max_queue_size=10) + errors: queue.Queue[dict[str, object]] = queue.Queue() + captured_worker: LoggingWorker | None = None + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.set_exception_handler(lambda _loop, context: errors.put(context)) + started = asyncio.Event() + + async def wait_forever() -> None: + started.set() + await asyncio.Event().wait() + + async def run() -> None: + nonlocal captured_worker + worker.ensure_initialized_and_enqueue(wait_forever()) + captured_worker = worker._workers[id(asyncio.get_running_loop())][1] + await asyncio.wait_for(started.wait(), timeout=2.0) + + try: + loop.run_until_complete(run()) + finally: + loop.close() + + next_loop = asyncio.new_event_loop() + asyncio.set_event_loop(next_loop) + + async def trigger_cleanup() -> None: + worker.start() + + try: + next_loop.run_until_complete(trigger_cleanup()) + finally: + next_loop.run_until_complete(worker.stop()) + next_loop.close() + asyncio.set_event_loop(None) + + assert captured_worker is not None + assert captured_worker._running_tasks == set() + assert captured_worker._task_coroutines == {} + assert captured_worker._queue is None + assert worker._workers == {} + assert errors.empty()