diff --git a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py index a9fb7669785..73068381dab 100644 --- a/litellm/proxy/spend_tracking/daily_global_spend_rollup.py +++ b/litellm/proxy/spend_tracking/daily_global_spend_rollup.py @@ -2,9 +2,12 @@ Only days that are over get rolled up, so a pod still flushing per-key spend for the current day can never leave the global table short; usage reads serve days through the recorded -marker from the global table and later days live from the per-key table. The marker lives in -``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on a -large deployment the first backfill is minutes of work. +marker from the global table and later days live from the per-key table. Per-key rows are +dated by request start, so spend can land on a day that was already rolled up (a flush +straddling midnight, a retry after an outage). Each run therefore also rewrites every closed +day that has rows touched since the previous run's scan, whatever the date. The marker lives +in ``LiteLLM_Config``. This runs as a background cron, never in a Prisma migration, since on +a large deployment the first backfill is minutes of work. """ from collections.abc import Awaitable, Callable @@ -27,7 +30,6 @@ if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient -_REPLAY_DAYS: Final = 1 GLOBAL_SPEND_TABLE_NAME: Final = "LiteLLM_DailyGlobalSpend" # The unique constraint, in constraint order. NULL never matches itself in a unique index, so # every column is normalized to '' or the same group would be inserted again on every run. @@ -69,15 +71,26 @@ def _reconcile_day_sql() -> str: RECONCILE_DAY_SQL: Final = _reconcile_day_sql() +_DB_NOW_SQL: Final = "SELECT (NOW() AT TIME ZONE 'UTC')::text AS now" +_ALL_CLOSED_DAYS_SQL: Final = 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ORDER BY "date"' +# Pod clocks drift from the database clock and from each other, so rows are picked up from a +# little before the previous scan; rewriting a day twice is idempotent. _PENDING_DAYS_SQL: Final = ( - 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" >= $1 AND "date" <= $2 ORDER BY "date"' + 'SELECT DISTINCT "date" FROM "LiteLLM_DailyUserSpend" WHERE "date" <= $1 ' + 'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') ' + 'ORDER BY "date"' ) class ReconciledThrough(BaseModel): + """``reconciled_through`` is the last closed UTC day the global table covers. ``scanned_at`` is + the database clock when the scan behind the last fully successful run started: every per-key + row written before it, on any day through the marker, is in the global table.""" + model_config = ConfigDict(frozen=True, extra="ignore") reconciled_through: str + scanned_at: str | None = None class _MarkerRow(BaseModel): @@ -92,6 +105,12 @@ class _DateRow(BaseModel): date: str +class _NowRow(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + now: str + + @dataclass(frozen=True, slots=True) class ReconcileResult: days_reconciled: tuple[str, ...] @@ -99,49 +118,70 @@ class ReconcileResult: failed_day: str | None = None -def _marker_from_param_value(value: object) -> str | None: +@dataclass(frozen=True, slots=True) +class _PendingScan: + marker: ReconciledThrough | None + scanned_at: str + days: tuple[str, ...] + + +def _marker_from_param_value(value: object) -> ReconciledThrough | None: try: - parsed: Final = ( + return ( ReconciledThrough.model_validate_json(value) if isinstance(value, str) else ReconciledThrough.model_validate(value) ) except ValidationError: return None - return parsed.reconciled_through -async def reconciled_through(prisma_client: "PrismaClient") -> str | None: - """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" +async def read_marker(prisma_client: "PrismaClient") -> ReconciledThrough | None: from litellm.proxy.utils import get_config_param row: Final = await get_config_param(prisma_client, DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) return None if row is None else _marker_from_param_value(_MarkerRow.model_validate(row).param_value) -async def _record_reconciled_through(prisma_client: "PrismaClient", day: str) -> None: +async def reconciled_through(prisma_client: "PrismaClient") -> str | None: + """The last UTC day ``LiteLLM_DailyGlobalSpend`` is known to cover, or None before the first run.""" + marker: Final = await read_marker(prisma_client) + return None if marker is None else marker.reconciled_through + + +async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThrough) -> None: from litellm.proxy.utils import invalidate_config_param await ConfigRepository(prisma_client).set_param( - DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, ReconciledThrough(reconciled_through=day).model_dump_json() + DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, marker.model_dump_json() ) await invalidate_config_param(DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM) -def _first_pending_day(marker: str | None) -> str: - if marker is None: - return "" - return (date.fromisoformat(marker) - timedelta(days=_REPLAY_DAYS)).isoformat() +async def _db_now(prisma_client: "PrismaClient") -> str: + rows: Final = await prisma_client.db.query_raw(_DB_NOW_SQL) + return _NowRow.model_validate(rows[0]).now + + +async def _scan_pending(prisma_client: "PrismaClient", today: date) -> _PendingScan: + """Every closed UTC day (strictly before today) still to roll up, oldest first: days past the + marker, plus any day with per-key rows written since the scan behind the marker. Before a + run has fully succeeded there is no such scan, so every closed day is rolled up.""" + marker: Final = await read_marker(prisma_client) + scanned_at: Final = await _db_now(prisma_client) + last_closed_day: Final = (today - timedelta(days=1)).isoformat() + rows: Final = ( + await prisma_client.db.query_raw(_ALL_CLOSED_DAYS_SQL, last_closed_day) + if marker is None or marker.scanned_at is None + else await prisma_client.db.query_raw( + _PENDING_DAYS_SQL, last_closed_day, marker.reconciled_through, marker.scanned_at + ) + ) + return _PendingScan(marker, scanned_at, tuple(_DateRow.model_validate(row).date for row in rows)) async def pending_days(prisma_client: "PrismaClient", today: date) -> tuple[str, ...]: - """Every closed UTC day (strictly before today) still to roll up, oldest first. The marker - day and the one before it are replayed so per-key rows that landed after their day was - rolled up (a flush straddling midnight, a late retry) are folded in.""" - marker: Final = await reconciled_through(prisma_client) - last_closed_day: Final = (today - timedelta(days=1)).isoformat() - rows: Final = await prisma_client.db.query_raw(_PENDING_DAYS_SQL, _first_pending_day(marker), last_closed_day) - return tuple(_DateRow.model_validate(row).date for row in rows) + return (await _scan_pending(prisma_client, today)).days async def reconcile_day(prisma_client: "PrismaClient", day: str) -> None: @@ -155,26 +195,42 @@ async def run_daily_global_spend_reconcile( today: date | None = None, ) -> ReconcileResult: """Roll up every pending day, advancing the marker after each; a failing day stops the run - with the marker on the last good day so the next run resumes there.""" + with the marker on the last good day so the next run resumes there. The scan time is only + recorded once every pending day is done, so late rows a failed run saw are found again.""" effective_today: Final = today or datetime.now(timezone.utc).date() - days: Final = await pending_days(prisma_client, effective_today) - done: Final = await _reconcile_until_failure(prisma_client, days) - failed: Final = days[len(done)] if len(done) < len(days) else None - marker: Final = done[-1] if done else await reconciled_through(prisma_client) - return ReconcileResult(days_reconciled=done, reconciled_through=marker, failed_day=failed) + scan: Final = await _scan_pending(prisma_client, effective_today) + done: Final = await _reconcile_until_failure(prisma_client, scan) + if len(done) < len(scan.days): + 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)) + return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client)) -async def _reconcile_until_failure(prisma_client: "PrismaClient", days: tuple[str, ...]) -> tuple[str, ...]: - for index, day in enumerate(days): - if not await _reconcile_and_record(prisma_client, day): - return days[:index] - return days +def _advanced(marker: ReconciledThrough | None, days: tuple[str, ...], *, scanned_at: str | None) -> ReconciledThrough: + """The marker after ``days`` were rewritten: a late old day never moves it back.""" + through: Final = max((marker.reconciled_through if marker is not None else "", *days)) + return ReconciledThrough(reconciled_through=through, scanned_at=scanned_at) -async def _reconcile_and_record(prisma_client: "PrismaClient", day: str) -> bool: +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]): + return scan.days[:index] + return scan.days + + +async def _reconcile_and_record( + prisma_client: "PrismaClient", marker: ReconciledThrough | None, done_with_this: tuple[str, ...] +) -> bool: + day: Final = done_with_this[-1] try: await reconcile_day(prisma_client, day) - await _record_reconciled_through(prisma_client, day) + await _record_marker( + prisma_client, + _advanced(marker, done_with_this, scanned_at=None if marker is None else marker.scanned_at), + ) 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 11ca72e7b3d..9a098744f08 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 @@ -15,6 +15,7 @@ from litellm.constants import DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_bulk_upsert, merge_by_conflict_key from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( RECONCILE_DAY_SQL, + read_marker, reconciled_through, run_daily_global_spend_reconcile, run_scheduled_daily_global_spend_reconcile, @@ -41,13 +42,25 @@ class _FakeConfigTable: class _FakeDb: + """Per-key rows are ``{date: updated_at}`` with a fake database clock that ticks per query, + so "rows written since the last scan" behaves like Postgres would.""" + def __init__(self, prisma: "_FakePrisma") -> None: self._prisma = prisma self.litellm_config = _FakeConfigTable() async def query_raw(self, sql: str, *params: str) -> list[dict[str, str]]: - first, last = params - return [{"date": d} for d in sorted(self._prisma.user_days) if first <= d <= last] + if sql.startswith("SELECT (NOW()"): + self._prisma.clock += 1 + return [{"now": f"clock-{self._prisma.clock:04d}"}] + rows = self._prisma.user_rows + if len(params) == 1: + (last,) = params + return [{"date": d} for d in sorted(rows) if d <= last] + last, marker, scanned_at = params + return [ + {"date": d} for d, written in sorted(rows.items()) if d <= last and (d > marker or written >= scanned_at) + ] async def execute_raw(self, sql: str, *params: str) -> int: (day,) = params @@ -61,11 +74,17 @@ class _FakePrisma: """Enough of PrismaClient for the reconcile: per-key dates, a config table, and execute_raw.""" def __init__(self, user_days: tuple[str, ...], failing_days: frozenset[str] = frozenset()) -> None: - self.user_days = user_days + self.clock = 0 + self.user_rows: dict[str, str] = {d: "clock-0000" for d in user_days} self.failing_days = failing_days self.reconciled: list[str] = [] self.db = _FakeDb(self) + def write_late_row(self, day: str) -> None: + """A per-key row for ``day`` lands now, after whatever scans already happened.""" + self.clock += 1 + self.user_rows[day] = f"clock-{self.clock:04d}" + async def get_generic_data(self, key: str, value: str, table_name: str) -> _FakeConfigRow | None: stored = self.db.litellm_config.rows.get(value) return None if stored is None else _FakeConfigRow(value, stored) @@ -94,20 +113,66 @@ async def test_first_run_rolls_up_every_closed_day_and_never_today(): @pytest.mark.asyncio -async def test_later_run_replays_the_marker_day_and_the_day_before_only(): - """Days older than marker-1 are settled; the marker day and its predecessor are replayed so - per-key rows that landed after their day was rolled up get folded in.""" +async def test_later_run_rolls_up_only_new_days_when_nothing_old_changed(): prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-12", "2026-09-13", "2026-09-14")) await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) prisma.reconciled.clear() result = await run_daily_global_spend_reconcile(prisma, today=TODAY) - assert result.days_reconciled == ("2026-09-12", "2026-09-13", "2026-09-14") - assert "2026-09-01" not in prisma.reconciled + assert result.days_reconciled == ("2026-09-14",) assert await reconciled_through(prisma) == "2026-09-14" +@pytest.mark.asyncio +async def test_spend_landing_on_an_old_rolled_up_day_is_folded_in_by_the_next_run(): + """Per-key rows carry the request start date, so a delayed flush or retry can add spend to a + day far behind the marker. That day is rewritten, and the marker never moves back for it.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-05", "2026-09-13")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.reconciled.clear() + prisma.write_late_row("2026-09-01") + prisma.write_late_row("2026-09-03") + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-03") + assert "2026-09-05" not in prisma.reconciled + assert await reconciled_through(prisma) == "2026-09-13" + + +@pytest.mark.asyncio +async def test_a_late_row_seen_by_a_failed_run_is_seen_again_by_the_next_one(): + """The scan time only advances when every pending day was rewritten, otherwise a late row + found by the failed run would be counted as handled.""" + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) + await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) + prisma.write_late_row("2026-09-01") + prisma.failing_days = frozenset({"2026-09-01"}) + failed = await run_daily_global_spend_reconcile(prisma, today=TODAY) + prisma.failing_days = frozenset() + prisma.reconciled.clear() + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert failed.failed_day == "2026-09-01" + assert failed.reconciled_through == "2026-09-13" + assert result.days_reconciled == ("2026-09-01",) + assert result.failed_day is None + + +@pytest.mark.asyncio +async def test_a_marker_without_a_scan_time_rolls_every_closed_day_up_again(): + prisma = _FakePrisma(user_days=("2026-09-01", "2026-09-13")) + prisma.db.litellm_config.rows[DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM] = '{"reconciled_through": "2026-09-13"}' + + result = await run_daily_global_spend_reconcile(prisma, today=TODAY) + + assert result.days_reconciled == ("2026-09-01", "2026-09-13") + marker = await read_marker(prisma) + assert marker is not None and marker.reconciled_through == "2026-09-13" and marker.scanned_at is not None + + @pytest.mark.asyncio async def test_a_run_with_no_new_closed_days_keeps_the_marker(): prisma = _FakePrisma(user_days=("2026-09-13",)) @@ -116,7 +181,7 @@ async def test_a_run_with_no_new_closed_days_keeps_the_marker(): result = await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) - assert result.days_reconciled == ("2026-09-13",) + assert result.days_reconciled == () assert result.reconciled_through == "2026-09-13" @@ -149,11 +214,10 @@ async def test_the_next_run_resumes_from_the_failed_day(): @pytest.mark.asyncio async def test_a_failure_with_nothing_done_reports_the_previous_marker_and_alerts(): - """A late flush for the day before the marker is exactly the replay case; when that replay - fails the marker must stay put and the operator must hear about it.""" + """When the rewrite of a late day fails the marker must stay put and the operator must hear about it.""" prisma = _FakePrisma(user_days=("2026-09-13",)) await run_daily_global_spend_reconcile(prisma, today=date(2026, 9, 14)) - prisma.user_days = ("2026-09-12", "2026-09-13") + prisma.write_late_row("2026-09-12") prisma.failing_days = frozenset({"2026-09-12"}) alert = AsyncMock()