scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(logging): isolate workers per event loop

Browse source 
Use a worker registry keyed by event loop so queue state and callback tasks never cross loop boundaries. Idle runners exit when their queue is drained, and shutdown no longer strands pending tasks.

Signed-off-by: King Star <mcxin.y@gmail.com>
This commit is contained in:
King Star 2026-08-13 20:46:28 +08:00
parent 09889e1986
commit 1d4a0c9248
No known key found for this signature in database
2 changed files with 273 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

204
litellm/litellm_core_utils/logging_worker.py
View file

@ -5,6 +5,8 @@ import asyncio
import atexit
import contextvars
import logging
import threading
import weakref
from collections.abc import Coroutine
from typing import Final
@ -28,7 +30,7 @@ class LoggingTask(TypedDict):
the original task's context.
"""
coroutine: Coroutine
coroutine: Coroutine[object, object, object]
context: contextvars.Context
@ -46,20 +48,21 @@ class LoggingWorker:
timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
concurrency: int = LOGGING_WORKER_CONCURRENCY,
register_atexit: bool = True,
):
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self._queue: asyncio.Queue[LoggingTask] | None = None
self._worker_task: asyncio.Task | None = None
self._running_tasks: set[asyncio.Task] = set()
self._sem: asyncio.Semaphore | 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._bound_loop: asyncio.AbstractEventLoop | None = None
self._stopping: bool = False
self._last_aggressive_clear_time: float = 0.0
self._aggressive_clear_in_progress: bool = False
# Register cleanup handler to flush remaining events on exit
atexit.register(self._flush_on_exit)
if register_atexit:
atexit.register(self._flush_on_exit)
def _ensure_queue(self) -> None:
"""Initialize the queue if it doesn't exist or if event loop has changed."""
@ -74,7 +77,6 @@ class LoggingWorker:
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._worker_task = None
self._running_tasks.clear()
@ -85,56 +87,53 @@ class LoggingWorker:
def start(self) -> None:
"""Start the logging worker. Idempotent - safe to call multiple times."""
self._ensure_queue()
if self._sem is None:
self._sem = asyncio.Semaphore(self.concurrency)
if self._worker_task is None or self._worker_task.done():
self._worker_task = asyncio.create_task(self._worker_loop())
async def _process_log_task(self, task: LoggingTask, sem: asyncio.Semaphore):
"""Runs the logging task and handles cleanup. Releases semaphore when done."""
async def _process_log_task(self, task: LoggingTask) -> None:
"""Run one logging task and update the queue completion counter."""
try:
if self._queue is not None:
try:
# Run the coroutine in its original context
await asyncio.wait_for(
task["context"].run(asyncio.create_task, task["coroutine"]),
timeout=self.timeout,
)
except Exception as e:
verbose_logger.exception("LoggingWorker error: %s", e)
finally:
self._queue.task_done()
await asyncio.wait_for(task["coroutine"], timeout=self.timeout)
except Exception as e:
verbose_logger.exception("LoggingWorker error: %s", e)
finally:
# Always release semaphore, even if queue is None
sem.release()
if self._queue is not None:
self._queue.task_done()
async def _worker_loop(self) -> None:
"""Main worker loop that gets tasks and schedules them to run concurrently."""
def _start_queued_tasks(self) -> None:
if self._queue is None or self._stopping:
return
try:
if self._queue is None or self._sem is None:
asyncio.get_running_loop()
except RuntimeError:
return
while len(self._running_tasks) < self.concurrency:
if not self._start_one_queued_task():
return
while True:
# Acquire semaphore before removing task from queue to prevent
# unbounded growth of waiting tasks
await self._sem.acquire()
try:
task = await self._queue.get()
# Track each spawned coroutine so we can cancel on shutdown.
processing_task = asyncio.create_task(self._process_log_task(task, self._sem))
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()
raise
def _start_one_queued_task(self) -> bool:
if self._queue is None:
return False
try:
logging_task: Final = self._queue.get_nowait()
except asyncio.QueueEmpty:
return False
except asyncio.CancelledError:
verbose_logger.debug("LoggingWorker cancelled during shutdown")
# Attempt to clear remaining items to prevent "never awaited" warnings
await self.clear_queue()
processing_task: Final = logging_task["context"].run(
asyncio.create_task,
self._process_log_task(logging_task),
)
self._running_tasks.add(processing_task)
self._worker_task = processing_task
processing_task.add_done_callback(self._processing_task_done)
return True
def enqueue(self, coroutine: Coroutine) -> None:
def _processing_task_done(self, task: asyncio.Task[None]) -> None:
self._running_tasks.discard(task)
if self._worker_task is task:
self._worker_task = next(iter(self._running_tasks), None)
self._start_queued_tasks()
def enqueue(self, coroutine: Coroutine[object, object, object]) -> None:
"""
Add a coroutine to the logging queue.
Hot path: never blocks, aggressively clears queue if full.
@ -147,6 +146,7 @@ class LoggingWorker:
try:
self._queue.put_nowait(task)
self._start_queued_tasks()
except asyncio.QueueFull:
# Queue is full - handle it appropriately
verbose_logger.exception("LoggingWorker queue is full")
@ -330,7 +330,7 @@ class LoggingWorker:
# Process all tasks concurrently for maximum speed
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine):
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, object]) -> None:
"""
Ensure the logging worker is initialized and enqueue the coroutine.
"""
@ -340,13 +340,11 @@ class LoggingWorker:
async def stop(self) -> None:
"""Stop the logging worker and clean up resources."""
if self._worker_task is None and not self._running_tasks:
# No worker launched and no in-flight tasks to drain.
await self.clear_queue()
return
tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks)
if self._worker_task:
# Include the main worker loop so it stops fetching work.
tasks_to_cancel.append(self._worker_task)
self._stopping = True
tasks_to_cancel: Final[list[asyncio.Task[None]]] = list(self._running_tasks)
for task in tasks_to_cancel:
# Propagate cancellation to every pending task.
@ -358,6 +356,8 @@ class LoggingWorker:
self._worker_task = None
# Drop references to completed tasks so we can restart cleanly.
self._running_tasks.clear()
await self.clear_queue()
self._stopping = False
async def flush(self) -> None:
"""Flush the logging queue.
@ -519,6 +519,100 @@ class LoggingWorker:
finally:
loop.close()
def flush_on_exit(self) -> None:
self._flush_on_exit()
class LoopAwareLoggingWorker:
def __init__(
self,
timeout: float = LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
concurrency: int = LOGGING_WORKER_CONCURRENCY,
) -> None:
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self._workers: dict[ # mutable-ok: registry is updated as event loops appear and close
int,
tuple[weakref.ReferenceType[asyncio.AbstractEventLoop], LoggingWorker],
] = {}
self._lock = threading.Lock()
atexit.register(self._flush_on_exit)
def _worker_for_current_loop(self, *, create: bool = True) -> LoggingWorker | None:
try:
loop: Final = asyncio.get_running_loop()
except RuntimeError:
return None
with self._lock:
stale_loop_ids: Final = tuple(
loop_id
for loop_id, (loop_ref, _worker) in self._workers.items()
if (tracked_loop := loop_ref()) is None or tracked_loop.is_closed()
)
stale_workers: Final = tuple(self._workers.pop(loop_id)[1] for loop_id in stale_loop_ids)
worker_entry: Final = self._workers.get(id(loop))
worker: LoggingWorker | None = (
worker_entry[1] if worker_entry is not None and worker_entry[0]() is loop else None
)
if worker is None and create:
worker = LoggingWorker(
timeout=self.timeout,
max_queue_size=self.max_queue_size,
concurrency=self.concurrency,
register_atexit=False,
)
self._workers[id(loop)] = (weakref.ref(loop), worker)
for stale_worker in stale_workers:
stale_worker.flush_on_exit()
return worker
def start(self) -> None:
worker: Final = self._worker_for_current_loop()
if worker is not None:
worker.start()
def enqueue(self, coroutine: Coroutine[object, object, object]) -> None:
worker: Final = self._worker_for_current_loop()
if worker is not None:
worker.enqueue(coroutine)
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, object]) -> None:
worker: Final = self._worker_for_current_loop()
if worker is not None:
worker.ensure_initialized_and_enqueue(async_coroutine)
async def stop(self) -> None:
loop: Final = asyncio.get_running_loop()
worker: Final = self._worker_for_current_loop(create=False)
if worker is None:
return
await worker.stop()
with self._lock:
worker_entry: Final = self._workers.get(id(loop))
if worker_entry is not None and worker_entry[0]() is loop:
self._workers.pop(id(loop))
async def flush(self) -> None:
worker: Final = self._worker_for_current_loop(create=False)
if worker is not None:
await worker.flush()
async def clear_queue(self) -> None:
worker: Final = self._worker_for_current_loop(create=False)
if worker is not None:
await worker.clear_queue()
def _flush_on_exit(self) -> None:
with self._lock:
workers: Final = tuple(worker for _loop_ref, worker in self._workers.values())
for worker in workers:
worker.flush_on_exit()
# Global instance for backward compatibility
GLOBAL_LOGGING_WORKER: Final = LoggingWorker()
GLOBAL_LOGGING_WORKER: Final = LoopAwareLoggingWorker()

