fix: bound CustomBatchLogger queue and call super().__init__ in ContextCachingEndpoints

Co-authored-by: Yassin Kortam <yassin@berri.ai>

litellm/integrations/custom_batch_logger.py

@@ -16,22 +16,36 @@ from litellm.integrations.custom_logger import CustomLogger
 class CustomBatchLogger(CustomLogger):
     preserve_events_added_during_flush = False
 
+    # Default cap on the in-memory log queue. Prevents unbounded memory growth
+    # if ``async_send_batch`` consistently fails (e.g. the destination is
+    # unreachable) and events are preserved across flush attempts. Subclasses
+    # may override by passing ``max_queue_size`` or by setting the attribute
+    # directly (see ``RubrikLogger`` for an example).
+    DEFAULT_MAX_QUEUE_SIZE = 50_000
+
     def __init__(
         self,
         flush_lock: Optional[asyncio.Lock] = None,
         batch_size: Optional[int] = None,
         flush_interval: Optional[int] = None,
+        max_queue_size: Optional[int] = None,
         **kwargs,
     ) -> None:
         """
         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 number of events to retain in ``log_queue``. When the limit is exceeded (e.g. because the send destination is unreachable and events are preserved for retry), the oldest events are dropped. Defaults to ``DEFAULT_MAX_QUEUE_SIZE``.
         """
         self.log_queue: List = []
         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()
         self.flush_lock = flush_lock
+        self.max_queue_size: int = (
+            max_queue_size
+            if max_queue_size is not None
+            else self.DEFAULT_MAX_QUEUE_SIZE
+        )
 
         super().__init__(**kwargs)
 
@@ -66,6 +80,18 @@ class CustomBatchLogger(CustomLogger):
                         "%s events in queue for retry",
                         log_queue_length,
                     )
+                    # Guard against unbounded queue growth if the destination
+                    # is persistently unreachable. Drop the oldest events
+                    # beyond ``max_queue_size``.
+                    overflow = len(self.log_queue) - self.max_queue_size
+                    if overflow > 0:
+                        del self.log_queue[:overflow]
+                        verbose_logger.warning(
+                            "CustomLogger: log queue exceeded max_queue_size=%s; "
+                            "dropped %s oldest events.",
+                            self.max_queue_size,
+                            overflow,
+                        )
                     return
                 if self.preserve_events_added_during_flush:
                     del self.log_queue[:log_queue_length]

litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py

@@ -41,7 +41,7 @@ class ContextCachingEndpoints(VertexBase):
     """
 
     def __init__(self) -> None:
-        pass
+        super().__init__()
 
     def _get_token_and_url_context_caching(
         self,