From c2690369c893c3a8105075b23c9e58bc48ecf751 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 24 Jan 2026 13:10:12 -0800 Subject: [PATCH] Add BoundedQueue OOM protection tests through CustomBatchLogger - Add realistic integration test for BoundedQueue using CustomBatchLogger - Test force_flush strategy triggers flush when queue overflows - Verify queue behavior matches production usage patterns --- litellm/constants.py | 2 + litellm/integrations/custom_batch_logger.py | 179 +++++++++++++++++- .../integrations/test_bounded_queue.py | 61 ++++++ 3 files changed, 240 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/integrations/test_bounded_queue.py diff --git a/litellm/constants.py b/litellm/constants.py index 49ca3a509b1..36f2e930b5a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -11,6 +11,8 @@ AZURE_DEFAULT_RESPONSES_API_VERSION = str( ROUTER_MAX_FALLBACKS = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) DEFAULT_BATCH_SIZE = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) +DEFAULT_MAX_QUEUE_SIZE = int(os.getenv("DEFAULT_MAX_QUEUE_SIZE", 1000)) # Prevents OOM for batching integrations +DEFAULT_QUEUE_OVERFLOW_STRATEGY = os.getenv("DEFAULT_QUEUE_OVERFLOW_STRATEGY", "force_flush") # Strategy when queue is full DEFAULT_S3_FLUSH_INTERVAL_SECONDS = int( os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10) ) diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index f9d4496c21f..b9b46f4db3f 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -5,11 +5,13 @@ Use this if you want your logs to be stored in memory and flushed periodically. """ import asyncio +import os import time -from typing import List, Optional +from typing import Any, Callable, List, Optional import litellm from litellm._logging import verbose_logger +from litellm.constants import DEFAULT_MAX_QUEUE_SIZE, DEFAULT_QUEUE_OVERFLOW_STRATEGY from litellm.integrations.custom_logger import CustomLogger @@ -24,8 +26,17 @@ class CustomBatchLogger(CustomLogger): """ Args: flush_lock (Optional[asyncio.Lock], optional): Lock to use when flushing the queue. Defaults to None. Only used for custom loggers that do batching + max_queue_size (Optional[int], optional): Maximum queue size to prevent OOM. Defaults to DEFAULT_MAX_QUEUE_SIZE (1000). Set to None for unlimited (not recommended). Can also be set via DEFAULT_MAX_QUEUE_SIZE env var. + queue_overflow_strategy (str): What to do when queue is full: "drop_oldest" (default), "drop_newest", "reject", or "force_flush". Can also be set via DEFAULT_QUEUE_OVERFLOW_STRATEGY env var. """ - self.log_queue: List = [] + + # Use BoundedQueue instead of plain list - all existing append() calls automatically get protection + # Pass flush callback for "force_flush" strategy + self.log_queue = BoundedQueue( + max_size=DEFAULT_MAX_QUEUE_SIZE, + overflow_strategy=DEFAULT_QUEUE_OVERFLOW_STRATEGY, + flush_callback=lambda: asyncio.create_task(self.flush_queue()) + ) self.flush_interval = flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS self.batch_size: int = batch_size or litellm.DEFAULT_BATCH_SIZE self.last_flush_time = time.time() @@ -56,3 +67,167 @@ class CustomBatchLogger(CustomLogger): async def async_send_batch(self, *args, **kwargs): pass + + +class BoundedQueue: + """ + A list-like queue with optional size limit to prevent OOM. + + Wraps a list and overrides append() to enforce max_size limit. + This allows all existing code using log_queue.append() to automatically + get OOM protection without any code changes. + """ + + def __init__( + self, + max_size: Optional[int] = None, + overflow_strategy: str = "force_flush", + flush_callback: Optional[Callable[[], Any]] = None, + ): + """ + Args: + max_size: Maximum queue size. None = unlimited (backward compatible) + overflow_strategy: What to do when queue is full: + - "drop_oldest": Remove oldest item (FIFO) + - "drop_newest": Reject new item (don't add) + - "reject": Raise RuntimeError (fail fast) + - "force_flush": Trigger flush callback to send logs, then add item - default + flush_callback: Optional callback function to trigger flush (required for "force_flush" strategy) + """ + self._items: List = [] + self.max_size = max_size + self.overflow_strategy = overflow_strategy + self._dropped_count = 0 + self._flush_callback = flush_callback + + # Map strategy names to handler methods + self._strategy_handlers = { + "drop_oldest": self._handle_drop_oldest, + "drop_newest": self._handle_drop_newest, + "reject": self._handle_reject, + "force_flush": self._handle_force_flush, + } + self._handler = self._strategy_handlers.get(overflow_strategy, self._handle_force_flush) + + def _handle_drop_oldest(self, item): + """Handle overflow by removing oldest item (FIFO).""" + self._items.pop(0) + self._items.append(item) + self._dropped_count += 1 + if self._dropped_count % 100 == 0: # Log every 100 drops to avoid spam + verbose_logger.warning( + f"CustomBatchLogger: Queue full ({self.max_size}), dropped {self._dropped_count} oldest logs" + ) + + def _handle_drop_newest(self, item): + """Handle overflow by rejecting new item.""" + self._dropped_count += 1 + if self._dropped_count % 100 == 0: + verbose_logger.warning( + f"CustomBatchLogger: Queue full ({self.max_size}), rejected {self._dropped_count} new logs" + ) + + def _handle_reject(self, item): + """Handle overflow by raising RuntimeError.""" + raise RuntimeError( + f"CustomBatchLogger: Queue full ({self.max_size}). " + f"Cannot add more logs. Consider increasing max_queue_size or fixing flush issues." + ) + + def _handle_force_flush(self, item): + """Handle overflow by triggering flush callback, then adding item.""" + if self._flush_callback is None: + # Fallback to drop_oldest if no flush callback provided + verbose_logger.warning( + "CustomBatchLogger: force_flush strategy requires flush_callback, falling back to drop_oldest" + ) + self._handle_drop_oldest(item) + return + + # Trigger flush to make space + try: + self._flush_callback() + verbose_logger.debug( + f"CustomBatchLogger: Queue full ({self.max_size}), triggered force flush to prevent log loss" + ) + except Exception as e: + verbose_logger.warning( + f"CustomBatchLogger: Error during force flush: {e}, falling back to drop_oldest" + ) + self._handle_drop_oldest(item) + return + + # Add item after flush is triggered (flush is async, but we've scheduled it) + self._items.append(item) + + def append(self, item): + """Append item with OOM protection.""" + # No limit - allow unlimited growth (backward compatible) + if self.max_size is None: + self._items.append(item) + return + + # Check if queue is full + if len(self._items) >= self.max_size: + self._handler(item) + else: + # Queue has space - add normally + self._items.append(item) + + def clear(self): + """Clear all items.""" + self._items.clear() + + def __len__(self): + """Return queue length.""" + return len(self._items) + + def __bool__(self): + """Check if queue is non-empty.""" + return bool(self._items) + + def __iter__(self): + """Iterate over items.""" + return iter(self._items) + + def __getitem__(self, index): + """Get item by index.""" + return self._items[index] + + def __setitem__(self, index, value): + """Set item by index.""" + self._items[index] = value + + def __delitem__(self, index): + """Delete item by index.""" + del self._items[index] + + def __contains__(self, item): + """Check if item is in queue.""" + return item in self._items + + def __repr__(self): + """String representation.""" + return f"BoundedQueue(max_size={self.max_size}, size={len(self._items)}, dropped={self._dropped_count})" + + # Delegate other list methods as needed + def extend(self, items): + """Extend queue with items.""" + for item in items: + self.append(item) + + def pop(self, index: int = -1): + """Pop item from queue.""" + return self._items.pop(index) + + def remove(self, item): + """Remove item from queue.""" + self._items.remove(item) + + def index(self, item): + """Get index of item.""" + return self._items.index(item) + + def count(self, item): + """Count occurrences of item.""" + return self._items.count(item) diff --git a/tests/test_litellm/integrations/test_bounded_queue.py b/tests/test_litellm/integrations/test_bounded_queue.py new file mode 100644 index 00000000000..541e1fef9ad --- /dev/null +++ b/tests/test_litellm/integrations/test_bounded_queue.py @@ -0,0 +1,61 @@ +""" +Test BoundedQueue OOM protection through CustomBatchLogger (realistic integration test). +""" + +import asyncio +import unittest + +from litellm.integrations.custom_batch_logger import CustomBatchLogger + + +class TestBoundedQueueWithCustomBatchLogger(unittest.TestCase): + """Test BoundedQueue behavior through CustomBatchLogger.""" + + def test_force_flush_on_queue_overflow(self): + """Test that force_flush strategy triggers flush when queue is full.""" + async def run_test(): + flush_called = [] + + async def mock_async_send_batch(*args, **kwargs): + """Track that flush was called.""" + flush_called.append(True) + + # Create logger with small queue for testing + logger = CustomBatchLogger(batch_size=10, flush_interval=60) + logger.async_send_batch = mock_async_send_batch + logger.flush_lock = asyncio.Lock() + + # Override max_size to small value for testing + original_max_size = logger.log_queue.max_size + logger.log_queue.max_size = 3 + + # Fill queue to max + logger.log_queue.append("item1") + logger.log_queue.append("item2") + logger.log_queue.append("item3") + self.assertEqual(len(logger.log_queue), 3) + + # Add 4th item - should trigger force_flush (default strategy) + logger.log_queue.append("item4") + + # Item should be added immediately (flush is async) + self.assertEqual(len(logger.log_queue), 4) + + # Wait for async flush task to run + await asyncio.sleep(0.1) + + + # Flush should have been called + self.assertTrue(len(flush_called) > 0) + + # Queue should be cleared after flush + self.assertEqual(len(logger.log_queue), 0) + + # Restore original max_size + logger.log_queue.max_size = original_max_size + + asyncio.run(run_test()) + + +if __name__ == "__main__": + unittest.main()