Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions litellm/litellm_core_utils/logging_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -484,9 +484,13 @@ async def flush(self) -> None:
so it correctly handles items that have been dequeued but whose
callback hasn't finished yet — ``queue.empty()`` would return True in
that window and cause us to skip the wait.

``start()`` runs first so a queue left behind by a previous event loop
is carried onto this one and drained here instead of joined forever.
"""
if self._queue is None:
return
self.start()
await self._queue.join()

async def clear_queue(self):
Expand Down
47 changes: 47 additions & 0 deletions tests/test_litellm/litellm_core_utils/test_logging_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,53 @@ async def second_script():

assert sorted(fired) == ["first", "second"]

@pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"])
def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded):
"""
Regression: ``flush()`` from a new event loop used to ``join()`` the queue bound to the
previous loop, whose unfinished counter nothing on the new loop ever decrements. The first
such flush hung until pytest-timeout killed it and every later one raised
``RuntimeError: ... is bound to a different event loop`` from the queue's Event.
"""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()

async def enqueue_on_first_loop():
if stranded == "still_queued":
worker._ensure_queue()
worker.enqueue(callback())
return
worker.ensure_initialized_and_enqueue(callback())

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 callback.await_count == 0, "precondition: the callback never ran before the first loop closed"

async def flush_twice_on_second_loop():
await asyncio.wait_for(worker.flush(), timeout=5)
await asyncio.wait_for(worker.flush(), timeout=5)

asyncio.run(flush_twice_on_second_loop())

assert callback.await_count == 1

def test_flush_starts_a_worker_when_the_queue_has_none(self):
"""``flush()`` must drain a queue that exists on the current loop without a running worker."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()

async def enqueue_then_flush():
worker._ensure_queue()
worker.enqueue(callback())
assert worker._worker_task is None, "precondition: nothing is draining the queue yet"
await asyncio.wait_for(worker.flush(), timeout=3)

asyncio.run(enqueue_then_flush())

assert callback.await_count == 1

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)
Expand Down
Loading