mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(proxy): fold late-arriving per-key spend into already rolled-up global days
The reconcile now records the database clock of the scan behind the last complete run and, on the next run, rewrites every closed day with per-key rows updated since then, however old the day is. Replaying only the marker day and the one before it missed a delayed flush or retry that landed on an older date, and reads through the marker come from the global table alone, so that spend was never counted. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
ad8de0e192
commit
84c098df92
2 changed files with 168 additions and 48 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue