Merge pull request #37387 from BerriAI/litellm_guardrail_usage_requeue

fix(guardrails): requeue usage rollup rows dropped after retry exhaustion

litellm/proxy/guardrails/usage_tracking.py

@@ -28,6 +28,7 @@ if TYPE_CHECKING:
 
 
 _UPSERT_RETRY_TIMES: Final = 3
+_MAX_PENDING_ROWS: Final = 10_000
 
 _RowKey = TypeVar("_RowKey")
 _RowValue = TypeVar("_RowValue")
@@ -46,6 +47,58 @@ class _MetricsKey(NamedTuple):
     date: str
 
 
+class PendingRollups:
+    """Rollup rows whose connection-error retries exhausted, held for the next flush."""
+
+    def __init__(self) -> None:
+        self.lock: Final = asyncio.Lock()
+        self.metrics: Mapping[_MetricsKey, Mapping[str, int]] = MappingProxyType({})
+        self.units: Mapping[_UsageUnitKey, int] = MappingProxyType({})
+
+
+_PENDING_ROLLUPS: Final = PendingRollups()
+
+_NO_COUNTERS: Final[Mapping[str, int]] = MappingProxyType({})
+
+
+def _merged_keys(base: Mapping[_RowKey, object], extra: Mapping[_RowKey, object]) -> tuple[_RowKey, ...]:
+    return (*base, *(key for key in extra if key not in base))
+
+
+def _merged_unit_rows(
+    base: Mapping[_UsageUnitKey, int], extra: Mapping[_UsageUnitKey, int]
+) -> Mapping[_UsageUnitKey, int]:
+    return MappingProxyType({key: base.get(key, 0) + extra.get(key, 0) for key in _merged_keys(base, extra)})
+
+
+def _merged_metric_rows(
+    base: Mapping[_MetricsKey, Mapping[str, int]], extra: Mapping[_MetricsKey, Mapping[str, int]]
+) -> Mapping[_MetricsKey, Mapping[str, int]]:
+    def merged_counters(key: _MetricsKey) -> Mapping[str, int]:
+        base_counters: Final = base.get(key, _NO_COUNTERS)
+        extra_counters: Final = extra.get(key, _NO_COUNTERS)
+        return MappingProxyType(
+            {
+                counter: int(base_counters.get(counter, 0)) + int(extra_counters.get(counter, 0))
+                for counter in _merged_keys(base_counters, extra_counters)
+            }
+        )
+
+    return MappingProxyType({key: merged_counters(key) for key in _merged_keys(base, extra)})
+
+
+def _capped(rows: Mapping[_RowKey, _RowValue], label: str) -> Mapping[_RowKey, _RowValue]:
+    if len(rows) <= _MAX_PENDING_ROWS:
+        return rows
+    verbose_proxy_logger.warning(
+        "Guardrail usage tracking: pending %s requeue exceeds %d rows; dropping the %d oldest (non-fatal)",
+        label,
+        _MAX_PENDING_ROWS,
+        len(rows) - _MAX_PENDING_ROWS,
+    )
+    return MappingProxyType(dict(tuple(rows.items())[len(rows) - _MAX_PENDING_ROWS :]))
+
+
 async def _attempt_upsert(
     upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], key: _RowKey, value: _RowValue
 ) -> Exception | None:
@@ -62,7 +115,8 @@ async def _upsert_rows_with_retry(
     label: str,
     sleep: Callable[[float], Awaitable[None]],
     retries_left: int = _UPSERT_RETRY_TIMES,
-) -> None:
+) -> Mapping[_RowKey, _RowValue]:
+    """Returns the rows still failing with connection errors once retries exhaust, for requeueing."""
     outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()}
     for key, error in outcomes.items():
         if error is not None and not isinstance(error, DB_RETRY_SAFE_ERROR_TYPES):
@@ -76,19 +130,20 @@ async def _upsert_rows_with_retry(
         {key: rows[key] for key, error in outcomes.items() if isinstance(error, DB_RETRY_SAFE_ERROR_TYPES)}
     )
     if not retryable:
-        return
+        return MappingProxyType({})
     if retries_left == 0:
         for key in retryable:
             verbose_proxy_logger.warning(
-                "Guardrail usage tracking: %s upsert failed for %s after %d retries (non-fatal): %s",
+                "Guardrail usage tracking: %s upsert failed for %s after %d retries; requeued for the next flush "
+                "(non-fatal): %s",
                 label,
                 key,
                 _UPSERT_RETRY_TIMES,
                 outcomes[key],
             )
-        return
+        return retryable
     await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left))
-    await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1)
+    return await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1)
 
 
 def _guardrail_status_to_action(status: str | None) -> str:
@@ -217,6 +272,7 @@ async def process_spend_logs_guardrail_usage(
     prisma_client: PrismaClient,
     logs_to_process: list[dict[str, Any]],
     sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    pending: PendingRollups = _PENDING_ROLLUPS,
 ) -> None:
     """
     After spend logs are written: update DailyGuardrailMetrics and insert
@@ -265,9 +321,20 @@ async def process_spend_logs_guardrail_usage(
                 }
             )
 
-    usage_unit_totals: Final = _sum_usage_unit_increments(logs_to_process)
+    async with pending.lock:
+        pending_metrics: Final = pending.metrics
+        pending_units: Final = pending.units
+        pending.metrics = MappingProxyType({})
+        pending.units = MappingProxyType({})
 
-    if not daily_guardrail and not index_rows and not usage_unit_totals:
+    # Upsert daily guardrail metrics (counts only; latency/score dropped)
+    evaluated_metrics: Final = MappingProxyType(
+        {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0}
+    )
+    metrics_rows: Final = _merged_metric_rows(pending_metrics, evaluated_metrics)
+    unit_rows: Final = _merged_unit_rows(pending_units, _sum_usage_unit_increments(logs_to_process))
+
+    if not metrics_rows and not index_rows and not unit_rows:
         return
 
     try:
@@ -281,13 +348,15 @@ async def process_spend_logs_guardrail_usage(
             except Exception as e:
                 verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e)
 
-        # Upsert daily guardrail metrics (counts only; latency/score dropped)
-        metrics_rows: Final = MappingProxyType(
-            {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0}
+        failed_metrics: Final = await _upsert_rows_with_retry(
+            metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep
         )
-        await _upsert_rows_with_retry(metrics_rows, partial(_upsert_metrics_row, prisma_client), "daily metrics", sleep)
-        await _upsert_rows_with_retry(
-            usage_unit_totals, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep
+        failed_units: Final = await _upsert_rows_with_retry(
+            unit_rows, partial(_upsert_usage_unit_row, prisma_client), "usage unit", sleep
         )
+        if failed_metrics or failed_units:
+            async with pending.lock:
+                pending.metrics = _capped(_merged_metric_rows(pending.metrics, failed_metrics), "daily metrics")
+                pending.units = _capped(_merged_unit_rows(pending.units, failed_units), "usage unit")
     except Exception as e:
         verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e)

tests/test_litellm/proxy/guardrails/test_usage_tracking.py

@@ -6,7 +6,12 @@ from unittest.mock import AsyncMock, MagicMock
 import httpx
 import pytest
 
-from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage
+from litellm.proxy.guardrails.usage_tracking import (
+    _MAX_PENDING_ROWS,
+    PendingRollups,
+    _capped,
+    process_spend_logs_guardrail_usage,
+)
 
 
 def _prisma() -> MagicMock:
@@ -103,7 +108,7 @@ async def test_one_failing_upsert_does_not_drop_remaining_writes():
         _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
     ]
 
-    await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
+    await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep, pending=PendingRollups())
 
     assert _units_upserts(prisma) == {
         ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
@@ -140,12 +145,86 @@ async def test_persistent_upsert_failure_stops_after_three_retries():
     prisma = _prisma()
     prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
     sleep, delays = _fake_sleep()
+    pending = PendingRollups()
 
-    await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep)
+    await process_spend_logs_guardrail_usage(
+        prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep, pending=pending
+    )
 
     assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4
     assert delays == [1, 2, 4]
     assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1
+    assert dict(pending.metrics) == {
+        ("bedrock-guard", "2026-08-17"): {
+            "requests_evaluated": 1,
+            "passed_count": 1,
+            "blocked_count": 0,
+            "flagged_count": 0,
+        }
+    }
+
+
+@pytest.mark.asyncio
+async def test_retry_exhausted_rows_are_requeued_and_land_on_the_next_flush():
+    """
+    LIT-5761: rollup rows whose connection-error retries exhaust must not be
+    silently lost. They are requeued and merged into the next flushed batch,
+    so the aggregates catch up once the database is reachable again.
+    """
+    pending = PendingRollups()
+    down = _prisma()
+    down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
+    down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down")
+    sleep, _ = _fake_sleep()
+
+    await process_spend_logs_guardrail_usage(
+        down, [_payload("r1", usage={"topicPolicyUnits": 2})], sleep=sleep, pending=pending
+    )
+
+    assert dict(pending.units) == {("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 2}
+
+    recovered = _prisma()
+    await process_spend_logs_guardrail_usage(
+        recovered, [_payload("r2", usage={"topicPolicyUnits": 3})], sleep=sleep, pending=pending
+    )
+
+    assert _units_upserts(recovered) == {
+        ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 5,
+    }
+    metrics_create = recovered.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]
+    assert metrics_create["requests_evaluated"] == 2
+    assert not pending.units
+    assert not pending.metrics
+
+
+@pytest.mark.asyncio
+async def test_ambiguous_failures_are_never_requeued():
+    """
+    A post-send failure (the increment may have committed) must stay dropped:
+    requeueing it would re-send a possibly applied increment and double-count.
+    """
+    pending = PendingRollups()
+    prisma = _prisma()
+    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ReadTimeout("maybe committed")
+    sleep, delays = _fake_sleep()
+
+    await process_spend_logs_guardrail_usage(
+        prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep, pending=pending
+    )
+
+    assert delays == []
+    assert not pending.units
+    assert not pending.metrics
+
+
+def test_pending_requeue_is_capped_dropping_oldest_rows():
+    rows = {index: index for index in range(_MAX_PENDING_ROWS + 5)}
+
+    capped = _capped(rows, "usage unit")
+
+    assert len(capped) == _MAX_PENDING_ROWS
+    assert 4 not in capped
+    assert _MAX_PENDING_ROWS + 4 in capped
 
 
 def _units_upsert_wheres(prisma: MagicMock) -> list[tuple]: