From 96080bdd24800f048e2e9358bb43adb60a549b57 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Wed, 19 Aug 2026 16:49:58 -0700 Subject: [PATCH] feat(ptu): accrue flat cost for PTU deployments declared in config.yaml The flat-cost rollup reads deployments from LiteLLM_ProxyModelTable, and config.yaml models never reach that table by design, so a PTU deployment declared there accrued no flat cost at all while still billing its traffic per token. The provider bills the reservation whichever file declared it. The rollup now also reads the deployments the router holds that no database row owns, identified by db_model, skipping the per-request credential clones that carry original_model_id and reuse their source's PTU config under a fresh id. Registering such a deployment zeroes its pricing, since reserved capacity already pays for the traffic it serves, and leaving a rate unset falls back to the public cost map, which makes the double charge the default rather than an opt-in. The rules both halves apply now live in one module. The rollup's test for what it will charge and the router's test for what to zero have to agree, or a deployment one accepts and the other declines serves its traffic for free. That module also owns the fields the write endpoints already zero, so the two paths cannot drift: tiered_pricing is emptied rather than zeroed because its tiers outrank the rates beside them, the search context table is written zeroed because an absent one means the provider default, and any further rate the deployment itself declares is zeroed alongside the standing set. The prune is bounded to the deployments a run scanned, but only for a run that priced a config-declared deployment. Deciding a row is garbage on staleness alone stays correct while every run derives its charges from the same table, so a database-only run sweeps exactly as it did before; once one host's charges come from a file the others cannot read, a row it never considered is not evidence of anything. Behaviour change worth calling out: a zeroed deployment sorts ahead of an unpriced sibling in QualityRouter's cost tiebreak, where an unset rate previously sorted last. Reserved capacity really is the cheaper choice, but the ordering moves. --- litellm/litellm_core_utils/ptu_pricing.py | 146 +++++++++ .../model_management_endpoints.py | 22 +- .../proxy/spend_tracking/ptu_feature_flag.py | 22 +- .../spend_tracking/ptu_flat_cost_rollup.py | 220 ++++++++----- litellm/router.py | 12 +- .../litellm_core_utils/test_ptu_pricing.py | 163 ++++++++++ .../test_ptu_flat_cost_rollup.py | 305 +++++++++++++++++- .../test_router_model_cost_isolation.py | 113 +++++++ 8 files changed, 887 insertions(+), 116 deletions(-) create mode 100644 litellm/litellm_core_utils/ptu_pricing.py create mode 100644 tests/test_litellm/litellm_core_utils/test_ptu_pricing.py diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py new file mode 100644 index 00000000000..fa14d43ea24 --- /dev/null +++ b/litellm/litellm_core_utils/ptu_pricing.py @@ -0,0 +1,146 @@ +"""Which deployments accrue PTU flat cost, and what that costs them per token. + +Reserved provisioned throughput is billed by the hour whether or not requests are sent, so +a deployment that accrues flat cost must not also bill per token. The two halves live here +together because they have to agree: a deployment the rollup declines to charge but the +router prices at zero serves its traffic for free. +""" + +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime, timezone +from types import MappingProxyType +from typing import Final + +from litellm.secret_managers.main import get_secret_bool +from litellm.types.router import ModelInfo +from litellm.types.utils import CustomPricingLiteLLMParams, MirroredPricingParams + +PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" + + +def is_ptu_cost_attribution_enabled() -> bool: + """Whether PTU flat-cost attribution is turned on for this process.""" + return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True + + +PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in MirroredPricingParams.model_fields if f != "tiered_pricing") + ( + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost_above_200k_tokens", +) +# tiered_pricing is emptied rather than zeroed: its tiers outrank the zeros written beside +# them, so a zero here would leave the cost map's tiers billing the traffic the reserved +# capacity already covers. +PTU_EMPTIED_PRICING_FIELDS: Final = frozenset(("tiered_pricing",)) +# search_context_cost_per_query holds its rates in a table keyed by context size, and an +# absent table means the provider's own default rather than free, so it is zeroed in place +# and written on every PTU deployment rather than only where a table is already stored. +PTU_ZEROED_TABLE_FIELDS: Final = frozenset(("search_context_cost_per_query",)) +SEARCH_CONTEXT_SIZES: Final = ("search_context_size_low", "search_context_size_medium", "search_context_size_high") +# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges, +# and zeroing one of those would destroy the deployment's configuration rather than stop a +# charge. +CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f) +PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()] | Mapping[str, float]]] = MappingProxyType( + { + **dict.fromkeys(PTU_ZEROED_PRICING_FIELDS, 0.0), + **dict.fromkeys(PTU_EMPTIED_PRICING_FIELDS, ()), + **dict.fromkeys(PTU_ZEROED_TABLE_FIELDS, MappingProxyType(dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0))), + } +) + + +@dataclass(frozen=True, slots=True) +class PTUTerms: + """The reservation a deployment declares, once every field has been validated.""" + + team_id: str + ptu_count: int + cost_per_ptu_per_hour: float + effective_from: datetime + effective_to: datetime | None + + +def _as_utc(value: object) -> datetime | None: + """A model_info datetime as UTC, parsing an ISO string, else None.""" + if isinstance(value, datetime): + parsed: Final = value + elif isinstance(value, str): + try: + parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) # rebind-ok: one parsed value, two sources + except ValueError: + return None + else: + return None + return parsed.replace(tzinfo=timezone.utc) if parsed.tzinfo is None else parsed.astimezone(timezone.utc) + + +def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None: + """The reservation this deployment accrues flat cost for, else None. + + A start is required rather than inferred because flat cost accrues from it, and a + present but unparseable bound would read as no bound and widen the window to the whole + day, so either one leaves the deployment unpriced until the config is fixed. + """ + ptu_count: Final = model_info.get("ptu_count") + cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour") + team_id: Final = model_info.get("team_id") + if ptu_count is None or cost_per_hour is None or not team_id: + return None + try: + ptu_count_int: Final = int(ptu_count) + cost_per_hour_float: Final = float(cost_per_hour) + except (TypeError, ValueError, OverflowError): + return None + if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT: + return None + if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR: + return None + + raw_from: Final = model_info.get("ptu_effective_from") + raw_to: Final = model_info.get("ptu_effective_to") + effective_from: Final = _as_utc(raw_from) + effective_to: Final = _as_utc(raw_to) + if effective_from is None or (raw_to is not None and effective_to is None): + return None + if effective_to is not None and effective_to <= effective_from: + return None + return PTUTerms( + team_id=str(team_id), + ptu_count=ptu_count_int, + cost_per_ptu_per_hour=cost_per_hour_float, + effective_from=effective_from, + effective_to=effective_to, + ) + + +def zeroed_ptu_pricing( + model_info: Mapping[str, object], declared: Mapping[str, object] +) -> Mapping[str, float | tuple[()] | Mapping[str, float]] | None: + """The pricing a deployment accruing flat cost must carry, else None. + + Both conditions hold or nothing is zeroed. Without the flag no flat cost accrues, so + zeroing would leave the deployment serving for free with nothing charged in its place, + which is what an SDK user who happens to carry ptu_count would otherwise get. The terms + are checked first only because they are a few dict reads, while the flag can resolve + through a configured secret manager, and this runs for every deployment registered. + + Any further rate the deployment itself declares is zeroed alongside the standing set, + since one left standing bills the traffic the reserved capacity already paid for. + """ + if ptu_terms(model_info) is None: + return None + if not is_ptu_cost_attribution_enabled(): + return None + return MappingProxyType( + { + **PTU_ZEROED_PRICING, + **dict.fromkeys( + CUSTOM_PRICING_FIELDS.intersection(declared) + .difference(PTU_ZEROED_TABLE_FIELDS) + .difference(PTU_EMPTIED_PRICING_FIELDS), + 0.0, + ), + } + ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ade24d194d2..1b49e2455e4 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -24,6 +24,13 @@ from pydantic import BaseModel, ConfigDict, Field, ValidationError from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.litellm_core_utils.ptu_pricing import ( + CUSTOM_PRICING_FIELDS, + PTU_EMPTIED_PRICING_FIELDS, + PTU_ZEROED_PRICING_FIELDS, + PTU_ZEROED_TABLE_FIELDS, + SEARCH_CONTEXT_SIZES, +) from litellm.proxy._types import ( BlockModelRequest, CommonProxyErrors, @@ -89,7 +96,6 @@ from litellm.types.router import ( ModelInfo, updateDeployment, ) -from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import get_utc_datetime router: Final = APIRouter() @@ -346,12 +352,8 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: # tiered_pricing is the one mirrored field that is a table of ranges, not a rate, so it is stored # empty (see _PTU_EMPTIED_PRICING_FIELDS): its tiers outrank the zeros written beside them, so # dropping it would leave the cost map's tiers billing the traffic the reserved capacity covers. -_PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in SPECIAL_MODEL_INFO_PARAMS if f != "tiered_pricing") + ( - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost_above_200k_tokens", -) -_PTU_EMPTIED_PRICING_FIELDS: Final = frozenset({"tiered_pricing"}) +_PTU_ZEROED_PRICING_FIELDS: Final = PTU_ZEROED_PRICING_FIELDS +_PTU_EMPTIED_PRICING_FIELDS: Final = PTU_EMPTIED_PRICING_FIELDS _PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()]]] = MappingProxyType( { **dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0), @@ -363,13 +365,13 @@ _EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE # Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges # (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of # those would destroy the deployment's configuration rather than stop a charge. -_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f) +_CUSTOM_PRICING_FIELDS: Final = CUSTOM_PRICING_FIELDS # search_context_cost_per_query holds its rates in a table keyed by context size, and an absent # table means the provider's own default rate rather than free (litellm/llms/gemini/cost_calculator # falls back to $0.035), so it is zeroed in place rather than emptied like tiered_pricing, and # written on every PTU deployment rather than only where a table is already stored. -_PTU_ZEROED_TABLE_FIELDS: Final = frozenset({"search_context_cost_per_query"}) -_SEARCH_CONTEXT_SIZES: Final = ("search_context_size_low", "search_context_size_medium", "search_context_size_high") +_PTU_ZEROED_TABLE_FIELDS: Final = PTU_ZEROED_TABLE_FIELDS +_SEARCH_CONTEXT_SIZES: Final = SEARCH_CONTEXT_SIZES def _is_nonzero_rate(value: object) -> bool: diff --git a/litellm/proxy/spend_tracking/ptu_feature_flag.py b/litellm/proxy/spend_tracking/ptu_feature_flag.py index 9078079b676..7f52dfa155d 100644 --- a/litellm/proxy/spend_tracking/ptu_feature_flag.py +++ b/litellm/proxy/spend_tracking/ptu_feature_flag.py @@ -1,18 +1,12 @@ -"""Opt-in flag for PTU (provisioned throughput unit) flat-cost attribution. +"""Re-exported from ``litellm.litellm_core_utils.ptu_pricing``. -The whole feature is inert unless an operator sets -``LITELLM_ENABLE_PTU_COST_ATTRIBUTION``: the daily rollup is not scheduled, the -model endpoints reject PTU config, the daily activity read path reports zero flat -cost, and the model form hides the PTU inputs. +The flag lives in core because the router reads it while registering a deployment, and +router code cannot import from the proxy. """ -from typing import Final +from litellm.litellm_core_utils.ptu_pricing import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) -from litellm.secret_managers.main import get_secret_bool - -PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" - - -def is_ptu_cost_attribution_enabled() -> bool: - """Report whether this deployment opted into PTU flat-cost attribution.""" - return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True +__all__ = ("PTU_COST_ATTRIBUTION_ENV_VAR", "is_ptu_cost_attribution_enabled") diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 25de9f6d065..eb5ac72de89 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -14,6 +14,7 @@ and share the existing unique constraint. import asyncio import json +import sys from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone @@ -29,14 +30,15 @@ from litellm.constants import ( PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY, ) +from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled -from litellm.types.router import ModelInfo if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager 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 @@ -72,28 +74,6 @@ class PTUModel: effective_to: datetime | None = None -def _parse_utc_datetime(value: object) -> datetime | None: - """Parse a model_info datetime (ISO string or datetime) into a UTC-aware datetime, else None.""" - parsed: Final = _coerce_datetime(value) - if parsed is None: - return None - if parsed.tzinfo is None: - return parsed.replace(tzinfo=timezone.utc) - return parsed.astimezone(timezone.utc) - - -def _coerce_datetime(value: object) -> datetime | None: - """``value`` as a datetime, parsing an ISO string, else None.""" - if isinstance(value, datetime): - return value - if not isinstance(value, str): - return None - try: - return datetime.fromisoformat(value.replace("Z", "+00:00")) - except ValueError: - return None - - def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: """The name an operator recognises for this deployment. @@ -167,46 +147,20 @@ def _parse_ptu_model(row: object) -> PTUModel | None: Valid means model_info has a positive ptu_count, a non-negative cost_per_ptu_per_hour, and a team_id (1 model -> 1 team). """ - raw_model_info: Final = getattr(row, "model_info", None) - model_info: Final = _decode_model_info(raw_model_info) + model_info: Final = _decode_model_info(getattr(row, "model_info", None)) if model_info is None: return None - ptu_count: Final = model_info.get("ptu_count") - cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour") - team_id: Final = model_info.get("team_id") - if ptu_count is None or cost_per_hour is None or not team_id: - return None - try: - ptu_count_int: Final = int(ptu_count) - cost_per_hour_float: Final = float(cost_per_hour) - except (TypeError, ValueError, OverflowError): - return None - if not 0 < ptu_count_int <= ModelInfo.MAX_PTU_COUNT: - return None - if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR: - return None - if model_info.get("ptu_effective_from") is None: - # The endpoints require a start; a row without one predates that rule or was - # written around them, and inferring one would bill days the deployment did not exist - return None - raw_from: Final = model_info.get("ptu_effective_from") - raw_to: Final = model_info.get("ptu_effective_to") - effective_from: Final = _parse_utc_datetime(raw_from) - effective_to: Final = _parse_utc_datetime(raw_to) - # A present-but-unparseable bound would read as "no bound" and silently widen the - # window to the whole day, so the deployment is skipped until the config is fixed - if (raw_from is not None and effective_from is None) or (raw_to is not None and effective_to is None): - return None - if effective_from is not None and effective_to is not None and effective_to <= effective_from: + terms: Final = ptu_terms(model_info) + if terms is None: return None return PTUModel( model_id=str(getattr(row, "model_id", "") or ""), model_name=_public_model_name(row, model_info), - team_id=str(team_id), - ptu_count=ptu_count_int, - cost_per_ptu_per_hour=cost_per_hour_float, - effective_from=effective_from, - effective_to=effective_to, + team_id=terms.team_id, + ptu_count=terms.ptu_count, + cost_per_ptu_per_hour=terms.cost_per_ptu_per_hour, + effective_from=terms.effective_from, + effective_to=terms.effective_to, ) @@ -358,10 +312,71 @@ 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] + config_sourced: bool + + +def _running_router() -> object | None: + """The proxy's router, or None outside a running proxy. + + Read out of ``sys.modules`` rather than imported, so a rollup driven from a test or a + script does not pull the whole proxy server in behind it. + """ + proxy_server: Final = sys.modules.get("litellm.proxy.proxy_server") + return getattr(proxy_server, "llm_router", None) if proxy_server is not None else None + + +def _config_deployments(router: object | None, *, owned_by_db: frozenset[str]) -> tuple[_PTUDeployment, ...]: + """Deployments the router holds that no ``LiteLLM_ProxyModelTable`` row owns. + + ``db_model`` is forced True on every deployment loaded from that table and defaults to + False on ModelInfo, so the complement is what config.yaml declared. A per-request + credential clone carries ``original_model_id`` and reuses its source's PTU config under + a fresh id, so pricing it would bill one reservation once per distinct client key. + """ + entries: Final = tuple(getattr(router, "model_list", None) or ()) + return tuple( + record + for entry in entries + if isinstance(entry, Mapping) + and isinstance(entry.get("model_info"), Mapping) + and entry["model_info"].get("db_model") is not True + and entry["model_info"].get("original_model_id") is None + for record in (_router_deployment(entry),) + if record is not None and record.model_id not in owned_by_db + ) + + +async def _load_ptu_models(prisma_client: "PrismaClient") -> _LoadedDeployments: + """Every deployment carrying valid manual PTU config, and every id the scan saw. + + Reserved capacity is billed by the provider whichever file declared it, so a + deployment the proxy only knows from config.yaml accrues alongside the stored ones. + """ 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) + db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or ""))) + config_records: Final = _config_deployments(_running_router(), owned_by_db=db_ids) + models: Final = tuple( + parsed for parsed in (_parse_ptu_model(row) for row in (*rows, *config_records)) if parsed is not None + ) + 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), + ) async def run_ptu_flat_cost_rollup( @@ -378,8 +393,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 +407,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 +433,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 if loaded.config_sourced else None, + ) verbose_proxy_logger.info( "PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed", @@ -524,7 +547,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 +730,61 @@ async def _prune_unrefreshed_sentinel_rows( *, date_str: str, run_started: datetime, + scanned_ids: frozenset[str] | None, ) -> 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. + + 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. + + 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. + """ 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 - } + 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,) + 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) + ) ) + deletions: Final = tuple( + [await prisma_client.db.litellm_dailyteamspend.delete_many(where=where) for where in filters] + ) + deleted: Final = sum(deletions) + if deleted: + verbose_proxy_logger.info( + "PTU rollup for %s: pruned %s stale sentinel row(s) across %s deployment(s)", + date_str, + deleted, + "every" if scanned_ids is None else len(scanned_ids), + ) __all__ = ( diff --git a/litellm/router.py b/litellm/router.py index b25b4f92467..91c8c123ba4 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -64,6 +64,7 @@ from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.ptu_pricing import zeroed_ptu_pricing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) @@ -7695,7 +7696,16 @@ class Router: - None: If the deployment is not active for the current environment (if 'supported_environments' is set in litellm_params) """ try: - litellm_params: Final[LiteLLM_Params] = LiteLLM_Params(**_litellm_params) + zeroed_pricing: Final = ( + zeroed_ptu_pricing(_model_info, _litellm_params) if _model_info.get("db_model") is not True else None + ) + litellm_params: Final[LiteLLM_Params] = LiteLLM_Params( + **( + _litellm_params + if zeroed_pricing is None + else MappingProxyType({**_litellm_params, **zeroed_pricing}) + ) + ) warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=_litellm_params) deployment = Deployment( **deployment_info, diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py new file mode 100644 index 00000000000..270c59f595f --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py @@ -0,0 +1,163 @@ +"""Tests for the shared PTU rules: which deployments accrue flat cost, and what that zeroes.""" + +import os +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest + +from litellm.litellm_core_utils.ptu_pricing import ( + CUSTOM_PRICING_FIELDS, + PTU_EMPTIED_PRICING_FIELDS, + PTU_ZEROED_PRICING_FIELDS, + PTU_ZEROED_TABLE_FIELDS, + SEARCH_CONTEXT_SIZES, + ptu_terms, + zeroed_ptu_pricing, +) +from litellm.types.router import ModelInfo + +_VALID = { + "team_id": "team-alpha", + "ptu_count": 100, + "cost_per_ptu_per_hour": 0.02, + "ptu_effective_from": "2026-01-01T00:00:00Z", +} + + +def _with_flag(model_info, declared=None, enabled=True): + with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True" if enabled else ""}, clear=False): + return zeroed_ptu_pricing(model_info, declared or {}) + + +def test_a_complete_reservation_is_accepted(): + terms = ptu_terms(_VALID) + + assert terms is not None + assert terms.team_id == "team-alpha" + assert terms.ptu_count == 100 + assert terms.effective_from == datetime(2026, 1, 1, tzinfo=timezone.utc) + assert terms.effective_to is None + + +@pytest.mark.parametrize( + "override", + [ + {"team_id": None}, + {"team_id": ""}, + {"ptu_count": None}, + {"cost_per_ptu_per_hour": None}, + {"ptu_count": 0}, + {"ptu_count": -1}, + {"ptu_count": ModelInfo.MAX_PTU_COUNT + 1}, + {"cost_per_ptu_per_hour": -0.01}, + {"cost_per_ptu_per_hour": ModelInfo.MAX_COST_PER_PTU_PER_HOUR + 1}, + {"ptu_count": "not-a-number"}, + {"ptu_effective_from": None}, + {"ptu_effective_from": "not-a-date"}, + {"ptu_effective_to": "not-a-date"}, + {"ptu_effective_to": "2025-01-01T00:00:00Z"}, + {"ptu_effective_to": "2026-01-01T00:00:00Z"}, + ], + ids=[ + "no team", + "blank team", + "no count", + "no rate", + "zero count", + "negative count", + "count over the cap", + "negative rate", + "rate over the cap", + "count not a number", + "no start", + "unparseable start", + "unparseable end", + "end before start", + "end equal to start", + ], +) +def test_an_incomplete_reservation_accrues_nothing(override): + """Anything the rollup declines to charge must also decline to be zeroed, or the + deployment serves its traffic for free with nothing charged in its place.""" + assert ptu_terms({**_VALID, **override}) is None + assert _with_flag({**_VALID, **override}) is None + + +def test_a_naive_start_is_read_as_utc(): + """config.yaml is hand-typed, and pydantic hands back a naive datetime for a date with + no offset.""" + terms = ptu_terms({**_VALID, "ptu_effective_from": datetime(2026, 5, 1, 12, 0)}) + + assert terms is not None + assert terms.effective_from == datetime(2026, 5, 1, 12, 0, tzinfo=timezone.utc) + + +def test_an_offset_start_is_converted_rather_than_relabelled(): + terms = ptu_terms({**_VALID, "ptu_effective_from": "2026-05-01T12:00:00-05:00"}) + + assert terms is not None + assert terms.effective_from == datetime(2026, 5, 1, 17, 0, tzinfo=timezone.utc) + + +def test_nothing_is_zeroed_while_the_feature_is_off(): + """No flat cost accrues with the flag off, so zeroing would serve the traffic free.""" + assert _with_flag(_VALID, enabled=False) is None + + +def test_the_standing_rates_are_all_zeroed(): + override = _with_flag(_VALID) + + assert override is not None + assert [field for field in PTU_ZEROED_PRICING_FIELDS if override[field] != 0.0] == [] + + +def test_tiered_pricing_is_emptied_rather_than_zeroed(): + """A tier outranks the flat rates written beside it, so a zero there would leave the + cost map's tiers billing the traffic the reserved capacity already covers.""" + override = _with_flag(_VALID, declared={"tiered_pricing": [{"range": [0, 1000], "input_cost_per_token": 0.003}]}) + + assert override is not None + for field in PTU_EMPTIED_PRICING_FIELDS: + assert override[field] == () + + +def test_the_search_context_table_is_zeroed_in_place_on_every_deployment(): + """An absent table means the provider's own default rather than free, so it is written + even when the deployment never declared one.""" + override = _with_flag(_VALID) + + assert override is not None + for field in PTU_ZEROED_TABLE_FIELDS: + assert dict(override[field]) == dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0) + + +def test_a_declared_table_does_not_become_a_scalar(): + """Zeroing it as a plain 0.0 would leave the provider's reader without a table to + consult, which is the same as absent.""" + override = _with_flag(_VALID, declared={"search_context_cost_per_query": {"search_context_size_medium": 0.05}}) + + assert override is not None + assert dict(override["search_context_cost_per_query"]) == dict.fromkeys(SEARCH_CONTEXT_SIZES, 0.0) + + +def test_a_rate_the_deployment_declares_itself_is_zeroed_too(): + """The standing set covers the mirrored rates. Anything else the operator wrote would + otherwise survive and bill the traffic the hourly charge already paid for.""" + extra = "input_cost_per_token_above_200k_tokens" + assert extra in CUSTOM_PRICING_FIELDS + assert extra not in PTU_ZEROED_PRICING_FIELDS + + override = _with_flag(_VALID, declared={extra: 9e-06}) + + assert override is not None + assert override[extra] == 0.0 + + +def test_a_setting_that_is_not_a_charge_is_left_alone(): + """CustomPricingLiteLLMParams also carries configuration, and zeroing one of those + would break the deployment rather than stop a charge.""" + override = _with_flag(_VALID, declared={"output_vector_size": 1536}) + + assert override is not None + assert "output_vector_size" not in override 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..333fd597b49 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 "model" not in where, "a database-only run has no reason to bound the sweep" @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,127 @@ 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_pricing_config_cannot_prune_a_row_it_did_not_scan(monkeypatch): + """Staleness alone stops being evidence once two hosts hold different configuration: a + row this run never considered belongs to a deployment another host is pricing from its + own file, 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)) + entry = _router_entry(model_id="cfg-here", model_info=dict(_VALID_PTU)) + 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), target_date=DAY) + + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-elsewhere") in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "cfg-here") in table.rows + assert table.delete_many_calls[-1]["model"]["in"] == ("cfg-here",) + + +@pytest.mark.asyncio +async def test_a_deployment_deleted_from_the_table_keeps_the_day_it_was_charged(monkeypatch): + """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") + monkeypatch.setattr( + ptu_rollup, "_running_router", lambda: _router_holding(_router_entry(model_id="cfg", model_info=dict(ptu))) + ) + + 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 "dep-doomed" not in table.delete_many_calls[-1]["model"]["in"] + + +@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.""" + table = _FakeSentinelTable() + table.seed("t", DAY, "dep-gone", 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-gone") not in table.rows + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-live") 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(4)] + monkeypatch.setattr( + ptu_rollup, "_running_router", lambda: _router_holding(_router_entry(model_id="dep-4", model_info=dict(ptu))) + ) + 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 @@ -1760,3 +1900,148 @@ async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled( assert result is None assert table.rows == {} assert table.upsert_keys == [] + + +# --- config.yaml deployments reach the rollup through the router ---------------- + + +def _router_holding(*entries): + """A stand-in for the proxy's router, carrying whatever model_list is passed.""" + return types.SimpleNamespace(model_list=list(entries)) + + +@pytest.mark.asyncio +async def test_a_config_declared_deployment_is_priced(monkeypatch): + """The whole point. A PTU deployment the proxy only knows from config.yaml is not in + LiteLLM_ProxyModelTable, so a DB-only scan bills the provider's reservation to nobody.""" + entry = _router_entry(model_id="cfg-1", model_name="gpt-4o-ptu", model_info=dict(_VALID_PTU)) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert [(m.model_id, m.model_name, m.team_id) for m in loaded.models] == [("cfg-1", "gpt-4o-ptu", "t")] + assert "cfg-1" in loaded.scanned_ids + + +@pytest.mark.asyncio +async def test_a_database_backed_router_entry_is_not_counted_twice(monkeypatch): + """Every deployment loaded from the table is also in the router, flagged db_model. Pricing + both copies would write two charges for one reservation.""" + row = _model_row(model_id="db-1", model_info=dict(_VALID_PTU)) + mirrored = _router_entry(model_id="db-1", model_info={**_VALID_PTU, "db_model": True}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(mirrored)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([row], _FakeSentinelTable())) + + assert [m.model_id for m in loaded.models] == ["db-1"] + + +@pytest.mark.asyncio +async def test_a_router_entry_sharing_an_id_with_the_table_is_priced_once(monkeypatch): + """db_model is data the router carries rather than something this module controls, so the + id anti-join is what actually maps onto the failure: two charges under one id.""" + row = _model_row(model_id="both-1", model_info=dict(_VALID_PTU)) + unflagged = _router_entry(model_id="both-1", model_info=dict(_VALID_PTU)) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(unflagged)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([row], _FakeSentinelTable())) + + assert [m.model_id for m in loaded.models] == ["both-1"] + + +@pytest.mark.asyncio +async def test_a_client_credential_clone_is_not_priced(monkeypatch): + """Supplying an api_key on a request mints a clone of the deployment under a fresh id, + carrying the source's PTU config. Pricing it bills one reservation per distinct caller key.""" + source = _router_entry(model_id="cfg-1", model_info=dict(_VALID_PTU)) + clone = _router_entry(model_id="cfg-1-clone", model_info={**_VALID_PTU, "original_model_id": "cfg-1"}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(source, clone)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert [m.model_id for m in loaded.models] == ["cfg-1"] + + +@pytest.mark.asyncio +async def test_a_config_deployment_without_ptu_config_is_scanned_but_not_priced(monkeypatch): + """It has to stay in the scanned set or its leftover sentinel rows become unprunable.""" + entry = _router_entry(model_id="cfg-plain", model_info={"team_id": "t"}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(entry)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert loaded.models == () + assert "cfg-plain" in loaded.scanned_ids + + +@pytest.mark.asyncio +async def test_no_router_in_the_process_prices_the_database_alone(monkeypatch): + """The rollup is importable and callable outside a running proxy.""" + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: None) + + loaded = await ptu_rollup._load_ptu_models( + _prisma_for([_model_row(model_id="db-1", model_info=dict(_VALID_PTU))], _FakeSentinelTable()) + ) + + assert [m.model_id for m in loaded.models] == ["db-1"] + + +@pytest.mark.asyncio +async def test_a_config_deployment_is_charged_end_to_end(monkeypatch): + """Through the scheduled entry point, so the charge lands in a sentinel row rather than + stopping at the loader.""" + table = _FakeSentinelTable() + entry = _router_entry(model_id="cfg-1", model_name="gpt-4o-ptu", model_info=dict(_VALID_PTU)) + 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), target_date=DAY) + + assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "cfg-1") in table.rows + + +@pytest.mark.asyncio +async def test_a_stale_database_backed_router_entry_is_not_treated_as_config(monkeypatch): + """The reconcile can leave a deployment on the router after its row is gone. The id + anti-join cannot see that one, so the flag is what keeps it from being priced as though + config.yaml had declared it.""" + stale = _router_entry(model_id="db-gone", model_info={**_VALID_PTU, "db_model": True}) + monkeypatch.setattr(ptu_rollup, "_running_router", lambda: _router_holding(stale)) + + loaded = await ptu_rollup._load_ptu_models(_prisma_for([], _FakeSentinelTable())) + + assert loaded.models == () + + +def test_the_router_lookup_reads_the_proxys_own_global(): + """Every other config test replaces this helper, so without one test driving the real + body a typo in the module path or the attribute name leaves the whole feature dead in + production with the suite still green.""" + import sys + import types as _types + + assert ptu_rollup._running_router() is None or "litellm.proxy.proxy_server" in sys.modules + + sentinel = object() + stub = _types.SimpleNamespace(llm_router=sentinel) + real = sys.modules.get("litellm.proxy.proxy_server") + sys.modules["litellm.proxy.proxy_server"] = stub + try: + assert ptu_rollup._running_router() is sentinel + del stub.llm_router + assert ptu_rollup._running_router() is None + finally: + if real is None: + del sys.modules["litellm.proxy.proxy_server"] + else: + sys.modules["litellm.proxy.proxy_server"] = real + + +def test_the_router_lookup_returns_none_outside_a_proxy(): + import sys + + real = sys.modules.pop("litellm.proxy.proxy_server", None) + try: + assert ptu_rollup._running_router() is None + finally: + if real is not None: + sys.modules["litellm.proxy.proxy_server"] = real diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index dfe46d54ab8..4674a8b1dfa 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -1584,3 +1584,116 @@ def test_inherit_builtin_tiered_output_rate_leaves_a_user_rate_alone(): ) assert model_info["output_cost_per_token"] == 9e-07 + + +# --- a config.yaml PTU deployment must not also bill per token ------------------ + +_PTU_MODEL_INFO = { + "team_id": "team-alpha", + "ptu_count": 100, + "cost_per_ptu_per_hour": 0.02, + "ptu_effective_from": "2026-01-01T00:00:00Z", +} + + +def _ptu_router(model_info=None, litellm_params=None, ptu_enabled=True): + """A router built the way loading config.yaml builds one.""" + with patch.dict(os.environ, {"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True" if ptu_enabled else ""}, clear=False): + return Router( + model_list=[ + { + "model_name": "gpt-4o-ptu", + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5-20250929", + "api_key": "sk-not-used", + **(litellm_params or {}), + }, + "model_info": dict(_PTU_MODEL_INFO if model_info is None else model_info), + } + ] + ) + + +def test_a_config_ptu_deployment_bills_nothing_per_token(): + """Reserved capacity is already billed by the hour, so charging its traffic bills the + same tokens twice. Left unset the rate falls back to the public cost map, which makes + the double charge the default rather than an opt-in.""" + router = _ptu_router(litellm_params={"input_cost_per_token": 5e-06, "output_cost_per_token": 1.5e-05}) + entry = router.model_list[0] + + assert entry["litellm_params"]["input_cost_per_token"] == 0.0 + assert entry["litellm_params"]["output_cost_per_token"] == 0.0 + assert entry["model_info"]["input_cost_per_token"] == 0.0 + assert litellm.model_cost[entry["model_info"]["id"]]["input_cost_per_token"] == 0.0 + + +@pytest.mark.parametrize( + "backend", + ["anthropic/claude-sonnet-4-5-20250929", "azure/gpt-4o", "gemini/gemini-2.5-flash"], +) +def test_a_config_ptu_deployment_imports_no_cache_rate_from_its_backend(backend): + """The cache back-fill runs whenever input_cost_per_token is set, and 0.0 is set, so a + partially zeroed deployment would silently inherit the backend model's real cache rates. + Every backend here publishes non-zero ones, which is what makes the assertion mean + something.""" + cache_fields = ( + "cache_creation_input_token_cost", + "cache_creation_input_token_cost_above_1hr", + "cache_creation_input_token_cost_above_200k_tokens", + "cache_read_input_token_cost", + "cache_read_input_token_cost_above_200k_tokens", + ) + builtin = litellm.get_model_info(model=backend) + assert any(builtin.get(field) for field in cache_fields), "backend publishes no cache pricing to leak" + + router = _ptu_router(litellm_params={"model": backend}) + priced = litellm.model_cost[router.model_list[0]["model_info"]["id"]] + + assert [field for field in cache_fields if priced.get(field)] == [] + + +def test_zeroing_a_ptu_deployment_leaves_its_backend_model_priced(): + """A sibling deployment on the same backend must keep billing normally.""" + backend = "anthropic/claude-sonnet-4-5-20250929" + builtin = litellm.get_model_info(model=backend)["input_cost_per_token"] + assert builtin > 0 + + _ptu_router(litellm_params={"model": backend}) + + assert litellm.get_model_info(model=backend)["input_cost_per_token"] == builtin + + +def test_zeroing_does_not_change_the_deployment_id(): + """The id is a hash of the deployment's params and keys its cooldowns, its budget, and + every spend row already written against it.""" + params = {"input_cost_per_token": 5e-06} + priced = _ptu_router(litellm_params=params, ptu_enabled=False).model_list[0]["model_info"]["id"] + zeroed = _ptu_router(litellm_params=params).model_list[0]["model_info"]["id"] + + assert priced == zeroed + + +def test_a_database_backed_deployment_is_left_alone(): + """The write endpoints already zero those, and they answer 400 rather than silently + rewriting a rate the caller sent.""" + entry = _ptu_router(model_info={**_PTU_MODEL_INFO, "db_model": True}).model_list[0] + + assert entry["litellm_params"].get("input_cost_per_token") is None + + +def test_nothing_is_zeroed_while_the_feature_is_off(): + """No flat cost accrues with the flag off, so zeroing would serve the traffic free.""" + entry = _ptu_router(litellm_params={"input_cost_per_token": 5e-06}, ptu_enabled=False).model_list[0] + + assert entry["litellm_params"]["input_cost_per_token"] == 5e-06 + + +@pytest.mark.parametrize("dropped", ["team_id", "ptu_effective_from"], ids=["no team_id", "no ptu_effective_from"]) +def test_a_deployment_the_rollup_will_not_charge_is_not_zeroed(dropped): + """The rollup refuses to price a reservation missing either field, so zeroing on the + looser count-and-rate test alone would leave the deployment serving for free with + nothing charged in its place.""" + incomplete = {k: v for k, v in _PTU_MODEL_INFO.items() if k != dropped} + entry = _ptu_router(model_info=incomplete, litellm_params={"input_cost_per_token": 5e-06}).model_list[0] + + assert entry["litellm_params"]["input_cost_per_token"] == 5e-06