mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(rubrik): override flush_queue so a single snapshot drives send and drain
Previously RubrikLogger relied on CustomBatchLogger.flush_queue, which captured len(self.log_queue) separately from the snapshot taken inside async_send_batch. Although both happen without an intervening await today (so they agree in practice), they are semantically disconnected: a future refactor that adds an await between the two captures, or that changes the async_send_batch contract, could cause the parent to delete a different number of items than were actually sent and trigger duplicate deliveries to Rubrik. Override flush_queue on RubrikLogger so a single snapshot drives both the HTTP POST and the queue truncation. async_send_batch is preserved for direct callers/tests but no longer participates in the canonical flush path. Existing tests (including the one that explicitly invokes the base CustomBatchLogger.flush_queue path) still pass. Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
c60b7b0a33
commit
f3635c86bf
1 changed files with 34 additions and 1 deletions
|
|
@ -387,7 +387,13 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
raise
|
||||
|
||||
async def async_send_batch(self):
|
||||
"""Handles sending batches of responses to Rubrik."""
|
||||
"""Handles sending batches of responses to Rubrik.
|
||||
|
||||
Note: the canonical flush path is :meth:`flush_queue`, which takes a
|
||||
single snapshot used for both sending and queue draining. This method
|
||||
is kept for direct callers / tests; it intentionally does NOT remove
|
||||
events from the queue.
|
||||
"""
|
||||
if not self.log_queue:
|
||||
return
|
||||
|
||||
|
|
@ -399,6 +405,33 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
data=log_queue_snapshot,
|
||||
)
|
||||
|
||||
async def flush_queue(self):
|
||||
"""Snapshot, send, and drain in one consistent step.
|
||||
|
||||
Overrides the base implementation so the same snapshot drives both
|
||||
the HTTP send and the queue truncation. This avoids the subtle
|
||||
coupling where the base class captures `len(self.log_queue)`
|
||||
separately from the snapshot taken inside `async_send_batch`,
|
||||
which could otherwise drift in a future refactor and cause
|
||||
duplicate deliveries to Rubrik.
|
||||
"""
|
||||
if self.flush_lock is None:
|
||||
return
|
||||
|
||||
async with self.flush_lock:
|
||||
if not self.log_queue:
|
||||
return
|
||||
snapshot = list(self.log_queue)
|
||||
verbose_logger.debug("Rubrik: Flushing batch of %s events", len(snapshot))
|
||||
try:
|
||||
await self._log_batch_to_rubrik(data=snapshot)
|
||||
except Exception:
|
||||
# Already logged with traceback inside _log_batch_to_rubrik.
|
||||
# Preserve the in-flight events for retry on the next flush.
|
||||
return
|
||||
del self.log_queue[: len(snapshot)]
|
||||
self.last_flush_time = time.time()
|
||||
|
||||
# -- Tool blocking service -------------------------------------------------
|
||||
|
||||
async def _post_to_tool_blocking_service(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue