Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/revert-pr-37554-migration-5af8ea

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),
         )
 
 

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