From 5910750ed6d9da402d1d57b9505909fb362ade5c Mon Sep 17 00:00:00 2001 From: AlexsanderHamir Date: Wed, 12 Nov 2025 17:15:54 -0800 Subject: [PATCH] fix: prevent memory blowout in LoggingWorker 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. --- docs/my-website/docs/proxy/config_settings.md | 7 + litellm/constants.py | 9 + litellm/litellm_core_utils/logging_worker.py | 344 ++++++++++++++++-- .../litellm_core_utils/test_logging_worker.py | 56 +++ 4 files changed, 386 insertions(+), 30 deletions(-) diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index df2b7fc803c..981b09f6024 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -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 diff --git a/litellm/constants.py b/litellm/constants.py index 64e92e382f8..f82c007f945 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" ) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 3c475f133a8..5f1bb5f51fa 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -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) 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 a35af315322..0fe15168467 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -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)}"