diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 381641be96d..f1f7248c064 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -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] ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index 333fd597b49..e039455607d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -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())