fix(ptu): hand the prune a plain delete filter the query builder can serialise (#37571)

* fix(ptu): hand the prune a plain delete filter the query builder can serialise

The bounded sweep built its predicate as a read-only mapping view, which the query
builder refuses to serialise, so the nightly job raised as soon as a config-declared
deployment was priced. The charges were already written by then, which is why the run
looked like it had produced its rows.

The in-memory table these tests run against accepts any mapping, so only a live run
caught it. A predicate builder now returns a plain dict and is asserted as one, and the
catch-up pass has a test covering a config-declared reservation.

* refactor(ptu): build the prune predicate in one shot

Both filter shapes are known upfront, so the bounded one is constructed
directly rather than by mutating a value already declared Final.

The catch-up test took two independent clock reads, which disagree across
UTC midnight; it now derives both the reservation start and the expected
last charged day from a single read, matching the three sibling tests.
This commit is contained in:
yucheng-berri 2026-08-19 22:18:58 -07:00 committed by GitHub
parent 57b328ff96
commit 8cf0b50125
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 55 additions and 16 deletions

View file

@ -724,6 +724,20 @@ async def _deliver_alert(alert: "Callable[[str], Awaitable[None]] | None", messa
verbose_proxy_logger.error("PTU rollup: could not deliver the failed-charge alert: %s", exc)
def _prune_filter(*, date_str: str, cutoff: datetime, chunk: "tuple[str, ...] | None") -> "Mapping[str, object]":
"""One delete statement's predicate. An absent chunk leaves the sweep unbounded.
Returns a plain dict because the query builder serialises the mapping it is handed and
rejects a read-only view of one.
"""
return { # mutable-ok: prisma delete filter
"date": date_str,
"api_key": PTU_SENTINEL_API_KEY,
"updated_at": {"lt": cutoff}, # mutable-ok: prisma comparison filter
**({} if chunk is None else {"model": {"in": chunk}}), # mutable-ok: prisma membership filter
}
async def _prune_unrefreshed_sentinel_rows(
prisma_client: "PrismaClient",
*,
@ -752,27 +766,15 @@ async def _prune_unrefreshed_sentinel_rows(
that many deployments would otherwise hit every night with no handler above here.
"""
cutoff: Final = run_started - timedelta(seconds=PTU_PRUNE_SKEW_GRACE_SECONDS)
unbounded: Final = { # 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 = () if scanned_ids is None else tuple(sorted(scanned_ids))
filters: Final = (
(unbounded,)
chunks: Final = (
(None,)
if scanned_ids is None
else tuple(
MappingProxyType(
{
**unbounded,
"model": { # mutable-ok: prisma membership filter
"in": ordered[start : start + _PRUNE_ID_CHUNK_SIZE]
},
}
)
for start in range(0, len(ordered), _PRUNE_ID_CHUNK_SIZE)
ordered[start : start + _PRUNE_ID_CHUNK_SIZE] for start in range(0, len(ordered), _PRUNE_ID_CHUNK_SIZE)
)
)
filters: Final = tuple(_prune_filter(date_str=date_str, cutoff=cutoff, chunk=chunk) for chunk in chunks)
deletions: Final = tuple(
[await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters]
)

View file

@ -2045,3 +2045,40 @@ def test_the_router_lookup_returns_none_outside_a_proxy():
finally:
if real is not None:
sys.modules["litellm.proxy.proxy_server"] = real
@pytest.mark.parametrize("chunk", [None, ("dep-a", "dep-b")], ids=["unbounded", "bounded"])
def test_the_prune_filter_is_a_plain_dict(chunk):
"""The query builder serialises the mapping it is handed and rejects a read-only view of
one, which the in-memory table in these tests accepts happily. Only a live run caught it."""
predicate = ptu_rollup._prune_filter(date_str=DAY.isoformat(), cutoff=datetime.now(timezone.utc), chunk=chunk)
assert type(predicate) is dict
assert type(predicate["updated_at"]) is dict
if chunk is None:
assert "model" not in predicate
else:
assert type(predicate["model"]) is dict
assert predicate["model"]["in"] == chunk
@pytest.mark.asyncio
async def test_the_catch_up_pass_reaches_a_config_declared_deployment(monkeypatch):
"""The catch-up shares the loader, so config deployments join it without being wired in.
That is what prices the elapsed days of a reservation declared before today."""
table = _FakeSentinelTable()
now = datetime.now(timezone.utc)
started = (now - timedelta(days=3)).strftime("%Y-%m-%dT00:00:00Z")
entry = _router_entry(
model_id="cfg-back",
model_info={"ptu_count": 100, "cost_per_ptu_per_hour": 0.02, "team_id": "t", "ptu_effective_from": started},
)
monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry))
await run_scheduled_ptu_rollup(_prisma_for([], table), pod_lock_manager=_pod_lock(acquired=True))
charged = sorted(day for (_, day, _, model) in table.rows if model == "cfg-back")
yesterday = (now.date() - timedelta(days=1)).isoformat()
assert len(charged) == 3, charged
assert charged[-1] == yesterday
assert all(row["ptu_flat_cost"] == pytest.approx(48.0) for row in table.rows.values())