From 0a5fa4fdc6599eb236c8116f9ce77e1ec29f83ae Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 21 Aug 2026 11:31:45 -0700 Subject: [PATCH] fix(ptu): never retract a flat charge for a deployment the run cannot see (#37793) The sweep ran unbounded whenever no config.yaml deployment was present, deleting the day's sentinel rows for deployments absent from the run's own view. A written charge records capacity that was reserved, so the only rows a run may retract are the ones it can reassess: a deployment it scanned and then declined to charge, because the window closed or the PTU config was removed. It is now always bounded to the ids it scanned --- .../spend_tracking/ptu_flat_cost_rollup.py | 53 +++++++++---------- .../test_ptu_flat_cost_rollup.py | 45 ++++++++++------ 2 files changed, 53 insertions(+), 45 deletions(-) diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index f1f7248c064..6f1bbaa722b 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -324,7 +324,6 @@ class _LoadedDeployments: models: tuple[PTUModel, ...] scanned_ids: frozenset[str] - config_sourced: bool def _running_router() -> object | None: @@ -371,7 +370,6 @@ async def _load_ptu_models(prisma_client: "PrismaClient") -> _LoadedDeployments: ) return _LoadedDeployments( models=models, - config_sourced=bool(config_records), scanned_ids=db_ids | frozenset(record.model_id for record in config_records) | frozenset(model.model_id for model in models), @@ -385,9 +383,11 @@ async def run_ptu_flat_cost_rollup( ) -> RollupResult: """Rollup one UTC day of flat PTU cost across all PTU-configured model deployments. - Defaults to yesterday UTC. Authoritative for the day: it upserts the current charges - first, then deletes the day's sentinel rows this run did not refresh, so a - since-removed, invalidated, or now-out-of-window deployment leaves no stale charge. + Defaults to yesterday UTC. It upserts the current charges first, then deletes the + day's sentinel rows it scanned and did not refresh, so an invalidated or + now-out-of-window deployment leaves no stale charge. A deployment it cannot see is + left alone, since its charge records capacity that was reserved and this run has no + grounds to retract it. 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 @@ -436,7 +436,7 @@ async def run_ptu_flat_cost_rollup( prisma_client, date_str=date_str, run_started=run_started, - scanned_ids=loaded.scanned_ids if loaded.config_sourced else None, + scanned_ids=loaded.scanned_ids, ) verbose_proxy_logger.info( @@ -724,8 +724,8 @@ 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. +def _prune_filter(*, date_str: str, cutoff: datetime, chunk: "tuple[str, ...]") -> "Mapping[str, object]": + """One delete statement's predicate, bounded to the deployments in ``chunk``. Returns a plain dict because the query builder serialises the mapping it is handed and rejects a read-only view of one. @@ -734,7 +734,7 @@ def _prune_filter(*, date_str: str, cutoff: datetime, chunk: "tuple[str, ...] | "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 + "model": {"in": chunk}, # mutable-ok: prisma membership filter } @@ -743,7 +743,7 @@ async def _prune_unrefreshed_sentinel_rows( *, date_str: str, run_started: datetime, - scanned_ids: frozenset[str] | None, + scanned_ids: frozenset[str], ) -> None: """Delete the day's PTU sentinel rows this run looked at and did not refresh. @@ -754,25 +754,22 @@ async def _prune_unrefreshed_sentinel_rows( different hosts, and the grace separates a row that is hours old from one written seconds ago without waiting on clocks agreeing. - A run that priced a deployment only its own host declares must also name the - deployments it scanned. Staleness alone is sufficient while every run derives its - charges from the same table, because then any two runs compute the same set, so a - database-only run still sweeps by timestamp exactly as it always has. Once one host's - charges come from a file the others cannot read, a row it never considered is not - evidence of anything, and deleting it drops a charge that host is responsible for. + It must also be a deployment this run could see. A charge already written is a record + of capacity that was reserved, so the only rows a run may retract are the ones it can + reassess: a deployment it scanned and then declined to charge, because the window + closed or the PTU config was removed. A row whose deployment is absent from every + source the run reads is not evidence that the reservation never happened, only that + this host cannot account for it. A deployment the router refused to register is in that + same bucket as one that was removed, because neither reaches the scan. - Where the bound applies the ids go out in chunks, because each is one bind variable 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. + The ids go out in chunks, because each is one bind variable 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) - ordered: Final = () if scanned_ids is None else tuple(sorted(scanned_ids)) - chunks: Final = ( - (None,) - if scanned_ids is None - else tuple( - ordered[start : start + _PRUNE_ID_CHUNK_SIZE] for start in range(0, len(ordered), _PRUNE_ID_CHUNK_SIZE) - ) + 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) ) filters: Final = tuple(_prune_filter(date_str=date_str, cutoff=cutoff, chunk=chunk) for chunk in chunks) deletions: Final = tuple( @@ -781,10 +778,10 @@ async def _prune_unrefreshed_sentinel_rows( deleted: Final = sum(deletions) if deleted: verbose_proxy_logger.info( - "PTU rollup for %s: pruned %s stale sentinel row(s) across %s deployment(s)", + "PTU rollup for %s: pruned %s stale sentinel row(s) of %s deployment(s) considered", date_str, deleted, - "every" if scanned_ids is None else len(scanned_ids), + len(ordered), ) 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 e039455607d..dda2f5a4d73 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 @@ -196,10 +196,12 @@ async def test_rollup_writes_sentinel_row_with_hourly_cost(): @pytest.mark.asyncio -async def test_rollup_prunes_stale_row_when_config_is_gone(): +async def test_rollup_prunes_a_scanned_deployment_whose_ptu_config_is_gone(): + """A deployment the run can still see, and can therefore judge, is the one case where + retracting the charge is justified.""" prisma, table = _prisma_with_models( - [_model_row(model_info={"team_id": "team_x"})], - existing_sentinel_rows=[_sentinel_row("stale-1", "team_x", "gpt-4o-mini-ptu")], + [_model_row(model_id="m1", model_info={"team_id": "team_x"})], + existing_sentinel_rows=[_sentinel_row("stale-1", "team_x", "m1")], ) result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) @@ -210,10 +212,8 @@ 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, and it is reachable at all - # because the run scanned the deployment it belongs to assert "lt" in where["updated_at"] - assert "model" not in where, "a database-only run has no reason to bound the sweep" + assert where["model"]["in"] == ("m1",) @pytest.mark.asyncio @@ -1812,9 +1812,10 @@ async def test_a_deployment_deleted_from_the_table_keeps_the_day_it_was_charged( @pytest.mark.asyncio -async def test_a_database_only_run_sweeps_exactly_as_it_did_before(): - """The bound exists for charges another host declares. A deployment nobody declares any - more still has its leftover row swept, which is what the table-only sweep always did.""" +async def test_a_charge_the_run_cannot_reassess_is_left_alone(): + """A written charge records capacity that was reserved. A deployment absent from every + source this run reads cannot be reassessed, and another host may be the one declaring + it, so retracting the charge would drop money the provider still invoiced.""" table = _FakeSentinelTable() table.seed("t", DAY, "dep-gone", 480.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) prisma = _prisma_for( @@ -1824,7 +1825,7 @@ async def test_a_database_only_run_sweeps_exactly_as_it_did_before(): 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-gone") not in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-gone") in table.rows assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") in table.rows @@ -2047,19 +2048,29 @@ def test_the_router_lookup_returns_none_outside_a_proxy(): 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): +def test_the_prune_filter_is_a_plain_dict(): """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.""" + chunk = ("dep-a", "dep-b") 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 + assert type(predicate["model"]) is dict + assert predicate["model"]["in"] == chunk + + +@pytest.mark.asyncio +async def test_a_run_that_scanned_nothing_issues_no_delete_statements(): + """The window where a master-key rotation wipes and recreates the model table. A run that + can see no deployment can reassess none of them, so it must not reach for the day's rows.""" + table = _FakeSentinelTable() + table.seed("t", DAY, "dep-orphan", 240.0, updated_at=datetime(2020, 1, 1, tzinfo=timezone.utc)) + + await run_ptu_flat_cost_rollup(_prisma_for([], table), target_date=DAY) + + assert table.delete_many_calls == [] + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-orphan") in table.rows @pytest.mark.asyncio