scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(proxy): advance the daily global spend marker in one conditional upsert so overlapping runs cannot rewind it

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-18 21:15:32 +00:00
parent abf530fbeb
commit 3449ae9d0d
2 changed files with 74 additions and 38 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/spend_tracking/daily_global_spend_rollup.py
View file

@ -23,7 +23,6 @@ from litellm.constants import (
DAILY_GLOBAL_SPEND_RECONCILE_LOCK_TTL_SECONDS,
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
)
from litellm.repositories.config_repository import ConfigRepository
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
@ -82,6 +81,17 @@ _PENDING_DAYS_SQL: Final = (
'AND ("date" > $2 OR "updated_at" >= $3::timestamp - INTERVAL \'1 hour\') '
'ORDER BY "date"'
)
# Runs can overlap (Redis unreachable, lock expired on a long backfill), so the database keeps the
# later of the stored and the incoming day and scan time in one statement; GREATEST skips NULL.
_ADVANCE_MARKER_SQL: Final = (
'INSERT INTO "LiteLLM_Config" ("param_name", "param_value") '
"VALUES ($1, jsonb_build_object('reconciled_through', $2::text, 'scanned_at', $3::text)) "
'ON CONFLICT ("param_name") DO UPDATE SET "param_value" = jsonb_build_object('
"'reconciled_through', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'reconciled_through', "
"EXCLUDED.\"param_value\" ->> 'reconciled_through'), "
"'scanned_at', GREATEST(\"LiteLLM_Config\".\"param_value\" ->> 'scanned_at', "
"EXCLUDED.\"param_value\" ->> 'scanned_at'))"
)
class ReconciledThrough(BaseModel):
@ -152,33 +162,20 @@ async def reconciled_through(prisma_client: "PrismaClient") -> str | None:
return None if marker is None else marker.reconciled_through
async def _record_marker(prisma_client: "PrismaClient", marker: ReconciledThrough) -> None:
async def _advance_marker(prisma_client: "PrismaClient", days: tuple[str, ...], *, scanned_at: str | None) -> None:
"""Move the stored marker to the last of ``days`` and to ``scanned_at`` where those are later
than what is stored, so a slower overlapping run can only add to a faster run's marker."""
from litellm.proxy.utils import invalidate_config_param
await ConfigRepository(prisma_client).set_param(
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM, marker.model_dump_json()
await prisma_client.db.execute_raw(
_ADVANCE_MARKER_SQL,
DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM,
max(days) if days else None,
scanned_at,
)
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])
@ -221,16 +218,10 @@ 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_advanced(prisma_client, done, scanned_at=scan.scanned_at)
await _advance_marker(prisma_client, done, scanned_at=scan.scanned_at)
return ReconcileResult(days_reconciled=done, reconciled_through=await reconciled_through(prisma_client))
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_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.days[: index + 1]):
@ -242,7 +233,7 @@ async def _reconcile_and_record(prisma_client: "PrismaClient", done_with_this: t
day: Final = done_with_this[-1]
try:
await reconcile_day(prisma_client, day)
await _record_advanced(prisma_client, done_with_this, scanned_at=None)
await _advance_marker(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

61
tests/test_litellm/proxy/spend_tracking/test_daily_global_spend_rollup.py
View file

@ -1,5 +1,6 @@
"""Tests for the LiteLLM_DailyGlobalSpend reconcile job (LIT-7818)."""
import json
import pathlib
import re
from datetime import date
@ -14,6 +15,7 @@ from pytest_postgresql import factories
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 (
_ADVANCE_MARKER_SQL,
RECONCILE_DAY_SQL,
read_marker,
reconciled_through,
@ -36,13 +38,21 @@ class _FakeConfigTable:
def __init__(self) -> None:
self.rows: dict[str, object] = {}
async def upsert(self, *, where: dict[str, str], data: dict[str, dict[str, str]]) -> _FakeConfigRow:
self.rows[where["param_name"]] = data["update"]["param_value"]
return _FakeConfigRow(where["param_name"], data["update"]["param_value"])
def advance(self, param_name: str, through: str | None, scanned_at: str | None) -> None:
"""What ``_ADVANCE_MARKER_SQL`` does in Postgres: keep the later of stored and incoming per field."""
stored = self.rows.get(param_name)
current: dict[str, str | None] = json.loads(stored) if isinstance(stored, str) else {}
self.rows[param_name] = json.dumps(
{
"reconciled_through": _greatest(current.get("reconciled_through"), through),
"scanned_at": _greatest(current.get("scanned_at"), scanned_at),
}
)
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)
def _greatest(stored: str | None, incoming: str | None) -> str | None:
present = [value for value in (stored, incoming) if value is not None]
return max(present) if present else None
class _FakeDb:
@ -67,9 +77,14 @@ class _FakeDb:
{"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:
async def execute_raw(self, sql: str, *params: str | None) -> int:
if sql == _ADVANCE_MARKER_SQL:
param_name, through, scanned_at = params
assert param_name is not None
self.litellm_config.advance(param_name, through, scanned_at)
return 1
(day,) = params
if day in self._prisma.failing_days:
if day is None or 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)
@ -485,3 +500,33 @@ def test_reconcile_day_sql_makes_the_global_day_equal_the_per_key_sums(_rollup_p
assert sum(int(r["total_response_time_ms"]) for r in global_rows) == 1600 # pyright: ignore[reportArgumentType] # dict_row values are untyped
assert [(r["model"], r["model_group"]) for r in global_rows] == [("gpt-5", ""), ("gpt-5", "gpt-5")]
assert untouched == []
_CONFIG_DDL: Final = 'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB)'
_MARKER_SQL: Final = 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s'
def test_advance_marker_sql_only_ever_moves_the_stored_marker_forward(_rollup_postgresql: psycopg.Connection):
"""Against real Postgres: the statement a slower overlapping run issues after the faster run
already stored a later marker leaves that marker alone, whether it carries an older scan time or
none at all, while a run that is further along moves both fields on."""
conn: Final = _rollup_postgresql
conn.execute(_CONFIG_DDL) # pyright: ignore[reportArgumentType] # DDL literal
conn.commit()
param: Final = DAILY_GLOBAL_SPEND_RECONCILED_THROUGH_PARAM
def stored() -> object:
with conn.cursor(row_factory=dict_row) as cur:
row = cur.execute(_MARKER_SQL, (param,)).fetchone()
return None if row is None else row["param_value"]
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-01", None))
assert stored() == {"reconciled_through": "2026-09-01", "scanned_at": None}
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-14", "2026-09-15 00:30:02.5"))
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-02", None))
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-03", "2026-09-15 00:30:01.25"))
assert stored() == {"reconciled_through": "2026-09-14", "scanned_at": "2026-09-15 00:30:02.5"}
_execute_dollar_sql(conn, _ADVANCE_MARKER_SQL, (param, "2026-09-15", "2026-09-16 00:30:00.75"))
assert stored() == {"reconciled_through": "2026-09-15", "scanned_at": "2026-09-16 00:30:00.75"}