mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
57b328ff96
commit
8cf0b50125
2 changed files with 55 additions and 16 deletions
|
|
@ -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]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue