scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

fix(ptu): bound the sentinel prune to the deployments a run scanned

Browse source 
The prune decides a sentinel row is garbage on staleness alone, which holds only while
every run derives its charges from the same table. Once a run's charges can depend on
configuration that only one host holds, a row it never considered is not evidence of
anything, and deleting it drops a charge another host is responsible for.

The loader now returns the deployment ids it looked at alongside the models it priced, and
the delete is bounded to that set. The set is wider than the priced set on purpose, so a
deployment that still exists but no longer carries PTU config stays prunable, and it is a
guaranteed superset of the priced set so a run cannot write a charge outside its own
filter. The ids go out in chunks, since each one is a bind variable and the server refuses
a statement carrying more than 32767 of them.

Behaviour change worth calling out: a deployment deleted from the model table keeps the
day it was already charged for, because nothing scans it any more and so nothing may judge
its row. That day the reservation did exist, and no later run revisits the date regardless.
This commit is contained in:
Yucheng Zhu 2026-08-19 13:42:02 -07:00
parent a1afc2f433
commit f82ef6adf3
2 changed files with 222 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

103
litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
View file

@ -37,6 +37,7 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_HOURS_PER_DAY: Final = 24
_PRUNE_ID_CHUNK_SIZE: Final = 5_000
_UPSERT_ATTEMPTS: Final = 3
_UPSERT_RETRY_BACKOFF_SECONDS: Final = 0.5
@ -358,10 +359,29 @@ async def _upsert_charge_with_retry(
return False
async def _load_ptu_models(prisma_client: "PrismaClient") -> tuple[PTUModel, ...]:
"""Every model deployment currently carrying valid manual PTU config."""
@dataclass(frozen=True, slots=True)
class _LoadedDeployments:
"""The deployments a run will price, and every deployment id it looked at.
The id set is deliberately wider than the priced set. A deployment whose PTU config
was removed produces no charge and still has to be prunable, so bounding the prune on
what priced would strand its old rows forever. It is also a guaranteed superset of the
priced set, or a run could write a charge that falls outside its own delete filter.
"""
models: tuple[PTUModel, ...]
scanned_ids: frozenset[str]
async def _load_ptu_models(prisma_client: "PrismaClient") -> _LoadedDeployments:
"""Every deployment carrying valid manual PTU config, and every id the scan saw."""
rows: Final = await prisma_client.db.litellm_proxymodeltable.find_many()
return tuple(parsed for parsed in (_parse_ptu_model(row) for row in rows) if parsed is not None)
models: Final = tuple(parsed for parsed in (_parse_ptu_model(row) for row in rows) if parsed is not None)
return _LoadedDeployments(
models=models,
scanned_ids=frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or "")))
| frozenset(model.model_id for model in models),
)
async def run_ptu_flat_cost_rollup(
@ -378,8 +398,10 @@ async def run_ptu_flat_cost_rollup(
The prune predicate is ``updated_at < run_started`` rather than "not in the charge
set I computed", which matters under concurrency: whether a row is garbage becomes a
property of the row instead of one run's in-memory config snapshot, so a run can
never delete a row a concurrent run just wrote. It is still skipped when any charge
failed to write, since a row whose replacement never landed would look unrefreshed.
never delete a row a concurrent run just wrote. It is bounded to the deployments this
run looked at, so a row it cannot account for is out of reach either way. It is still
skipped when any charge failed to write, since a row whose replacement never landed
would look unrefreshed.
"""
day: Final = target_date or (datetime.now(timezone.utc).date() - timedelta(days=1))
@ -390,7 +412,8 @@ async def run_ptu_flat_cost_rollup(
date_str: Final = day.isoformat()
run_started: Final = datetime.now(timezone.utc)
ptu_models: Final = await _load_ptu_models(prisma_client)
loaded: Final = await _load_ptu_models(prisma_client)
ptu_models: Final = loaded.models
charges: Final = _aggregate_charges(ptu_models, day)
landed: Final = tuple(
@ -415,7 +438,12 @@ async def run_ptu_flat_cost_rollup(
date_str,
)
else:
await _prune_unrefreshed_sentinel_rows(prisma_client, date_str=date_str, run_started=run_started)
await _prune_unrefreshed_sentinel_rows(
prisma_client,
date_str=date_str,
run_started=run_started,
scanned_ids=loaded.scanned_ids,
)
verbose_proxy_logger.info(
"PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed",
@ -524,7 +552,7 @@ async def run_ptu_flat_cost_backfill(
verbose_proxy_logger.warning("PTU backfill: prisma_client is None, skipping")
return BackfillResult(start=end, end=end, days_scanned=0, rows_written=0)
ptu_models: Final = await _load_ptu_models(prisma_client)
ptu_models: Final = (await _load_ptu_models(prisma_client)).models
days: Final = _backfill_window(ptu_models, end)
if not days:
@ -707,26 +735,53 @@ async def _prune_unrefreshed_sentinel_rows(
*,
date_str: str,
run_started: datetime,
scanned_ids: frozenset[str],
) -> None:
"""Delete the day's PTU sentinel rows this run did not refresh.
"""Delete the day's PTU sentinel rows this run looked at and did not refresh.
Every charge the run wrote bumps ``updated_at`` past ``run_started``, so anything
left below that mark is a (team, model) the current config no longer prices. The mark
is pulled back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come
from different hosts: a stale row is hours old, a concurrently written one is seconds
old, and the grace separates them without waiting on clocks agreeing. The
predicate reads only the row, never the caller's config snapshot, which is what
makes it safe to run twice, out of order, or beside another pod: a row written
after this run began is out of reach of its delete. Mirrors the retention predicate
``SpendLogCleanup`` deletes by."""
Two conditions, and a row survives unless it meets both. It must be stale: every
charge the run wrote bumps ``updated_at`` past ``run_started``, so anything left below
that mark is a (team, model) the current config no longer prices. The mark is pulled
back by ``PTU_PRUNE_SKEW_GRACE_SECONDS`` because the two timestamps come from
different hosts, and the grace separates a row that is hours old from one written
seconds ago without waiting on clocks agreeing.
It must also name a deployment this run scanned. Staleness alone was sufficient while
every run derived its charges from the same table, because then any two runs computed
the same set. Once a run's charges depend on configuration only that host holds, a row
it never considered is not evidence of anything, and deleting it drops a charge some
other host is responsible for.
The ids go out in chunks because the whole set is one bind variable each and the
server rejects a statement carrying more than 32767 of them, which a proxy holding
that many deployments would otherwise hit every night with no handler above here.
"""
cutoff: Final = run_started - timedelta(seconds=PTU_PRUNE_SKEW_GRACE_SECONDS)
await prisma_client.db.litellm_dailyteamspend.delete_many(
where={ # mutable-ok: prisma delete filter
"date": date_str,
"api_key": PTU_SENTINEL_API_KEY,
"updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter
}
ordered: Final = tuple(sorted(scanned_ids))
chunks: Final = tuple(
ordered[start : start + _PRUNE_ID_CHUNK_SIZE] for start in range(0, len(ordered), _PRUNE_ID_CHUNK_SIZE)
)
deletions: Final = tuple(
[
await prisma_client.db.litellm_dailyteamspend.delete_many(
where={ # mutable-ok: prisma delete filter
"date": date_str,
"api_key": PTU_SENTINEL_API_KEY,
"updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter
"model": {"in": chunk}, # mutable-ok: prisma membership filter
}
)
for chunk in chunks
]
)
deleted: Final = sum(deletions)
if deleted:
verbose_proxy_logger.info(
"PTU rollup for %s: pruned %s stale sentinel row(s) across %d scanned deployment(s)",
date_str,
deleted,
len(scanned_ids),
)
__all__ = (

153
tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py
View file

@ -210,8 +210,10 @@ async def test_rollup_prunes_stale_row_when_config_is_gone():
where = table.delete_many.await_args.kwargs["where"]
assert where["date"] == DAY.isoformat()
assert where["api_key"] == PTU_SENTINEL_API_KEY
# the row is garbage because this run did not refresh it, not because of a key list
# the row is garbage because this run did not refresh it, and it is reachable at all
# because the run scanned the deployment it belongs to
assert "lt" in where["updated_at"]
assert where["model"]["in"] == ("m1",)
@pytest.mark.asyncio
@ -705,13 +707,20 @@ class _FakeSentinelTable:
async def delete_many(self, where):
self.delete_many_calls.append(where)
cutoff = where["updated_at"]["lt"]
# honouring "model" matters: a fake that ignored an unknown clause would delete
# the row the prune-scoping test exists to protect and still report a pass
allowed = where.get("model", {}).get("in")
doomed = [
k
for k, v in self.rows.items()
if k[1] == where["date"] and k[2] == where["api_key"] and v["updated_at"] < cutoff
if k[1] == where["date"]
and k[2] == where["api_key"]
and v["updated_at"] < cutoff
and (allowed is None or k[3] in allowed)
]
for k in doomed:
del self.rows[k]
return len(doomed)
async def find_many(self, where=None):
"""Read back sentinel rows the way prisma would, honouring api_key and a date range."""
@ -785,11 +794,11 @@ async def test_an_older_run_cannot_delete_a_newer_runs_row():
@pytest.mark.asyncio
async def test_a_later_clean_run_clears_the_row_the_race_left_behind():
"""The race can leave a charge for a since-removed deployment in place for a day; the
next run, seeing only the current config, must sweep it."""
"""The race can leave a charge for a no-longer-priced deployment in place for a day;
the next run, seeing only the current config, must sweep it."""
table = _FakeSentinelTable()
ptu = {"ptu_count": 10, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}
stale_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-removed")
stale_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-retired")
table.rows[stale_key] = {
"ptu_flat_cost": 480.0,
"model_group": "retired",
@ -797,7 +806,14 @@ async def test_a_later_clean_run_clears_the_row_the_race_left_behind():
}
await run_ptu_flat_cost_rollup(
_prisma_for([_model_row(model_id="dep-live", model_info=ptu)], table), target_date=DAY
_prisma_for(
[
_model_row(model_id="dep-live", model_info=ptu),
_model_row(model_id="dep-retired", model_info={"team_id": "t"}),
],
table,
),
target_date=DAY,
)
assert stale_key not in table.rows
@ -1715,16 +1731,19 @@ async def test_a_run_holding_the_lock_still_prunes():
"""Losing the sweep entirely would leave stale charges forever, so the guarded path,
which is the normal one, keeps it."""
table = _FakeSentinelTable()
table.seed("t", DAY, "dep-gone", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc))
table.seed("t", DAY, "dep-unpriced", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc))
prisma = _prisma_for(
[_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})],
[
_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}),
_model_row(model_id="dep-unpriced", model_info={"team_id": "t"}),
],
table,
)
await run_scheduled_ptu_rollup(prisma, pod_lock_manager=_pod_lock(acquired=True), target_date=DAY)
assert table.delete_many_calls != []
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-gone") not in table.rows
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-unpriced") not in table.rows
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows
@ -1737,7 +1756,7 @@ async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts():
just_written = datetime.now(timezone.utc) - timedelta(seconds=30)
table.seed("t", DAY, "dep-concurrent", 480.0, updated_at=just_written)
table.seed("t", DAY, "dep-stale", 480.0, updated_at=datetime.now(timezone.utc) - timedelta(hours=6))
prisma = _prisma_for([], table)
prisma = _prisma_for([_model_row(model_id="dep-concurrent"), _model_row(model_id="dep-stale")], table)
await run_ptu_flat_cost_rollup(prisma, target_date=DAY)
@ -1747,6 +1766,120 @@ async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts():
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-stale") not in table.rows
@pytest.mark.asyncio
async def test_a_run_cannot_prune_a_row_for_a_deployment_it_did_not_scan():
"""Staleness alone stops being evidence once two hosts can hold different configuration:
a row this run never considered belongs to a deployment some other host is pricing, and
sweeping it drops that charge."""
table = _FakeSentinelTable()
table.seed("t", DAY, "dep-elsewhere", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc))
prisma = _prisma_for(
[_model_row(model_id="dep-live", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})],
table,
)
await run_scheduled_ptu_rollup(prisma, pod_lock_manager=_pod_lock(acquired=True), target_date=DAY)
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-elsewhere") in table.rows
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows
assert table.delete_many_calls[-1]["model"]["in"] == ("dep-live",)
@pytest.mark.asyncio
async def test_a_deployment_deleted_from_the_table_keeps_the_day_it_was_charged():
"""The accepted cost of bounding the prune, driven through the sequence that produces
it: charge the day while the deployment exists, remove it, run the day again. Nothing
scans it now, so nothing may judge its row, and the amount it was billed stands."""
table = _FakeSentinelTable()
ptu = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}
live_row = _model_row(model_id="dep-live", model_info=ptu)
doomed_row = _model_row(model_id="dep-doomed", model_info=ptu)
charged_key = ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-doomed")
await run_scheduled_ptu_rollup(
_prisma_for([live_row, doomed_row], table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY
)
billed = table.rows[charged_key]["ptu_flat_cost"]
table.rows[charged_key]["updated_at"] = datetime(2020, 1, 1, tzinfo=timezone.utc)
await run_scheduled_ptu_rollup(
_prisma_for([live_row], table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY
)
assert table.rows[charged_key]["ptu_flat_cost"] == billed
assert table.delete_many_calls[-1]["model"]["in"] == ("dep-live",)
@pytest.mark.asyncio
async def test_a_run_that_scanned_nothing_prunes_nothing():
"""An empty scan has nothing it can account for, so it issues no delete at all rather
than one that can match nothing."""
table = _FakeSentinelTable()
table.seed("t", DAY, "dep-stale", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc))
prisma = _prisma_for([], table)
await run_scheduled_ptu_rollup(prisma, pod_lock_manager=_pod_lock(acquired=True), target_date=DAY)
assert table.delete_many_calls == []
assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-stale") in table.rows
@pytest.mark.asyncio
async def test_every_deployment_that_prices_is_inside_the_set_that_bounds_the_prune():
"""The bound has to be a superset of what the same run wrote, or a run's own charge
could fall outside its own delete filter and never be reconciled."""
table = _FakeSentinelTable()
prisma = _prisma_for(
[
_model_row(model_id="dep-a", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}),
_model_row(model_id="dep-b", model_info={"ptu_count": 9, "cost_per_ptu_per_hour": 1.0, "team_id": "u"}),
_model_row(model_id="dep-unpriced", model_info={"team_id": "t"}),
],
table,
)
loaded = await ptu_rollup._load_ptu_models(prisma)
assert {model.model_id for model in loaded.models} <= loaded.scanned_ids
assert loaded.scanned_ids == {"dep-a", "dep-b", "dep-unpriced"}
@pytest.mark.asyncio
async def test_a_priced_deployment_is_in_the_bound_even_with_an_id_the_scan_skips():
"""The bound is built by construction rather than by coincidence. The row scan drops a
falsy id while the parser still prices one, and a charge outside its own run's delete
filter could never be reconciled by any later run."""
prisma = _prisma_for(
[_model_row(model_id="", model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})],
_FakeSentinelTable(),
)
loaded = await ptu_rollup._load_ptu_models(prisma)
assert {model.model_id for model in loaded.models} <= loaded.scanned_ids
@pytest.mark.asyncio
async def test_the_prune_splits_the_id_set_across_statements(monkeypatch):
"""Every id is one bind variable and the server refuses a statement carrying more than
32767, so a proxy with that many deployments would fail the prune outright, and with it
the rest of the scheduled run."""
monkeypatch.setattr(ptu_rollup, "_PRUNE_ID_CHUNK_SIZE", 2)
table = _FakeSentinelTable()
ptu = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"}
deployments = [_model_row(model_id=f"dep-{n}", model_info=ptu) for n in range(5)]
table.seed("t", DAY, "dep-3", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc))
await run_scheduled_ptu_rollup(
_prisma_for(deployments, table), pod_lock_manager=_pod_lock(acquired=True), target_date=DAY
)
chunks = [call["model"]["in"] for call in table.delete_many_calls]
assert len(chunks) == 3
assert all(len(chunk) <= 2 for chunk in chunks)
assert sorted(i for chunk in chunks for i in chunk) == [f"dep-{n}" for n in range(5)]
@pytest.mark.asyncio
async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled(monkeypatch):
"""Startup already skips scheduling the cron, so this guards the function itself: a