diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index c135c7d1d9c..d50feb6f16b 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -161,6 +161,24 @@ async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThroug await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) +async def _stored_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: + """The marker as another pod may have just written it, bypassing this pod's config cache.""" + param: Final = await ConfigRepository(prisma_client).get_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) + return None if param is None else _marker_from_param_value(param.param_value) + + +async def _record_advanced(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None: + """Advance the stored marker by ``days``. Two runs can overlap (Redis unreachable, lock expired + on a long backfill), so the base is what is stored now, not the snapshot this run scanned from: + a slower run may then only add to the faster run's marker, never rewind it. Without a new scan + time the stored one is kept.""" + stored: Final = await _stored_marker(prisma_client) + kept_scanned_at: Final = None if stored is None else stored.scanned_at + await _record_marker( + prisma_client, _advanced(stored, days, scanned_at=scanned_at if scanned_at is not None else kept_scanned_at) + ) + + async def _db_now(prisma_client: "PrismaClient") -> _NowRow: rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) return _NowRow.model_validate(rows[0]) @@ -203,7 +221,7 @@ async def run_daily_global_spend_reconcile(prisma_client: "PrismaClient") -> Rec marker: Final = await reconciled_through(prisma_client) return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=scan.days[len(done)]) if scan.marker is not None or done: - await _record_marker(prisma_client, _advanced(scan.marker, done, scanned_at=scan.scanned_at)) + await _record_advanced(prisma_client, done, scanned_at=scan.scanned_at) return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) @@ -215,21 +233,16 @@ def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanne async def _reconcile_until_failure(prisma_client: "PrismaClient", scan: _PendingScan) -> tuple[str, ...]: for index, day in enumerate(scan.days): - if not await _reconcile_and_record(prisma_client, scan.marker, scan.days[: index + 1]): + if not await _reconcile_and_record(prisma_client, scan.days[: index + 1]): return scan.days[:index] return scan.days -async def _reconcile_and_record( - prisma_client: "PrismaClient", marker: ReconciledThrough | None, done_with_this: tuple[str, ...] -) -> bool: +async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: tuple[str, ...]) -> bool: day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_marker( - prisma_client, - _advanced(marker, done_with_this, scanned_at=None if marker is None else marker.scanned_at), - ) + await _record_advanced(prisma_client, done_with_this, scanned_at=None) except Exception as exc: # noqa: BLE001 # one bad day must not lose the days already done verbose_proxy_logger.exception("Daily global spend reconcile: day %s failed: %s", day, exc) return False diff --git a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py index 9655953134a..69f06fad081 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py @@ -40,6 +40,10 @@ class _FakeConfigTable: self.rows[where["param_name"]] = data["update"]["param_value"] return _FakeConfigRow(where["param_name"], data["update"]["param_value"]) + async def find_unique(self, *, where: dict[str, str]) -> _FakeConfigRow | None: + stored = self.rows.get(where["param_name"]) + return None if stored is None else _FakeConfigRow(where["param_name"], stored) + class _FakeDb: """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, @@ -68,11 +72,15 @@ class _FakeDb: if day in self._prisma.failing_days: raise RuntimeError(f"day {day} exploded") self._prisma.reconciled.append(day) + landing = self._prisma.marker_landing_on_day.get(day) + if landing is not None: + self.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = landing return 1 class _FakePrisma: - """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" + """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw. + ``marker_landing_on_day`` stores another pod's marker the moment this run rewrites that day.""" def __init__( self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset(), today: date = TODAY @@ -81,6 +89,7 @@ class _FakePrisma: self.today = today self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days + self.marker_landing_on_day: dict[str, str] = {} self.reconciled: list[str] = [] self.db = _FakeDb(self) @@ -221,6 +230,25 @@ async def test_the_next_run_resumes_from_the_failed_day(): assert await reconciled_through(prisma) == "2026-09-03" +@pytest.mark.asyncio +async def test_a_slower_overlapping_run_never_rewinds_the_marker_a_faster_run_stored(): + """Two pods can reconcile at once (Redis unreachable, or the lock expired on a long backfill). + When the faster one has already stored a later marker, the slower one may only add to it. Putting + its own older prefix back, or dropping the scan time, would send usage reads for every day in + between back to the per-key table until the next run.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-02", "2026-09-03"), failing_days=frozenset({"2026-09-03"})) + prisma.marker_landing_on_day = { + "2026-09-02": '{"reconciled_through": "2026-09-14", "scanned_at": "clock-0009"}', + } + + result = await run_daily_global_spend_reconcile(prisma) + + assert result.days_reconciled == ("2026-09-01", "2026-09-02") + assert result.reconciled_through == "2026-09-14" + marker = await read_marker(prisma) + assert marker is not None and (marker.reconciled_through, marker.scanned_at) == ("2026-09-14", "clock-0009") + + @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): """When the rewrite of a late day fails the marker must stay put and the operator must hear about it."""