fix(guardrails): retry failed daily metrics and usage unit upserts with backoff

A transient DB error during the spend log flush dropped that batch's guardrail
metrics and usage unit rows for good. Retry only the rows that failed, up to 3
times with 1s/2s/4s backoff, mirroring the daily spend writer, and inject the
sleep so tests stay fast. Lowers the lint budgets the refactor freed up

basedpyright-code-budget.json

@@ -57,7 +57,7 @@
     "limit": 5681
   },
   "reportMissingTypeArgument": {
-    "limit": 15609
+    "limit": 15608
   },
   "reportMissingTypeStubs": {
     "limit": 40

litellm/proxy/guardrails/usage_tracking.py

@@ -3,14 +3,16 @@ Track guardrail and policy usage for the dashboard: upsert daily metrics and
 insert into SpendLogGuardrailIndex when spend logs are written.
 """
 
+import asyncio
 import json
 from collections import defaultdict
-from collections.abc import Iterator, Mapping, Sequence
+from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
 from datetime import datetime, timezone
+from functools import partial
 from itertools import groupby
 from operator import itemgetter
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, NamedTuple
+from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy.utils import PrismaClient
@@ -24,6 +26,12 @@ if TYPE_CHECKING:
     from prisma import types as prisma_types
 
 
+_UPSERT_RETRY_TIMES: Final = 3
+
+_RowKey = TypeVar("_RowKey")
+_RowValue = TypeVar("_RowValue")
+
+
 class _UsageUnitKey(NamedTuple):
     guardrail_id: str
     date: str
@@ -32,6 +40,46 @@ class _UsageUnitKey(NamedTuple):
     usage_unit: str
 
 
+class _MetricsKey(NamedTuple):
+    guardrail_id: str
+    date: str
+
+
+async def _attempt_upsert(
+    upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]], key: _RowKey, value: _RowValue
+) -> Exception | None:
+    try:
+        await upsert_row(key, value)
+    except Exception as error:
+        return error
+    return None
+
+
+async def _upsert_rows_with_retry(
+    rows: Mapping[_RowKey, _RowValue],
+    upsert_row: Callable[[_RowKey, _RowValue], Awaitable[None]],
+    label: str,
+    sleep: Callable[[float], Awaitable[None]],
+    retries_left: int = _UPSERT_RETRY_TIMES,
+) -> None:
+    outcomes: Final = {key: await _attempt_upsert(upsert_row, key, value) for key, value in rows.items()}
+    failed: Final = MappingProxyType({key: rows[key] for key, error in outcomes.items() if error is not None})
+    if not failed:
+        return
+    if retries_left == 0:
+        for key in failed:
+            verbose_proxy_logger.warning(
+                "Guardrail usage tracking: %s upsert failed for %s after %d retries (non-fatal): %s",
+                label,
+                key,
+                _UPSERT_RETRY_TIMES,
+                outcomes[key],
+            )
+        return
+    await sleep(2 ** (_UPSERT_RETRY_TIMES - retries_left))
+    await _upsert_rows_with_retry(failed, upsert_row, label, sleep, retries_left - 1)
+
+
 def _guardrail_status_to_action(status: str | None) -> str:
     """Map StandardLogging guardrail_status to blocked/passed/flagged."""
     if not status:
@@ -131,9 +179,33 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey
     await DailyGuardrailUsageUnitsRepository(prisma_client).table.upsert(where=where, data=data)
 
 
+async def _upsert_metrics_row(prisma_client: PrismaClient, key: _MetricsKey, agg: Mapping[str, int]) -> None:
+    n: Final = int(agg["requests_evaluated"])
+    await DailyGuardrailMetricsRepository(prisma_client).table.upsert(
+        where={"guardrail_id_date": {"guardrail_id": key.guardrail_id, "date": key.date}},
+        data={
+            "create": {
+                "guardrail_id": key.guardrail_id,
+                "date": key.date,
+                "requests_evaluated": n,
+                "passed_count": int(agg["passed_count"]),
+                "blocked_count": int(agg["blocked_count"]),
+                "flagged_count": int(agg["flagged_count"]),
+            },
+            "update": {
+                "requests_evaluated": {"increment": n},
+                "passed_count": {"increment": int(agg["passed_count"])},
+                "blocked_count": {"increment": int(agg["blocked_count"])},
+                "flagged_count": {"increment": int(agg["flagged_count"])},
+            },
+        },
+    )
+
+
 async def process_spend_logs_guardrail_usage(
     prisma_client: PrismaClient,
     logs_to_process: list[dict[str, Any]],
+    sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
 ) -> None:
     """
     After spend logs are written: update DailyGuardrailMetrics and insert
@@ -142,7 +214,7 @@ async def process_spend_logs_guardrail_usage(
     if not logs_to_process:
         return
     # Aggregate daily metrics by (guardrail_id, date). Latency/score metrics dropped.
-    daily_guardrail: Final[dict[tuple, dict[str, Any]]] = defaultdict(
+    daily_guardrail: Final[dict[_MetricsKey, dict[str, Any]]] = defaultdict(
         lambda: {
             "requests_evaluated": 0,
             "passed_count": 0,
@@ -163,7 +235,7 @@ async def process_spend_logs_guardrail_usage(
             guardrail_id = entry.get("guardrail_id") or entry.get("guardrail_name") or ""
             if not guardrail_id:
                 continue
-            key = (guardrail_id, date_key)
+            key = _MetricsKey(guardrail_id, date_key)
             daily_guardrail[key]["requests_evaluated"] += 1
             action = _guardrail_status_to_action(entry.get("guardrail_status"))
             if action == "passed":
@@ -199,51 +271,12 @@ async def process_spend_logs_guardrail_usage(
                 verbose_proxy_logger.debug("Guardrail usage tracking: index create_many skipped: %s", e)
 
         # Upsert daily guardrail metrics (counts only; latency/score dropped)
-        for (guardrail_id, date_key), agg in daily_guardrail.items():
-            n = int(agg["requests_evaluated"])
-            if n == 0:
-                continue
-            try:
-                await DailyGuardrailMetricsRepository(prisma_client).table.upsert(
-                    where={
-                        "guardrail_id_date": {
-                            "guardrail_id": guardrail_id,
-                            "date": date_key,
-                        }
-                    },
-                    data={
-                        "create": {
-                            "guardrail_id": guardrail_id,
-                            "date": date_key,
-                            "requests_evaluated": n,
-                            "passed_count": int(agg["passed_count"]),
-                            "blocked_count": int(agg["blocked_count"]),
-                            "flagged_count": int(agg["flagged_count"]),
-                        },
-                        "update": {
-                            "requests_evaluated": {"increment": n},
-                            "passed_count": {"increment": int(agg["passed_count"])},
-                            "blocked_count": {"increment": int(agg["blocked_count"])},
-                            "flagged_count": {"increment": int(agg["flagged_count"])},
-                        },
-                    },
-                )
-            except Exception as metrics_error:
-                verbose_proxy_logger.warning(
-                    "Guardrail usage tracking: daily metrics upsert failed for %s on %s (non-fatal): %s",
-                    guardrail_id,
-                    date_key,
-                    metrics_error,
-                )
-
-        for unit_key, units in usage_unit_totals.items():
-            try:
-                await _upsert_usage_unit_row(prisma_client, unit_key, units)
-            except Exception as unit_error:
-                verbose_proxy_logger.warning(
-                    "Guardrail usage tracking: usage unit upsert failed for %s (non-fatal): %s",
-                    unit_key,
-                    unit_error,
-                )
+        metrics_rows: Final = MappingProxyType(
+            {key: agg for key, agg in daily_guardrail.items() if int(agg["requests_evaluated"]) > 0}
+        )
+        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
+        )
     except Exception as e:
         verbose_proxy_logger.warning("Guardrail usage tracking failed (non-fatal): %s", e)

ruff-strict-budget.json

@@ -78,7 +78,7 @@
     "limit": 1
   },
   "C901": {
-    "limit": 313
+    "limit": 312
   },
   "D419": {
     "limit": 6

tests/test_litellm/proxy/guardrails/test_usage_tracking.py

@@ -80,24 +80,70 @@ async def test_usage_units_rolled_up_by_guardrail_team_key_and_date():
     }
 
 
+def _fake_sleep() -> tuple[AsyncMock, list[float]]:
+    delays: list[float] = []
+    sleep = AsyncMock(side_effect=lambda delay: delays.append(delay))
+    return sleep, delays
+
+
 @pytest.mark.asyncio
 async def test_one_failing_upsert_does_not_drop_remaining_writes():
     """
     A DB error on one daily-metrics or usage-unit upsert must not cancel the
     remaining upserts in the flushed batch, or the usage endpoints would
-    permanently under-report billable counters (batches are never retried).
+    permanently under-report billable counters.
     """
     prisma = _prisma()
     prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down")
-    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None]
+    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("db down"), None, None]
+    sleep, _ = _fake_sleep()
     logs = [
         _payload("r1", usage={"topicPolicyUnits": 1}),
         _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
     ]
 
-    await process_spend_logs_guardrail_usage(prisma, logs)
+    await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
 
-    assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 2
+    assert _units_upserts(prisma) == {
+        ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
+        ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1,
+    }
+
+
+@pytest.mark.asyncio
+async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only():
+    """
+    A transient DB error must not permanently drop billed units from the
+    aggregates: only the rows that failed are re-sent, after exponential
+    backoff, and the batch ends once every row has landed.
+    """
+    prisma = _prisma()
+    prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [RuntimeError("blip"), None, None]
+    sleep, delays = _fake_sleep()
+    logs = [
+        _payload("r1", usage={"topicPolicyUnits": 1}),
+        _payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
+    ]
+
+    await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
+
+    calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
+    assert len(calls) == 3
+    assert calls[2].kwargs["where"] == calls[0].kwargs["where"]
+    assert delays == [1]
+
+
+@pytest.mark.asyncio
+async def test_persistent_upsert_failure_stops_after_three_retries():
+    prisma = _prisma()
+    prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("db down")
+    sleep, delays = _fake_sleep()
+
+    await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep)
+
+    assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4
+    assert delays == [1, 2, 4]
+    assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1
 
 
 @pytest.mark.asyncio

type-discipline-budget.json

@@ -1,9 +1,9 @@
 {
   "LIT001": {
-    "limit": 22903
+    "limit": 22900
   },
   "LIT002": {
-    "limit": 26894
+    "limit": 26893
   },
   "LIT003": {
     "limit": 269