diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 25de9f6d065..7332b7f273d 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -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__ = ( 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 f41756d2b87..94f7ed42acc 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 @@ -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