mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(logging_worker): carry queued tasks across event-loop change instead of dropping them
LoggingWorker._ensure_queue nulled self._queue on a loop change, discarding every pending LoggingTask (each an un-awaited spend-logging coroutine) with no counter and only a debug log. SDK callers using asyncio.run() per request and mixed sync/async processes rebind the queue's loop and silently lose spend rows and observability events. Drain the stale queue and move the pending tasks onto a fresh queue bound to the new loop, warn with the carried-over count, and keep flush()/join() honest since the queue is no longer thrown away. Adds a regression test that fills the queue before the loop change and asserts every task survives and still executes.
This commit is contained in:
parent
85d5ac2b5c
commit
12a34a10d8
2 changed files with 70 additions and 5 deletions
|
|
@ -5,7 +5,7 @@ import asyncio
|
||||||
import atexit
|
import atexit
|
||||||
import contextvars
|
import contextvars
|
||||||
import logging
|
import logging
|
||||||
from collections.abc import Coroutine
|
from collections.abc import Coroutine, Iterator
|
||||||
from typing import Final
|
from typing import Final
|
||||||
|
|
||||||
from typing_extensions import TypedDict
|
from typing_extensions import TypedDict
|
||||||
|
|
@ -61,6 +61,19 @@ class LoggingWorker:
|
||||||
# Register cleanup handler to flush remaining events on exit
|
# Register cleanup handler to flush remaining events on exit
|
||||||
atexit.register(self._flush_on_exit)
|
atexit.register(self._flush_on_exit)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _drain_pending(queue: "asyncio.Queue[LoggingTask]") -> tuple[LoggingTask, ...]:
|
||||||
|
"""Pop every task still queued, without awaiting them, so they can be moved to another queue."""
|
||||||
|
|
||||||
|
def _pop_until_empty() -> Iterator[LoggingTask]:
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
yield queue.get_nowait()
|
||||||
|
except asyncio.QueueEmpty:
|
||||||
|
return
|
||||||
|
|
||||||
|
return tuple(_pop_until_empty())
|
||||||
|
|
||||||
def _ensure_queue(self) -> None:
|
def _ensure_queue(self) -> None:
|
||||||
"""Initialize the queue if it doesn't exist or if event loop has changed."""
|
"""Initialize the queue if it doesn't exist or if event loop has changed."""
|
||||||
try:
|
try:
|
||||||
|
|
@ -69,14 +82,27 @@ class LoggingWorker:
|
||||||
# No running loop, can't initialize
|
# No running loop, can't initialize
|
||||||
return
|
return
|
||||||
|
|
||||||
# Check if we need to reinitialize due to event loop change
|
# The queue, semaphore and worker task are all bound to the loop that created them. On a
|
||||||
|
# loop change we hand the still-pending tasks to a fresh queue instead of dropping them,
|
||||||
|
# so queued spend-logging coroutines are not silently discarded (and never left un-awaited).
|
||||||
if self._queue is not None and self._bound_loop is not current_loop:
|
if self._queue is not None and self._bound_loop is not current_loop:
|
||||||
|
carried_over: Final = self._drain_pending(self._queue)
|
||||||
|
new_queue: Final[asyncio.Queue[LoggingTask]] = asyncio.Queue(maxsize=self.max_queue_size)
|
||||||
|
for carried_task in carried_over:
|
||||||
|
new_queue.put_nowait(carried_task)
|
||||||
|
if carried_over:
|
||||||
|
verbose_logger.warning(
|
||||||
|
"LoggingWorker: event loop changed; carried %d pending logging task(s) onto the new loop",
|
||||||
|
len(carried_over),
|
||||||
|
)
|
||||||
|
else:
|
||||||
verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
|
verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
|
||||||
# Clear old state - these are bound to the old loop
|
|
||||||
self._queue = None
|
|
||||||
self._sem = None
|
self._sem = None
|
||||||
self._worker_task = None
|
self._worker_task = None
|
||||||
self._running_tasks.clear()
|
self._running_tasks.clear()
|
||||||
|
self._queue = new_queue
|
||||||
|
self._bound_loop = current_loop
|
||||||
|
return
|
||||||
|
|
||||||
if self._queue is None:
|
if self._queue is None:
|
||||||
self._queue = asyncio.Queue(maxsize=self.max_queue_size)
|
self._queue = asyncio.Queue(maxsize=self.max_queue_size)
|
||||||
|
|
|
||||||
|
|
@ -413,3 +413,42 @@ class TestLoggingWorker:
|
||||||
assert worker2._bound_loop is not None
|
assert worker2._bound_loop is not None
|
||||||
|
|
||||||
await worker2.stop()
|
await worker2.stop()
|
||||||
|
|
||||||
|
def test_event_loop_change_carries_pending_tasks_over(self):
|
||||||
|
"""Regression (LIT-6028): a loop change must not silently drop queued coroutines.
|
||||||
|
|
||||||
|
Before the fix ``_ensure_queue`` nulled ``self._queue`` on a loop change, discarding
|
||||||
|
every pending ``LoggingTask`` (each an un-awaited spend-logging coroutine). The tasks
|
||||||
|
must instead be moved onto the queue bound to the new loop and still execute there.
|
||||||
|
"""
|
||||||
|
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
|
||||||
|
executed: list[int] = []
|
||||||
|
|
||||||
|
async def spend_log(index: int) -> None:
|
||||||
|
executed.append(index)
|
||||||
|
|
||||||
|
async def enqueue_on_first_loop() -> None:
|
||||||
|
worker._ensure_queue()
|
||||||
|
for i in range(5):
|
||||||
|
worker.enqueue(spend_log(i))
|
||||||
|
assert worker._queue is not None
|
||||||
|
assert worker._queue.qsize() == 5
|
||||||
|
|
||||||
|
asyncio.run(enqueue_on_first_loop())
|
||||||
|
|
||||||
|
stale_queue = worker._queue
|
||||||
|
assert stale_queue is not None
|
||||||
|
|
||||||
|
async def rebind_on_second_loop() -> None:
|
||||||
|
worker._ensure_queue()
|
||||||
|
assert worker._queue is not None
|
||||||
|
# A fresh queue bound to the new loop, holding every carried-over task (not dropped).
|
||||||
|
assert worker._queue is not stale_queue
|
||||||
|
assert worker._queue.qsize() == 5
|
||||||
|
while not worker._queue.empty():
|
||||||
|
task = worker._queue.get_nowait()
|
||||||
|
await task["context"].run(asyncio.create_task, task["coroutine"])
|
||||||
|
|
||||||
|
asyncio.run(rebind_on_second_loop())
|
||||||
|
|
||||||
|
assert sorted(executed) == [0, 1, 2, 3, 4]
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue