scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix: prevent memory blowout in LoggingWorker

Browse source 
Previously, tasks were executed sequentially, awaiting each task before
starting the next. With large queues (10k+ tasks), this caused objects to
accumulate in memory, holding references to heavy resources and leading
to memory blowouts.

This update:

1. Introduces a semaphore to allow a configurable number of concurrent tasks,
   improving throughput and preventing queue buildup.
2. Implements a configurable cleaning mechanism for when the queue reaches its
   limit, ensuring tasks are not dropped.
This commit is contained in:
AlexsanderHamir 2025-11-12 17:15:54 -08:00
parent c471bf1f16
commit 5910750ed6
4 changed files with 386 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
docs/my-website/docs/proxy/config_settings.md
View file

@ -643,7 +643,14 @@ router_settings:
| LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD | If true, prints the standard logging payload to the console - useful for debugging
| LITELM_ENVIRONMENT | Environment for LiteLLM Instance. This is currently only logged to DeepEval to determine the environment for DeepEval integration.
| LOGFIRE_TOKEN | Token for Logfire logging service
| LOGGING_WORKER_CONCURRENCY | Maximum number of concurrent coroutine slots for the logging worker on the asyncio event loop. **Default is 100**. Setting too high will flood the event loop with logging tasks which will lower the overall latency of the requests. |
| LOGGING_WORKER_MAX_QUEUE_SIZE | Maximum size of the logging worker queue. When the queue is full, the worker aggressively clears tasks to make room instead of dropping logs. **Default is 50,000** |
| LOGGING_WORKER_MAX_TIME_PER_COROUTINE | Maximum time in seconds allowed for each coroutine in the logging worker before timing out. **Default is 20.0** |
| LOGGING_WORKER_CLEAR_PERCENTAGE | Percentage of the queue to extract when clearing. **Default is 50** (50%) |
| MAX_EXCEPTION_MESSAGE_LENGTH | Maximum length for exception messages. Default is 2000
| MAX_ITERATIONS_TO_CLEAR_QUEUE | Maximum number of iterations to attempt when clearing the logging worker queue during shutdown. **Default is 200** |
| MAX_TIME_TO_CLEAR_QUEUE | Maximum time in seconds to spend clearing the logging worker queue during shutdown. **Default is 5.0** |
| LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS | Cooldown time in seconds before allowing another aggressive clear operation when the queue is full. **Default is 0.5** |
| MAX_STRING_LENGTH_PROMPT_IN_DB | Maximum length for strings in spend logs when sanitizing request bodies. Strings longer than this will be truncated. Default is 1000
| MAX_IN_MEMORY_QUEUE_FLUSH_COUNT | Maximum count for in-memory queue flush operations. Default is 1000
| MAX_LONG_SIDE_FOR_IMAGE_HIGH_RES | Maximum length for the long side of high-resolution images. Default is 2000

9
litellm/constants.py
View file

@ -253,6 +253,15 @@ REDACTED_BY_LITELM_STRING = "REDACTED_BY_LITELM"
MAX_LANGFUSE_INITIALIZED_CLIENTS = int(
os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)
)
LOGGING_WORKER_CONCURRENCY = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
LOGGING_WORKER_MAX_QUEUE_SIZE = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
LOGGING_WORKER_MAX_TIME_PER_COROUTINE = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
LOGGING_WORKER_CLEAR_PERCENTAGE = int(os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50)) # Percentage of queue to clear (default: 50%)
MAX_ITERATIONS_TO_CLEAR_QUEUE = int(os.getenv("MAX_ITERATIONS_TO_CLEAR_QUEUE", 200))
MAX_TIME_TO_CLEAR_QUEUE = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0))
LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS = float(
os.getenv("LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS", 0.5)
) # Cooldown time in seconds before allowing another aggressive clear (default: 0.5s)
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv(
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
)

344
litellm/litellm_core_utils/logging_worker.py
View file

@ -1,11 +1,22 @@
# This file may be a good candidate to be the first one to be refactored into a separate process,
# for the sake of performance and scalability.
import asyncio
import contextlib
import contextvars
from typing import Coroutine, Optional
import atexit
from typing_extensions import TypedDict
from litellm._logging import verbose_logger
from litellm.constants import (
LOGGING_WORKER_CONCURRENCY,
LOGGING_WORKER_MAX_QUEUE_SIZE,
LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
LOGGING_WORKER_CLEAR_PERCENTAGE,
LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS,
MAX_ITERATIONS_TO_CLEAR_QUEUE,
MAX_TIME_TO_CLEAR_QUEUE,
)
class LoggingTask(TypedDict):
@ -27,21 +38,21 @@ class LoggingWorker:
- Use this to queue coroutine tasks that are not critical to the main flow of the application. e.g Success/Error callbacks, logging, etc.
"""
LOGGING_WORKER_MAX_QUEUE_SIZE = 50_000
LOGGING_WORKER_MAX_TIME_PER_COROUTINE = 20.0
MAX_ITERATIONS_TO_CLEAR_QUEUE = 200
MAX_TIME_TO_CLEAR_QUEUE = 5.0
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,
):
self.timeout = timeout
self.max_queue_size = max_queue_size
self.concurrency = concurrency
self._queue: Optional[asyncio.Queue[LoggingTask]] = None
self._worker_task: Optional[asyncio.Task] = None
self._running_tasks: set[asyncio.Task] = set()
self._sem: Optional[asyncio.Semaphore] = None
self._last_aggressive_clear_time: float = 0.0
self._aggressive_clear_in_progress: bool = False
def _ensure_queue(self) -> None:
"""Initialize the queue if it doesn't exist."""
@ -51,18 +62,15 @@ 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 _worker_loop(self) -> None:
"""Main worker loop that processes log coroutines sequentially."""
async def _process_log_task(self, task: LoggingTask, sem: asyncio.Semaphore):
"""Runs the logging task and handles cleanup. Releases semaphore when done."""
try:
if self._queue is None:
return
while True:
# Process one coroutine at a time to keep event loop load predictable
task = await self._queue.get()
if self._queue is not None:
try:
# Run the coroutine in its original context
await asyncio.wait_for(
@ -71,9 +79,34 @@ class LoggingWorker:
)
except Exception as e:
verbose_logger.exception(f"LoggingWorker error: {e}")
pass
finally:
self._queue.task_done()
finally:
# Always release semaphore, even if queue is None
sem.release()
async def _worker_loop(self) -> None:
"""Main worker loop that gets tasks and schedules them to run concurrently."""
try:
if self._queue is None or self._sem is None:
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
except asyncio.CancelledError:
verbose_logger.debug("LoggingWorker cancelled during shutdown")
@ -83,20 +116,199 @@ class LoggingWorker:
def enqueue(self, coroutine: Coroutine) -> None:
"""
Add a coroutine to the logging queue.
Hot path: never blocks, drops logs if queue is full.
Hot path: never blocks, aggressively clears queue if full.
"""
if self._queue is None:
return
# Capture the current context when enqueueing
task = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
try:
# Capture the current context when enqueueing
task = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
self._queue.put_nowait(task)
except asyncio.QueueFull as e:
verbose_logger.exception(f"LoggingWorker queue is full: {e}")
# Drop logs on overload to protect request throughput
except asyncio.QueueFull:
# Queue is full - handle it appropriately
self._handle_queue_full(task)
def _should_start_aggressive_clear(self) -> bool:
"""
Check if we should start a new aggressive clear operation.
Returns True if cooldown period has passed and no clear is in progress.
"""
if self._aggressive_clear_in_progress:
return False
try:
loop = asyncio.get_running_loop()
current_time = loop.time()
time_since_last_clear = current_time - self._last_aggressive_clear_time
if time_since_last_clear < LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS:
return False
return True
except RuntimeError:
# No event loop running, drop the task
return False
def _mark_aggressive_clear_started(self) -> None:
"""
Mark that an aggressive clear operation has started.
Note: This should only be called after _should_start_aggressive_clear()
returns True, which guarantees an event loop exists.
"""
loop = asyncio.get_running_loop()
self._last_aggressive_clear_time = loop.time()
self._aggressive_clear_in_progress = True
def _handle_queue_full(self, task: LoggingTask) -> None:
"""
Handle queue full condition by either starting an aggressive clear
or scheduling a delayed retry.
"""
if self._should_start_aggressive_clear():
self._mark_aggressive_clear_started()
# Schedule clearing as async task so enqueue returns immediately (non-blocking)
asyncio.create_task(self._aggressively_clear_queue_async(task))
else:
# Cooldown active or clear in progress, schedule a delayed retry
self._schedule_delayed_enqueue_retry(task)
def _calculate_retry_delay(self) -> float:
"""
Calculate the delay before retrying an enqueue operation.
Returns the delay in seconds.
"""
try:
loop = asyncio.get_running_loop()
current_time = loop.time()
time_since_last_clear = current_time - self._last_aggressive_clear_time
remaining_cooldown = max(
0.0,
LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS - time_since_last_clear
)
# Add a small buffer (10% of cooldown or 50ms, whichever is larger) to ensure
# cooldown has expired and aggressive clear has completed
return remaining_cooldown + max(
0.05, LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS * 0.1
)
except RuntimeError:
# No event loop, return minimum delay
return 0.1
def _schedule_delayed_enqueue_retry(self, task: LoggingTask) -> None:
"""
Schedule a delayed retry to enqueue the task after cooldown expires.
This prevents dropping tasks when the queue is full during cooldown.
Preserves the original task context.
"""
try:
# Check that we have a running event loop (will raise RuntimeError if not)
asyncio.get_running_loop()
delay = self._calculate_retry_delay()
# Schedule the retry as a background task
asyncio.create_task(self._retry_enqueue_task(task, delay))
except RuntimeError:
# No event loop, drop the task as we can't schedule a retry
pass
async def _retry_enqueue_task(self, task: LoggingTask, delay: float) -> None:
"""
Retry enqueueing the task after delay, preserving original context.
This is called as a background task from _schedule_delayed_enqueue_retry.
"""
await asyncio.sleep(delay)
# Try to enqueue the task directly, preserving its original context
if self._queue is None:
return
try:
self._queue.put_nowait(task)
except asyncio.QueueFull:
# Still full - handle it appropriately (clear or retry again)
self._handle_queue_full(task)
def _extract_tasks_from_queue(self) -> list[LoggingTask]:
"""
Extract tasks from the queue to make room.
Returns a list of extracted tasks based on percentage of queue size.
"""
if self._queue is None:
return []
# Calculate items based on percentage of queue size
items_to_extract = (self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE) // 100
# Use actual queue size to avoid unnecessary iterations
actual_size = self._queue.qsize()
if actual_size == 0:
return []
items_to_extract = min(items_to_extract, actual_size)
# Extract tasks from queue (using list comprehension would require wrapping in try/except)
extracted_tasks = []
for _ in range(items_to_extract):
try:
extracted_tasks.append(self._queue.get_nowait())
except asyncio.QueueEmpty:
break
return extracted_tasks
async def _aggressively_clear_queue_async(self, new_task: Optional[LoggingTask] = None) -> None:
"""
Aggressively clear the queue by extracting and processing items.
This is called when the queue is full to prevent dropping logs.
Fully async and non-blocking - runs in background task.
"""
try:
if self._queue is None:
return
extracted_tasks = self._extract_tasks_from_queue()
# Add new task to extracted tasks to process directly
if new_task is not None:
extracted_tasks.append(new_task)
# Process extracted tasks directly
if extracted_tasks:
await self._process_extracted_tasks(extracted_tasks)
except Exception as e:
verbose_logger.exception(f"LoggingWorker error during aggressive clear: {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 single task and mark it done."""
if self._queue is None:
return
try:
await asyncio.wait_for(
task["context"].run(asyncio.create_task, task["coroutine"]),
timeout=self.timeout,
)
except Exception:
# Suppress errors during processing to ensure we keep going
pass
finally:
self._queue.task_done()
async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None:
"""
Process tasks that were extracted from the queue to make room.
Processes them concurrently without semaphore limits for maximum speed.
"""
if not tasks or self._queue is None:
return
# 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):
"""
Ensure the logging worker is initialized and enqueue the coroutine.
@ -106,11 +318,25 @@ 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.
return
tasks_to_cancel: list[asyncio.Task] = list(self._running_tasks)
if self._worker_task:
self._worker_task.cancel()
with contextlib.suppress(Exception):
await self._worker_task
self._worker_task = None
# Include the main worker loop so it stops fetching work.
tasks_to_cancel.append(self._worker_task)
for task in tasks_to_cancel:
# Propagate cancellation to every pending task.
task.cancel()
# 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()
async def flush(self) -> None:
"""Flush the logging queue."""
@ -128,14 +354,14 @@ class LoggingWorker:
start_time = asyncio.get_event_loop().time()
for _ in range(self.MAX_ITERATIONS_TO_CLEAR_QUEUE):
for _ in range(MAX_ITERATIONS_TO_CLEAR_QUEUE):
# Check if we've exceeded the maximum time
if (
asyncio.get_event_loop().time() - start_time
>= self.MAX_TIME_TO_CLEAR_QUEUE
>= MAX_TIME_TO_CLEAR_QUEUE
):
verbose_logger.warning(
f"clear_queue exceeded max_time of {self.MAX_TIME_TO_CLEAR_QUEUE}s, stopping early"
f"clear_queue exceeded max_time of {MAX_TIME_TO_CLEAR_QUEUE}s, stopping early"
)
break
@ -154,6 +380,64 @@ class LoggingWorker:
except asyncio.QueueEmpty:
break
def _flush_on_exit(self):
"""
Flush remaining events synchronously before process exit.
Called automatically via atexit handler.
This ensures callbacks queued by async completions are processed
even when the script exits before the worker loop can handle them.
"""
if self._queue is None:
verbose_logger.debug("[LoggingWorker] atexit: No queue initialized")
return
if self._queue.empty():
verbose_logger.debug("[LoggingWorker] atexit: Queue is empty")
return
queue_size = self._queue.qsize()
verbose_logger.info(f"[LoggingWorker] atexit: Flushing {queue_size} remaining events...")
# Create a new event loop since the original is closed
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
# Process remaining queue items with time limit
processed = 0
start_time = loop.time()
while not self._queue.empty() and processed < MAX_ITERATIONS_TO_CLEAR_QUEUE:
if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE:
verbose_logger.warning(
f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush"
)
break
try:
task = self._queue.get_nowait()
except asyncio.QueueEmpty:
break
# Run the coroutine synchronously in new loop
# Note: We run the coroutine directly, not via create_task,
# since we're in a new event loop context
try:
loop.run_until_complete(task["coroutine"])
processed += 1
except Exception as e:
# Silent failure to not break user's program
verbose_logger.debug(f"[LoggingWorker] atexit: Error flushing callback: {e}")
verbose_logger.info(f"[LoggingWorker] atexit: Successfully flushed {processed} events!")
finally:
loop.close()
# Global instance for backward compatibility
GLOBAL_LOGGING_WORKER = LoggingWorker()
# Register flush on exit to ensure remaining events are processed
atexit.register(GLOBAL_LOGGING_WORKER._flush_on_exit)

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

@ -8,6 +8,7 @@ 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
@ -267,3 +268,58 @@ class TestLoggingWorker:
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):
"""Test that the worker respects the semaphore concurrency limit."""
worker = LoggingWorker(timeout=5.0, max_queue_size=20, concurrency=2)
worker.start()
running_tasks, max_concurrent, lock = set(), 0, asyncio.Lock()
completed = asyncio.Event()
async def tracked_task(task_id: int):
async with lock:
running_tasks.add(task_id)
nonlocal max_concurrent
max_concurrent = max(max_concurrent, len(running_tasks))
await asyncio.sleep(0.2)
async with lock:
running_tasks.remove(task_id)
if not running_tasks:
completed.set()
for i in range(5):
worker.enqueue(tracked_task(i))
await asyncio.wait_for(completed.wait(), timeout=5.0)
await worker.stop()
assert max_concurrent <= 2, f"Max {max_concurrent} exceeded limit 2"
assert max_concurrent >= 2, f"Expected 2+ concurrent, got {max_concurrent}"
@pytest.mark.asyncio
async def test_aggressive_queue_clearing(self):
"""Test that aggressive queue clearing processes tasks when queue is full."""
worker = LoggingWorker(timeout=2.0, max_queue_size=4, concurrency=1)
worker.start()
processed, lock = [], asyncio.Lock()
async def tracked_task(task_id: int):
async with lock:
processed.append(task_id)
await asyncio.sleep(0.01)
for i in range(4):
worker.enqueue(tracked_task(i))
await asyncio.sleep(0.1)
for i in range(4, 8):
worker.enqueue(tracked_task(i))
await asyncio.sleep(LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS + 0.3)
await worker.stop()
await worker.clear_queue()
assert len(processed) >= 4, f"Expected 4+ tasks processed, got {len(processed)}"