171
tests/test_litellm/litellm_core_utils/test_logging_worker.py
View file

@ -6,12 +6,17 @@ import asyncio
import contextvars
import io
import logging
import queue
import threading
from unittest.mock import AsyncMock, patch
import pytest
from litellm.constants import LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS
from litellm.litellm_core_utils.logging_worker import LoggingWorker
from litellm.litellm_core_utils.logging_worker import (
LoggingWorker,
LoopAwareLoggingWorker,
)
class TestLoggingWorker:
@ -26,27 +31,19 @@ class TestLoggingWorker:
async def test_graceful_shutdown_with_clear_queue(self, logging_worker):
"""Test that cancellation triggers clear_queue to prevent 'never awaited' warnings."""
# Mock the clear_queue method to verify it's called during cancellation
with patch.object(
logging_worker, "clear_queue", new_callable=AsyncMock
) as mock_clear_queue:
# Start the worker
logging_worker.start()
with patch.object(logging_worker, "clear_queue", new_callable=AsyncMock) as mock_clear_queue:
started = asyncio.Event()
# Give it a moment to start
await asyncio.sleep(0.1)
async def wait_until_cancelled():
started.set()
await asyncio.Event().wait()
# Cancel the worker task to simulate shutdown
if logging_worker._worker_task:
logging_worker._worker_task.cancel()
logging_worker.ensure_initialized_and_enqueue(wait_until_cancelled())
# Wait for the task to handle the cancellation
try:
await logging_worker._worker_task
except asyncio.CancelledError:
# Expected during cancellation
pass
await started.wait()
await logging_worker.stop()
# Verify that clear_queue was called during cancellation
mock_clear_queue.assert_called_once()
@pytest.mark.asyncio
@ -122,9 +119,7 @@ class TestLoggingWorker:
async def test_worker_handles_cancellation_gracefully(self, logging_worker):
"""Test that the worker handles cancellation without throwing exceptions."""
# Mock verbose_logger to capture debug messages
with patch(
"litellm.litellm_core_utils.logging_worker.verbose_logger"
) as mock_logger:
with patch("litellm.litellm_core_utils.logging_worker.verbose_logger") as mock_logger:
# Start the worker
logging_worker.start()
@ -142,6 +137,29 @@ class TestLoggingWorker:
]
assert len(debug_calls) >= 0 # May be 0 if no cancellation occurred
@pytest.mark.asyncio
async def test_stop_drains_queue_without_leaving_runners(self):
worker = LoggingWorker(timeout=1.0, max_queue_size=10, concurrency=1)
started = asyncio.Event()
queued_ran = False
async def wait_until_cancelled() -> None:
started.set()
await asyncio.Event().wait()
async def queued_task() -> None:
nonlocal queued_ran
queued_ran = True
worker.ensure_initialized_and_enqueue(wait_until_cancelled())
worker.ensure_initialized_and_enqueue(queued_task())
await started.wait()
await worker.stop()
assert queued_ran is True
assert worker._running_tasks == set()
@pytest.mark.asyncio
async def test_enqueue_and_process_single_item(self, logging_worker):
"""Test basic enqueue and process functionality."""
@ -187,33 +205,34 @@ class TestLoggingWorker:
async def test_queue_full_handling(self, logging_worker):
"""Test that queue full condition is handled gracefully."""
# Create a worker with very small queue size
small_worker = LoggingWorker(timeout=1.0, max_queue_size=2)
small_worker = LoggingWorker(timeout=1.0, max_queue_size=2, concurrency=0)
small_worker._ensure_queue()
# Mock verbose_logger to capture exception messages
with patch(
"litellm.litellm_core_utils.logging_worker.verbose_logger"
) as mock_logger:
with patch("litellm.litellm_core_utils.logging_worker.verbose_logger") as mock_logger:
# Fill the queue beyond capacity
mock_coro = AsyncMock()
async def mock_coro() -> None:
pass
pending_coroutines = []
for _ in range(5): # More than max_queue_size of 2
small_worker.enqueue(mock_coro())
coroutine = mock_coro()
pending_coroutines.append(coroutine)
small_worker.enqueue(coroutine)
# Should have logged queue full exceptions
exception_calls = [
call
for call in mock_logger.exception.call_args_list
if "queue is full" in str(call)
]
exception_calls = [call for call in mock_logger.exception.call_args_list if "queue is full" in str(call)]
assert len(exception_calls) > 0
await small_worker.clear_queue()
for coroutine in pending_coroutines:
coroutine.close()
@pytest.mark.asyncio
async def test_context_propagation(self, logging_worker):
"""Test that enqueued tasks execute in their original context."""
# Create a context variable for testing
test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar(
"test_context_var"
)
test_context_var: contextvars.ContextVar[str] = contextvars.ContextVar("test_context_var")
# Track results from multiple tasks using asyncio.Event for synchronization
task_results = []
@ -291,16 +310,12 @@ class TestLoggingWorker:
task_results.sort(key=lambda x: x["task_id"])
# Verify that each task saw its own context
assert (
len(task_results) == 3
), f"Expected 3 results, got {len(task_results)}: {task_results}"
assert len(task_results) == 3, f"Expected 3 results, got {len(task_results)}: {task_results}"
# Task 1 should see "context_1"
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_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']}"
@ -308,9 +323,7 @@ class TestLoggingWorker:
# 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_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']}"
@ -318,9 +331,7 @@ class TestLoggingWorker:
# 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)
assert task3_result is not None, "Task 3 result not found"
assert (
task3_result["context_accessible"] is False
), "Task 3 should not have access to context variable"
assert task3_result["context_accessible"] is False, "Task 3 should not have access to context variable"
@pytest.mark.asyncio
async def test_semaphore_concurrency_limit(self):
@ -413,3 +424,69 @@ class TestLoggingWorker:
assert worker2._bound_loop is not None
await worker2.stop()
def test_loop_aware_worker_keeps_event_loop_queues_isolated():
worker = LoopAwareLoggingWorker(timeout=1.0, max_queue_size=10)
ready = threading.Barrier(2)
processed: queue.Queue[int] = queue.Queue()
errors: queue.Queue[BaseException | dict[str, object]] = queue.Queue()
def run_worker(index: int) -> None:
loop = asyncio.new_event_loop()
loop.set_exception_handler(lambda _loop, context: errors.put(context))
ready.wait()
async def run() -> None:
completed = asyncio.Event()
async def record_and_complete() -> None:
processed.put(index)
completed.set()
worker.ensure_initialized_and_enqueue(record_and_complete())
await asyncio.wait_for(completed.wait(), timeout=2.0)
await worker.flush()
await worker.stop()
try:
loop.run_until_complete(run())
except BaseException as exc:
errors.put(exc)
finally:
loop.close()
threads = tuple(threading.Thread(target=run_worker, args=(index,)) for index in range(2))
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout=5.0)
assert all(not thread.is_alive() for thread in threads)
collected_errors = tuple(errors.get_nowait() for _ in range(errors.qsize()))
processed_indexes = tuple(processed.get_nowait() for _ in range(processed.qsize()))
assert collected_errors == ()
assert sorted(processed_indexes) == [0, 1]
assert worker._workers == {}
def test_loop_aware_worker_does_not_strand_tasks_when_event_loops_close():
worker = LoopAwareLoggingWorker(timeout=1.0, max_queue_size=10)
errors: queue.Queue[dict[str, object]] = queue.Queue()
async def noop() -> None:
pass
for _ in range(2):
async def run() -> None:
worker.ensure_initialized_and_enqueue(noop())
await asyncio.sleep(0)
loop = asyncio.new_event_loop()
loop.set_exception_handler(lambda _loop, context: errors.put(context))
asyncio.set_event_loop(loop)
loop.run_until_complete(run())
loop.close()
assert errors.empty()