From 2d49656a2ac6e6dfc14831757933c184b3ff7081 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 13:37:08 -0700 Subject: [PATCH 01/18] feat(proxy): split Azure PTU deployments across teams by share A PTU deployment's model_info takes ptu_shares (team_id -> whole PTUs adding up to ptu_count) in place of a single team_id. With LITELLM_ENABLE_PTU_COST_ATTRIBUTION on, routing hides the deployment from teams holding no share, the v3 limiter enforces each share as a per-team TPM ceiling from Azure's per-model TPM-per-PTU sizing, the daily flat-cost rollup splits the hourly cost by share, and /team/daily/activity plus the Team Usage page report PTU-hours next to tokens. --- .../litellm_core_utils/azure_ptu_capacity.py | 104 ++++++++++++ litellm/litellm_core_utils/ptu_pricing.py | 88 +++++++++- litellm/proxy/_lazy_openapi_snapshot.json | 10 ++ .../hooks/parallel_request_limiter_v3.py | 90 +++++++++++ .../management_endpoints/ptu_consumption.py | 95 +++++++++++ .../management_endpoints/team_endpoints.py | 22 ++- .../spend_tracking/ptu_flat_cost_rollup.py | 61 +++---- litellm/router.py | 21 ++- litellm/router_utils/ptu_shares.py | 101 ++++++++++++ .../common_daily_activity.py | 2 + litellm/types/router.py | 1 + .../test_azure_ptu_capacity.py | 88 ++++++++++ .../litellm_core_utils/test_ptu_pricing.py | 81 +++++++++- .../hooks/test_parallel_request_limiter_v3.py | 151 ++++++++++++++++++ .../test_ptu_consumption.py | 130 +++++++++++++++ .../test_ptu_model_settings.py | 46 ++++++ .../test_team_endpoints.py | 87 ++++++++++ .../test_ptu_flat_cost_rollup.py | 58 ++++++- .../router_utils/test_ptu_shares.py | 115 +++++++++++++ tests/test_litellm/test_router.py | 69 ++++++++ .../components/EntityUsage/EntityUsage.tsx | 16 +- .../EntityUsage/entityUsageSummary.test.ts | 30 +++- .../EntityUsage/entityUsageSummary.ts | 22 ++- .../hooks/usePaginatedDailyActivity.test.ts | 1 + .../hooks/usePaginatedDailyActivity.ts | 1 + .../src/components/EntityUsageExport/types.ts | 2 + .../EntityUsageExport/utils.test.ts | 17 ++ .../src/components/EntityUsageExport/utils.ts | 3 + .../src/components/UsagePage/types.ts | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 14 ++ 30 files changed, 1479 insertions(+), 48 deletions(-) create mode 100644 litellm/litellm_core_utils/azure_ptu_capacity.py create mode 100644 litellm/proxy/management_endpoints/ptu_consumption.py create mode 100644 litellm/router_utils/ptu_shares.py create mode 100644 tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py create mode 100644 tests/test_litellm/router_utils/test_ptu_shares.py diff --git a/litellm/litellm_core_utils/azure_ptu_capacity.py b/litellm/litellm_core_utils/azure_ptu_capacity.py new file mode 100644 index 00000000000..3ba4b9b6abf --- /dev/null +++ b/litellm/litellm_core_utils/azure_ptu_capacity.py @@ -0,0 +1,104 @@ +"""How much throughput one Azure OpenAI provisioned throughput unit (PTU) serves per model. + +Azure sizes a provisioned deployment in normalized tokens per minute: +``input TPM x (1 - cache hit rate) + output-to-input ratio x output TPM``, divided by the +model's "Input TPM per PTU" to get the PTUs required. The same two numbers turn a team's +PTU share into a per-minute token ceiling and a request's usage back into PTU-hours. + +Table read from +https://learn.microsoft.com/en-us/azure/foundry/openai/how-to/provisioned-throughput-sizing#deployment-parameters-and-throughput-values-by-model +on 2026-09-24 (page dated 2026-09-23). Azure deducts cached input tokens in full for every +model except the GPT-6 family, where a cached input token costs a tenth of an uncached one. +""" + +import re +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + + +@dataclass(frozen=True, slots=True) +class PTUCapacity: + input_tpm_per_ptu: int + output_to_input_ratio: float + cached_input_ratio: float = 0.0 + + @property + def normalized_tokens_per_ptu_hour(self) -> int: + return self.input_tpm_per_ptu * 60 + + +AZURE_PTU_CAPACITY: Final[Mapping[str, PTUCapacity]] = MappingProxyType( + { + "gpt-6-sol": PTUCapacity(3_000, 5.0, cached_input_ratio=0.1), + "gpt-6-astra": PTUCapacity(600, 5.0, cached_input_ratio=0.1), + "gpt-5.6-luna": PTUCapacity(30_000, 6.0), + "gpt-5.6-terra": PTUCapacity(3_000, 6.0), + "gpt-5.6-sol": PTUCapacity(1_200, 6.0), + "gpt-5.5": PTUCapacity(1_200, 6.0), + "gpt-5.4": PTUCapacity(2_400, 6.0), + "gpt-5.4-mini": PTUCapacity(7_900, 6.0), + "gpt-5.3-codex": PTUCapacity(3_400, 8.0), + "gpt-5.2": PTUCapacity(3_400, 8.0), + "gpt-5.2-codex": PTUCapacity(3_400, 8.0), + "gpt-5.1": PTUCapacity(4_750, 8.0), + "gpt-5.1-codex": PTUCapacity(4_750, 8.0), + "gpt-5": PTUCapacity(4_750, 8.0), + "gpt-5-mini": PTUCapacity(23_750, 8.0), + "gpt-4.1": PTUCapacity(3_000, 4.0), + "gpt-4.1-mini": PTUCapacity(14_900, 4.0), + "gpt-4.1-nano": PTUCapacity(59_400, 4.0), + "o3": PTUCapacity(3_000, 4.0), + "o4-mini": PTUCapacity(5_400, 4.0), + "gpt-4o": PTUCapacity(2_500, 4.0), + "gpt-4o-mini": PTUCapacity(37_000, 4.0), + "o3-mini": PTUCapacity(2_500, 4.0), + "o1": PTUCapacity(230, 4.0), + "llama-3.3-70b-instruct": PTUCapacity(8_450, 4.0), + } +) + +_VERSION_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$") + + +def azure_ptu_capacity(model: str) -> PTUCapacity | None: + """The sizing row for ``model``, read as its last path segment with a dated version dropped. + + ``azure/gpt-4.1-2025-04-14`` and ``gpt-4.1`` both resolve to the ``gpt-4.1`` row; a + deployment name that is not a model name resolves to nothing, which is why callers + prefer ``model_info.base_model``. + """ + name: Final = model.rsplit("/", 1)[-1].strip().lower() + return AZURE_PTU_CAPACITY.get(name) or AZURE_PTU_CAPACITY.get(_VERSION_SUFFIX.sub("", name)) + + +def deployment_ptu_capacity(deployment: Mapping[str, object]) -> PTUCapacity | None: + """The sizing row a deployment resolves to: ``model_info.base_model`` first, since an + Azure deployment name is arbitrary, then ``litellm_params.model``.""" + model_info: Final = deployment.get("model_info") + litellm_params: Final = deployment.get("litellm_params") + candidates: Final = tuple( + value + for value in ( + model_info.get("base_model") if isinstance(model_info, Mapping) else None, + litellm_params.get("model") if isinstance(litellm_params, Mapping) else None, + ) + if isinstance(value, str) and value + ) + return next((capacity for capacity in map(azure_ptu_capacity, candidates) if capacity is not None), None) + + +def normalized_tokens( + capacity: PTUCapacity, *, prompt_tokens: int, completion_tokens: int, cache_read_tokens: int = 0 +) -> float: + """Azure's normalized token count for one request: uncached input in full, cached input + at the model's cached ratio, output weighted by the output-to-input ratio.""" + cached: Final = min(max(cache_read_tokens, 0), max(prompt_tokens, 0)) + uncached: Final = max(prompt_tokens, 0) - cached + return uncached + capacity.cached_input_ratio * cached + capacity.output_to_input_ratio * max(completion_tokens, 0) + + +def ptu_hours(capacity: PTUCapacity, normalized: float) -> float: + """PTU-hours ``normalized`` tokens amount to: one PTU serves its input TPM for sixty minutes.""" + return normalized / capacity.normalized_tokens_per_ptu_hour diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py index 80f7a822b96..b68b5cc93c2 100644 --- a/litellm/litellm_core_utils/ptu_pricing.py +++ b/litellm/litellm_core_utils/ptu_pricing.py @@ -56,9 +56,13 @@ PTU_ZEROED_PRICING: Final[Mapping[str, float | tuple[()] | Mapping[str, float]]] @dataclass(frozen=True, slots=True) class PTUTerms: - """The reservation a deployment declares, once every field has been validated.""" + """The reservation a deployment declares, once every field has been validated. - team_id: str + ``shares`` maps every team the capacity is attributed to onto its PTUs and adds up to + ``ptu_count``: a deployment declaring a single ``team_id`` holds the whole count under it. + """ + + shares: Mapping[str, int] ptu_count: int cost_per_ptu_per_hour: float effective_from: datetime @@ -127,7 +131,44 @@ def ptu_identity_error( return None -PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to") +PTU_MODEL_INFO_FIELDS: Final = ( + "ptu_count", + "cost_per_ptu_per_hour", + "ptu_effective_from", + "ptu_effective_to", + "ptu_shares", +) + + +def parsed_ptu_shares(raw: object) -> Mapping[str, int] | None: + """``ptu_shares`` as team id -> whole PTUs, else None when empty or any entry is unusable. + + A share is a count of reserved units, so it has to be a positive integer; ``bool`` is + excluded because it is an ``int`` subclass and ``True`` would read as one PTU. + """ + if not isinstance(raw, Mapping) or not raw: + return None + entries: Final = tuple((str(team_id), share) for team_id, share in raw.items()) + if any( + not team_id or isinstance(share, bool) or not isinstance(share, int) or share <= 0 for team_id, share in entries + ): + return None + return MappingProxyType(dict(entries)) + + +def _declared_shares(model_info: Mapping[str, object], ptu_count: int) -> Mapping[str, int] | None: + """Who holds the capacity: the single ``team_id`` holding all of it, or the ``ptu_shares`` + that add up to it, else None.""" + team_id: Final = model_info.get("team_id") + raw_shares: Final = model_info.get("ptu_shares") + if team_id and raw_shares is None: + return MappingProxyType({str(team_id): ptu_count}) + if team_id or raw_shares is None: + return None + shares: Final = parsed_ptu_shares(raw_shares) + if shares is None or sum(shares.values()) != ptu_count: + return None + return shares def declares_ptu(model_info: Mapping[str, object]) -> bool: @@ -164,8 +205,37 @@ def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None "today could be billed for days it did not exist", model_name, ) - if not model_info.get("team_id"): - return _named("team_id is required when PTU fields are set (one model maps to one team)", model_name) + return _ptu_holder_error(model_info, model_name) + + +def _ptu_holder_error(model_info: Mapping[str, object], model_name: str | None) -> str | None: + """Why the teams this capacity is attributed to cannot be read, else None. + + The shares have to add up to the count exactly: a shortfall would leave PTU-hours the + provider bills attributed to nobody, and a surplus would attribute capacity that was + never reserved. + """ + team_id: Final = model_info.get("team_id") + raw_shares: Final = model_info.get("ptu_shares") + if team_id and raw_shares is not None: + return _named( + "team_id and ptu_shares cannot both be set; ptu_shares lists every team the capacity is split across", + model_name, + ) + if not team_id and raw_shares is None: + return _named("team_id or ptu_shares is required when PTU fields are set", model_name) + if raw_shares is None: + return None + shares: Final = parsed_ptu_shares(raw_shares) + if shares is None: + return _named("ptu_shares must map at least one team_id to a positive whole number of PTUs", model_name) + try: + ptu_count: Final = int(str(model_info.get("ptu_count"))) + except ValueError: + return None + allocated: Final = sum(shares.values()) + if allocated != ptu_count: + return _named(f"ptu_shares must add up to ptu_count ({allocated} of {ptu_count} allocated)", model_name) return None @@ -178,8 +248,7 @@ def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | 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: + if ptu_count is None or cost_per_hour is None: return None try: ptu_count_int: Final = int(ptu_count) @@ -199,8 +268,11 @@ def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None: return None if effective_to is not None and effective_to <= effective_from: return None + shares: Final = _declared_shares(model_info, ptu_count_int) + if shares is None: + return None return PTUTerms( - team_id=str(team_id), + shares=shares, ptu_count=ptu_count_int, cost_per_ptu_per_hour=cost_per_hour_float, effective_from=effective_from, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 0b43c3864ab..f767cf7d5f0 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3177,6 +3177,11 @@ "title": "Total Prompt Tokens", "type": "integer" }, + "total_ptu_hours": { + "default": 0.0, + "title": "Total Ptu Hours", + "type": "number" + }, "total_response_time_ms": { "default": 0, "title": "Total Response Time Ms", @@ -3720,6 +3725,11 @@ "title": "Prompt Tokens", "type": "integer" }, + "ptu_hours": { + "default": 0.0, + "title": "Ptu Hours", + "type": "number" + }, "spend": { "default": 0.0, "title": "Spend", diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 17cf7382246..c9f8e27aa5b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -36,6 +36,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) +from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -64,6 +65,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( response_has_hidden_params, ) from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.router_utils.ptu_shares import PTUTeamCeiling, team_ptu_ceiling from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -112,6 +114,14 @@ def _resolve_model_group_alias_via_proxy_router(model: str) -> str | None: return resolve_model_group_alias(llm_router.model_group_alias, model) +def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) -> PTUTeamCeiling | None: + from litellm.proxy.proxy_server import llm_router + + if llm_router is None or not is_ptu_cost_attribution_enabled(): + return None + return team_ptu_ceiling(llm_router.get_model_list(model_name=model_group) or (), team_id) + + def _sibling_counter_keys(window_key: str) -> tuple[str, str]: prefix: Final = window_key.removesuffix(":window") return f"{prefix}:requests", f"{prefix}:tokens" @@ -420,6 +430,10 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600 # on the same project+model simultaneously without colliding on cache keys. PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" +# Descriptor "key" for a team's PTU share of a shared Azure provisioned deployment, +# counted in Azure normalized tokens (output weighted by the model's ratio) so it +# never collides with the raw-token "model_per_team" counter on the same team+model. +PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu" # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request @@ -648,10 +662,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): internal_usage_cache: InternalUsageCache, time_provider: Callable[[], datetime] | None = None, model_group_resolver: Callable[[str], str | None] = _resolve_model_group_alias_via_proxy_router, + ptu_team_ceiling_resolver: Callable[ + [str, str], PTUTeamCeiling | None + ] = _resolve_ptu_team_ceiling_via_proxy_router, ): self.internal_usage_cache = internal_usage_cache self._time_provider = time_provider or datetime.now self._model_group_resolver = model_group_resolver + self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver if self.internal_usage_cache.dual_cache.redis_cache is not None: self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( BATCH_RATE_LIMITER_SCRIPT @@ -2879,6 +2897,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model if isinstance(requested_model, str) else None, descriptors=descriptors, ) + self._add_team_ptu_rate_limit_descriptor( + user_api_key_dict=user_api_key_dict, + requested_model=requested_model if isinstance(requested_model, str) else None, + descriptors=descriptors, + ) # Agent-level and session-level rate limits resolved_agent_id: Final = self._get_resolved_agent_id(user_api_key_dict, data) @@ -3017,6 +3040,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + def _add_team_ptu_rate_limit_descriptor( + self, + user_api_key_dict: UserAPIKeyAuth, + requested_model: str | None, + descriptors: list[RateLimitDescriptor], + ) -> None: + model: Final = self._rate_limited_model(requested_model) + if model is None or not user_api_key_dict.team_id: + return + ceiling: Final = self._ptu_team_ceiling_resolver(user_api_key_dict.team_id, model.group) + if ceiling is None: + return + descriptors.append( + RateLimitDescriptor( + key=PTU_TEAM_DESCRIPTOR_KEY, + value=f"{user_api_key_dict.team_id}:{model.group}", + rate_limit=RateLimitDescriptorRateLimitObject( + requests_per_unit=None, tokens_per_unit=ceiling.tpm_limit, window_size=self.window_size + ), + ) + ) + def _add_project_model_rate_limit_descriptor_from_metadata( self, user_api_key_dict: UserAPIKeyAuth, @@ -4550,9 +4595,54 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens=reserved_tokens, ) ) + pipeline_operations.extend( + self._build_team_ptu_tpm_ops( + standard_logging_metadata=standard_logging_metadata, + response_obj=response_obj, + reconcile_model=reconcile_model, + reserved_scopes=reserved_scopes, + reserved_tokens=reserved_tokens, + total_tokens=total_tokens, + ) + ) return pipeline_operations + def _build_team_ptu_tpm_ops( + self, + standard_logging_metadata: Mapping[str, object], + response_obj: object, + reconcile_model: RateLimitedModel | None, + reserved_scopes: Set[tuple[str, str]], + reserved_tokens: int, + total_tokens: int, + ) -> Sequence[RedisPipelineIncrementOperation]: + """Settle the team's PTU counter in Azure normalized tokens: uncached input in full plus + output weighted by the model's output-to-input ratio, the way Azure sizes a PTU. + + The pre-call reservation was raw estimated tokens, so this is the same reconcile as the + other TPM scopes with a weighted actual; when usage cannot be resolved it charges the raw + total the other scopes charge. + """ + team_id: Final = standard_logging_metadata.get("user_api_key_team_id") + if reconcile_model is None or not isinstance(team_id, str) or not team_id: + return () + ceiling: Final = self._ptu_team_ceiling_resolver(team_id, reconcile_model.group) + if ceiling is None: + return () + billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage(response_obj) + normalized: Final = ( + billable_input + round(ceiling.output_to_input_ratio * completion_tokens) + if usage_resolved + else total_tokens + ) + return self._build_reservation_aware_tpm_ops( + targets=((PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{reconcile_model.group}"),), + reserved_scopes=reserved_scopes, + actual_tokens=normalized, + reserved_tokens=reserved_tokens, + ) + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Update TPM usage on successful API calls by incrementing counters using pipeline diff --git a/litellm/proxy/management_endpoints/ptu_consumption.py b/litellm/proxy/management_endpoints/ptu_consumption.py new file mode 100644 index 00000000000..12736b312a8 --- /dev/null +++ b/litellm/proxy/management_endpoints/ptu_consumption.py @@ -0,0 +1,95 @@ +"""PTU-hours a team's tokens amount to, attached to a daily activity response. + +Azure sizes a provisioned deployment in normalized tokens per minute per PTU, so the prompt, +cached, and completion tokens a team sent to a PTU model group convert back to the share of +a PTU-hour it consumed, reported next to the raw token counts. +""" + +from collections.abc import Callable +from types import MappingProxyType +from typing import Final + +from litellm.litellm_core_utils.azure_ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + DailySpendData, + MetricWithMetadata, + SpendAnalyticsPaginatedResponse, + SpendMetrics, +) + + +def _with_ptu_hours(metrics: SpendMetrics, capacity: PTUCapacity) -> SpendMetrics: + consumed: Final = ptu_hours( + capacity, + normalized_tokens( + capacity, + prompt_tokens=metrics.prompt_tokens, + completion_tokens=metrics.completion_tokens, + cache_read_tokens=metrics.cache_read_input_tokens, + ), + ) + return metrics.model_copy(update=MappingProxyType({"ptu_hours": consumed})) + + +def _model_group_with_ptu_hours(bucket: MetricWithMetadata, capacity: PTUCapacity) -> MetricWithMetadata: + api_key_breakdown: Final = { # mutable-ok: pydantic serializes a dict[...] field only from a plain dict + api_key: key_bucket.model_copy( + update=MappingProxyType({"metrics": _with_ptu_hours(key_bucket.metrics, capacity)}) + ) + for api_key, key_bucket in bucket.api_key_breakdown.items() + } + return bucket.model_copy( + update=MappingProxyType( + {"metrics": _with_ptu_hours(bucket.metrics, capacity), "api_key_breakdown": api_key_breakdown} + ) + ) + + +def _day_with_ptu_hours( + day: DailySpendData, capacity_for_model_group: Callable[[str], PTUCapacity | None] +) -> DailySpendData: + priced: Final = MappingProxyType( + { + model_group: _model_group_with_ptu_hours(bucket, capacity) + for model_group, bucket in day.breakdown.model_groups.items() + if (capacity := capacity_for_model_group(model_group)) is not None + } + ) + if not priced: + return day + model_groups: Final = { # mutable-ok: pydantic serializes a dict[...] field only from a plain dict + **day.breakdown.model_groups, + **priced, + } + return day.model_copy( + update=MappingProxyType( + { + "metrics": day.metrics.model_copy( + update=MappingProxyType({"ptu_hours": sum(bucket.metrics.ptu_hours for bucket in priced.values())}) + ), + "breakdown": day.breakdown.model_copy(update=MappingProxyType({"model_groups": model_groups})), + } + ) + ) + + +def attach_ptu_hours( + response: SpendAnalyticsPaginatedResponse, capacity_for_model_group: Callable[[str], PTUCapacity | None] +) -> SpendAnalyticsPaginatedResponse: + """The response with ``ptu_hours`` filled in on every PTU model group and its api keys, + on each day, and on the total, from the tokens already on the page. + + A model group the resolver has no sizing row for keeps ``ptu_hours`` at zero. + """ + days: Final = tuple(_day_with_ptu_hours(day, capacity_for_model_group) for day in response.results) + results: Final = list(days) # mutable-ok: pydantic serializes a list[...] field only from a plain list + return response.model_copy( + update=MappingProxyType( + { + "results": results, + "metadata": response.metadata.model_copy( + update=MappingProxyType({"total_ptu_hours": sum(day.metrics.ptu_hours for day in days)}) + ), + } + ) + ) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 6cec3e714ec..d32e63b3a6b 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -41,6 +41,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.integrations.prometheus import PrometheusLogger +from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( UI_TEAM_ID, @@ -143,6 +144,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, ) +from litellm.proxy.management_endpoints.ptu_consumption import attach_ptu_hours from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, @@ -194,6 +196,7 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.router import Router +from litellm.router_utils.ptu_shares import model_group_ptu_capacity from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, @@ -6647,6 +6650,17 @@ async def _resolve_team_daily_activity_scope( ) +def _with_ptu_consumption( + activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None +) -> SpendAnalyticsPaginatedResponse: + if llm_router is None or not is_ptu_cost_attribution_enabled(): + return activity + return attach_ptu_hours( + activity, + lambda model_group: model_group_ptu_capacity(llm_router.get_model_list(model_name=model_group) or ()), + ) + + @router.get( "/team/daily/activity", response_model=SpendAnalyticsPaginatedResponse, @@ -6679,6 +6693,7 @@ async def get_team_daily_activity( SpendAnalyticsPaginatedResponse: Paginated response containing daily activity data. """ from litellm.proxy.proxy_server import ( + llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -6697,7 +6712,7 @@ async def get_team_daily_activity( proxy_logging_obj=proxy_logging_obj, ) - return await get_daily_activity( + activity: Final = await get_daily_activity( prisma_client=prisma_client, table_name="litellm_dailyteamspend", entity_id_field="team_id", @@ -6711,6 +6726,7 @@ async def get_team_daily_activity( page=page, page_size=page_size, ) + return _with_ptu_consumption(activity, llm_router) _MAX_AGGREGATED_RANGE_DAYS: Final = 400 @@ -6767,6 +6783,7 @@ async def get_team_daily_activity_aggregated( SpendAnalyticsPaginatedResponse: Response containing all daily activity data for the range. """ from litellm.proxy.proxy_server import ( + llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -6789,7 +6806,7 @@ async def get_team_daily_activity_aggregated( proxy_logging_obj=proxy_logging_obj, ) - return await get_daily_activity_aggregated( + activity: Final = await get_daily_activity_aggregated( prisma_client=prisma_client, table_name="litellm_dailyteamspend", entity_id_field="team_id", @@ -6803,6 +6820,7 @@ async def get_team_daily_activity_aggregated( timezone_offset_minutes=timezone, include_entity_breakdown=True, ) + return _with_ptu_consumption(activity, llm_router) def _team_user_spend_sql(*, team_count: int, restrict_to_user: bool) -> str: diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index b8af432029f..e2c0ccb0327 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -82,7 +82,8 @@ class BackfillResult: @dataclass(frozen=True, slots=True) class PTUModel: - """A model deployment carrying valid manual PTU config.""" + """One team's holding on a deployment carrying valid manual PTU config: the whole count + under a single team_id, or that team's share of a split deployment.""" model_id: str model_name: str @@ -98,10 +99,10 @@ def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: Creating a team-scoped deployment rewrites model_name to a synthetic routing key (``model_name__``) and keeps the chosen name in - ``model_info.team_public_model_name``. PTU config is only accepted alongside a - team_id, so every PTU deployment carries that synthetic name; keying the sentinel - row on it would file each charge under a UUID that no usage view can resolve and - that never lines up with the same model's request rows. + ``model_info.team_public_model_name``. A PTU deployment held by one team_id carries + that synthetic name; keying the sentinel row on it would file each charge under a + UUID that no usage view can resolve and that never lines up with the same model's + request rows. A deployment split into ptu_shares keeps the name it was declared with. """ public_name: Final = model_info.get("team_public_model_name") if isinstance(public_name, str) and public_name: @@ -128,7 +129,7 @@ def _decode_model_info(raw: object) -> "Mapping[str, object] | None": @dataclass(frozen=True, slots=True) class _PTUDeployment: - """A deployment in the shape ``_parse_ptu_model`` reads, whatever declared it. + """A deployment in the shape ``_parse_ptu_models`` reads, whatever declared it. A ``LiteLLM_ProxyModelTable`` row already has it. A router entry does not: its id lives in ``model_info.id`` rather than on the entry itself. @@ -160,26 +161,30 @@ def _router_deployment(deployment: Mapping[str, object]) -> _PTUDeployment | Non ) -def _parse_ptu_model(row: object) -> PTUModel | None: - """Return a PTUModel when the deployment carries valid manual PTU config, else None. +def _parse_ptu_models(row: object) -> tuple[PTUModel, ...]: + """One PTUModel per team the deployment's valid PTU config attributes capacity to, else empty. - 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). + A deployment held by one team_id yields one entry carrying the whole count; a deployment + split into ptu_shares yields one per team carrying that team's share, so each team's row + accrues its share of the flat cost. """ model_info: Final = _decode_model_info(getattr(row, "model_info", None)) if model_info is None: - return None + return () 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=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, + return () + return tuple( + PTUModel( + model_id=str(getattr(row, "model_id", "") or ""), + model_name=_public_model_name(row, model_info), + team_id=team_id, + ptu_count=share, + cost_per_ptu_per_hour=terms.cost_per_ptu_per_hour, + effective_from=terms.effective_from, + effective_to=terms.effective_to, + ) + for team_id, share in sorted(terms.shares.items()) ) @@ -376,9 +381,7 @@ async def _load_ptu_models(prisma_client: "PrismaClient", *, router: object | No rows: Final = await _proxy_model_table(prisma_client).find_many() db_ids: Final = frozenset(model_id for row in rows if (model_id := str(getattr(row, "model_id", "") or ""))) config_records: Final = _config_deployments(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 - ) + models: Final = tuple(parsed for row in (*rows, *config_records) for parsed in _parse_ptu_models(row)) return _LoadedDeployments( models=models, scanned_ids=db_ids @@ -485,11 +488,13 @@ def _lapsed_models(ptu_models: tuple[PTUModel, ...], now: datetime) -> tuple[str charge the provider does not make for reserved capacity. """ return tuple( - _slack_safe(model.model_name) - for model in sorted( - (m for m in ptu_models if m.effective_to is not None and m.effective_to <= now), - key=lambda m: m.effective_to, - reverse=True, + dict.fromkeys( + _slack_safe(model.model_name) + for model in sorted( + (m for m in ptu_models if m.effective_to is not None and m.effective_to <= now), + key=lambda m: m.effective_to, + reverse=True, + ) ) ) diff --git a/litellm/router.py b/litellm/router.py index d328fbbb12f..6313dd3594c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -238,6 +238,7 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, ) +from litellm.router_utils.ptu_shares import filter_ptu_shared_deployments, ptu_capacity_warning from litellm.router_utils.reasoning_effort_capability import ( deployment_is_catalog_mapped, intersect_supported_reasoning_efforts, @@ -8824,6 +8825,15 @@ class Router: access_windows_error: Final = access_windows_config_error(_model_info, model_name=_model_name) if access_windows_error is not None: raise ValueError(access_windows_error) + capacity_warning: Final = ( + ptu_capacity_warning( + _model_name, MappingProxyType({"model_info": _model_info, "litellm_params": _litellm_params}) + ) + if is_ptu_cost_attribution_enabled() + else None + ) + if capacity_warning is not None: + verbose_router_logger.warning(capacity_warning) zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None litellm_params: Final[LiteLLM_Params] = LiteLLM_Params( **( # pyright: ignore[reportArgumentType] # untyped merged dict; already true for every field here @@ -12778,7 +12788,16 @@ class Router: model=model, llm_provider="", ) - return result.deployments + if not is_ptu_cost_attribution_enabled(): + return result.deployments + shared: Final = filter_ptu_shared_deployments(result.deployments, request_team_id) + if shared.withheld and len(shared.deployments) == 0: + raise litellm.BadRequestError( + message=f"Deployment {model} is reserved for the teams holding a PTU share of it", + model=model, + llm_provider="", + ) + return shared.deployments def _filter_deployments_by_model_access_groups( self, diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py new file mode 100644 index 00000000000..04e54395b5a --- /dev/null +++ b/litellm/router_utils/ptu_shares.py @@ -0,0 +1,101 @@ +"""Per-team PTU shares on a shared Azure provisioned deployment. + +A deployment declaring ``model_info.ptu_shares`` is served only to the teams named in it, +and each team's share converts to a normalized-tokens-per-minute ceiling through the +model's Azure sizing row, so the proxy enforces the split instead of an operator hand +converting PTUs to ``model_tpm_limit``. +""" + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Final, Generic, TypeVar + +from litellm.litellm_core_utils.azure_ptu_capacity import PTUCapacity, deployment_ptu_capacity +from litellm.litellm_core_utils.ptu_pricing import parsed_ptu_shares, ptu_terms + +_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) + + +@dataclass(frozen=True, slots=True) +class PTUTeamCeiling: + tpm_limit: int + output_to_input_ratio: float + + +@dataclass(frozen=True, slots=True) +class PTUShareFilterResult(Generic[_DeploymentT]): + deployments: tuple[_DeploymentT, ...] + withheld: bool + + +def _deployment_shares(deployment: Mapping[str, object]) -> Mapping[str, int] | None: + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping): + return None + return parsed_ptu_shares(model_info.get("ptu_shares")) + + +def filter_ptu_shared_deployments( + healthy_deployments: Sequence[_DeploymentT], request_team_id: str | None +) -> PTUShareFilterResult[_DeploymentT]: + """Drop every deployment split into PTU shares that ``request_team_id`` holds none of. + + A caller with no team, the master key included, holds no share, the same way an access + window reserves a deployment away from it. + """ + checks: Final = tuple((deployment, _deployment_shares(deployment)) for deployment in healthy_deployments) + kept: Final = tuple( + deployment for deployment, shares in checks if shares is None or (request_team_id or "") in shares + ) + return PTUShareFilterResult(deployments=kept, withheld=len(kept) < len(checks)) + + +def team_ptu_ceiling(deployments: Sequence[Mapping[str, object]], team_id: str) -> PTUTeamCeiling | None: + """The per-minute normalized-token ceiling ``team_id``'s shares across ``deployments`` add + up to, else None when the team holds no share on a deployment with a known sizing row. + + Two shared deployments of different models in one group are weighted by the larger + output ratio, which over-counts output on the cheaper one rather than under-counting it + on the dearer one. + """ + priced: Final = tuple( + (shares[team_id], capacity) + for deployment in deployments + if (shares := _deployment_shares(deployment)) is not None + and team_id in shares + and (capacity := deployment_ptu_capacity(deployment)) is not None + ) + if not priced: + return None + return PTUTeamCeiling( + tpm_limit=sum(share * capacity.input_tpm_per_ptu for share, capacity in priced), + output_to_input_ratio=max(capacity.output_to_input_ratio for _, capacity in priced), + ) + + +def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]]) -> PTUCapacity | None: + """The sizing row of the group's first reserved deployment, single-team or shared, so a + team's tokens on that group convert to PTU-hours.""" + return next( + ( + capacity + for deployment in deployments + if isinstance(model_info := deployment.get("model_info"), Mapping) + and ptu_terms(model_info) is not None + and (capacity := deployment_ptu_capacity(deployment)) is not None + ), + None, + ) + + +def ptu_capacity_warning(model_name: str, deployment: Mapping[str, object]) -> str | None: + """Why this reserved deployment's tokens cannot be converted to PTUs, else None.""" + model_info: Final = deployment.get("model_info") + if not isinstance(model_info, Mapping) or ptu_terms(model_info) is None: + return None + if deployment_ptu_capacity(deployment) is not None: + return None + return ( + f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its PTU shares set no " + "team ceiling and its usage reports no PTU hours; set model_info.base_model to the Azure model name" + ) diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 2a4f6b2944a..b4a04950969 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -19,6 +19,7 @@ class GroupByDimension(str, Enum): class SpendMetrics(BaseModel): spend: float = Field(default=0.0) flat_cost: float = Field(default=0.0) + ptu_hours: float = Field(default=0.0) prompt_tokens: int = Field(default=0) completion_tokens: int = Field(default=0) cache_read_input_tokens: int = Field(default=0) @@ -83,6 +84,7 @@ class DailySpendData(BaseModel): class DailySpendMetadata(BaseModel): total_spend: float = Field(default=0.0) total_flat_cost: float = Field(default=0.0) + total_ptu_hours: float = Field(default=0.0) total_prompt_tokens: int = Field(default=0) total_completion_tokens: int = Field(default=0) total_tokens: int = Field(default=0) diff --git a/litellm/types/router.py b/litellm/types/router.py index c0f724584fd..5ead0dd7626 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -260,6 +260,7 @@ class ModelInfo(MirroredPricingParams): cost_per_ptu_per_hour: float | None = None ptu_effective_from: datetime.datetime | None = None ptu_effective_to: datetime.datetime | None = None + ptu_shares: Mapping[str, int] | None = None # when tag-based routing's "!" or "&" constraints eliminate every deployment # in this model group, fall back to the default-tagged pool instead of diff --git a/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py b/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py new file mode 100644 index 00000000000..d07d71f0d82 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py @@ -0,0 +1,88 @@ +"""Tests for the Azure PTU sizing table and the conversions built on it.""" + +from typing import Final + +import pytest + +from litellm.litellm_core_utils.azure_ptu_capacity import ( + AZURE_PTU_CAPACITY, + PTUCapacity, + azure_ptu_capacity, + deployment_ptu_capacity, + normalized_tokens, + ptu_hours, +) + +_GPT41: Final = AZURE_PTU_CAPACITY["gpt-4.1"] +_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0) +_CACHED_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=5.0, cached_input_ratio=0.1) + + +@pytest.mark.parametrize( + "model", + ["gpt-4.1", "azure/gpt-4.1", "azure/gpt-4.1-2025-04-14", "GPT-4.1", "azure/eastus/gpt-4.1-2025-04-14"], +) +def test_a_model_resolves_to_its_row_through_a_provider_prefix_and_a_dated_version(model): + assert azure_ptu_capacity(model) is _GPT41 + + +def test_a_deployment_name_that_is_not_a_model_has_no_row(): + assert azure_ptu_capacity("azure/my-ptu-deployment") is None + assert azure_ptu_capacity("") is None + + +def test_a_dated_version_only_drops_a_full_date_suffix(): + """``gpt-4o-mini`` must not lose its ``-mini`` the way ``-2024-07-18`` is dropped.""" + assert azure_ptu_capacity("azure/gpt-4o-mini-2024-07-18") is AZURE_PTU_CAPACITY["gpt-4o-mini"] + assert azure_ptu_capacity("gpt-4o-mini") is not AZURE_PTU_CAPACITY["gpt-4o"] + + +def test_every_row_serves_its_input_tpm_for_an_hour(): + for capacity in AZURE_PTU_CAPACITY.values(): + assert capacity.normalized_tokens_per_ptu_hour == capacity.input_tpm_per_ptu * 60 + assert capacity.input_tpm_per_ptu > 0 + assert capacity.output_to_input_ratio >= 1.0 + assert 0.0 <= capacity.cached_input_ratio < 1.0 + + +def test_a_deployment_prefers_its_declared_base_model_over_its_deployment_name(): + deployment: Final = { + "model_info": {"base_model": "azure/gpt-4.1"}, + "litellm_params": {"model": "azure/gpt-4o"}, + } + assert deployment_ptu_capacity(deployment) is _GPT41 + + +def test_a_deployment_falls_back_to_its_litellm_model_when_no_base_model_is_declared(): + assert deployment_ptu_capacity({"litellm_params": {"model": "azure/gpt-4o"}}) is AZURE_PTU_CAPACITY["gpt-4o"] + assert deployment_ptu_capacity({"model_info": {"base_model": ""}, "litellm_params": {"model": "azure/gpt-4o"}}) is ( + AZURE_PTU_CAPACITY["gpt-4o"] + ) + + +def test_a_deployment_with_no_recognisable_model_has_no_row(): + assert deployment_ptu_capacity({"litellm_params": {"model": "azure/team-a-ptu"}}) is None + assert deployment_ptu_capacity({"model_info": None, "litellm_params": None}) is None + assert deployment_ptu_capacity({}) is None + + +def test_output_is_weighted_by_the_models_ratio_and_uncached_input_counts_in_full(): + assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=10) == pytest.approx(140.0) + + +def test_cached_input_is_free_unless_the_row_prices_it(): + assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60) == pytest.approx(40.0) + assert normalized_tokens(_CACHED_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60) == pytest.approx( + 46.0 + ) + + +def test_cached_input_never_exceeds_the_prompt_and_negatives_count_as_zero(): + assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=250) == 0.0 + assert normalized_tokens(_ROW, prompt_tokens=-5, completion_tokens=-5, cache_read_tokens=-5) == 0.0 + + +def test_one_ptu_hour_is_one_ptus_input_tpm_served_for_sixty_minutes(): + assert ptu_hours(_ROW, _ROW.input_tpm_per_ptu * 60) == pytest.approx(1.0) + assert ptu_hours(_ROW, _ROW.input_tpm_per_ptu * 30) == pytest.approx(0.5) + assert ptu_hours(_ROW, 0.0) == 0.0 diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py index 1689da2696f..ef8d7703924 100644 --- a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py +++ b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py @@ -38,7 +38,7 @@ def test_a_complete_reservation_is_accepted(): terms = ptu_terms(_VALID) assert terms is not None - assert terms.team_id == "team-alpha" + assert terms.shares == {"team-alpha": 100} assert terms.ptu_count == 100 assert terms.effective_from == datetime(2026, 1, 1, tzinfo=timezone.utc) assert terms.effective_to is None @@ -193,8 +193,8 @@ def test_a_deployment_with_no_ptu_fields_is_not_a_ptu_deployment(): @pytest.mark.parametrize( "override, expected", [ - ({"team_id": None}, "team_id is required when PTU fields are set (one model maps to one team)"), - ({"team_id": ""}, "team_id is required when PTU fields are set (one model maps to one team)"), + ({"team_id": None}, "team_id or ptu_shares is required when PTU fields are set"), + ({"team_id": ""}, "team_id or ptu_shares is required when PTU fields are set"), ({"cost_per_ptu_per_hour": None}, "ptu_count and cost_per_ptu_per_hour must be set together"), ({"ptu_count": None}, "ptu_count and cost_per_ptu_per_hour must be set together"), ({"ptu_effective_to": "2025-01-01T00:00:00Z"}, "ptu_effective_to must be after ptu_effective_from"), @@ -356,3 +356,78 @@ def test_no_spillover_marker_returns_none(): is None ) assert azure_spillover(response_headers=None, additional_headers=None) is None + + +# --- a reservation split into per-team shares ------------------------------------------ + +_SHARED = {**{k: v for k, v in _VALID.items() if k != "team_id"}, "ptu_shares": {"team-a": 60, "team-b": 40}} + + +def test_a_split_reservation_attributes_each_share_to_its_team(): + terms = ptu_terms(_SHARED) + + assert terms is not None + assert terms.shares == {"team-a": 60, "team-b": 40} + assert terms.ptu_count == 100 + assert ptu_config_error(_SHARED) is None + + +def test_a_single_team_reservation_holds_the_whole_count_under_that_team(): + terms = ptu_terms(_VALID) + + assert terms is not None + assert sum(terms.shares.values()) == terms.ptu_count + + +@pytest.mark.parametrize( + "override, expected", + [ + ( + {"team_id": "team-alpha"}, + "team_id and ptu_shares cannot both be set; ptu_shares lists every team the capacity is split across", + ), + ({"ptu_shares": {}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ( + {"ptu_shares": {"team-a": 0, "team-b": 100}}, + "ptu_shares must map at least one team_id to a positive whole number of PTUs", + ), + ( + {"ptu_shares": {"team-a": 50.5, "team-b": 49.5}}, + "ptu_shares must map at least one team_id to a positive whole number of PTUs", + ), + ({"ptu_shares": {"team-a": True}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ({"ptu_shares": {"": 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ({"ptu_shares": ["team-a"]}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ({"ptu_shares": {"team-a": 60, "team-b": 30}}, "ptu_shares must add up to ptu_count (90 of 100 allocated)"), + ({"ptu_shares": {"team-a": 60, "team-b": 50}}, "ptu_shares must add up to ptu_count (110 of 100 allocated)"), + ], + ids=[ + "team_id and shares", + "empty shares", + "zero share", + "fractional share", + "boolean share", + "blank team", + "not a mapping", + "shares short of the count", + "shares over the count", + ], +) +def test_an_incoherent_split_names_its_reason_and_reserves_nothing(override, expected): + assert ptu_config_error({**_SHARED, **override}) == expected + assert ptu_terms({**_SHARED, **override}) is None + + +def test_the_split_is_named_after_the_deployment_when_the_caller_supplies_one(): + error = ptu_config_error({**_SHARED, "ptu_shares": {"team-a": 1}}, model_name="gpt-4.1-ptu") + + assert error is not None + assert "'gpt-4.1-ptu'" in error + assert "1 of 100 allocated" in error + + +def test_a_split_reservation_is_zeroed_like_a_single_team_one(): + zeroed = _with_flag(_SHARED) + + assert zeroed["input_cost_per_token"] == 0.0 + assert zeroed["output_cost_per_token"] == 0.0 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 8c0dcd3383c..a18b38453b0 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -37,11 +37,13 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token +from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, ModelResponse, + PromptTokensDetailsWrapper, TextCompletionResponse, Usage, ) @@ -7150,3 +7152,152 @@ async def test_success_tpm_accounting_keeps_the_admission_target_after_an_alias_ charged: Final = {op["key"]: op["increment_value"] for op in ops} assert charged[admission_bucket] == 150 - stash.reserved_tokens assert not any(":target-b" in key for key in charged) + + +# --- a team's PTU share on a shared Azure provisioned deployment --------------------------- + + +def _ptu_ceiling_for(team_id: str, model_group: str, tpm_limit: int, ratio: float): + calls: list[tuple[str, str]] = [] + + def resolve(requested_team: str, requested_group: str) -> PTUTeamCeiling | None: + calls.append((requested_team, requested_group)) + if (requested_team, requested_group) != (team_id, model_group): + return None + return PTUTeamCeiling(tpm_limit=tpm_limit, output_to_input_ratio=ratio) + + return resolve, calls + + +def _ptu_request(model: str = "test-model") -> dict: + return {"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 300} + + +@pytest.mark.asyncio +async def test_a_teams_ptu_share_is_a_hard_tpm_ceiling_on_the_shared_model(): + cache = DualCache() + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion") + with pytest.raises(HTTPException) as exc: + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion") + + assert exc.value.status_code == 429 + assert "model_per_team_ptu" in str(exc.value.detail) + assert exc.value.headers["rate_limit_type"] == "tokens" + + +@pytest.mark.asyncio +async def test_a_ptu_ceiling_on_one_model_leaves_the_teams_other_models_alone(): + cache = DualCache() + resolve, calls = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + for _ in range(3): + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_ptu_request("other-model"), call_type="acompletion" + ) + + assert calls == [("t", "other-model")] * 3 + + +@pytest.mark.asyncio +async def test_a_team_without_a_share_and_a_key_without_a_team_get_no_ptu_ceiling(): + cache = DualCache() + resolve, calls = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve) + other_team = UserAPIKeyAuth(api_key=hash_token("sk-other"), team_id="u") + no_team = UserAPIKeyAuth(api_key=hash_token("sk-no-team")) + + for _ in range(3): + await handler.async_pre_call_hook(user_api_key_dict=other_team, cache=cache, data=_ptu_request(), call_type="acompletion") + await handler.async_pre_call_hook(user_api_key_dict=no_team, cache=cache, data=_ptu_request(), call_type="acompletion") + + assert set(calls) == {("u", "test-model")} + + +def _ptu_success_kwargs() -> dict: + return { + "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"}}, + "litellm_params": {"metadata": {"model_group": "test-model", "user_api_key_metadata": {}, "user_api_key_team_metadata": {}}}, + "model": "test-model", + } + + +def _ptu_response(usage: Usage) -> ModelResponse: + return ModelResponse( + id="ptu-share", object="chat.completion", created=int(datetime.now().timestamp()), model="test-model", usage=usage, choices=[] + ) + + +def _ptu_increment(handler: _PROXY_MaxParallelRequestsHandler, ops: list) -> float | None: + ptu_key = handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens") + return next((op["increment_value"] for op in ops if op["key"] == ptu_key), None) + + +def test_success_accounting_settles_the_ptu_counter_in_azure_normalized_tokens(): + """100 uncached input tokens plus 50 output tokens at a 4:1 output ratio are 300 + normalized tokens, which is what the share was sized in.""" + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve + ) + response = _ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)) + + ops = handler._build_success_event_pipeline_operations( + kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output" + ) + + assert _ptu_increment(handler, ops) == 300 + + +def test_cached_input_is_not_charged_to_the_ptu_counter(): + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve + ) + response = _ptu_response( + Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=40), + ) + ) + + ops = handler._build_success_event_pipeline_operations( + kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output" + ) + + assert _ptu_increment(handler, ops) == 60 + 4 * 50 + + +def test_usage_that_only_reports_a_total_charges_that_total_to_the_ptu_counter(): + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve + ) + response = _ptu_response(Usage(prompt_tokens=0, completion_tokens=0, total_tokens=150)) + + ops = handler._build_success_event_pipeline_operations( + kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output" + ) + + assert _ptu_increment(handler, ops) == 150 + + +def test_success_accounting_charges_no_ptu_counter_without_a_ceiling(): + resolve, _ = _ptu_ceiling_for("u", "test-model", tpm_limit=500, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve + ) + response = _ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)) + + ops = handler._build_success_event_pipeline_operations( + kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output" + ) + + assert _ptu_increment(handler, ops) is None diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py new file mode 100644 index 00000000000..e9ef4ea2994 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py @@ -0,0 +1,130 @@ +"""Tests for attaching PTU-hours to a daily activity response.""" + +from typing import Final + +import pytest + +from litellm.litellm_core_utils.azure_ptu_capacity import PTUCapacity +from litellm.proxy.management_endpoints.ptu_consumption import attach_ptu_hours +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + BreakdownMetrics, + DailySpendData, + DailySpendMetadata, + KeyMetricWithMetadata, + MetricWithMetadata, + SpendAnalyticsPaginatedResponse, + SpendMetrics, +) + +_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0) +_CACHED_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0, cached_input_ratio=0.1) +_CAPACITY: Final = {"gpt-4.1-ptu": _ROW, "gpt-6-ptu": _CACHED_ROW} + + +def _metrics(prompt: int, completion: int, cached: int = 0) -> SpendMetrics: + return SpendMetrics( + prompt_tokens=prompt, + completion_tokens=completion, + cache_read_input_tokens=cached, + total_tokens=prompt + completion, + api_requests=1, + successful_requests=1, + ) + + +def _bucket(metrics: SpendMetrics, keys: dict[str, SpendMetrics] | None = None) -> MetricWithMetadata: + return MetricWithMetadata( + metrics=metrics, + metadata={}, + api_key_breakdown={key: KeyMetricWithMetadata(metrics=m, metadata={}) for key, m in (keys or {}).items()}, + ) + + +def _day(date: str, model_groups: dict[str, MetricWithMetadata]) -> DailySpendData: + total: Final = _metrics( + sum(b.metrics.prompt_tokens for b in model_groups.values()), + sum(b.metrics.completion_tokens for b in model_groups.values()), + ) + return DailySpendData(date=date, metrics=total, breakdown=BreakdownMetrics(model_groups=model_groups)) + + +def _response(*days: DailySpendData) -> SpendAnalyticsPaginatedResponse: + return SpendAnalyticsPaginatedResponse(results=list(days), metadata=DailySpendMetadata()) + + +def test_a_ptu_model_groups_tokens_become_ptu_hours_on_the_group_the_day_and_the_total(): + """60,000 normalized tokens on a 1,000 input-TPM row is one PTU-hour.""" + day: Final = _day("2026-09-23", {"gpt-4.1-ptu": _bucket(_metrics(prompt=40_000, completion=5_000))}) + + attached: Final = attach_ptu_hours(_response(day), _CAPACITY.get) + + group: Final = attached.results[0].breakdown.model_groups["gpt-4.1-ptu"] + assert group.metrics.ptu_hours == pytest.approx(1.0) + assert attached.results[0].metrics.ptu_hours == pytest.approx(1.0) + assert attached.metadata.total_ptu_hours == pytest.approx(1.0) + + +def test_a_model_group_without_a_sizing_row_keeps_zero_ptu_hours_and_the_rest_of_the_day_intact(): + day: Final = _day( + "2026-09-23", + { + "gpt-4.1-ptu": _bucket(_metrics(prompt=30_000, completion=0)), + "gpt-4o-mini": _bucket(_metrics(prompt=1_000_000, completion=1_000_000)), + }, + ) + + attached: Final = attach_ptu_hours(_response(day), _CAPACITY.get) + + groups: Final = attached.results[0].breakdown.model_groups + assert groups["gpt-4o-mini"].metrics.ptu_hours == 0.0 + assert groups["gpt-4o-mini"].metrics.prompt_tokens == 1_000_000 + assert groups["gpt-4.1-ptu"].metrics.ptu_hours == pytest.approx(0.5) + assert attached.results[0].metrics.ptu_hours == pytest.approx(0.5) + assert attached.results[0].metrics.total_tokens == day.metrics.total_tokens + + +def test_cached_input_is_charged_at_the_rows_cached_ratio(): + day: Final = _day("2026-09-23", {"gpt-6-ptu": _bucket(_metrics(prompt=60_000, completion=0, cached=60_000))}) + + attached: Final = attach_ptu_hours(_response(day), _CAPACITY.get) + + assert attached.results[0].breakdown.model_groups["gpt-6-ptu"].metrics.ptu_hours == pytest.approx(0.1) + + +def test_each_api_key_under_a_ptu_group_gets_its_own_ptu_hours(): + day: Final = _day( + "2026-09-23", + { + "gpt-4.1-ptu": _bucket( + _metrics(prompt=60_000, completion=0), + keys={"key-a": _metrics(prompt=45_000, completion=0), "key-b": _metrics(prompt=15_000, completion=0)}, + ) + }, + ) + + attached: Final = attach_ptu_hours(_response(day), _CAPACITY.get) + + breakdown: Final = attached.results[0].breakdown.model_groups["gpt-4.1-ptu"].api_key_breakdown + assert breakdown["key-a"].metrics.ptu_hours == pytest.approx(0.75) + assert breakdown["key-b"].metrics.ptu_hours == pytest.approx(0.25) + + +def test_the_total_sums_every_day_on_the_page(): + first: Final = _day("2026-09-22", {"gpt-4.1-ptu": _bucket(_metrics(prompt=60_000, completion=0))}) + second: Final = _day("2026-09-23", {"gpt-4.1-ptu": _bucket(_metrics(prompt=0, completion=15_000))}) + + attached: Final = attach_ptu_hours(_response(first, second), _CAPACITY.get) + + assert [day.metrics.ptu_hours for day in attached.results] == [pytest.approx(1.0), pytest.approx(1.0)] + assert attached.metadata.total_ptu_hours == pytest.approx(2.0) + + +def test_a_page_with_no_ptu_group_is_returned_unchanged(): + day: Final = _day("2026-09-23", {"gpt-4o-mini": _bucket(_metrics(prompt=100, completion=100))}) + response: Final = _response(day) + + attached: Final = attach_ptu_hours(response, _CAPACITY.get) + + assert attached.results[0] is day + assert attached.metadata.total_ptu_hours == 0.0 + assert response.metadata.total_ptu_hours == 0.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index d35b77f732c..a7c21d74f1e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -16,6 +16,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.auth.auth_checks import _is_model_cost_zero from litellm.llms.gemini.cost_calculator import cost_per_web_search_request from litellm.proxy.management_endpoints.model_management_endpoints import ( @@ -1284,3 +1285,48 @@ class TestPtuDeploymentsAreNotBilledPerToken: assert "input_cost_per_token" in str(exc.value) add_team_model_to_db.assert_not_called() + + +_SHARED_START = "2026-08-01T00:00:00Z" + + +def test_validate_helper_accepts_shares_in_place_of_a_team_id(): + shared = { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": _SHARED_START, + "ptu_shares": {"team-a": 3, "team-b": 2}, + } + + _validate_ptu_model_info(shared) + + assert ptu_terms(shared) is not None + + +def test_validate_helper_refuses_shares_that_do_not_add_up_to_the_count(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": _SHARED_START, + "ptu_shares": {"team-a": 3, "team-b": 1}, + } + ) + assert exc.value.status_code == 400 + assert "4 of 5 allocated" in exc.value.detail + + +def test_validate_helper_refuses_a_team_id_beside_shares(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info( + { + "team_id": "team-a", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": _SHARED_START, + "ptu_shares": {"team-a": 5}, + } + ) + assert exc.value.status_code == 400 + assert "cannot both be set" in exc.value.detail diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b066b3b80e6..379bbe5d032 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -16889,3 +16889,90 @@ def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hi assert response.status_code == 503, response.text assert response.json() == _DB_OUTAGE_503_BODY + + +# --- PTU-equivalent consumption on /team/daily/activity ------------------------------------ + + +def _ptu_activity_page(): + from litellm.types.proxy.management_endpoints.common_daily_activity import ( + BreakdownMetrics, + DailySpendData, + DailySpendMetadata, + MetricWithMetadata, + SpendAnalyticsPaginatedResponse, + SpendMetrics, + ) + + metrics = SpendMetrics( + prompt_tokens=180_000, completion_tokens=0, total_tokens=180_000, api_requests=3, successful_requests=3 + ) + return SpendAnalyticsPaginatedResponse( + results=[ + DailySpendData( + date=datetime(2026, 9, 23).date(), + metrics=metrics, + breakdown=BreakdownMetrics( + model_groups={"gpt-4.1-ptu": MetricWithMetadata(metrics=metrics, metadata={}, api_key_breakdown={})} + ), + ) + ], + metadata=DailySpendMetadata(total_tokens=180_000, total_api_requests=3, total_successful_requests=3), + ) + + +def _shared_ptu_router() -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "shared-ptu", + "base_model": "azure/gpt-4.1", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 30, "team-b": 20}, + }, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("attribution_enabled", [True, False]) +async def test_team_daily_activity_reports_ptu_hours_only_while_attribution_is_on( + mock_db_client, mock_admin_auth, monkeypatch, attribution_enabled +): + """One PTU serves 3,000 input tokens per minute on gpt-4.1, so 180,000 uncached + input tokens are one PTU-hour; the figure appears beside tokens only once the + PTU flag is on, and the token totals are untouched either way.""" + from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity + + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True" if attribution_enabled else "False") + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + page = _ptu_activity_page() + + with ( + patch("litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", AsyncMock(return_value=page)), + patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router()), + ): + result = await get_team_daily_activity( + team_ids="team-a", + start_date="2026-09-23", + end_date="2026-09-24", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=mock_admin_auth, + ) + + expected_ptu_hours = 1.0 if attribution_enabled else 0.0 + assert result.metadata.total_ptu_hours == expected_ptu_hours + assert result.results[0].metrics.ptu_hours == expected_ptu_hours + assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == expected_ptu_hours + assert result.metadata.total_tokens == 180_000 + assert result.results[0].metrics.total_tokens == 180_000 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 8f25cffecf5..73e3e5f3ec6 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 @@ -15,7 +15,7 @@ from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( PTUModel, _active_hours_on_day, _compute_daily_flat_cost, - _parse_ptu_model, + _parse_ptu_models, run_ptu_flat_cost_backfill, run_ptu_flat_cost_rollup, run_scheduled_ptu_rollup, @@ -41,6 +41,14 @@ def _ptu_enabled(monkeypatch): _VALID_PTU = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} +def _parse_ptu_model(row): + """The single holding a team_id deployment parses to, else None; shared deployments + parse to one holding per team and have their own tests below.""" + parsed = _parse_ptu_models(row) + assert len(parsed) <= 1 + return parsed[0] if parsed else None + + def _model_row(model_id="m1", model_name="gpt-4o-mini-ptu", model_info=None, with_start=True): row = MagicMock() row.model_id = model_id @@ -2082,3 +2090,51 @@ async def test_the_catch_up_pass_reaches_a_config_declared_deployment(): assert len(charged) == 3, charged assert charged[-1] == yesterday assert all(row["ptu_flat_cost"] == pytest.approx(48.0) for row in table.rows.values()) + + +# --- a deployment split into per-team PTU shares ----------------------------------------- + +_SHARED_PTU = {"ptu_count": 50, "cost_per_ptu_per_hour": 1.0, "ptu_shares": {"team-b": 20, "team-a": 30}} + + +def test_a_shared_deployment_parses_to_one_holding_per_team_carrying_its_share(): + parsed = _parse_ptu_models(_model_row(model_info=dict(_SHARED_PTU))) + + assert [(m.team_id, m.ptu_count) for m in parsed] == [("team-a", 30), ("team-b", 20)] + assert {m.model_id for m in parsed} == {"m1"} + assert {m.model_name for m in parsed} == {"gpt-4o-mini-ptu"} + assert {m.cost_per_ptu_per_hour for m in parsed} == {1.0} + + +def test_a_shared_deployment_whose_shares_do_not_add_up_is_not_priced(): + assert _parse_ptu_models(_model_row(model_info={**_SHARED_PTU, "ptu_shares": {"team-a": 30}})) == () + + +@pytest.mark.asyncio +async def test_rollup_splits_a_shared_deployments_flat_cost_by_share(): + """50 PTUs at $1/hour for a day is $1,200; team-a's 30 PTUs are $720 of it and team-b's + 20 are $480, each keyed on the same deployment id and public name.""" + prisma, table = _prisma_with_models([_model_row(model_info=dict(_SHARED_PTU))]) + + result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) + + assert result.rows_written == 2 + created = {call.kwargs["data"]["create"]["team_id"]: call.kwargs["data"]["create"] for call in table.upsert.await_args_list} + assert created["team-a"]["ptu_flat_cost"] == pytest.approx(720.0) + assert created["team-b"]["ptu_flat_cost"] == pytest.approx(480.0) + assert sum(row["ptu_flat_cost"] for row in created.values()) == pytest.approx(50 * 1.0 * 24) + assert {row["model"] for row in created.values()} == {"m1"} + assert {row["model_group"] for row in created.values()} == {"gpt-4o-mini-ptu"} + assert {row["api_key"] for row in created.values()} == {PTU_SENTINEL_API_KEY} + + +@pytest.mark.asyncio +async def test_a_shared_deployment_with_a_closed_window_is_reported_lapsed_once(): + closed = {**_SHARED_PTU, "ptu_effective_from": "2020-01-01T00:00:00Z", "ptu_effective_to": "2020-02-01T00:00:00Z"} + prisma, _ = _prisma_with_models([_model_row(model_id="dep-shared", model_info=closed)]) + alert = AsyncMock() + + result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert) + + assert result.lapsed == ("gpt-4o-mini-ptu",) + alert.assert_awaited_once() diff --git a/tests/test_litellm/router_utils/test_ptu_shares.py b/tests/test_litellm/router_utils/test_ptu_shares.py new file mode 100644 index 00000000000..150cfd31b75 --- /dev/null +++ b/tests/test_litellm/router_utils/test_ptu_shares.py @@ -0,0 +1,115 @@ +"""Tests for per-team PTU shares: who a shared deployment is served to and what a share is worth.""" + +from typing import Final + +from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY +from litellm.router_utils.ptu_shares import ( + PTUTeamCeiling, + filter_ptu_shared_deployments, + model_group_ptu_capacity, + ptu_capacity_warning, + team_ptu_ceiling, +) + +_GPT41: Final = AZURE_PTU_CAPACITY["gpt-4.1"] +_GPT4O: Final = AZURE_PTU_CAPACITY["gpt-4o"] +_SHARES: Final = {"team-a": 30, "team-b": 20} + + +def _shared(model: str = "azure/gpt-4.1", shares: object = _SHARES, deployment_id: str = "shared") -> dict: + return { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": model}, + "model_info": { + "id": deployment_id, + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": shares, + }, + } + + +def _single_team(model: str = "azure/gpt-4.1") -> dict: + return { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": model}, + "model_info": { + "id": "single", + "team_id": "team-a", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + }, + } + + +_OPEN: Final = {"model_name": "gpt-4.1-ptu", "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {"id": "open"}} + + +def test_a_team_holding_a_share_keeps_the_shared_deployment(): + result: Final = filter_ptu_shared_deployments([_shared(), _OPEN], "team-a") + assert [d["model_info"]["id"] for d in result.deployments] == ["shared", "open"] + assert result.withheld is False + + +def test_a_team_without_a_share_only_sees_the_unshared_deployments(): + result: Final = filter_ptu_shared_deployments([_shared(), _OPEN], "team-c") + assert [d["model_info"]["id"] for d in result.deployments] == ["open"] + assert result.withheld is True + + +def test_a_caller_with_no_team_holds_no_share(): + for team_id in (None, ""): + result = filter_ptu_shared_deployments([_shared()], team_id) + assert result.deployments == () + assert result.withheld is True + + +def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_here(): + """A ``team_id`` deployment is scoped by the router's team filter, and an unusable + ``ptu_shares`` is refused at registration, so neither is withheld by the share filter.""" + result: Final = filter_ptu_shared_deployments([_single_team(), _shared(shares={"team-a": 0})], "team-z") + assert [d["model_info"]["id"] for d in result.deployments] == ["single", "shared"] + assert result.withheld is False + + +def test_a_share_converts_to_the_models_input_tpm_per_ptu(): + ceiling: Final = team_ptu_ceiling([_shared()], "team-a") + assert ceiling == PTUTeamCeiling( + tpm_limit=30 * _GPT41.input_tpm_per_ptu, output_to_input_ratio=_GPT41.output_to_input_ratio + ) + + +def test_shares_across_deployments_add_up_and_the_larger_output_ratio_wins(): + gpt4o: Final = _shared(model="azure/gpt-4o", shares={"team-a": 10}, deployment_id="shared-4o") + ceiling: Final = team_ptu_ceiling([_shared(), gpt4o, _OPEN], "team-a") + assert ceiling is not None + assert ceiling.tpm_limit == 30 * _GPT41.input_tpm_per_ptu + 10 * _GPT4O.input_tpm_per_ptu + assert ceiling.output_to_input_ratio == max(_GPT41.output_to_input_ratio, _GPT4O.output_to_input_ratio) + + +def test_no_share_or_no_sizing_row_sets_no_ceiling(): + assert team_ptu_ceiling([_shared()], "team-c") is None + assert team_ptu_ceiling([_shared(model="azure/unknown-deployment")], "team-a") is None + assert team_ptu_ceiling([_single_team(), _OPEN], "team-a") is None + + +def test_a_groups_capacity_comes_from_its_first_reserved_deployment_with_a_row(): + assert model_group_ptu_capacity([_OPEN, _single_team()]) is _GPT41 + assert model_group_ptu_capacity([_shared(model="azure/unknown"), _shared(model="azure/gpt-4o")]) is _GPT4O + assert model_group_ptu_capacity([_OPEN]) is None + assert model_group_ptu_capacity([]) is None + + +def test_a_reserved_deployment_without_a_sizing_row_is_warned_about_by_name(): + warning: Final = ptu_capacity_warning("gpt-4.1-ptu", _shared(model="azure/my-ptu-deployment")) + assert warning is not None + assert "gpt-4.1-ptu" in warning + assert "base_model" in warning + + +def test_a_sized_reservation_and_an_unreserved_deployment_raise_no_warning(): + assert ptu_capacity_warning("gpt-4.1-ptu", _shared()) is None + assert ptu_capacity_warning("gpt-4.1-ptu", _single_team()) is None + assert ptu_capacity_warning("gpt-4.1-ptu", {**_OPEN, "litellm_params": {"model": "azure/my-ptu-deployment"}}) is None diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 411377f29cf..ca10d8f4925 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -17829,3 +17829,72 @@ def test_access_windows_filter_reserved_deployments_method(): request_team_id="team-a", ) ] == ["reserved-deployment", "open-deployment"] + + +def _shared_ptu_model_list() -> list: + return [ + { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": "gpt-4.1", "mock_response": "shared"}, + "model_info": { + "id": "shared-deployment", + "base_model": "azure/gpt-4.1", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 30, "team-b": 20}, + }, + }, + { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": "gpt-4.1", "mock_response": "open"}, + "model_info": {"id": "open-deployment"}, + }, + ] + + +def test_ptu_shares_hide_the_shared_deployment_from_teams_holding_no_share(monkeypatch): + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + router = Router(model_list=_shared_ptu_model_list()) + for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-c"}}, {"metadata": {}}, {}): + _, deployments = router._common_checks_available_deployment(model="gpt-4.1-ptu", request_kwargs=request_kwargs) + assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] + for team_id in ("team-a", "team-b"): + _, deployments = router._common_checks_available_deployment( + model="gpt-4.1-ptu", + request_kwargs={"metadata": {"user_api_key_team_id": team_id}}, + ) + assert {d["model_info"]["id"] for d in deployments} == {"shared-deployment", "open-deployment"} + + +def test_ptu_shares_raise_when_only_shared_deployments_remain(monkeypatch): + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + router = Router(model_list=_shared_ptu_model_list()[:1]) + with pytest.raises(litellm.BadRequestError, match="reserved for the teams holding a PTU share"): + router._common_checks_available_deployment( + model="gpt-4.1-ptu", + request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, + ) + _, deployments = router._common_checks_available_deployment( + model="gpt-4.1-ptu", + request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, + ) + assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"] + + +def test_ptu_shares_do_not_filter_while_the_feature_is_off(monkeypatch): + monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False) + router = Router(model_list=_shared_ptu_model_list()) + _, deployments = router._common_checks_available_deployment( + model="gpt-4.1-ptu", + request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, + ) + assert {d["model_info"]["id"] for d in deployments} == {"shared-deployment", "open-deployment"} + + +def test_a_shared_ptu_deployment_whose_shares_do_not_add_up_is_refused_at_registration(monkeypatch): + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + model_list = _shared_ptu_model_list()[:1] + model_list[0]["model_info"]["ptu_shares"] = {"team-a": 30} + with pytest.raises(ValueError, match=r"gpt-4\.1-ptu.*30 of 50 allocated"): + Router(model_list=model_list) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 27460b21108..20fa128b764 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -9,7 +9,13 @@ import { type ExtendedDailyData, type ProviderSpendRow, } from "./entityUsageAggregations"; -import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost, type SummaryTile } from "./entityUsageSummary"; +import { + buildCostBreakdownTiles, + buildSummaryTiles, + hasFlatCost, + hasPtuHours, + type SummaryTile, +} from "./entityUsageSummary"; import { MoneyCell } from "@/components/shared/table_cells"; import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { hasCapability, type Capability } from "@/utils/capabilities"; @@ -67,6 +73,7 @@ interface EntitySpendData { metadata: { total_spend: number; total_flat_cost?: number; + total_ptu_hours?: number; total_api_requests: number; total_successful_requests: number; total_failed_requests: number; @@ -284,6 +291,7 @@ const EntityUsage: React.FC = ({ const capitalizedEntityLabel = entityType.charAt(0).toUpperCase() + entityType.slice(1); const showFlatCost = entityType === "team" && hasFlatCost(spendData.metadata); + const showPtuHours = entityType === "team" && hasPtuHours(spendData.metadata); const userSpendTeamIds = useMemo( () => selectedTags.length > 0 @@ -392,7 +400,7 @@ const EntityUsage: React.FC = ({ ); const breakdownTiles = showFlatCost && showCostBreakdown ? buildCostBreakdownTiles(spendData.metadata) : []; - const summaryTiles = [...buildSummaryTiles(spendData.metadata, showFlatCost), ...breakdownTiles]; + const summaryTiles = [...buildSummaryTiles(spendData.metadata, showFlatCost, showPtuHours), ...breakdownTiles]; const modelViewTitle = modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models"; @@ -402,7 +410,9 @@ const EntityUsage: React.FC = ({

{capitalizedEntityLabel} Spend Overview

-
{summaryTiles.map(renderSummaryTile)}
+
+ {summaryTiles.map(renderSummaryTile)} +
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts index 403f9473391..7798f90e86a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost } from "./entityUsageSummary"; +import { buildCostBreakdownTiles, buildSummaryTiles, hasFlatCost, hasPtuHours } from "./entityUsageSummary"; const metadata = { total_spend: 100, @@ -10,6 +10,17 @@ const metadata = { total_tokens: 3456, }; +describe("hasPtuHours", () => { + it("is false when no PTU hours were consumed or the backend predates the field", () => { + expect(hasPtuHours({ ...metadata, total_ptu_hours: 0 })).toBe(false); + expect(hasPtuHours(metadata)).toBe(false); + }); + + it("is true once a PTU model group has been used", () => { + expect(hasPtuHours({ ...metadata, total_ptu_hours: 0.25 })).toBe(true); + }); +}); + describe("hasFlatCost", () => { it("is false when there is no flat cost to report", () => { expect(hasFlatCost({ ...metadata, total_flat_cost: 0 })).toBe(false); @@ -28,6 +39,23 @@ describe("buildSummaryTiles", () => { expect(buildSummaryTiles(metadata, true)).toHaveLength(5); }); + it("adds a sixth PTU Hours tile only when asked, after Total Tokens", () => { + const withPtu = { ...metadata, total_ptu_hours: 1234.5678 }; + expect(buildSummaryTiles(withPtu, true)).toHaveLength(5); + const tiles = buildSummaryTiles(withPtu, true, true); + expect(tiles).toHaveLength(6); + const last = tiles[tiles.length - 1]; + expect(last.title).toBe("PTU Hours"); + expect(last.value).toBe("1,234.57"); + expect(last.tooltip).toContain("PTU-hours"); + expect(tiles[tiles.length - 2].title).toBe("Total Tokens"); + }); + + it("treats missing PTU hours as zero when the tile is forced on", () => { + const tiles = buildSummaryTiles(metadata, false, true); + expect(tiles[tiles.length - 1].value).toBe("0.00"); + }); + it("shows request-only spend under the original title when there is no flat cost", () => { const [first] = buildSummaryTiles(metadata, false); expect(first.title).toBe("Total Spend"); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts index 01d0cb6bda9..f27b9929368 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/entityUsageSummary.ts @@ -11,6 +11,7 @@ export interface SummaryTile { interface SpendSummaryMetadata { total_spend: number; total_flat_cost?: number; + total_ptu_hours?: number; total_api_requests: number; total_successful_requests: number; total_failed_requests: number; @@ -26,10 +27,28 @@ export const REQUEST_COST_TOOLTIP = export const FLAT_COST_TOOLTIP = "Reserved provisioned throughput, billed per hour whether or not requests are sent. Reported here only; it does not count toward team, key, user, or organization budgets."; +export const PTU_HOURS_TOOLTIP = + "Provisioned throughput consumed, in PTU-hours: the tokens sent to PTU deployments converted through the Azure sizing table for each model. Compare it with the PTU-hours reserved to see how much of the reservation was used."; + export const hasFlatCost = (metadata: SpendSummaryMetadata): boolean => (metadata.total_flat_cost ?? 0) > 0; -export const buildSummaryTiles = (metadata: SpendSummaryMetadata, showFlatCost: boolean): SummaryTile[] => { +export const hasPtuHours = (metadata: SpendSummaryMetadata): boolean => (metadata.total_ptu_hours ?? 0) > 0; + +export const buildSummaryTiles = ( + metadata: SpendSummaryMetadata, + showFlatCost: boolean, + showPtuHours: boolean = false, +): SummaryTile[] => { const flatCost = metadata.total_flat_cost ?? 0; + const ptuHoursTile: SummaryTile[] = showPtuHours + ? [ + { + title: "PTU Hours", + value: formatNumberWithCommas(metadata.total_ptu_hours ?? 0, 2), + tooltip: PTU_HOURS_TOOLTIP, + }, + ] + : []; return [ showFlatCost ? { @@ -47,6 +66,7 @@ export const buildSummaryTiles = (metadata: SpendSummaryMetadata, showFlatCost: }, { title: "Failed Requests", value: metadata.total_failed_requests.toLocaleString(), className: "text-destructive" }, { title: "Total Tokens", value: metadata.total_tokens.toLocaleString() }, + ...ptuHoursTile, ]; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts index 0b8cfecbfc6..816d7979f23 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.test.ts @@ -43,6 +43,7 @@ describe("sumMetadata", () => { total_cache_read_input_tokens: 1, total_cache_creation_input_tokens: 1, total_flat_cost: 1, + total_ptu_hours: 1, total_response_time_ms: 1, total_timed_requests: 1, }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts index 1f03f6a4fcb..37a3c1e2bee 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts @@ -30,6 +30,7 @@ const SUMMABLE_METADATA_KEYS = [ "total_cache_read_input_tokens", "total_cache_creation_input_tokens", "total_flat_cost", + "total_ptu_hours", "total_response_time_ms", "total_timed_requests", ] as const; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index 15f193ecc3f..55939e2cd95 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -10,6 +10,7 @@ export interface EntitySpendData { metadata: { total_spend: number; total_flat_cost?: number; + total_ptu_hours?: number; total_api_requests: number; total_successful_requests: number; total_failed_requests: number; @@ -41,6 +42,7 @@ export interface ExportMetadata { total_spend: number; total_flat_cost?: number; total_cost?: number; + total_ptu_hours?: number; total_requests: number; successful_requests: number; failed_requests: number; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index 3f9cf58ec20..e744f6deb68 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -2100,6 +2100,23 @@ describe("EntityUsageExport utils", () => { expect(result.summary.total_flat_cost).toBeUndefined(); expect(result.summary.total_cost).toBeUndefined(); }); + + it("should carry total_ptu_hours into the summary once a PTU model group was used", () => { + const withPtu: EntitySpendData = { + ...mockSpendData, + metadata: { ...mockSpendData.metadata, total_ptu_hours: 2.5 }, + }; + const result = generateMetadata("team", mockDateRange, [], "daily", withPtu); + expect(result.summary.total_ptu_hours).toBeCloseTo(2.5, 4); + }); + + it("should omit total_ptu_hours when no PTU hours were consumed", () => { + const zeroPtu = { ...mockSpendData, metadata: { ...mockSpendData.metadata, total_ptu_hours: 0 } }; + expect(generateMetadata("team", mockDateRange, [], "daily", zeroPtu).summary.total_ptu_hours).toBeUndefined(); + expect( + generateMetadata("team", mockDateRange, [], "daily", mockSpendData).summary.total_ptu_hours, + ).toBeUndefined(); + }); }); describe("generateDailyData PTU flat cost", () => { diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 95ce584cc89..f236ec5582c 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -436,6 +436,9 @@ export const generateMetadata = ( summary.total_flat_cost = flatCost; summary.total_cost = spendData.metadata.total_spend + flatCost; } + if ((spendData.metadata.total_ptu_hours ?? 0) > 0) { + summary.total_ptu_hours = spendData.metadata.total_ptu_hours; + } return { export_date: new Date().toISOString(), entity_type: entityType, diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index d53db68bb9f..7aa38cca38e 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -1,6 +1,7 @@ export interface SpendMetrics { spend: number; flat_cost?: number; + ptu_hours?: number; prompt_tokens: number; completion_tokens: number; total_tokens: number; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2938e65cfde..787b44d8477 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29311,6 +29311,11 @@ export interface components { * @default 0 */ total_prompt_tokens: number; + /** + * Total Ptu Hours + * @default 0 + */ + total_ptu_hours: number; /** * Total Response Time Ms * @default 0 @@ -42229,6 +42234,11 @@ export interface components { * @default 0 */ prompt_tokens: number; + /** + * Ptu Hours + * @default 0 + */ + ptu_hours: number; /** * Spend * @default 0 @@ -46156,6 +46166,10 @@ export interface components { ptu_effective_from?: string | null; /** Ptu Effective To */ ptu_effective_to?: string | null; + /** Ptu Shares */ + ptu_shares?: { + [key: string]: number; + } | null; /** Team Id */ team_id?: string | null; /** Team Public Model Name */ From 304e9d2ca0a0f6418b5b01045e52a92458b7ca0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:44:47 -0700 Subject: [PATCH 02/18] fix(proxy): weight cached input in PTU ceilings and size team-scoped deployments by public name The PTU flag is read from the process environment on every check instead of through the secret-manager path, so a proxy with a hosted secret manager does not pay a round trip per request. A team's ceiling now carries the model's cached-input ratio and settlement charges cache reads at that ratio, matching how Azure sizes a PTU. A deployment registered for one team is found by its team_public_model_name for PTU-hours and the ceiling, since its model_name is rewritten to a unique internal name. ptu_shares are StrictInt on ModelInfo and parsed_ptu_shares refuses booleans, strings, and non-string team ids, a fractional ptu_count reserves nothing, and models_processed counts distinct deployments rather than holdings. --- .../litellm_core_utils/azure_ptu_capacity.py | 14 +- litellm/litellm_core_utils/ptu_pricing.py | 49 ++++-- .../hooks/parallel_request_limiter_v3.py | 104 ++++++++---- .../management_endpoints/team_endpoints.py | 6 +- .../spend_tracking/ptu_flat_cost_rollup.py | 5 +- litellm/router_utils/ptu_shares.py | 21 ++- litellm/types/router.py | 4 +- .../test_azure_ptu_capacity.py | 11 +- .../litellm_core_utils/test_ptu_pricing.py | 23 ++- .../hooks/test_parallel_request_limiter_v3.py | 160 ++++++++++++++++-- .../test_ptu_model_settings.py | 23 ++- .../test_team_endpoints.py | 71 +++++++- .../spend_tracking/test_ptu_feature_flag.py | 12 ++ .../test_ptu_flat_cost_rollup.py | 5 +- .../router_utils/test_ptu_shares.py | 35 +++- 15 files changed, 448 insertions(+), 95 deletions(-) diff --git a/litellm/litellm_core_utils/azure_ptu_capacity.py b/litellm/litellm_core_utils/azure_ptu_capacity.py index 3ba4b9b6abf..c68429dc7b7 100644 --- a/litellm/litellm_core_utils/azure_ptu_capacity.py +++ b/litellm/litellm_core_utils/azure_ptu_capacity.py @@ -15,7 +15,15 @@ import re from collections.abc import Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import Final, Protocol + + +class NormalizedTokenWeights(Protocol): + @property + def output_to_input_ratio(self) -> float: ... + + @property + def cached_input_ratio(self) -> float: ... @dataclass(frozen=True, slots=True) @@ -90,13 +98,13 @@ def deployment_ptu_capacity(deployment: Mapping[str, object]) -> PTUCapacity | N def normalized_tokens( - capacity: PTUCapacity, *, prompt_tokens: int, completion_tokens: int, cache_read_tokens: int = 0 + weights: NormalizedTokenWeights, *, prompt_tokens: int, completion_tokens: int, cache_read_tokens: int = 0 ) -> float: """Azure's normalized token count for one request: uncached input in full, cached input at the model's cached ratio, output weighted by the output-to-input ratio.""" cached: Final = min(max(cache_read_tokens, 0), max(prompt_tokens, 0)) uncached: Final = max(prompt_tokens, 0) - cached - return uncached + capacity.cached_input_ratio * cached + capacity.output_to_input_ratio * max(completion_tokens, 0) + return uncached + weights.cached_input_ratio * cached + weights.output_to_input_ratio * max(completion_tokens, 0) def ptu_hours(capacity: PTUCapacity, normalized: float) -> float: diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py index b68b5cc93c2..e9cd57258ab 100644 --- a/litellm/litellm_core_utils/ptu_pricing.py +++ b/litellm/litellm_core_utils/ptu_pricing.py @@ -6,13 +6,14 @@ together because they have to agree: a deployment the rollup declines to charge router prices at zero serves its traffic for free. """ +import os from collections.abc import Mapping from dataclasses import dataclass from datetime import date, datetime, time, timezone from types import MappingProxyType from typing import Final -from litellm.secret_managers.main import get_secret_bool +from litellm.secret_managers.main import str_to_bool from litellm.types.router import ModelInfo from litellm.types.utils import AzureSpillover, CustomPricingLiteLLMParams, MirroredPricingParams @@ -22,8 +23,12 @@ AZURE_SPILLOVER_FROM_HEADER: Final = "x-ms-spillover-from-deployment" 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 + """Whether PTU flat-cost attribution is turned on for this process. + + Read from the environment alone: the router and the rate limiter ask on every request, + and ``get_secret`` would forward each of those reads to a configured secret manager. + """ + return str_to_bool(os.environ.get(PTU_COST_ATTRIBUTION_ENV_VAR)) is True PTU_ZEROED_PRICING_FIELDS: Final = tuple(f for f in MirroredPricingParams.model_fields if f != "tiered_pricing") + ( @@ -148,14 +153,30 @@ def parsed_ptu_shares(raw: object) -> Mapping[str, int] | None: """ if not isinstance(raw, Mapping) or not raw: return None - entries: Final = tuple((str(team_id), share) for team_id, share in raw.items()) - if any( - not team_id or isinstance(share, bool) or not isinstance(share, int) or share <= 0 for team_id, share in entries - ): + entries: Final = tuple( + (team_id, share) + for team_id, share in raw.items() + if isinstance(team_id, str) and team_id and isinstance(share, int) and not isinstance(share, bool) and share > 0 + ) + if len(entries) != len(raw): return None return MappingProxyType(dict(entries)) +def _parsed_ptu_count(model_info: Mapping[str, object]) -> int | None: + """``ptu_count`` as the whole number of reserved units within bounds, else None.""" + raw: Final = model_info.get("ptu_count") + if isinstance(raw, bool) or not isinstance(raw, (int, float, str)): + return None + if isinstance(raw, float) and not raw.is_integer(): + return None + try: + count: Final = int(raw) + except (ValueError, OverflowError): + return None + return count if 0 < count <= ModelInfo.MAX_PTU_COUNT else None + + def _declared_shares(model_info: Mapping[str, object], ptu_count: int) -> Mapping[str, int] | None: """Who holds the capacity: the single ``team_id`` holding all of it, or the ``ptu_shares`` that add up to it, else None.""" @@ -229,9 +250,8 @@ def _ptu_holder_error(model_info: Mapping[str, object], model_name: str | None) shares: Final = parsed_ptu_shares(raw_shares) if shares is None: return _named("ptu_shares must map at least one team_id to a positive whole number of PTUs", model_name) - try: - ptu_count: Final = int(str(model_info.get("ptu_count"))) - except ValueError: + ptu_count: Final = _parsed_ptu_count(model_info) + if ptu_count is None: return None allocated: Final = sum(shares.values()) if allocated != ptu_count: @@ -246,16 +266,13 @@ def ptu_terms(model_info: Mapping[str, object]) -> PTUTerms | None: 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") + ptu_count_int: Final = _parsed_ptu_count(model_info) cost_per_hour: Final = model_info.get("cost_per_ptu_per_hour") - if ptu_count is None or cost_per_hour is None: + if ptu_count_int is None or isinstance(cost_per_hour, bool) or not isinstance(cost_per_hour, (int, float, str)): 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: + except (ValueError, OverflowError): return None if not 0 <= cost_per_hour_float <= ModelInfo.MAX_COST_PER_PTU_PER_HOUR: return None diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c9f8e27aa5b..4fc994c5939 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -33,6 +33,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.azure_ptu_capacity import normalized_tokens from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -65,7 +66,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( response_has_hidden_params, ) from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.ptu_shares import PTUTeamCeiling, team_ptu_ceiling +from litellm.router_utils.ptu_shares import PTUTeamCeiling, model_group_deployments, team_ptu_ceiling from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -94,6 +95,17 @@ else: _REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object]) +@dataclass(frozen=True, slots=True) +class _ReconciledUsage: + prompt_tokens: int + completion_tokens: int + cached_tokens: int + + @property + def billable_input_tokens(self) -> int: + return max(0, self.prompt_tokens - self.cached_tokens) + + @dataclass(frozen=True, slots=True) class RateLimitedModel: requested: str @@ -119,7 +131,7 @@ def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) - if llm_router is None or not is_ptu_cost_attribution_enabled(): return None - return team_ptu_ceiling(llm_router.get_model_list(model_name=model_group) or (), team_id) + return team_ptu_ceiling(model_group_deployments(llm_router.get_model_list() or (), model_group), team_id) def _sibling_counter_keys(window_key: str) -> tuple[str, str]: @@ -430,9 +442,6 @@ _AUDIO_BYTES_PER_TOKEN: Final = 1600 # on the same project+model simultaneously without colliding on cache keys. PROJECT_ITPM_DESCRIPTOR_KEY: Final = "model_per_project_itpm" PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" -# Descriptor "key" for a team's PTU share of a shared Azure provisioned deployment, -# counted in Azure normalized tokens (output weighted by the model's ratio) so it -# never collides with the raw-token "model_per_team" counter on the same team+model. PTU_TEAM_DESCRIPTOR_KEY: Final = "model_per_team_ptu" # How long an acquired slot counts toward the in-flight total before it is # considered leaked (worker crashed without any release callback firing) and @@ -4200,20 +4209,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return max(0, total_tokens), 0, True return None - def _resolve_io_token_reconcile_usage( - self, - response_obj: object, - ) -> tuple[int, int, bool]: - """ - Resolve ``(billable_input_tokens, completion_tokens, usage_resolved)`` - for ITPM/OTPM reconciliation. Cache-read tokens are excluded from - billable input -- Bedrock Mantle doesn't count them toward ITPM -- - but they're untouched everywhere else (cost/usage logging still sees - the full prompt token count). - """ + def _resolve_reconciled_usage(self, response_obj: object) -> _ReconciledUsage | None: + """The prompt, completion, and cache-read token counts a response reports, else None + when it reports no usage at all. Cache-read tokens stay inside ``prompt_tokens`` here; + each consumer decides what they cost it.""" rerank_usage: Final = self._resolve_rerank_token_usage(response_obj) if rerank_usage is not None: - return rerank_usage + return _ReconciledUsage(prompt_tokens=rerank_usage[0], completion_tokens=rerank_usage[1], cached_tokens=0) usage: Final = self._response_usage(response_obj) @@ -4226,8 +4228,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else 0 ) if prompt_tokens == 0 and completion_tokens == 0: - return 0, 0, False - return max(0, prompt_tokens - cached_tokens), completion_tokens, True + return None + return _ReconciledUsage( + prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, cached_tokens=cached_tokens + ) if isinstance(usage, ResponseAPIUsage): response_input_tokens: Final = usage.input_tokens or 0 @@ -4236,8 +4240,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): usage.input_tokens_details.cached_tokens or 0 if usage.input_tokens_details is not None else 0 ) if response_input_tokens == 0 and response_output_tokens == 0: - return 0, 0, False - return max(0, response_input_tokens - response_cached_tokens), response_output_tokens, True + return None + return _ReconciledUsage( + prompt_tokens=response_input_tokens, + completion_tokens=response_output_tokens, + cached_tokens=response_cached_tokens, + ) if isinstance(usage, Mapping): raw_prompt_tokens: Final = usage.get("prompt_tokens") or usage.get("input_tokens") or 0 @@ -4252,10 +4260,30 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) mapped_cached_tokens: Final = raw_cached_tokens if isinstance(raw_cached_tokens, int) else 0 if mapped_prompt_tokens == 0 and mapped_completion_tokens == 0: - return 0, 0, False - return max(0, mapped_prompt_tokens - mapped_cached_tokens), mapped_completion_tokens, True + return None + return _ReconciledUsage( + prompt_tokens=mapped_prompt_tokens, + completion_tokens=mapped_completion_tokens, + cached_tokens=mapped_cached_tokens, + ) - return 0, 0, False + return None + + def _resolve_io_token_reconcile_usage( + self, + response_obj: object, + ) -> tuple[int, int, bool]: + """ + Resolve ``(billable_input_tokens, completion_tokens, usage_resolved)`` + for ITPM/OTPM reconciliation. Cache-read tokens are excluded from + billable input -- Bedrock Mantle doesn't count them toward ITPM -- + but they're untouched everywhere else (cost/usage logging still sees + the full prompt token count). + """ + usage: Final = self._resolve_reconciled_usage(response_obj) + if usage is None: + return 0, 0, False + return usage.billable_input_tokens, usage.completion_tokens, True def _build_io_token_reservation_ops( self, @@ -4617,27 +4645,37 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reserved_tokens: int, total_tokens: int, ) -> Sequence[RedisPipelineIncrementOperation]: - """Settle the team's PTU counter in Azure normalized tokens: uncached input in full plus - output weighted by the model's output-to-input ratio, the way Azure sizes a PTU. + """Settle the team's PTU counter in Azure normalized tokens: uncached input in full, + cached input at the model's cached ratio, output weighted by its output-to-input + ratio, the way Azure sizes a PTU. The pre-call reservation was raw estimated tokens, so this is the same reconcile as the - other TPM scopes with a weighted actual; when usage cannot be resolved it charges the raw - total the other scopes charge. + other TPM scopes with a weighted actual; when usage cannot be resolved, or the ceiling + is gone since the reservation was taken, it charges the raw total the other scopes + charge so the reservation is never left standing. """ team_id: Final = standard_logging_metadata.get("user_api_key_team_id") if reconcile_model is None or not isinstance(team_id, str) or not team_id: return () + scope: Final = (PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{reconcile_model.group}") ceiling: Final = self._ptu_team_ceiling_resolver(team_id, reconcile_model.group) - if ceiling is None: + if ceiling is None and scope not in reserved_scopes: return () - billable_input, completion_tokens, usage_resolved = self._resolve_io_token_reconcile_usage(response_obj) + usage: Final = self._resolve_reconciled_usage(response_obj) normalized: Final = ( - billable_input + round(ceiling.output_to_input_ratio * completion_tokens) - if usage_resolved + round( + normalized_tokens( + ceiling, + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + cache_read_tokens=usage.cached_tokens, + ) + ) + if ceiling is not None and usage is not None else total_tokens ) return self._build_reservation_aware_tpm_ops( - targets=((PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{reconcile_model.group}"),), + targets=(scope,), reserved_scopes=reserved_scopes, actual_tokens=normalized, reserved_tokens=reserved_tokens, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 34afa030287..b3dca8b127f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -201,7 +201,7 @@ from litellm.repositories.verification_token_repository import ( VerificationTokenRepository, ) from litellm.router import Router -from litellm.router_utils.ptu_shares import model_group_ptu_capacity +from litellm.router_utils.ptu_shares import model_group_deployments, model_group_ptu_capacity from litellm.types.proxy.auth.auth_checks import UserNotFoundError from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendMetadata, @@ -6670,7 +6670,9 @@ def _with_ptu_consumption( return activity return attach_ptu_hours( activity, - lambda model_group: model_group_ptu_capacity(llm_router.get_model_list(model_name=model_group) or ()), + lambda model_group: model_group_ptu_capacity( + model_group_deployments(llm_router.get_model_list() or (), model_group) + ), ) diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index e2c0ccb0327..1ca0abf19ed 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -454,16 +454,17 @@ async def run_ptu_flat_cost_rollup( scanned_ids=loaded.scanned_ids, ) + models_processed: Final = len(frozenset(model.model_id for model in ptu_models)) verbose_proxy_logger.info( "PTU rollup for %s: %d PTU models processed, %d rows written, %d rows failed", date_str, - len(ptu_models), + models_processed, rows_written, rows_failed, ) return RollupResult( day=day, - models_processed=len(ptu_models), + models_processed=models_processed, rows_written=rows_written, rows_failed=rows_failed, lapsed=_lapsed_models(ptu_models, run_started), diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 04e54395b5a..1401dd23d2e 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -20,6 +20,7 @@ _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) class PTUTeamCeiling: tpm_limit: int output_to_input_ratio: float + cached_input_ratio: float @dataclass(frozen=True, slots=True) @@ -55,8 +56,8 @@ def team_ptu_ceiling(deployments: Sequence[Mapping[str, object]], team_id: str) up to, else None when the team holds no share on a deployment with a known sizing row. Two shared deployments of different models in one group are weighted by the larger - output ratio, which over-counts output on the cheaper one rather than under-counting it - on the dearer one. + output and cached-input ratios, which over-counts those tokens on the cheaper one rather + than under-counting them on the dearer one. """ priced: Final = tuple( (shares[team_id], capacity) @@ -70,6 +71,22 @@ def team_ptu_ceiling(deployments: Sequence[Mapping[str, object]], team_id: str) return PTUTeamCeiling( tpm_limit=sum(share * capacity.input_tpm_per_ptu for share, capacity in priced), output_to_input_ratio=max(capacity.output_to_input_ratio for _, capacity in priced), + cached_input_ratio=max(capacity.cached_input_ratio for _, capacity in priced), + ) + + +def model_group_deployments(deployments: Sequence[_DeploymentT], model_group: str) -> tuple[_DeploymentT, ...]: + """Every deployment serving ``model_group``: by its own name, or by the public name a + team-scoped deployment keeps in ``model_info.team_public_model_name`` after the router + renames it to a unique internal one.""" + return tuple( + deployment + for deployment in deployments + if deployment.get("model_name") == model_group + or ( + isinstance(model_info := deployment.get("model_info"), Mapping) + and model_info.get("team_public_model_name") == model_group + ) ) diff --git a/litellm/types/router.py b/litellm/types/router.py index 5ead0dd7626..7c91b8f099f 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Final, Generic, Lite from zoneinfo import ZoneInfo, ZoneInfoNotFoundError import httpx -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, StrictInt, field_validator, model_validator from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable from litellm._logging import verbose_logger @@ -260,7 +260,7 @@ class ModelInfo(MirroredPricingParams): cost_per_ptu_per_hour: float | None = None ptu_effective_from: datetime.datetime | None = None ptu_effective_to: datetime.datetime | None = None - ptu_shares: Mapping[str, int] | None = None + ptu_shares: Mapping[str, StrictInt] | None = None # when tag-based routing's "!" or "&" constraints eliminate every deployment # in this model group, fall back to the default-tagged pool instead of diff --git a/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py b/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py index d07d71f0d82..4d96068cd8b 100644 --- a/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py +++ b/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py @@ -55,8 +55,9 @@ def test_a_deployment_prefers_its_declared_base_model_over_its_deployment_name() def test_a_deployment_falls_back_to_its_litellm_model_when_no_base_model_is_declared(): assert deployment_ptu_capacity({"litellm_params": {"model": "azure/gpt-4o"}}) is AZURE_PTU_CAPACITY["gpt-4o"] - assert deployment_ptu_capacity({"model_info": {"base_model": ""}, "litellm_params": {"model": "azure/gpt-4o"}}) is ( - AZURE_PTU_CAPACITY["gpt-4o"] + assert ( + deployment_ptu_capacity({"model_info": {"base_model": ""}, "litellm_params": {"model": "azure/gpt-4o"}}) + is (AZURE_PTU_CAPACITY["gpt-4o"]) ) @@ -72,9 +73,9 @@ def test_output_is_weighted_by_the_models_ratio_and_uncached_input_counts_in_ful def test_cached_input_is_free_unless_the_row_prices_it(): assert normalized_tokens(_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60) == pytest.approx(40.0) - assert normalized_tokens(_CACHED_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60) == pytest.approx( - 46.0 - ) + assert normalized_tokens( + _CACHED_ROW, prompt_tokens=100, completion_tokens=0, cache_read_tokens=60 + ) == pytest.approx(46.0) def test_cached_input_never_exceeds_the_prompt_and_negatives_count_as_zero(): diff --git a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py index ef8d7703924..7dda3432e43 100644 --- a/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py +++ b/tests/test_litellm/litellm_core_utils/test_ptu_pricing.py @@ -395,8 +395,13 @@ def test_a_single_team_reservation_holds_the_whole_count_under_that_team(): {"ptu_shares": {"team-a": 50.5, "team-b": 49.5}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs", ), - ({"ptu_shares": {"team-a": True}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ( + {"ptu_shares": {"team-a": True}}, + "ptu_shares must map at least one team_id to a positive whole number of PTUs", + ), ({"ptu_shares": {"": 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ({"ptu_shares": {None: 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), + ({"ptu_shares": {1: 100}}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), ({"ptu_shares": ["team-a"]}, "ptu_shares must map at least one team_id to a positive whole number of PTUs"), ({"ptu_shares": {"team-a": 60, "team-b": 30}}, "ptu_shares must add up to ptu_count (90 of 100 allocated)"), ({"ptu_shares": {"team-a": 60, "team-b": 50}}, "ptu_shares must add up to ptu_count (110 of 100 allocated)"), @@ -408,6 +413,8 @@ def test_a_single_team_reservation_holds_the_whole_count_under_that_team(): "fractional share", "boolean share", "blank team", + "null team", + "numeric team", "not a mapping", "shares short of the count", "shares over the count", @@ -418,6 +425,20 @@ def test_an_incoherent_split_names_its_reason_and_reserves_nothing(override, exp assert ptu_terms({**_SHARED, **override}) is None +def test_a_whole_count_written_as_a_float_is_checked_against_the_shares_all_the_same(): + assert ptu_config_error({**_SHARED, "ptu_count": 100.0, "ptu_shares": {"team-a": 60, "team-b": 30}}) == ( + "ptu_shares must add up to ptu_count (90 of 100 allocated)" + ) + terms = ptu_terms({**_SHARED, "ptu_count": 100.0}) + assert terms is not None + assert terms.ptu_count == 100 + + +def test_a_fractional_count_reserves_nothing(): + assert ptu_terms({**_VALID, "ptu_count": 100.5}) is None + assert ptu_terms({**_SHARED, "ptu_count": 100.5}) is None + + def test_the_split_is_named_after_the_deployment_when_the_caller_supplies_one(): error = ptu_config_error({**_SHARED, "ptu_shares": {"team-a": 1}}, model_name="gpt-4.1-ptu") diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index a18b38453b0..8f44945718d 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -11,6 +11,7 @@ from collections.abc import Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone from typing import Any, Dict, Final, List, Optional +from unittest.mock import patch import pytest from fastapi import HTTPException @@ -20,6 +21,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY +from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( @@ -36,6 +38,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler, ) +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.router_utils.ptu_shares import PTUTeamCeiling from litellm.types.caching import RedisPipelineIncrementOperation @@ -4726,8 +4729,8 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): from the outer generator: the counter returns to 0 (not -1), proving the nested hook does not also refund and there is no double decrement. """ - from litellm.integrations.custom_logger import CustomLogger import litellm.proxy.proxy_server as proxy_server + from litellm.integrations.custom_logger import CustomLogger class _PassthroughIteratorOverride(CustomLogger): async def async_post_call_streaming_iterator_hook( @@ -7157,14 +7160,14 @@ async def test_success_tpm_accounting_keeps_the_admission_target_after_an_alias_ # --- a team's PTU share on a shared Azure provisioned deployment --------------------------- -def _ptu_ceiling_for(team_id: str, model_group: str, tpm_limit: int, ratio: float): +def _ptu_ceiling_for(team_id: str, model_group: str, tpm_limit: int, ratio: float, cached_ratio: float = 0.0): calls: list[tuple[str, str]] = [] def resolve(requested_team: str, requested_group: str) -> PTUTeamCeiling | None: calls.append((requested_team, requested_group)) if (requested_team, requested_group) != (team_id, model_group): return None - return PTUTeamCeiling(tpm_limit=tpm_limit, output_to_input_ratio=ratio) + return PTUTeamCeiling(tpm_limit=tpm_limit, output_to_input_ratio=ratio, cached_input_ratio=cached_ratio) return resolve, calls @@ -7177,12 +7180,16 @@ def _ptu_request(model: str = "test-model") -> dict: async def test_a_teams_ptu_share_is_a_hard_tpm_ceiling_on_the_shared_model(): cache = DualCache() resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion") with pytest.raises(HTTPException) as exc: - await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion") + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion" + ) assert exc.value.status_code == 429 assert "model_per_team_ptu" in str(exc.value.detail) @@ -7193,7 +7200,9 @@ async def test_a_teams_ptu_share_is_a_hard_tpm_ceiling_on_the_shared_model(): async def test_a_ptu_ceiling_on_one_model_leaves_the_teams_other_models_alone(): cache = DualCache() resolve, calls = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") for _ in range(3): @@ -7208,28 +7217,102 @@ async def test_a_ptu_ceiling_on_one_model_leaves_the_teams_other_models_alone(): async def test_a_team_without_a_share_and_a_key_without_a_team_get_no_ptu_ceiling(): cache = DualCache() resolve, calls = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) other_team = UserAPIKeyAuth(api_key=hash_token("sk-other"), team_id="u") no_team = UserAPIKeyAuth(api_key=hash_token("sk-no-team")) for _ in range(3): - await handler.async_pre_call_hook(user_api_key_dict=other_team, cache=cache, data=_ptu_request(), call_type="acompletion") - await handler.async_pre_call_hook(user_api_key_dict=no_team, cache=cache, data=_ptu_request(), call_type="acompletion") + for caller in (other_team, no_team): + await handler.async_pre_call_hook( + user_api_key_dict=caller, cache=cache, data=_ptu_request(), call_type="acompletion" + ) assert set(calls) == {("u", "test-model")} +def _shared_ptu_router(model_group: str) -> Router: + return Router( + model_list=[ + { + "model_name": model_group, + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "shared-ptu", + "base_model": "azure/gpt-4.1", + "ptu_count": 1, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"t": 1}, + }, + } + ] + ) + + +def _two_thirds_of_a_ptu_minute() -> dict: + return {**_ptu_request(), "max_tokens": AZURE_PTU_CAPACITY["gpt-4.1"].input_tpm_per_ptu * 2 // 3} + + +@pytest.mark.asyncio +async def test_the_proxy_router_turns_a_teams_share_into_its_ceiling_when_attribution_is_on(monkeypatch): + """With no resolver injected the ceiling comes from the proxy router's own deployments: one + PTU of gpt-4.1 a minute, so two requests each reserving two thirds of it are one too many.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("test-model")): + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion" + ) + with pytest.raises(HTTPException) as exc: + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion" + ) + + assert exc.value.status_code == 429 + assert "model_per_team_ptu" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_the_proxy_router_sets_no_ceiling_while_attribution_is_off(monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("test-model")): + for _ in range(3): + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion" + ) + + assert not any("model_per_team_ptu" in cache_key for cache_key in cache.in_memory_cache.cache_dict) + + def _ptu_success_kwargs() -> dict: return { - "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"}}, - "litellm_params": {"metadata": {"model_group": "test-model", "user_api_key_metadata": {}, "user_api_key_team_metadata": {}}}, + "standard_logging_object": { + "metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"} + }, + "litellm_params": { + "metadata": {"model_group": "test-model", "user_api_key_metadata": {}, "user_api_key_team_metadata": {}} + }, "model": "test-model", } def _ptu_response(usage: Usage) -> ModelResponse: return ModelResponse( - id="ptu-share", object="chat.completion", created=int(datetime.now().timestamp()), model="test-model", usage=usage, choices=[] + id="ptu-share", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="test-model", + usage=usage, + choices=[], ) @@ -7275,6 +7358,59 @@ def test_cached_input_is_not_charged_to_the_ptu_counter(): assert _ptu_increment(handler, ops) == 60 + 4 * 50 +def test_cached_input_is_charged_at_the_models_cached_ratio(): + """40 of the 100 input tokens were cache reads; at a tenth each they are 4 normalized tokens + beside the 60 uncached ones and the 200 for 50 outputs at 4:1.""" + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0, cached_ratio=0.1) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()), ptu_team_ceiling_resolver=resolve + ) + response = _ptu_response( + Usage( + prompt_tokens=100, + completion_tokens=50, + total_tokens=150, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=40), + ) + ) + + ops = handler._build_success_event_pipeline_operations( + kwargs=_ptu_success_kwargs(), response_obj=response, rate_limit_type="output" + ) + + assert _ptu_increment(handler, ops) == 60 + 4 + 4 * 50 + + +@pytest.mark.asyncio +async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone(): + """The share can be removed between admission and completion; the reserved tokens still + come off the counter instead of standing in the window.""" + ceiling: dict[str, PTUTeamCeiling | None] = { + "current": PTUTeamCeiling(tpm_limit=500, output_to_input_ratio=4.0, cached_input_ratio=0.0) + } + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), + ptu_team_ceiling_resolver=lambda _team, _group: ceiling["current"], + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion") + stash = get_request_stash() + assert stash is not None + assert ("model_per_team_ptu", "t:test-model") in stash.reserved_scopes + assert stash.reserved_tokens > 150 + + ceiling["current"] = None + ops = handler._build_success_event_pipeline_operations( + kwargs=_ptu_success_kwargs(), + response_obj=_ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)), + rate_limit_type="total", + ) + + assert _ptu_increment(handler, ops) == 150 - stash.reserved_tokens + + def test_usage_that_only_reports_a_total_charges_that_total_to_the_ptu_counter(): resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) handler = _PROXY_MaxParallelRequestsHandler( diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index a7c21d74f1e..8ed3df44216 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -9,31 +9,30 @@ from unittest.mock import patch as patch_ctx import pytest from fastapi import HTTPException +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.ptu_pricing import ptu_terms +from litellm.llms.gemini.cost_calculator import cost_per_web_search_request from litellm.proxy._types import ( LiteLLM_ProxyModelTable, LitellmUserRoles, ReconcileOutcome, UserAPIKeyAuth, ) -from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.proxy.auth.auth_checks import _is_model_cost_zero -from litellm.llms.gemini.cost_calculator import cost_per_web_search_request from litellm.proxy.management_endpoints.model_management_endpoints import ( _PTU_ZEROED_PRICING_FIELDS, _SEARCH_CONTEXT_SIZES, _is_nonzero_price, _merged_ptu_model_info, - _update_team_model_in_db, _ptu_priced_deployment, _ptu_zeroed_pricing, _raise_if_ptu_cost_attribution_disabled, + _update_team_model_in_db, _validate_ptu_model_info, add_new_model, update_db_model, ) from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.types.utils import PromptTokensDetailsWrapper from litellm.router import Router from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, @@ -43,7 +42,7 @@ from litellm.types.router import ( updateDeployment, updateLiteLLMParams, ) -from litellm.types.utils import Usage +from litellm.types.utils import PromptTokensDetailsWrapper, Usage async def _passthrough_row(update_data): @@ -116,6 +115,18 @@ def test_model_info_allows_partial_delta_for_patch(): assert info.cost_per_ptu_per_hour is None +@pytest.mark.parametrize("share", [True, "2", 2.0]) +def test_model_info_rejects_a_share_that_is_not_a_whole_number(share): + with pytest.raises(ValueError, match="ptu_shares"): + ModelInfo(id="x", ptu_shares={"team-a": share}) + + +def test_model_info_keeps_whole_number_shares_and_refuses_a_fractional_count(): + assert ModelInfo(id="x", ptu_shares={"team-a": 2}).ptu_shares == {"team-a": 2} + with pytest.raises(ValueError, match="ptu_count"): + ModelInfo(id="x", ptu_count=100.5) + + def test_validate_helper_no_ptu_is_noop(): _validate_ptu_model_info({"team_id": "t"}) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 5a1b87b7a01..7e0792cdc2c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1,9 +1,9 @@ import asyncio import json +from collections.abc import Sequence from contextlib import asynccontextmanager, contextmanager from datetime import datetime, timezone from types import SimpleNamespace -from collections.abc import Sequence from typing import Final, Optional, cast from unittest.mock import AsyncMock, MagicMock, PropertyMock, call, patch @@ -15,6 +15,7 @@ from pydantic import ValidationError from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, @@ -17106,6 +17107,9 @@ def test_list_team_v2_answers_503_no_db_connection_when_the_callers_user_read_hi # --- PTU-equivalent consumption on /team/daily/activity ------------------------------------ +_ONE_PTU_HOUR_OF_INPUT: Final = AZURE_PTU_CAPACITY["gpt-4.1"].normalized_tokens_per_ptu_hour + + def _ptu_activity_page(): from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, @@ -17116,8 +17120,9 @@ def _ptu_activity_page(): SpendMetrics, ) + tokens: Final = _ONE_PTU_HOUR_OF_INPUT metrics = SpendMetrics( - prompt_tokens=180_000, completion_tokens=0, total_tokens=180_000, api_requests=3, successful_requests=3 + prompt_tokens=tokens, completion_tokens=0, total_tokens=tokens, api_requests=3, successful_requests=3 ) return SpendAnalyticsPaginatedResponse( results=[ @@ -17129,7 +17134,7 @@ def _ptu_activity_page(): ), ) ], - metadata=DailySpendMetadata(total_tokens=180_000, total_api_requests=3, total_successful_requests=3), + metadata=DailySpendMetadata(total_tokens=tokens, total_api_requests=3, total_successful_requests=3), ) @@ -17157,9 +17162,9 @@ def _shared_ptu_router() -> Router: async def test_team_daily_activity_reports_ptu_hours_only_while_attribution_is_on( mock_db_client, mock_admin_auth, monkeypatch, attribution_enabled ): - """One PTU serves 3,000 input tokens per minute on gpt-4.1, so 180,000 uncached - input tokens are one PTU-hour; the figure appears beside tokens only once the - PTU flag is on, and the token totals are untouched either way.""" + """An hour of one PTU's input rate on gpt-4.1, in uncached input tokens, is one PTU-hour; + the figure appears beside tokens only once the PTU flag is on, and the token totals are + untouched either way.""" from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True" if attribution_enabled else "False") @@ -17186,8 +17191,58 @@ async def test_team_daily_activity_reports_ptu_hours_only_while_attribution_is_o assert result.metadata.total_ptu_hours == expected_ptu_hours assert result.results[0].metrics.ptu_hours == expected_ptu_hours assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == expected_ptu_hours - assert result.metadata.total_tokens == 180_000 - assert result.results[0].metrics.total_tokens == 180_000 + assert result.metadata.total_tokens == _ONE_PTU_HOUR_OF_INPUT + assert result.results[0].metrics.total_tokens == _ONE_PTU_HOUR_OF_INPUT + + +@pytest.mark.asyncio +async def test_team_daily_activity_sizes_a_team_scoped_deployment_by_its_public_name( + mock_db_client, mock_admin_auth, monkeypatch +): + """A deployment registered for one team is renamed to a unique internal name and keeps the + name callers use in ``team_public_model_name``, which is the name the activity rows carry, + so its sizing row is still found.""" + from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity + + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + page = _ptu_activity_page() + team_scoped_router: Final = Router( + model_list=[ + { + "model_name": "gpt-4.1-ptu-3f9c1b", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "team-a-ptu", + "base_model": "azure/gpt-4.1", + "team_id": "team-a", + "team_public_model_name": "gpt-4.1-ptu", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + }, + } + ] + ) + + with ( + patch("litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", AsyncMock(return_value=page)), + patch("litellm.proxy.proxy_server.llm_router", team_scoped_router), + ): + result = await get_team_daily_activity( + team_ids="team-a", + start_date="2026-09-23", + end_date="2026-09-24", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=mock_admin_auth, + ) + + assert result.metadata.total_ptu_hours == 1.0 + assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 1.0 def test_team_export_csv_columns_match_the_dashboard_client_layout(): diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py index 7f4bd935a2b..23d0b4ed84f 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py @@ -31,3 +31,15 @@ def test_reads_the_env_var_on_every_call(monkeypatch): monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") assert is_ptu_cost_attribution_enabled() is True + + +def test_reads_the_process_environment_and_never_a_secret_manager(monkeypatch): + """The flag is checked on every request and every rollup, so it never costs a round trip to + a hosted secret manager even when one is configured for reads.""" + + def refuse(*args: object, **kwargs: object) -> object: + raise AssertionError("secret manager consulted") + + monkeypatch.setattr("litellm.secret_managers.main.get_secret", refuse) + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + assert is_ptu_cost_attribution_enabled() is True 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 73e3e5f3ec6..2bcb31c4db2 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 @@ -2119,7 +2119,10 @@ async def test_rollup_splits_a_shared_deployments_flat_cost_by_share(): result = await run_ptu_flat_cost_rollup(prisma, target_date=DAY) assert result.rows_written == 2 - created = {call.kwargs["data"]["create"]["team_id"]: call.kwargs["data"]["create"] for call in table.upsert.await_args_list} + assert result.models_processed == 1 + created = { + call.kwargs["data"]["create"]["team_id"]: call.kwargs["data"]["create"] for call in table.upsert.await_args_list + } assert created["team-a"]["ptu_flat_cost"] == pytest.approx(720.0) assert created["team-b"]["ptu_flat_cost"] == pytest.approx(480.0) assert sum(row["ptu_flat_cost"] for row in created.values()) == pytest.approx(50 * 1.0 * 24) diff --git a/tests/test_litellm/router_utils/test_ptu_shares.py b/tests/test_litellm/router_utils/test_ptu_shares.py index 150cfd31b75..110161bfa93 100644 --- a/tests/test_litellm/router_utils/test_ptu_shares.py +++ b/tests/test_litellm/router_utils/test_ptu_shares.py @@ -6,6 +6,7 @@ from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY from litellm.router_utils.ptu_shares import ( PTUTeamCeiling, filter_ptu_shared_deployments, + model_group_deployments, model_group_ptu_capacity, ptu_capacity_warning, team_ptu_ceiling, @@ -13,6 +14,7 @@ from litellm.router_utils.ptu_shares import ( _GPT41: Final = AZURE_PTU_CAPACITY["gpt-4.1"] _GPT4O: Final = AZURE_PTU_CAPACITY["gpt-4o"] +_GPT6SOL: Final = AZURE_PTU_CAPACITY["gpt-6-sol"] _SHARES: Final = {"team-a": 30, "team-b": 20} @@ -77,7 +79,9 @@ def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_her def test_a_share_converts_to_the_models_input_tpm_per_ptu(): ceiling: Final = team_ptu_ceiling([_shared()], "team-a") assert ceiling == PTUTeamCeiling( - tpm_limit=30 * _GPT41.input_tpm_per_ptu, output_to_input_ratio=_GPT41.output_to_input_ratio + tpm_limit=30 * _GPT41.input_tpm_per_ptu, + output_to_input_ratio=_GPT41.output_to_input_ratio, + cached_input_ratio=_GPT41.cached_input_ratio, ) @@ -89,6 +93,32 @@ def test_shares_across_deployments_add_up_and_the_larger_output_ratio_wins(): assert ceiling.output_to_input_ratio == max(_GPT41.output_to_input_ratio, _GPT4O.output_to_input_ratio) +def test_the_larger_cached_input_ratio_wins_across_deployments(): + """A team sharing two models is weighted by the one that charges more for cache reads, + whichever order the deployments come in.""" + gpt6sol: Final = _shared(model="azure/gpt-6-sol", shares={"team-a": 10}, deployment_id="shared-6") + ceiling: Final = team_ptu_ceiling([_shared(), gpt6sol], "team-a") + assert ceiling is not None + assert _GPT41.cached_input_ratio < _GPT6SOL.cached_input_ratio + assert ceiling.cached_input_ratio == _GPT6SOL.cached_input_ratio + + +def test_a_group_is_served_by_name_or_by_a_team_scoped_deployments_public_name(): + """A deployment registered for one team is renamed to a unique internal name and keeps + the name callers use in ``team_public_model_name``.""" + team_scoped: Final = { + "model_name": "gpt-4.1-ptu-3f9c1b", + "litellm_params": {"model": "azure/gpt-4.1"}, + "model_info": {"id": "team-scoped", "team_id": "team-a", "team_public_model_name": "gpt-4.1-ptu"}, + } + other: Final = {"model_name": "other", "litellm_params": {"model": "azure/gpt-4o"}, "model_info": {"id": "other"}} + deployments: Final = [team_scoped, _OPEN, other] + served: Final = model_group_deployments(deployments, "gpt-4.1-ptu") + assert [d["model_info"]["id"] for d in served] == ["team-scoped", "open"] + assert model_group_deployments(deployments, "gpt-4.1-ptu-3f9c1b") == (team_scoped,) + assert model_group_deployments(deployments, "missing") == () + + def test_no_share_or_no_sizing_row_sets_no_ceiling(): assert team_ptu_ceiling([_shared()], "team-c") is None assert team_ptu_ceiling([_shared(model="azure/unknown-deployment")], "team-a") is None @@ -112,4 +142,5 @@ def test_a_reserved_deployment_without_a_sizing_row_is_warned_about_by_name(): def test_a_sized_reservation_and_an_unreserved_deployment_raise_no_warning(): assert ptu_capacity_warning("gpt-4.1-ptu", _shared()) is None assert ptu_capacity_warning("gpt-4.1-ptu", _single_team()) is None - assert ptu_capacity_warning("gpt-4.1-ptu", {**_OPEN, "litellm_params": {"model": "azure/my-ptu-deployment"}}) is None + unsized_open: Final = {**_OPEN, "litellm_params": {"model": "azure/my-ptu-deployment"}} + assert ptu_capacity_warning("gpt-4.1-ptu", unsized_open) is None From 0d80930df80be57e27fafe67178df19ad06ad08f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 15:39:14 -0700 Subject: [PATCH 03/18] fix(proxy): reserve PTU ceilings in normalized tokens, keep share routing flag-independent, and warn only shared or Azure reservations --- .../azure/ptu_capacity.py} | 0 .../hooks/parallel_request_limiter_v3.py | 145 ++++++++++++++---- .../management_endpoints/ptu_consumption.py | 2 +- litellm/router.py | 2 - litellm/router_utils/ptu_shares.py | 39 ++++- .../azure}/test_azure_ptu_capacity.py | 2 +- .../hooks/test_parallel_request_limiter_v3.py | 116 +++++++++++++- .../test_ptu_consumption.py | 2 +- .../test_team_endpoints.py | 2 +- .../router_utils/test_ptu_shares.py | 23 ++- tests/test_litellm/test_router.py | 7 +- .../test_router_model_cost_isolation.py | 2 +- 12 files changed, 289 insertions(+), 53 deletions(-) rename litellm/{litellm_core_utils/azure_ptu_capacity.py => llms/azure/ptu_capacity.py} (100%) rename tests/test_litellm/{litellm_core_utils => llms/azure}/test_azure_ptu_capacity.py (98%) diff --git a/litellm/litellm_core_utils/azure_ptu_capacity.py b/litellm/llms/azure/ptu_capacity.py similarity index 100% rename from litellm/litellm_core_utils/azure_ptu_capacity.py rename to litellm/llms/azure/ptu_capacity.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 4fc994c5939..c61b35f5936 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -33,12 +33,12 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.azure_ptu_capacity import normalized_tokens from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( ESTIMATED_OUTPUT_TOKENS_FIELD, @@ -592,6 +592,8 @@ class RequestRateLimiterStash: reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) + ptu_ceiling: PTUTeamCeiling | None = None + ptu_reserved_tokens: int = 0 itpm_reserved_tokens: int = 0 itpm_reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset) itpm_reserved_window_identities: frozenset[tuple[str, str, Literal["redis", "local"]]] = field( @@ -973,6 +975,30 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return total_estimated + def _estimate_ptu_tokens_for_request( + self, + ceiling: PTUTeamCeiling | None, + data: dict, + min_configured_tpm_limit: int | None, + call_type: str | None, + configured_output_tokens: int | None, + raw_estimate: int, + ) -> int: + """The team PTU ceiling counts Azure normalized tokens, so its reservation weighs the + output budget the way the ceiling does instead of the raw sum the other scopes reserve.""" + if ceiling is None: + return raw_estimate + estimated_input_tokens, max_tokens_estimate = self._estimate_input_and_output_tokens( + data=data, + min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, + configured_output_tokens=configured_output_tokens, + ) + normalized: Final = normalized_tokens( + ceiling, prompt_tokens=estimated_input_tokens, completion_tokens=max_tokens_estimate + ) + return max(round(normalized), 1) + def _estimate_input_and_output_tokens( self, data: object, @@ -2184,11 +2210,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors: list[RateLimitDescriptor], estimated_tokens: int, parent_otel_span: Span | None = None, + scope_estimates: Mapping[str, int] = MappingProxyType({}), ) -> RateLimitResponse: """ Reserve ``estimated_tokens`` against every TPM-bearing descriptor BEFORE the upstream call, so concurrent requests cannot all observe "under limit" before any of them increments the counter. + ``scope_estimates`` replaces that amount per descriptor key for a + scope counted in other units, such as the team PTU ceiling. Thin wrapper around ``atomic_check_and_increment_by_n``: builds a TPM-only descriptor/increment list and delegates the all-or-nothing @@ -2215,7 +2244,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return RateLimitResponse(overall_code="OK", statuses=[]) increments: Final[list[dict[Literal["requests", "tokens"], int]]] = [ - {"tokens": estimated_tokens} for _ in tpm_descriptors + {"tokens": scope_estimates.get(d["key"], estimated_tokens)} for d in tpm_descriptors ] return await self.atomic_check_and_increment_by_n( descriptors=tpm_descriptors, @@ -3061,6 +3090,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ceiling: Final = self._ptu_team_ceiling_resolver(user_api_key_dict.team_id, model.group) if ceiling is None: return + stash: Final = get_request_stash() + if stash is not None: + stash.ptu_ceiling = ceiling descriptors.append( RateLimitDescriptor( key=PTU_TEAM_DESCRIPTOR_KEY, @@ -3482,10 +3514,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # (still-stashed) reservation and refunds it again. if tpm_reservation_amount > 0: await self._refund_reserved_tokens( - scopes=tpm_reservation_scopes, + scopes=tuple(scope for scope in tpm_reservation_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY), amount=tpm_reservation_amount, parent_otel_span=user_api_key_dict.parent_otel_span, ) + await self._refund_reserved_tokens( + scopes=tuple(scope for scope in tpm_reservation_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY), + amount=stash.ptu_reserved_tokens, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) stash.reservation_released = True await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( @@ -3795,10 +3832,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit, ) + ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request( + ceiling=stash.ptu_ceiling, + data=data, + min_configured_tpm_limit=min_configured_tpm_limit, + call_type=call_type, + configured_output_tokens=configured_output_tokens, + raw_estimate=estimated_tokens, + ) tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, estimated_tokens=estimated_tokens, parent_otel_span=user_api_key_dict.parent_otel_span, + scope_estimates=MappingProxyType({PTU_TEAM_DESCRIPTOR_KEY: ptu_estimated_tokens}), ) if tpm_response["overall_code"] == "OVER_LIMIT": @@ -3814,6 +3860,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # the (actual - reserved) delta to those — unreserved # scopes get charged the full actual usage instead. stash.reserved_tokens = estimated_tokens + stash.ptu_reserved_tokens = ptu_estimated_tokens stash.reserved_model = self._rate_limited_model(requested_model) stash.reserved_scopes = frozenset( (d["key"], d["value"]) @@ -4629,7 +4676,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_obj=response_obj, reconcile_model=reconcile_model, reserved_scopes=reserved_scopes, - reserved_tokens=reserved_tokens, + reserved_ceiling=stash.ptu_ceiling if stash is not None else None, + reserved_tokens=stash.ptu_reserved_tokens if stash is not None else 0, total_tokens=total_tokens, ) ) @@ -4642,6 +4690,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): response_obj: object, reconcile_model: RateLimitedModel | None, reserved_scopes: Set[tuple[str, str]], + reserved_ceiling: PTUTeamCeiling | None, reserved_tokens: int, total_tokens: int, ) -> Sequence[RedisPipelineIncrementOperation]: @@ -4649,38 +4698,44 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): cached input at the model's cached ratio, output weighted by its output-to-input ratio, the way Azure sizes a PTU. - The pre-call reservation was raw estimated tokens, so this is the same reconcile as the - other TPM scopes with a weighted actual; when usage cannot be resolved, or the ceiling - is gone since the reservation was taken, it charges the raw total the other scopes - charge so the reservation is never left standing. + The reservation was taken in those units against the ceiling admission resolved, so that + ceiling settles it even after the share changed or went away; when usage cannot be + resolved it charges the raw total the other scopes charge so the reservation is never + left standing. """ team_id: Final = standard_logging_metadata.get("user_api_key_team_id") if reconcile_model is None or not isinstance(team_id, str) or not team_id: return () scope: Final = (PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{reconcile_model.group}") - ceiling: Final = self._ptu_team_ceiling_resolver(team_id, reconcile_model.group) + ceiling: Final = ( + reserved_ceiling + if reserved_ceiling is not None + else self._ptu_team_ceiling_resolver(team_id, reconcile_model.group) + ) if ceiling is None and scope not in reserved_scopes: return () - usage: Final = self._resolve_reconciled_usage(response_obj) - normalized: Final = ( - round( - normalized_tokens( - ceiling, - prompt_tokens=usage.prompt_tokens, - completion_tokens=usage.completion_tokens, - cache_read_tokens=usage.cached_tokens, - ) - ) - if ceiling is not None and usage is not None - else total_tokens - ) return self._build_reservation_aware_tpm_ops( targets=(scope,), reserved_scopes=reserved_scopes, - actual_tokens=normalized, + actual_tokens=self._ptu_settlement_tokens( + ceiling, self._resolve_reconciled_usage(response_obj), total_tokens + ), reserved_tokens=reserved_tokens, ) + @staticmethod + def _ptu_settlement_tokens(ceiling: PTUTeamCeiling | None, usage: _ReconciledUsage | None, raw_total: int) -> int: + if ceiling is None or usage is None: + return raw_total + return round( + normalized_tokens( + ceiling, + prompt_tokens=usage.prompt_tokens, + completion_tokens=usage.completion_tokens, + cache_read_tokens=usage.cached_tokens, + ) + ) + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Update TPM usage on successful API calls by incrementing counters using pipeline @@ -4786,9 +4841,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses=statuses, ) - def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple[int, int, int]: + def _recovered_partial_usage(self, source: Mapping[str, object]) -> Usage | None: usage: Final = source.get("combined_usage_object") if not isinstance(usage, Usage) or (usage.completion_tokens or 0) <= 0: + return None + return usage + + def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple[int, int, int]: + usage: Final = self._recovered_partial_usage(source) + if usage is None: return 0, 0, 0 billable_input, completion_tokens, _ = self._resolve_io_token_reconcile_usage(usage) return ( @@ -4797,6 +4858,26 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): completion_tokens, ) + def _build_ptu_failure_settlement_ops( + self, stash: RequestRateLimiterStash, source: Mapping[str, object], raw_actual_tokens: int + ) -> Sequence[RedisPipelineIncrementOperation]: + """Settle the team PTU reservation on failure in the normalized tokens it was taken in: + at the recovered partial usage when there is one, else a full refund.""" + ptu_scopes: Final = tuple(scope for scope in stash.reserved_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY) + if not ptu_scopes: + return () + usage: Final = self._recovered_partial_usage(source) + return self._build_reservation_aware_tpm_ops( + targets=ptu_scopes, + reserved_scopes=stash.reserved_scopes, + actual_tokens=self._ptu_settlement_tokens( + stash.ptu_ceiling, + self._resolve_reconciled_usage(usage) if usage is not None else None, + raw_actual_tokens, + ), + reserved_tokens=stash.ptu_reserved_tokens, + ) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """ On failure: decrement max_parallel_requests and refund the upfront @@ -4838,12 +4919,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # refund there would drive their counter negative. pipeline_operations.extend( self._build_reservation_aware_tpm_ops( - targets=list(stash.reserved_scopes), + targets=tuple(scope for scope in stash.reserved_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY), reserved_scopes=stash.reserved_scopes, actual_tokens=tpm_actual, reserved_tokens=reserved_tokens, ) ) + pipeline_operations.extend(self._build_ptu_failure_settlement_ops(stash, kwargs, tpm_actual)) # Settle project ITPM/OTPM reservations the same way: at the # recovered partial usage, or a full refund when there is none. @@ -5034,11 +5116,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(request_data) combined_ops: Final = ( - self._build_reservation_aware_tpm_ops( - targets=tuple(stash.reserved_scopes), - reserved_scopes=stash.reserved_scopes, - actual_tokens=tpm_actual, - reserved_tokens=reserved_tokens, + ( + *self._build_reservation_aware_tpm_ops( + targets=tuple(scope for scope in stash.reserved_scopes if scope[0] != PTU_TEAM_DESCRIPTOR_KEY), + reserved_scopes=stash.reserved_scopes, + actual_tokens=tpm_actual, + reserved_tokens=reserved_tokens, + ), + *self._build_ptu_failure_settlement_ops(stash, request_data, tpm_actual), ) if reserved_tokens > 0 else () diff --git a/litellm/proxy/management_endpoints/ptu_consumption.py b/litellm/proxy/management_endpoints/ptu_consumption.py index 12736b312a8..2e125864c93 100644 --- a/litellm/proxy/management_endpoints/ptu_consumption.py +++ b/litellm/proxy/management_endpoints/ptu_consumption.py @@ -9,7 +9,7 @@ from collections.abc import Callable from types import MappingProxyType from typing import Final -from litellm.litellm_core_utils.azure_ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours +from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, MetricWithMetadata, diff --git a/litellm/router.py b/litellm/router.py index 79e9f0a2d4f..2275578c695 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12793,8 +12793,6 @@ class Router: model=model, llm_provider="", ) - if not is_ptu_cost_attribution_enabled(): - return result.deployments shared: Final = filter_ptu_shared_deployments(result.deployments, request_team_id) if shared.withheld and len(shared.deployments) == 0: raise litellm.BadRequestError( diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 1401dd23d2e..a797cd2d109 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -10,8 +10,11 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Final, Generic, TypeVar -from litellm.litellm_core_utils.azure_ptu_capacity import PTUCapacity, deployment_ptu_capacity from litellm.litellm_core_utils.ptu_pricing import parsed_ptu_shares, ptu_terms +from litellm.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity +from litellm.types.utils import LlmProviders + +_AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value}) _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) @@ -105,14 +108,38 @@ def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]]) -> PTU ) +def _is_azure_deployment(deployment: Mapping[str, object]) -> bool: + litellm_params: Final = deployment.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return False + provider: Final = litellm_params.get("custom_llm_provider") + if isinstance(provider, str): + return provider in _AZURE_PROVIDERS + model: Final = litellm_params.get("model") + return isinstance(model, str) and model.partition("/")[0] in _AZURE_PROVIDERS + + def ptu_capacity_warning(model_name: str, deployment: Mapping[str, object]) -> str | None: - """Why this reserved deployment's tokens cannot be converted to PTUs, else None.""" + """Why this reserved deployment's tokens cannot be converted to PTUs, else None. + + Only a deployment that declares shares (whose ceilings need a sizing row) or one served by + Azure (whose PTU hours need one) is worth warning about; a single-team reservation on another + provider only ever used the flat-cost rollup, which needs no sizing. + """ model_info: Final = deployment.get("model_info") if not isinstance(model_info, Mapping) or ptu_terms(model_info) is None: return None if deployment_ptu_capacity(deployment) is not None: return None - return ( - f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its PTU shares set no " - "team ceiling and its usage reports no PTU hours; set model_info.base_model to the Azure model name" - ) + has_shares: Final = model_info.get("ptu_shares") is not None + if has_shares: + return ( + f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its PTU shares set no " + "team ceiling and its usage reports no PTU hours; set model_info.base_model to the Azure model name" + ) + if _is_azure_deployment(deployment): + return ( + f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its usage reports no " + "PTU hours; set model_info.base_model to the Azure model name" + ) + return None diff --git a/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py b/tests/test_litellm/llms/azure/test_azure_ptu_capacity.py similarity index 98% rename from tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py rename to tests/test_litellm/llms/azure/test_azure_ptu_capacity.py index 4d96068cd8b..59e26ecb776 100644 --- a/tests/test_litellm/litellm_core_utils/test_azure_ptu_capacity.py +++ b/tests/test_litellm/llms/azure/test_azure_ptu_capacity.py @@ -4,7 +4,7 @@ from typing import Final import pytest -from litellm.litellm_core_utils.azure_ptu_capacity import ( +from litellm.llms.azure.ptu_capacity import ( AZURE_PTU_CAPACITY, PTUCapacity, azure_ptu_capacity, diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 8f44945718d..09a8314dc37 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -21,7 +21,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY -from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY +from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( @@ -4500,7 +4500,7 @@ async def test_tpm_over_limit_rejection_releases_parallel_slot_v3(monkeypatch): ) counter_key = f"{{api_key:{_api_key}}}:max_parallel_requests" - async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None): + async def over_limit_reservation(descriptors, estimated_tokens, parent_otel_span=None, **_kwargs): return { "overall_code": "OVER_LIMIT", "statuses": [ @@ -7179,7 +7179,7 @@ def _ptu_request(model: str = "test-model") -> dict: @pytest.mark.asyncio async def test_a_teams_ptu_share_is_a_hard_tpm_ceiling_on_the_shared_model(): cache = DualCache() - resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=500, ratio=4.0) + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0) handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve ) @@ -7252,7 +7252,55 @@ def _shared_ptu_router(model_group: str) -> Router: def _two_thirds_of_a_ptu_minute() -> dict: - return {**_ptu_request(), "max_tokens": AZURE_PTU_CAPACITY["gpt-4.1"].input_tpm_per_ptu * 2 // 3} + """An output budget worth two thirds of a gpt-4.1 PTU minute once weighted at the model's + output-to-input ratio, the way the ceiling counts it.""" + capacity = AZURE_PTU_CAPACITY["gpt-4.1"] + return {**_ptu_request(), "max_tokens": int(capacity.input_tpm_per_ptu * 2 / 3 / capacity.output_to_input_ratio)} + + +@pytest.mark.asyncio +async def test_the_reservation_weighs_output_the_way_the_ceiling_does(): + """300 output tokens are 1200 normalized tokens at 4:1, over a 1000-token ceiling the raw + 301-token estimate would clear; the same request at 1:1 is admitted.""" + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + unweighted_cache = DualCache() + unweighted = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(unweighted_cache), + ptu_team_ceiling_resolver=_ptu_ceiling_for("t", "test-model", tpm_limit=1000, ratio=1.0)[0], + ) + await unweighted.async_pre_call_hook( + user_api_key_dict=key, cache=unweighted_cache, data=_ptu_request(), call_type="acompletion" + ) + + weighted_cache = DualCache() + weighted = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(weighted_cache), + ptu_team_ceiling_resolver=_ptu_ceiling_for("t", "test-model", tpm_limit=1000, ratio=4.0)[0], + ) + with pytest.raises(HTTPException) as rejected: + await weighted.async_pre_call_hook( + user_api_key_dict=key, cache=weighted_cache, data=_ptu_request(), call_type="acompletion" + ) + assert rejected.value.status_code == 429 + assert "model_per_team_ptu" in str(rejected.value.detail) + + +@pytest.mark.asyncio +async def test_the_ptu_counter_holds_the_normalized_reservation_beside_the_raw_one(): + cache = DualCache() + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=_ptu_request(), call_type="acompletion") + + stash = get_request_stash() + assert stash is not None + assert stash.ptu_reserved_tokens == stash.reserved_tokens + 3 * 300 + ptu_key = handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens") + assert int(await cache.async_get_cache(key=ptu_key) or 0) == stash.ptu_reserved_tokens @pytest.mark.asyncio @@ -7386,7 +7434,7 @@ async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone(): """The share can be removed between admission and completion; the reserved tokens still come off the counter instead of standing in the window.""" ceiling: dict[str, PTUTeamCeiling | None] = { - "current": PTUTeamCeiling(tpm_limit=500, output_to_input_ratio=4.0, cached_input_ratio=0.0) + "current": PTUTeamCeiling(tpm_limit=2000, output_to_input_ratio=4.0, cached_input_ratio=0.0) } cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( @@ -7399,7 +7447,7 @@ async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone(): stash = get_request_stash() assert stash is not None assert ("model_per_team_ptu", "t:test-model") in stash.reserved_scopes - assert stash.reserved_tokens > 150 + assert stash.ptu_reserved_tokens > stash.reserved_tokens > 150 ceiling["current"] = None ops = handler._build_success_event_pipeline_operations( @@ -7408,7 +7456,61 @@ async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone(): rate_limit_type="total", ) - assert _ptu_increment(handler, ops) == 150 - stash.reserved_tokens + assert _ptu_increment(handler, ops) == 300 - stash.ptu_reserved_tokens + + +async def _reserve_a_ptu_minute(cache: DualCache, call_id: str) -> tuple[_PROXY_MaxParallelRequestsHandler, str]: + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + await handler.async_pre_call_hook( + user_api_key_dict=key, + cache=cache, + data={**_ptu_request(), "litellm_call_id": call_id}, + call_type="acompletion", + ) + stash = get_request_stash() + assert stash is not None and stash.ptu_reserved_tokens > 0 + return handler, handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens") + + +@pytest.mark.asyncio +async def test_a_failed_stream_settles_the_ptu_counter_in_normalized_tokens(): + """The partial usage a failed stream recovered is 20 input and 7 output tokens: 48 normalized + at 4:1, which is what stays in the window instead of the raw 27 or the whole reservation.""" + cache = DualCache() + handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-partial") + + await handler.async_log_failure_event( + kwargs={ + "litellm_call_id": "ptu-partial", + "standard_logging_object": { + "metadata": {"user_api_key_hash": hash_token("sk-ptu"), "user_api_key_team_id": "t"} + }, + "combined_usage_object": Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27), + }, + response_obj=None, + start_time=None, + end_time=None, + ) + + assert int(await cache.async_get_cache(key=ptu_key) or 0) == 20 + 4 * 7 + + +@pytest.mark.asyncio +async def test_a_proxy_side_rejection_refunds_the_whole_normalized_ptu_reservation(): + cache = DualCache() + handler, ptu_key = await _reserve_a_ptu_minute(cache, "ptu-rejected") + + await handler.async_post_call_failure_hook( + request_data={**_ptu_request(), "litellm_call_id": "ptu-rejected"}, + original_exception=Exception("guardrail rejected the request"), + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t"), + ) + + assert int(await cache.async_get_cache(key=ptu_key) or 0) == 0 def test_usage_that_only_reports_a_total_charges_that_total_to_the_ptu_counter(): diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py index e9ef4ea2994..432bf858899 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_consumption.py @@ -4,7 +4,7 @@ from typing import Final import pytest -from litellm.litellm_core_utils.azure_ptu_capacity import PTUCapacity +from litellm.llms.azure.ptu_capacity import PTUCapacity from litellm.proxy.management_endpoints.ptu_consumption import attach_ptu_hours from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 7e0792cdc2c..87fff1d5d6a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15,7 +15,7 @@ from pydantic import ValidationError from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY +from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_BudgetTableFull, diff --git a/tests/test_litellm/router_utils/test_ptu_shares.py b/tests/test_litellm/router_utils/test_ptu_shares.py index 110161bfa93..2c54c3a1c8d 100644 --- a/tests/test_litellm/router_utils/test_ptu_shares.py +++ b/tests/test_litellm/router_utils/test_ptu_shares.py @@ -2,7 +2,7 @@ from typing import Final -from litellm.litellm_core_utils.azure_ptu_capacity import AZURE_PTU_CAPACITY +from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.router_utils.ptu_shares import ( PTUTeamCeiling, filter_ptu_shared_deployments, @@ -144,3 +144,24 @@ def test_a_sized_reservation_and_an_unreserved_deployment_raise_no_warning(): assert ptu_capacity_warning("gpt-4.1-ptu", _single_team()) is None unsized_open: Final = {**_OPEN, "litellm_params": {"model": "azure/my-ptu-deployment"}} assert ptu_capacity_warning("gpt-4.1-ptu", unsized_open) is None + + +def test_an_unsized_single_team_azure_reservation_is_warned_about_its_ptu_hours_only(): + warning: Final = ptu_capacity_warning("gpt-4.1-ptu", _single_team(model="azure/my-ptu-deployment")) + assert warning is not None + assert "PTU hours" in warning + assert "ceiling" not in warning + + +def test_an_unsized_single_team_reservation_on_another_provider_is_not_warned_about(): + assert ptu_capacity_warning("claude-ptu", _single_team(model="anthropic/claude-sonnet-4-5")) is None + + +def test_a_bare_model_name_counts_as_azure_through_custom_llm_provider(): + deployment: Final = { + **_single_team(model="my-ptu-deployment"), + "litellm_params": {"model": "my-ptu-deployment", "custom_llm_provider": "azure"}, + } + warning: Final = ptu_capacity_warning("gpt-4.1-ptu", deployment) + assert warning is not None + assert "PTU hours" in warning diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ad9f6e9c4aa..a58292d6d22 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -17915,14 +17915,17 @@ def test_ptu_shares_raise_when_only_shared_deployments_remain(monkeypatch): assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"] -def test_ptu_shares_do_not_filter_while_the_feature_is_off(monkeypatch): +def test_ptu_shares_hide_the_shared_deployment_even_while_the_feature_is_off(monkeypatch): + """The flag switches cost attribution on; a declared split is an access rule and holds + without it, so a team never reaches a deployment reserved for others while the flag is + off.""" monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False) router = Router(model_list=_shared_ptu_model_list()) _, deployments = router._common_checks_available_deployment( model="gpt-4.1-ptu", request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, ) - assert {d["model_info"]["id"] for d in deployments} == {"shared-deployment", "open-deployment"} + assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] def test_a_shared_ptu_deployment_whose_shares_do_not_add_up_is_refused_at_registration(monkeypatch): diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index d73f5efa96b..3c42de71cf7 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -2142,7 +2142,7 @@ def test_an_incomplete_reservation_is_refused_rather_than_served(dropped): @pytest.mark.parametrize( "dropped, expected", [ - ("team_id", "team_id is required when PTU fields are set (one model maps to one team)"), + ("team_id", "team_id or ptu_shares is required when PTU fields are set"), ("cost_per_ptu_per_hour", "ptu_count and cost_per_ptu_per_hour must be set together"), ], ids=["no team_id", "count without rate"], From 4ad0db30b26c94a7d9295a99858ed8e3f62e53e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 16:06:01 -0700 Subject: [PATCH 04/18] fix(proxy): size the uncapped output floor against the PTU share in raw output tokens and escape the exported team id --- .../hooks/parallel_request_limiter_v3.py | 11 +++++-- .../management_endpoints/team_endpoints.py | 2 +- litellm/router_utils/ptu_shares.py | 4 +++ .../hooks/test_parallel_request_limiter_v3.py | 26 +++++++++++++++ .../test_team_endpoints.py | 32 +++++++++++++++++++ 5 files changed, 71 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index c61b35f5936..e0f0b9d714e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -3765,13 +3765,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # in-memory check otherwise — single-worker protection still holds # even without Redis. # ---------------------------------------------------------------- - configured_tpm_limits: Final = [ - int(v) + ptu_raw_output_limit: Final = ( + stash.ptu_ceiling.raw_output_token_limit() if stash.ptu_ceiling is not None else None + ) + configured_tpm_limits: Final = tuple( + ptu_raw_output_limit + if d["key"] == PTU_TEAM_DESCRIPTOR_KEY and ptu_raw_output_limit is not None + else int(v) for d in descriptors if d["key"] not in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY) for v in [(d.get("rate_limit") or {}).get("tokens_per_unit")] if v is not None - ] + ) has_tpm_limits: Final = bool(configured_tpm_limits) # Populated on a successful combined-TPM reservation below, so the diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b3dca8b127f..b0320cd6e6d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -6882,7 +6882,7 @@ def _export_csv_record(row: TeamDailyActivityExportRow) -> dict[str, object]: return { # mutable-ok: csv.DictWriter consumes a plain mapping per row "Date": row.date, "Team": _csv_safe(row.team_alias) if row.team_alias else "-", - "Team ID": row.team_id, + "Team ID": _csv_safe(row.team_id), "Key Alias": _csv_safe(row.key_alias) if row.key_alias else "-", "Key ID": row.api_key or "-", "User ID": _csv_safe(row.user_id) if row.user_id else "-", diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index a797cd2d109..3bc9d33349c 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -25,6 +25,10 @@ class PTUTeamCeiling: output_to_input_ratio: float cached_input_ratio: float + def raw_output_token_limit(self) -> int: + """The ceiling expressed in unweighted tokens: what fits under it when every token is output.""" + return max(1, int(self.tpm_limit / max(self.output_to_input_ratio, 1.0))) + @dataclass(frozen=True, slots=True) class PTUShareFilterResult(Generic[_DeploymentT]): diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 09a8314dc37..dee97db8472 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7285,6 +7285,32 @@ async def test_the_reservation_weighs_output_the_way_the_ceiling_does(): assert "model_per_team_ptu" in str(rejected.value.detail) +@pytest.mark.asyncio +async def test_a_one_ptu_share_admits_four_uncapped_requests_a_minute_and_rejects_the_fifth(): + """One PTU of gpt-4.1 is 3,000 normalized tokens a minute. A request without max_tokens gets the + proxy's output floor measured against the share in output tokens (750), so its cap costs at most + a quarter of the share once weighted 4:1 and four such requests fit before the fifth is refused.""" + cache = DualCache() + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=3000, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + uncapped = {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]} + + admitted = [dict(uncapped) for _ in range(4)] + for data in admitted: + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=data, call_type="acompletion") + with pytest.raises(HTTPException) as fifth: + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=dict(uncapped), call_type="acompletion" + ) + + assert fifth.value.status_code == 429 + assert "model_per_team_ptu" in str(fifth.value.detail) + assert all(data["max_tokens"] * 4 <= 3000 // 4 for data in admitted) + + @pytest.mark.asyncio async def test_the_ptu_counter_holds_the_normalized_reservation_beside_the_raw_one(): cache = DualCache() diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 87fff1d5d6a..6d87eda3d98 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -17295,6 +17295,38 @@ def test_team_export_csv_columns_match_the_dashboard_client_layout(): ] +def test_team_export_csv_escapes_a_formula_shaped_team_id_like_the_alias(): + import csv + import io + + from litellm.proxy.management_endpoints.team_endpoints import _team_export_csv + from litellm.types.proxy.management_endpoints.team_endpoints import TeamDailyActivityExportRow + + row: Final = TeamDailyActivityExportRow( + date="2026-06-01", + team_id="=HYPERLINK(\"https://example.com\")", + team_alias="=SUM(A1)", + api_key=None, + key_alias=None, + user_id=None, + user_email=None, + spend=0.0, + api_requests=0, + successful_requests=0, + failed_requests=0, + total_tokens=0, + prompt_tokens=0, + completion_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + ) + + record: Final = next(csv.DictReader(io.StringIO(_team_export_csv("daily", (row,))))) + + assert record["Team ID"] == "'=HYPERLINK(\"https://example.com\")" + assert record["Team"] == "'=SUM(A1)" + + def test_team_export_csv_omits_key_columns_for_the_plain_daily_scope(): import csv import io From 6a24e7f920582890ea74cc8f3a588ffcb866a17b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 17:53:51 -0700 Subject: [PATCH 05/18] fix(ptu): refuse ptu_shares naming a team that does not exist --- .../model_management_endpoints.py | 22 ++++- .../test_ptu_model_settings.py | 82 ++++++++++++++++++- 2 files changed, 101 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ae294871afc..ab33c35909a 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -17,6 +17,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequen from contextlib import AbstractAsyncContextManager, asynccontextmanager, suppress from dataclasses import dataclass from fnmatch import fnmatchcase +from functools import partial from json import JSONDecodeError from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable @@ -37,6 +38,7 @@ from litellm.litellm_core_utils.ptu_pricing import ( PTU_ZEROED_PRICING_FIELDS, PTU_ZEROED_TABLE_FIELDS, SEARCH_CONTEXT_SIZES, + parsed_ptu_shares, ptu_config_error, ) from litellm.proxy._types import ( @@ -779,6 +781,21 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: raise HTTPException(status_code=400, detail=error) +async def _raise_if_ptu_share_teams_missing( + model_info: Mapping[str, object], team_table: Callable[[], _TeamLookupTable] +) -> None: + """Hold every team named in ``ptu_shares`` to the existence check ``team_id`` already gets.""" + shares: Final = parsed_ptu_shares(model_info.get("ptu_shares")) + if shares is None: + return + table: Final = team_table() + rows: Final = await asyncio.gather(*(table.find_unique(where={"team_id": team_id}) for team_id in shares)) + missing: Final = tuple(team_id for team_id, row in zip(shares, rows, strict=True) if row is None) + if not missing: + return + raise HTTPException(status_code=400, detail={"error": f"Team id={', '.join(missing)} does not exist in db"}) + + # The mirrored per-token pricing fields plus the remaining rates the public cost map or a # provider default would otherwise supply (the cache back-fills, the Maps grounding rate). An # unset field falls back to those sources, so a field left out here is one a PTU deployment @@ -1614,8 +1631,10 @@ async def _update_team_model_in_db( # raising the rate on a configured model carries no ptu_effective_from, which the # stored row supplies. if patch_data.model_info is not None: - _raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True)) + incoming_model_info: Final = patch_data.model_info.model_dump(exclude_none=True) + _raise_if_ptu_cost_attribution_disabled(incoming_model_info) _validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data)) + await _raise_if_ptu_share_teams_missing(incoming_model_info, partial(_repo_team_table, prisma_client)) _raise_if_ptu_deployment_is_priced( model_info=_merged_ptu_model_info(db_model=db_model, patch_data=patch_data), supplied=( @@ -2456,6 +2475,7 @@ async def add_new_model( incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True) _raise_if_ptu_cost_attribution_disabled(incoming_model_info) _validate_ptu_model_info(incoming_model_info) + await _raise_if_ptu_share_teams_missing(incoming_model_info, partial(_repo_team_table, prisma_client)) priced_model_params: Final = _ptu_priced_deployment(model_params) if store_model_in_db is True: diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index 8ed3df44216..cc1ac071207 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -2,7 +2,9 @@ import datetime import json +from collections.abc import Mapping from contextlib import ExitStack +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import patch as patch_ctx @@ -14,7 +16,9 @@ from litellm.litellm_core_utils.ptu_pricing import ptu_terms from litellm.llms.gemini.cost_calculator import cost_per_web_search_request from litellm.proxy._types import ( LiteLLM_ProxyModelTable, + LiteLLM_TeamTable, LitellmUserRoles, + ProxyException, ReconcileOutcome, UserAPIKeyAuth, ) @@ -27,6 +31,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( _ptu_priced_deployment, _ptu_zeroed_pricing, _raise_if_ptu_cost_attribution_disabled, + _raise_if_ptu_share_teams_missing, _update_team_model_in_db, _validate_ptu_model_info, add_new_model, @@ -651,7 +656,7 @@ class TestAddNewModelPtuGate: monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) @staticmethod - def _patched_proxy(model_id: str): + def _patched_proxy(model_id: str, prisma_client: MagicMock | None = None): """Patch everything /model/new touches except the PTU gate, and hand back the DB writers.""" db_row = LiteLLM_ProxyModelTable( model_id=model_id, @@ -678,7 +683,7 @@ class TestAddNewModelPtuGate: proxy_server = "litellm.proxy.proxy_server" endpoints = "litellm.proxy.management_endpoints.model_management_endpoints" return (add_model_to_db, add_team_model_to_db), [ - patch(f"{proxy_server}.prisma_client", MagicMock()), + patch(f"{proxy_server}.prisma_client", prisma_client if prisma_client is not None else MagicMock()), patch(f"{proxy_server}.store_model_in_db", True), patch(f"{proxy_server}.proxy_config", mock_proxy_config), patch(f"{proxy_server}.proxy_logging_obj", MagicMock()), @@ -1328,6 +1333,79 @@ def test_validate_helper_refuses_shares_that_do_not_add_up_to_the_count(): assert "4 of 5 allocated" in exc.value.detail +class _TeamLookup: + def __init__(self, existing: frozenset[str]) -> None: + self.existing: Final = existing + self.looked_up: tuple[str, ...] = () + + async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_TeamTable | None: + team_id: Final = str(where["team_id"]) + self.looked_up = (*self.looked_up, team_id) + return LiteLLM_TeamTable(team_id=team_id) if team_id in self.existing else None + + +def _shared_model_info(shares: Mapping[str, int]) -> Mapping[str, object]: + return { + "ptu_count": sum(shares.values()), + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": _SHARED_START, + "ptu_shares": dict(shares), + } + + +@pytest.mark.asyncio +async def test_share_team_check_refuses_a_team_that_does_not_exist(): + with pytest.raises(HTTPException) as exc: + await _raise_if_ptu_share_teams_missing( + _shared_model_info({"team-a": 3, "ghost-team": 2}), lambda: _TeamLookup(frozenset({"team-a"})) + ) + assert exc.value.status_code == 400 + assert exc.value.detail == {"error": "Team id=ghost-team does not exist in db"} + + +@pytest.mark.asyncio +async def test_share_team_check_accepts_shares_naming_existing_teams(): + lookup: Final = _TeamLookup(frozenset({"team-a", "team-b"})) + await _raise_if_ptu_share_teams_missing(_shared_model_info({"team-a": 3, "team-b": 2}), lambda: lookup) + assert sorted(lookup.looked_up) == ["team-a", "team-b"] + + +@pytest.mark.asyncio +async def test_share_team_check_leaves_a_team_id_holder_to_the_team_model_check(): + lookup: Final = _TeamLookup(frozenset()) + await _raise_if_ptu_share_teams_missing( + {"team_id": "team-a", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "ptu_effective_from": _SHARED_START}, + lambda: lookup, + ) + assert lookup.looked_up == () + + +@pytest.mark.asyncio +async def test_model_new_refuses_shares_naming_a_team_that_does_not_exist(monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + prisma_client = MagicMock() + prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + (add_model_to_db, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy( + "ptu-shared-model", prisma_client=prisma_client + ) + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + base = TestAddNewModelPtuGate._ptu_deployment("ptu-shared-model") + shared = base.model_copy( + update={"model_info": base.model_info.model_copy(update={"team_id": None, "ptu_shares": {"ghost-team": 15}})} + ) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + with pytest.raises(ProxyException, match="Team id=ghost-team does not exist in db") as exc: + await add_new_model(model_params=shared, user_api_key_dict=admin) + + assert exc.value.code == "400" + prisma_client.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "ghost-team"}) + add_model_to_db.assert_not_called() + add_team_model_to_db.assert_not_called() + + def test_validate_helper_refuses_a_team_id_beside_shares(): with pytest.raises(HTTPException) as exc: _validate_ptu_model_info( From 00f6210f7fe6de21c441bd3bc910036d346b2ec6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 24 Sep 2026 18:20:34 -0700 Subject: [PATCH 06/18] refactor(ptu): keep Azure detection under llms/azure and check share teams with one query --- litellm/llms/azure/ptu_capacity.py | 19 +++++++++++++++++ .../model_management_endpoints.py | 18 +++++++++++++--- litellm/router_utils/ptu_shares.py | 20 +++--------------- .../test_ptu_model_settings.py | 21 ++++++++++++------- 4 files changed, 50 insertions(+), 28 deletions(-) diff --git a/litellm/llms/azure/ptu_capacity.py b/litellm/llms/azure/ptu_capacity.py index c68429dc7b7..ce50e885b66 100644 --- a/litellm/llms/azure/ptu_capacity.py +++ b/litellm/llms/azure/ptu_capacity.py @@ -17,6 +17,8 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final, Protocol +from litellm.types.utils import LlmProviders + class NormalizedTokenWeights(Protocol): @property @@ -36,6 +38,9 @@ class PTUCapacity: def normalized_tokens_per_ptu_hour(self) -> int: return self.input_tpm_per_ptu * 60 + def input_tpm_for(self, ptus: int) -> int: + return ptus * self.input_tpm_per_ptu + AZURE_PTU_CAPACITY: Final[Mapping[str, PTUCapacity]] = MappingProxyType( { @@ -68,6 +73,7 @@ AZURE_PTU_CAPACITY: Final[Mapping[str, PTUCapacity]] = MappingProxyType( ) _VERSION_SUFFIX: Final = re.compile(r"-\d{4}-\d{2}-\d{2}$") +_AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value}) def azure_ptu_capacity(model: str) -> PTUCapacity | None: @@ -97,6 +103,19 @@ def deployment_ptu_capacity(deployment: Mapping[str, object]) -> PTUCapacity | N return next((capacity for capacity in map(azure_ptu_capacity, candidates) if capacity is not None), None) +def is_azure_deployment(deployment: Mapping[str, object]) -> bool: + """Whether ``litellm_params`` route this deployment to Azure OpenAI or Azure AI, by + ``custom_llm_provider`` first and the ``model`` prefix otherwise.""" + litellm_params: Final = deployment.get("litellm_params") + if not isinstance(litellm_params, Mapping): + return False + provider: Final = litellm_params.get("custom_llm_provider") + if isinstance(provider, str): + return provider in _AZURE_PROVIDERS + model: Final = litellm_params.get("model") + return isinstance(model, str) and model.partition("/")[0] in _AZURE_PROVIDERS + + def normalized_tokens( weights: NormalizedTokenWeights, *, prompt_tokens: int, completion_tokens: int, cache_read_tokens: int = 0 ) -> float: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ab33c35909a..41df67557e7 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -24,6 +24,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -134,6 +135,7 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import AutoRouterClassifierDefaultPromptResponse, UpdateUsefulLinksRequest, ) +from litellm.types.proxy.management_endpoints.team_endpoints import TeamIdSearchFilter from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, @@ -237,15 +239,24 @@ class _ExistingModelRow(Protocol): class _TeamRow(Protocol): + @property + def team_id(self) -> str: ... + @property def models(self) -> Sequence[str]: ... def model_dump(self) -> Mapping[str, object]: ... +class _TeamIdsWhere(TypedDict): + team_id: ReadOnly[TeamIdSearchFilter] + + class _TeamLookupTable(Protocol): def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ... + def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_TeamRow]]: ... + class _TeamTable(_TeamLookupTable, Protocol): def update( @@ -788,9 +799,10 @@ async def _raise_if_ptu_share_teams_missing( shares: Final = parsed_ptu_shares(model_info.get("ptu_shares")) if shares is None: return - table: Final = team_table() - rows: Final = await asyncio.gather(*(table.find_unique(where={"team_id": team_id}) for team_id in shares)) - missing: Final = tuple(team_id for team_id, row in zip(shares, rows, strict=True) if row is None) + where: Final[_TeamIdsWhere] = {"team_id": {"in": tuple(shares)}} + rows: Final = await team_table().find_many(where=where) + found: Final = frozenset(row.team_id for row in rows) + missing: Final = tuple(team_id for team_id in shares if team_id not in found) if not missing: return raise HTTPException(status_code=400, detail={"error": f"Team id={', '.join(missing)} does not exist in db"}) diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 3bc9d33349c..11b61676ce5 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -11,10 +11,7 @@ from dataclasses import dataclass from typing import Final, Generic, TypeVar from litellm.litellm_core_utils.ptu_pricing import parsed_ptu_shares, ptu_terms -from litellm.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity -from litellm.types.utils import LlmProviders - -_AZURE_PROVIDERS: Final = frozenset({LlmProviders.AZURE.value, LlmProviders.AZURE_AI.value}) +from litellm.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity, is_azure_deployment _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) @@ -76,7 +73,7 @@ def team_ptu_ceiling(deployments: Sequence[Mapping[str, object]], team_id: str) if not priced: return None return PTUTeamCeiling( - tpm_limit=sum(share * capacity.input_tpm_per_ptu for share, capacity in priced), + tpm_limit=sum(capacity.input_tpm_for(share) for share, capacity in priced), output_to_input_ratio=max(capacity.output_to_input_ratio for _, capacity in priced), cached_input_ratio=max(capacity.cached_input_ratio for _, capacity in priced), ) @@ -112,17 +109,6 @@ def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]]) -> PTU ) -def _is_azure_deployment(deployment: Mapping[str, object]) -> bool: - litellm_params: Final = deployment.get("litellm_params") - if not isinstance(litellm_params, Mapping): - return False - provider: Final = litellm_params.get("custom_llm_provider") - if isinstance(provider, str): - return provider in _AZURE_PROVIDERS - model: Final = litellm_params.get("model") - return isinstance(model, str) and model.partition("/")[0] in _AZURE_PROVIDERS - - def ptu_capacity_warning(model_name: str, deployment: Mapping[str, object]) -> str | None: """Why this reserved deployment's tokens cannot be converted to PTUs, else None. @@ -141,7 +127,7 @@ def ptu_capacity_warning(model_name: str, deployment: Mapping[str, object]) -> s f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its PTU shares set no " "team ceiling and its usage reports no PTU hours; set model_info.base_model to the Azure model name" ) - if _is_azure_deployment(deployment): + if is_azure_deployment(deployment): return ( f"PTU deployment '{model_name}' has no Azure sizing row for its model, so its usage reports no " "PTU hours; set model_info.base_model to the Azure model name" diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index cc1ac071207..b81ae11ac5d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -2,9 +2,9 @@ import datetime import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from contextlib import ExitStack -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import patch as patch_ctx @@ -1339,9 +1339,14 @@ class _TeamLookup: self.looked_up: tuple[str, ...] = () async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_TeamTable | None: - team_id: Final = str(where["team_id"]) - self.looked_up = (*self.looked_up, team_id) - return LiteLLM_TeamTable(team_id=team_id) if team_id in self.existing else None + raise AssertionError(f"one lookup per team is what the review asked to avoid: {where}") + + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[LiteLLM_TeamTable]: + team_filter: Final = where["team_id"] + assert isinstance(team_filter, Mapping) + requested: Final = tuple(str(team_id) for team_id in cast(Sequence[object], team_filter["in"])) + self.looked_up = (*self.looked_up, *requested) + return tuple(LiteLLM_TeamTable(team_id=team_id) for team_id in requested if team_id in self.existing) def _shared_model_info(shares: Mapping[str, int]) -> Mapping[str, object]: @@ -1367,7 +1372,7 @@ async def test_share_team_check_refuses_a_team_that_does_not_exist(): async def test_share_team_check_accepts_shares_naming_existing_teams(): lookup: Final = _TeamLookup(frozenset({"team-a", "team-b"})) await _raise_if_ptu_share_teams_missing(_shared_model_info({"team-a": 3, "team-b": 2}), lambda: lookup) - assert sorted(lookup.looked_up) == ["team-a", "team-b"] + assert lookup.looked_up == ("team-a", "team-b") @pytest.mark.asyncio @@ -1384,7 +1389,7 @@ async def test_share_team_check_leaves_a_team_id_holder_to_the_team_model_check( async def test_model_new_refuses_shares_naming_a_team_that_does_not_exist(monkeypatch): monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") prisma_client = MagicMock() - prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) (add_model_to_db, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy( "ptu-shared-model", prisma_client=prisma_client ) @@ -1401,7 +1406,7 @@ async def test_model_new_refuses_shares_naming_a_team_that_does_not_exist(monkey await add_new_model(model_params=shared, user_api_key_dict=admin) assert exc.value.code == "400" - prisma_client.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "ghost-team"}) + prisma_client.db.litellm_teamtable.find_many.assert_awaited_once_with(where={"team_id": {"in": ("ghost-team",)}}) add_model_to_db.assert_not_called() add_team_model_to_db.assert_not_called() From 48c0be34aa49af7483343c2e821d6fb999aaa335 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 15:36:37 -0700 Subject: [PATCH 07/18] fix(proxy): drop the reverted TeamIdSearchFilter import after merging main Main reverted the team key search change that added TeamIdSearchFilter, so the PTU share team check types its where clause locally --- .../proxy/management_endpoints/model_management_endpoints.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 41df67557e7..b16ef81e6c2 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -135,7 +135,6 @@ from litellm.types.proxy.management_endpoints.model_management_endpoints import AutoRouterClassifierDefaultPromptResponse, UpdateUsefulLinksRequest, ) -from litellm.types.proxy.management_endpoints.team_endpoints import TeamIdSearchFilter from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, @@ -249,7 +248,7 @@ class _TeamRow(Protocol): class _TeamIdsWhere(TypedDict): - team_id: ReadOnly[TeamIdSearchFilter] + team_id: ReadOnly[Mapping[Literal["in"], Sequence[str]]] class _TeamLookupTable(Protocol): From d1f698fdd8cb82f4319c4e438746d647602a8a04 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 16:34:00 -0700 Subject: [PATCH 08/18] fix(proxy): chunk the ptu_shares team lookup under the IN-list bound --- .../model_management_endpoints.py | 11 +++-------- .../test_ptu_model_settings.py | 14 +++++++++++++- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b16ef81e6c2..0449ccccee2 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -24,7 +24,6 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias from fastapi import APIRouter, Depends, Header, HTTPException, Request, status from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator -from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -99,6 +98,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, ) from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.chunked_in import find_many_in from litellm.repositories.credentials_repository import CredentialsRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions @@ -247,14 +247,10 @@ class _TeamRow(Protocol): def model_dump(self) -> Mapping[str, object]: ... -class _TeamIdsWhere(TypedDict): - team_id: ReadOnly[Mapping[Literal["in"], Sequence[str]]] - - class _TeamLookupTable(Protocol): def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ... - def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_TeamRow]]: ... + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRow]: ... class _TeamTable(_TeamLookupTable, Protocol): @@ -798,8 +794,7 @@ async def _raise_if_ptu_share_teams_missing( shares: Final = parsed_ptu_shares(model_info.get("ptu_shares")) if shares is None: return - where: Final[_TeamIdsWhere] = {"team_id": {"in": tuple(shares)}} - rows: Final = await team_table().find_many(where=where) + rows: Final = await find_many_in(team_table(), "team_id", shares.keys()) found: Final = frozenset(row.team_id for row in rows) missing: Final = tuple(team_id for team_id in shares if team_id not in found) if not missing: diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index b81ae11ac5d..ce2e865d69d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -38,6 +38,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( update_db_model, ) from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR +from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE from litellm.router import Router from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, @@ -1337,6 +1338,7 @@ class _TeamLookup: def __init__(self, existing: frozenset[str]) -> None: self.existing: Final = existing self.looked_up: tuple[str, ...] = () + self.batch_sizes: tuple[int, ...] = () async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_TeamTable | None: raise AssertionError(f"one lookup per team is what the review asked to avoid: {where}") @@ -1346,6 +1348,7 @@ class _TeamLookup: assert isinstance(team_filter, Mapping) requested: Final = tuple(str(team_id) for team_id in cast(Sequence[object], team_filter["in"])) self.looked_up = (*self.looked_up, *requested) + self.batch_sizes = (*self.batch_sizes, len(requested)) return tuple(LiteLLM_TeamTable(team_id=team_id) for team_id in requested if team_id in self.existing) @@ -1375,6 +1378,15 @@ async def test_share_team_check_accepts_shares_naming_existing_teams(): assert lookup.looked_up == ("team-a", "team-b") +@pytest.mark.asyncio +async def test_share_team_check_splits_a_share_list_longer_than_one_in_list_chunk(): + team_ids: Final = tuple(f"team-{index}" for index in range(IN_LIST_CHUNK_SIZE + 1)) + lookup: Final = _TeamLookup(frozenset(team_ids)) + await _raise_if_ptu_share_teams_missing(_shared_model_info(dict.fromkeys(team_ids, 1)), lambda: lookup) + assert lookup.looked_up == team_ids + assert max(lookup.batch_sizes) <= IN_LIST_CHUNK_SIZE + + @pytest.mark.asyncio async def test_share_team_check_leaves_a_team_id_holder_to_the_team_model_check(): lookup: Final = _TeamLookup(frozenset()) @@ -1406,7 +1418,7 @@ async def test_model_new_refuses_shares_naming_a_team_that_does_not_exist(monkey await add_new_model(model_params=shared, user_api_key_dict=admin) assert exc.value.code == "400" - prisma_client.db.litellm_teamtable.find_many.assert_awaited_once_with(where={"team_id": {"in": ("ghost-team",)}}) + prisma_client.db.litellm_teamtable.find_many.assert_awaited_once_with(where={"team_id": {"in": ["ghost-team"]}}) add_model_to_db.assert_not_called() add_team_model_to_db.assert_not_called() From ea8b63d90fcc2b0e44e811fd523559653c2d6ec8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 17:52:01 -0700 Subject: [PATCH 09/18] fix(proxy): count a request naming a shared PTU deployment directly against its group's ceiling --- .../hooks/parallel_request_limiter_v3.py | 11 +++-- litellm/router_utils/ptu_shares.py | 47 ++++++++++++++++-- .../hooks/test_parallel_request_limiter_v3.py | 39 ++++++++++++++- tests/unit/router_utils/test_ptu_shares.py | 49 ++++++++++++++++--- 4 files changed, 130 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d10e8aabdd0..9fd40115e10 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -69,7 +69,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( response_has_hidden_params, ) from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.ptu_shares import PTUTeamCeiling, model_group_deployments, team_ptu_ceiling +from litellm.router_utils.ptu_shares import PTUTeamCeiling, team_ptu_ceiling from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.utils import ( @@ -134,7 +134,7 @@ def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) - if llm_router is None or not is_ptu_cost_attribution_enabled(): return None - return team_ptu_ceiling(model_group_deployments(llm_router.get_model_list() or (), model_group), team_id) + return team_ptu_ceiling(llm_router.get_model_list() or (), team_id, model_group) FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_SETTING: Final = "fail_closed_rate_limit_enforcement" @@ -3301,7 +3301,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors.append( RateLimitDescriptor( key=PTU_TEAM_DESCRIPTOR_KEY, - value=f"{user_api_key_dict.team_id}:{model.group}", + value=f"{user_api_key_dict.team_id}:{ceiling.model_group}", rate_limit=RateLimitDescriptorRateLimitObject( requests_per_unit=None, tokens_per_unit=ceiling.tpm_limit, window_size=self.window_size ), @@ -4930,12 +4930,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): team_id: Final = standard_logging_metadata.get("user_api_key_team_id") if reconcile_model is None or not isinstance(team_id, str) or not team_id: return () - scope: Final = (PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{reconcile_model.group}") ceiling: Final = ( reserved_ceiling if reserved_ceiling is not None else self._ptu_team_ceiling_resolver(team_id, reconcile_model.group) ) + scope: Final = ( + PTU_TEAM_DESCRIPTOR_KEY, + f"{team_id}:{ceiling.model_group if ceiling is not None else reconcile_model.group}", + ) if ceiling is None and scope not in reserved_scopes: return () return self._build_reservation_aware_tpm_ops( diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 11b61676ce5..91ccc585ba7 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -18,6 +18,7 @@ _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) @dataclass(frozen=True, slots=True) class PTUTeamCeiling: + model_group: str tpm_limit: int output_to_input_ratio: float cached_input_ratio: float @@ -55,17 +56,24 @@ def filter_ptu_shared_deployments( return PTUShareFilterResult(deployments=kept, withheld=len(kept) < len(checks)) -def team_ptu_ceiling(deployments: Sequence[Mapping[str, object]], team_id: str) -> PTUTeamCeiling | None: - """The per-minute normalized-token ceiling ``team_id``'s shares across ``deployments`` add - up to, else None when the team holds no share on a deployment with a known sizing row. +def team_ptu_ceiling( + deployments: Sequence[Mapping[str, object]], team_id: str, requested_model: str +) -> PTUTeamCeiling | None: + """The per-minute normalized-token ceiling ``team_id``'s shares on the group serving + ``requested_model`` add up to, else None when the team holds no share on a deployment with + a known sizing row. + + A request naming one deployment by its id or provider model counts against that deployment's + group, so every name the router serves it under shares one ceiling. Two shared deployments of different models in one group are weighted by the larger output and cached-input ratios, which over-counts those tokens on the cheaper one rather than under-counting them on the dearer one. """ + model_group: Final = _model_group_of(deployments, requested_model) priced: Final = tuple( (shares[team_id], capacity) - for deployment in deployments + for deployment in model_group_deployments(deployments, model_group) if (shares := _deployment_shares(deployment)) is not None and team_id in shares and (capacity := deployment_ptu_capacity(deployment)) is not None @@ -73,6 +81,7 @@ def team_ptu_ceiling(deployments: Sequence[Mapping[str, object]], team_id: str) if not priced: return None return PTUTeamCeiling( + model_group=model_group, tpm_limit=sum(capacity.input_tpm_for(share) for share, capacity in priced), output_to_input_ratio=max(capacity.output_to_input_ratio for _, capacity in priced), cached_input_ratio=max(capacity.cached_input_ratio for _, capacity in priced), @@ -94,6 +103,36 @@ def model_group_deployments(deployments: Sequence[_DeploymentT], model_group: st ) +def _names_deployment(deployment: Mapping[str, object], name: str) -> bool: + model_info: Final = deployment.get("model_info") + litellm_params: Final = deployment.get("litellm_params") + return (isinstance(model_info, Mapping) and model_info.get("id") == name) or ( + isinstance(litellm_params, Mapping) and litellm_params.get("model") == name + ) + + +def _deployment_model_group(deployment: Mapping[str, object]) -> str | None: + model_info: Final = deployment.get("model_info") + public_name: Final = model_info.get("team_public_model_name") if isinstance(model_info, Mapping) else None + if isinstance(public_name, str): + return public_name + model_name: Final = deployment.get("model_name") + return model_name if isinstance(model_name, str) else None + + +def _model_group_of(deployments: Sequence[Mapping[str, object]], requested_model: str) -> str: + """The group ``requested_model`` routes to: itself when it names a group, else the group of + the deployment it names by id or by provider model, the way the router falls back to them.""" + if model_group_deployments(deployments, requested_model): + return requested_model + named: Final = tuple(deployment for deployment in deployments if _names_deployment(deployment, requested_model)) + shared_first: Final = sorted(named, key=lambda deployment: _deployment_shares(deployment) is None) + return next( + (group for deployment in shared_first if (group := _deployment_model_group(deployment)) is not None), + requested_model, + ) + + def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]]) -> PTUCapacity | None: """The sizing row of the group's first reserved deployment, single-team or shared, so a team's tokens on that group convert to PTU-hours.""" diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index dfabeceb94c..0b3e54a8358 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7576,7 +7576,9 @@ def _ptu_ceiling_for(team_id: str, model_group: str, tpm_limit: int, ratio: floa calls.append((requested_team, requested_group)) if (requested_team, requested_group) != (team_id, model_group): return None - return PTUTeamCeiling(tpm_limit=tpm_limit, output_to_input_ratio=ratio, cached_input_ratio=cached_ratio) + return PTUTeamCeiling( + model_group=model_group, tpm_limit=tpm_limit, output_to_input_ratio=ratio, cached_input_ratio=cached_ratio + ) return resolve, calls @@ -7760,6 +7762,37 @@ async def test_the_proxy_router_turns_a_teams_share_into_its_ceiling_when_attrib assert "model_per_team_ptu" in str(exc.value.detail) +@pytest.mark.asyncio +@pytest.mark.parametrize("deployment_name", ["shared-ptu", "azure/gpt-4.1"]) +async def test_naming_the_shared_deployment_directly_draws_on_the_same_ceiling_as_its_group( + monkeypatch, deployment_name +): + """The router also serves a deployment named by its id or its provider model, so a team that + spent its share by group name cannot keep going under the deployment's other names.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + with patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router("test-model")): + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_two_thirds_of_a_ptu_minute(), call_type="acompletion" + ) + with pytest.raises(HTTPException) as exc: + await handler.async_pre_call_hook( + user_api_key_dict=key, + cache=cache, + data={**_two_thirds_of_a_ptu_minute(), "model": deployment_name}, + call_type="acompletion", + ) + + assert exc.value.status_code == 429 + assert "model_per_team_ptu" in str(exc.value.detail) + ptu_keys = [cache_key for cache_key in cache.in_memory_cache.cache_dict if "model_per_team_ptu" in cache_key] + assert handler.create_rate_limit_keys("model_per_team_ptu", "t:test-model", "tokens") in ptu_keys + assert all(cache_key.startswith("{model_per_team_ptu:t:test-model}") for cache_key in ptu_keys) + + @pytest.mark.asyncio async def test_the_proxy_router_sets_no_ceiling_while_attribution_is_off(monkeypatch): monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) @@ -7869,7 +7902,9 @@ async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone(): """The share can be removed between admission and completion; the reserved tokens still come off the counter instead of standing in the window.""" ceiling: dict[str, PTUTeamCeiling | None] = { - "current": PTUTeamCeiling(tpm_limit=2000, output_to_input_ratio=4.0, cached_input_ratio=0.0) + "current": PTUTeamCeiling( + model_group="test-model", tpm_limit=2000, output_to_input_ratio=4.0, cached_input_ratio=0.0 + ) } cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler( diff --git a/tests/unit/router_utils/test_ptu_shares.py b/tests/unit/router_utils/test_ptu_shares.py index 2c54c3a1c8d..0012b5628c2 100644 --- a/tests/unit/router_utils/test_ptu_shares.py +++ b/tests/unit/router_utils/test_ptu_shares.py @@ -77,8 +77,9 @@ def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_her def test_a_share_converts_to_the_models_input_tpm_per_ptu(): - ceiling: Final = team_ptu_ceiling([_shared()], "team-a") + ceiling: Final = team_ptu_ceiling([_shared()], "team-a", "gpt-4.1-ptu") assert ceiling == PTUTeamCeiling( + model_group="gpt-4.1-ptu", tpm_limit=30 * _GPT41.input_tpm_per_ptu, output_to_input_ratio=_GPT41.output_to_input_ratio, cached_input_ratio=_GPT41.cached_input_ratio, @@ -87,7 +88,7 @@ def test_a_share_converts_to_the_models_input_tpm_per_ptu(): def test_shares_across_deployments_add_up_and_the_larger_output_ratio_wins(): gpt4o: Final = _shared(model="azure/gpt-4o", shares={"team-a": 10}, deployment_id="shared-4o") - ceiling: Final = team_ptu_ceiling([_shared(), gpt4o, _OPEN], "team-a") + ceiling: Final = team_ptu_ceiling([_shared(), gpt4o, _OPEN], "team-a", "gpt-4.1-ptu") assert ceiling is not None assert ceiling.tpm_limit == 30 * _GPT41.input_tpm_per_ptu + 10 * _GPT4O.input_tpm_per_ptu assert ceiling.output_to_input_ratio == max(_GPT41.output_to_input_ratio, _GPT4O.output_to_input_ratio) @@ -97,7 +98,7 @@ def test_the_larger_cached_input_ratio_wins_across_deployments(): """A team sharing two models is weighted by the one that charges more for cache reads, whichever order the deployments come in.""" gpt6sol: Final = _shared(model="azure/gpt-6-sol", shares={"team-a": 10}, deployment_id="shared-6") - ceiling: Final = team_ptu_ceiling([_shared(), gpt6sol], "team-a") + ceiling: Final = team_ptu_ceiling([_shared(), gpt6sol], "team-a", "gpt-4.1-ptu") assert ceiling is not None assert _GPT41.cached_input_ratio < _GPT6SOL.cached_input_ratio assert ceiling.cached_input_ratio == _GPT6SOL.cached_input_ratio @@ -120,9 +121,45 @@ def test_a_group_is_served_by_name_or_by_a_team_scoped_deployments_public_name() def test_no_share_or_no_sizing_row_sets_no_ceiling(): - assert team_ptu_ceiling([_shared()], "team-c") is None - assert team_ptu_ceiling([_shared(model="azure/unknown-deployment")], "team-a") is None - assert team_ptu_ceiling([_single_team(), _OPEN], "team-a") is None + assert team_ptu_ceiling([_shared()], "team-c", "gpt-4.1-ptu") is None + assert team_ptu_ceiling([_shared(model="azure/unknown-deployment")], "team-a", "gpt-4.1-ptu") is None + assert team_ptu_ceiling([_single_team(), _OPEN], "team-a", "gpt-4.1-ptu") is None + + +def test_naming_a_shared_deployment_by_id_or_provider_model_draws_on_its_groups_ceiling(): + """The router serves a deployment id or a provider model string when no group has that + name, so those names share the group's ceiling instead of bypassing it.""" + payg: Final = {"model_name": "gpt-4.1-payg", "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {"id": "payg"}} + deployments: Final = [payg, _shared(), _OPEN] + by_group: Final = team_ptu_ceiling(deployments, "team-a", "gpt-4.1-ptu") + assert by_group is not None + assert by_group.model_group == "gpt-4.1-ptu" + assert team_ptu_ceiling(deployments, "team-a", "shared") == by_group + assert team_ptu_ceiling(deployments, "team-a", "azure/gpt-4.1") == by_group + assert team_ptu_ceiling(deployments, "team-a", "payg") is None + assert team_ptu_ceiling(deployments, "team-a", "missing") is None + + +def test_a_group_name_wins_over_a_deployment_id_it_collides_with(): + """The router routes a name that is both a group and a deployment id to the group.""" + colliding: Final = { + "model_name": "shared", + "litellm_params": {"model": "azure/gpt-4o"}, + "model_info": {"id": "colliding"}, + } + assert team_ptu_ceiling([_shared(), colliding], "team-a", "shared") is None + + +def test_a_team_scoped_deployment_named_by_id_draws_on_its_public_groups_ceiling(): + team_scoped: Final = { + "model_name": "gpt-4.1-ptu-3f9c1b", + "litellm_params": {"model": "azure/gpt-4.1"}, + "model_info": {**_shared()["model_info"], "id": "team-scoped", "team_public_model_name": "gpt-4.1-ptu"}, + } + ceiling: Final = team_ptu_ceiling([team_scoped], "team-a", "team-scoped") + assert ceiling is not None + assert ceiling.model_group == "gpt-4.1-ptu" + assert ceiling == team_ptu_ceiling([team_scoped], "team-a", "gpt-4.1-ptu") def test_a_groups_capacity_comes_from_its_first_reserved_deployment_with_a_row(): From b41eed4347c169a68de7fc0f9fa25dd0f8bcd450 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 28 Sep 2026 18:22:58 -0700 Subject: [PATCH 10/18] fix(proxy): resolve every name a shared PTU deployment is served under to its group's ceiling --- .../hooks/parallel_request_limiter_v3.py | 12 ++-- litellm/router_utils/ptu_shares.py | 52 ++++++++++++----- .../hooks/test_parallel_request_limiter_v3.py | 42 +++++++++++++- tests/unit/router_utils/test_ptu_shares.py | 58 ++++++++++++++----- 4 files changed, 129 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9fd40115e10..74daca823ac 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -134,7 +134,7 @@ def _resolve_ptu_team_ceiling_via_proxy_router(team_id: str, model_group: str) - if llm_router is None or not is_ptu_cost_attribution_enabled(): return None - return team_ptu_ceiling(llm_router.get_model_list() or (), team_id, model_group) + return team_ptu_ceiling(llm_router.get_model_list() or (), llm_router.model_list, team_id, model_group) FAIL_CLOSED_RATE_LIMIT_ENFORCEMENT_SETTING: Final = "fail_closed_rate_limit_enforcement" @@ -4935,14 +4935,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if reserved_ceiling is not None else self._ptu_team_ceiling_resolver(team_id, reconcile_model.group) ) - scope: Final = ( - PTU_TEAM_DESCRIPTOR_KEY, - f"{team_id}:{ceiling.model_group if ceiling is not None else reconcile_model.group}", + reserved_ptu_scopes: Final = tuple(scope for scope in reserved_scopes if scope[0] == PTU_TEAM_DESCRIPTOR_KEY) + targets: Final = reserved_ptu_scopes or ( + ((PTU_TEAM_DESCRIPTOR_KEY, f"{team_id}:{ceiling.model_group}"),) if ceiling is not None else () ) - if ceiling is None and scope not in reserved_scopes: + if not targets: return () return self._build_reservation_aware_tpm_ops( - targets=(scope,), + targets=targets, reserved_scopes=reserved_scopes, actual_tokens=self._ptu_settlement_tokens( ceiling, self._resolve_reconciled_usage(response_obj), total_tokens diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 91ccc585ba7..14ae11dc298 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -57,20 +57,26 @@ def filter_ptu_shared_deployments( def team_ptu_ceiling( - deployments: Sequence[Mapping[str, object]], team_id: str, requested_model: str + listed_rows: Sequence[Mapping[str, object]], + deployments: Sequence[Mapping[str, object]], + team_id: str, + requested_model: str, ) -> PTUTeamCeiling | None: """The per-minute normalized-token ceiling ``team_id``'s shares on the group serving ``requested_model`` add up to, else None when the team holds no share on a deployment with a known sizing row. - A request naming one deployment by its id or provider model counts against that deployment's - group, so every name the router serves it under shares one ceiling. + ``listed_rows`` is every row the router lists a name under, alias and routing-group copies + included, and ``deployments`` is the router's own deployments. A name resolves to the + deployments behind it, so a group, a routing group, a deployment id, and a provider model + all count against the one ceiling of the group whose shared deployment the team can be + served from. Two shared deployments of different models in one group are weighted by the larger output and cached-input ratios, which over-counts those tokens on the cheaper one rather than under-counting them on the dearer one. """ - model_group: Final = _model_group_of(deployments, requested_model) + model_group: Final = _model_group_of(listed_rows, deployments, team_id, requested_model) priced: Final = tuple( (shares[team_id], capacity) for deployment in model_group_deployments(deployments, model_group) @@ -103,10 +109,14 @@ def model_group_deployments(deployments: Sequence[_DeploymentT], model_group: st ) -def _names_deployment(deployment: Mapping[str, object], name: str) -> bool: +def _deployment_id(deployment: Mapping[str, object]) -> object: model_info: Final = deployment.get("model_info") + return model_info.get("id") if isinstance(model_info, Mapping) else None + + +def _names_deployment(deployment: Mapping[str, object], name: str) -> bool: litellm_params: Final = deployment.get("litellm_params") - return (isinstance(model_info, Mapping) and model_info.get("id") == name) or ( + return _deployment_id(deployment) == name or ( isinstance(litellm_params, Mapping) and litellm_params.get("model") == name ) @@ -120,13 +130,29 @@ def _deployment_model_group(deployment: Mapping[str, object]) -> str | None: return model_name if isinstance(model_name, str) else None -def _model_group_of(deployments: Sequence[Mapping[str, object]], requested_model: str) -> str: - """The group ``requested_model`` routes to: itself when it names a group, else the group of - the deployment it names by id or by provider model, the way the router falls back to them.""" - if model_group_deployments(deployments, requested_model): - return requested_model - named: Final = tuple(deployment for deployment in deployments if _names_deployment(deployment, requested_model)) - shared_first: Final = sorted(named, key=lambda deployment: _deployment_shares(deployment) is None) +def _routed_deployments( + listed_rows: Sequence[Mapping[str, object]], deployments: Sequence[Mapping[str, object]], requested_model: str +) -> tuple[Mapping[str, object], ...]: + """The router's own deployments behind ``requested_model``: those of the rows listed under it + when it names a group, else the one it names by id or by provider model, the way the router + falls back to them.""" + listed_ids: Final = frozenset(_deployment_id(row) for row in model_group_deployments(listed_rows, requested_model)) + if listed_ids: + return tuple(deployment for deployment in deployments if _deployment_id(deployment) in listed_ids) + return tuple(deployment for deployment in deployments if _names_deployment(deployment, requested_model)) + + +def _model_group_of( + listed_rows: Sequence[Mapping[str, object]], + deployments: Sequence[Mapping[str, object]], + team_id: str, + requested_model: str, +) -> str: + """The group of the deployment behind ``requested_model`` this team can be served from, one + holding its share first.""" + routed: Final = _routed_deployments(listed_rows, deployments, requested_model) + servable: Final = filter_ptu_shared_deployments(routed, team_id).deployments + shared_first: Final = sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None) return next( (group for deployment in shared_first if (group := _deployment_model_group(deployment)) is not None), requested_model, diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 0b3e54a8358..d22f595b634 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7658,7 +7658,8 @@ def _shared_ptu_router(model_group: str) -> Router: "ptu_shares": {"t": 1}, }, } - ] + ], + model_group_alias={f"{model_group}-alias": model_group}, ) @@ -7768,7 +7769,8 @@ async def test_naming_the_shared_deployment_directly_draws_on_the_same_ceiling_a monkeypatch, deployment_name ): """The router also serves a deployment named by its id or its provider model, so a team that - spent its share by group name cannot keep going under the deployment's other names.""" + spent its share by group name cannot keep going under the deployment's other names, even + though the router lists the deployment's alias copy ahead of it.""" monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") cache = DualCache() handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache)) @@ -7929,6 +7931,42 @@ async def test_a_reservation_is_settled_even_after_the_teams_share_is_gone(): assert _ptu_increment(handler, ops) == 300 - stash.ptu_reserved_tokens +@pytest.mark.asyncio +async def test_success_settles_the_scope_the_reservation_was_taken_on(): + """Settlement credits the scope admission reserved, not one rebuilt from the name the + request used, so a request naming the deployment by id cannot leave its reservation standing.""" + ceiling: dict[str, PTUTeamCeiling | None] = { + "current": PTUTeamCeiling( + model_group="test-model", tpm_limit=2000, output_to_input_ratio=4.0, cached_input_ratio=0.0 + ) + } + cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), + ptu_team_ceiling_resolver=lambda _team, _group: ceiling["current"], + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + + await handler.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=_ptu_request("shared-ptu"), call_type="acompletion" + ) + stash = get_request_stash() + assert stash is not None + assert ("model_per_team_ptu", "t:test-model") in stash.reserved_scopes + + ceiling["current"] = None + stash.ptu_ceiling = None + kwargs = _ptu_success_kwargs() + kwargs["litellm_params"]["metadata"]["model_group"] = "shared-ptu" + ops = handler._build_success_event_pipeline_operations( + kwargs=kwargs, + response_obj=_ptu_response(Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)), + rate_limit_type="total", + ) + + assert _ptu_increment(handler, ops) == 150 - stash.ptu_reserved_tokens + + async def _reserve_a_ptu_minute(cache: DualCache, call_id: str) -> tuple[_PROXY_MaxParallelRequestsHandler, str]: resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=2000, ratio=4.0) handler = _PROXY_MaxParallelRequestsHandler( diff --git a/tests/unit/router_utils/test_ptu_shares.py b/tests/unit/router_utils/test_ptu_shares.py index 0012b5628c2..f0cd0cb20bd 100644 --- a/tests/unit/router_utils/test_ptu_shares.py +++ b/tests/unit/router_utils/test_ptu_shares.py @@ -49,6 +49,10 @@ def _single_team(model: str = "azure/gpt-4.1") -> dict: _OPEN: Final = {"model_name": "gpt-4.1-ptu", "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {"id": "open"}} +def _unaliased_ceiling(deployments: list[dict], team_id: str, requested_model: str) -> PTUTeamCeiling | None: + return team_ptu_ceiling(deployments, deployments, team_id, requested_model) + + def test_a_team_holding_a_share_keeps_the_shared_deployment(): result: Final = filter_ptu_shared_deployments([_shared(), _OPEN], "team-a") assert [d["model_info"]["id"] for d in result.deployments] == ["shared", "open"] @@ -77,7 +81,7 @@ def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_her def test_a_share_converts_to_the_models_input_tpm_per_ptu(): - ceiling: Final = team_ptu_ceiling([_shared()], "team-a", "gpt-4.1-ptu") + ceiling: Final = _unaliased_ceiling([_shared()], "team-a", "gpt-4.1-ptu") assert ceiling == PTUTeamCeiling( model_group="gpt-4.1-ptu", tpm_limit=30 * _GPT41.input_tpm_per_ptu, @@ -88,7 +92,7 @@ def test_a_share_converts_to_the_models_input_tpm_per_ptu(): def test_shares_across_deployments_add_up_and_the_larger_output_ratio_wins(): gpt4o: Final = _shared(model="azure/gpt-4o", shares={"team-a": 10}, deployment_id="shared-4o") - ceiling: Final = team_ptu_ceiling([_shared(), gpt4o, _OPEN], "team-a", "gpt-4.1-ptu") + ceiling: Final = _unaliased_ceiling([_shared(), gpt4o, _OPEN], "team-a", "gpt-4.1-ptu") assert ceiling is not None assert ceiling.tpm_limit == 30 * _GPT41.input_tpm_per_ptu + 10 * _GPT4O.input_tpm_per_ptu assert ceiling.output_to_input_ratio == max(_GPT41.output_to_input_ratio, _GPT4O.output_to_input_ratio) @@ -98,7 +102,7 @@ def test_the_larger_cached_input_ratio_wins_across_deployments(): """A team sharing two models is weighted by the one that charges more for cache reads, whichever order the deployments come in.""" gpt6sol: Final = _shared(model="azure/gpt-6-sol", shares={"team-a": 10}, deployment_id="shared-6") - ceiling: Final = team_ptu_ceiling([_shared(), gpt6sol], "team-a", "gpt-4.1-ptu") + ceiling: Final = _unaliased_ceiling([_shared(), gpt6sol], "team-a", "gpt-4.1-ptu") assert ceiling is not None assert _GPT41.cached_input_ratio < _GPT6SOL.cached_input_ratio assert ceiling.cached_input_ratio == _GPT6SOL.cached_input_ratio @@ -121,9 +125,9 @@ def test_a_group_is_served_by_name_or_by_a_team_scoped_deployments_public_name() def test_no_share_or_no_sizing_row_sets_no_ceiling(): - assert team_ptu_ceiling([_shared()], "team-c", "gpt-4.1-ptu") is None - assert team_ptu_ceiling([_shared(model="azure/unknown-deployment")], "team-a", "gpt-4.1-ptu") is None - assert team_ptu_ceiling([_single_team(), _OPEN], "team-a", "gpt-4.1-ptu") is None + assert _unaliased_ceiling([_shared()], "team-c", "gpt-4.1-ptu") is None + assert _unaliased_ceiling([_shared(model="azure/unknown-deployment")], "team-a", "gpt-4.1-ptu") is None + assert _unaliased_ceiling([_single_team(), _OPEN], "team-a", "gpt-4.1-ptu") is None def test_naming_a_shared_deployment_by_id_or_provider_model_draws_on_its_groups_ceiling(): @@ -131,13 +135,13 @@ def test_naming_a_shared_deployment_by_id_or_provider_model_draws_on_its_groups_ name, so those names share the group's ceiling instead of bypassing it.""" payg: Final = {"model_name": "gpt-4.1-payg", "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {"id": "payg"}} deployments: Final = [payg, _shared(), _OPEN] - by_group: Final = team_ptu_ceiling(deployments, "team-a", "gpt-4.1-ptu") + by_group: Final = _unaliased_ceiling(deployments, "team-a", "gpt-4.1-ptu") assert by_group is not None assert by_group.model_group == "gpt-4.1-ptu" - assert team_ptu_ceiling(deployments, "team-a", "shared") == by_group - assert team_ptu_ceiling(deployments, "team-a", "azure/gpt-4.1") == by_group - assert team_ptu_ceiling(deployments, "team-a", "payg") is None - assert team_ptu_ceiling(deployments, "team-a", "missing") is None + assert _unaliased_ceiling(deployments, "team-a", "shared") == by_group + assert _unaliased_ceiling(deployments, "team-a", "azure/gpt-4.1") == by_group + assert _unaliased_ceiling(deployments, "team-a", "payg") is None + assert _unaliased_ceiling(deployments, "team-a", "missing") is None def test_a_group_name_wins_over_a_deployment_id_it_collides_with(): @@ -147,7 +151,7 @@ def test_a_group_name_wins_over_a_deployment_id_it_collides_with(): "litellm_params": {"model": "azure/gpt-4o"}, "model_info": {"id": "colliding"}, } - assert team_ptu_ceiling([_shared(), colliding], "team-a", "shared") is None + assert _unaliased_ceiling([_shared(), colliding], "team-a", "shared") is None def test_a_team_scoped_deployment_named_by_id_draws_on_its_public_groups_ceiling(): @@ -156,10 +160,36 @@ def test_a_team_scoped_deployment_named_by_id_draws_on_its_public_groups_ceiling "litellm_params": {"model": "azure/gpt-4.1"}, "model_info": {**_shared()["model_info"], "id": "team-scoped", "team_public_model_name": "gpt-4.1-ptu"}, } - ceiling: Final = team_ptu_ceiling([team_scoped], "team-a", "team-scoped") + ceiling: Final = _unaliased_ceiling([team_scoped], "team-a", "team-scoped") assert ceiling is not None assert ceiling.model_group == "gpt-4.1-ptu" - assert ceiling == team_ptu_ceiling([team_scoped], "team-a", "gpt-4.1-ptu") + assert ceiling == _unaliased_ceiling([team_scoped], "team-a", "gpt-4.1-ptu") + + +def test_alias_and_routing_group_copies_do_not_split_a_deployments_ceiling(): + """The router lists alias and routing-group copies of a deployment under their own names + ahead of its own rows, so every name still resolves to the deployment's group.""" + shared: Final = _shared() + listed: Final = [{**shared, "model_name": "ptu-alias"}, {**shared, "model_name": "ptu-routing-group"}, shared] + by_group: Final = team_ptu_ceiling(listed, [shared], "team-a", "gpt-4.1-ptu") + assert by_group is not None + assert by_group.model_group == "gpt-4.1-ptu" + for name in ("shared", "azure/gpt-4.1", "ptu-routing-group"): + assert team_ptu_ceiling(listed, [shared], "team-a", name) == by_group + + +def test_a_provider_model_draws_on_the_group_where_the_team_holds_its_share(): + """Two groups share deployments of one provider model among different teams, and the router + serves each team only the one it holds a share of.""" + east: Final = _shared(shares={"team-a": 30}, deployment_id="east") + west: Final = {**_shared(shares={"team-b": 20}, deployment_id="west"), "model_name": "gpt-4.1-ptu-west"} + by_provider_model: Final = _unaliased_ceiling([east, west], "team-b", "azure/gpt-4.1") + assert by_provider_model is not None + assert by_provider_model.model_group == "gpt-4.1-ptu-west" + assert by_provider_model == _unaliased_ceiling([east, west], "team-b", "gpt-4.1-ptu-west") + assert _unaliased_ceiling([east, west], "team-a", "azure/gpt-4.1") == _unaliased_ceiling( + [east, west], "team-a", "gpt-4.1-ptu" + ) def test_a_groups_capacity_comes_from_its_first_reserved_deployment_with_a_row(): From d2eda3adc01553d16b54055c5023a6f455417df8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:57:07 -0700 Subject: [PATCH 11/18] fix(ptu): refuse ptu_shares declared without ptu_count and cost_per_ptu_per_hour --- litellm/litellm_core_utils/ptu_pricing.py | 5 ++++- tests/unit/litellm_core_utils/test_ptu_pricing.py | 15 +++++++++++++++ .../test_ptu_model_settings.py | 7 +++++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/ptu_pricing.py b/litellm/litellm_core_utils/ptu_pricing.py index f9a4335bb4c..683ef9a5ba4 100644 --- a/litellm/litellm_core_utils/ptu_pricing.py +++ b/litellm/litellm_core_utils/ptu_pricing.py @@ -229,8 +229,11 @@ def ptu_config_error(model_info: Mapping[str, object], *, model_name: str | None has_count: Final = model_info.get("ptu_count") is not None has_rate: Final = model_info.get("cost_per_ptu_per_hour") is not None - if not has_count and not has_rate: + has_shares: Final = model_info.get("ptu_shares") is not None + if not has_count and not has_rate and not has_shares: return None + if not has_count and not has_rate: + return _named("ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set", model_name) if has_count != has_rate: return _named("ptu_count and cost_per_ptu_per_hour must be set together", model_name) if effective_from is None: diff --git a/tests/unit/litellm_core_utils/test_ptu_pricing.py b/tests/unit/litellm_core_utils/test_ptu_pricing.py index 7dda3432e43..28b57fb328d 100644 --- a/tests/unit/litellm_core_utils/test_ptu_pricing.py +++ b/tests/unit/litellm_core_utils/test_ptu_pricing.py @@ -439,6 +439,21 @@ def test_a_fractional_count_reserves_nothing(): assert ptu_terms({**_SHARED, "ptu_count": 100.5}) is None +@pytest.mark.parametrize( + "shares_only", + [ + {"ptu_shares": {"team-a": 60, "team-b": 40}}, + {"ptu_shares": {"team-a": 60, "team-b": 40}, "ptu_effective_from": "2026-01-01T00:00:00Z"}, + ], + ids=["shares alone", "shares with a window"], +) +def test_shares_without_the_count_and_rate_are_refused_instead_of_registering_unpriced(shares_only): + """Shares split a count, so a split with no count has nothing to add up to and would register a + deployment the share filter and ceiling honour while flat cost never rolls up.""" + assert ptu_config_error(shares_only) == "ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set" + assert ptu_terms(shares_only) is None + + def test_the_split_is_named_after_the_deployment_when_the_caller_supplies_one(): error = ptu_config_error({**_SHARED, "ptu_shares": {"team-a": 1}}, model_name="gpt-4.1-ptu") diff --git a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py index ce2e865d69d..8bbde32e756 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py @@ -1334,6 +1334,13 @@ def test_validate_helper_refuses_shares_that_do_not_add_up_to_the_count(): assert "4 of 5 allocated" in exc.value.detail +def test_validate_helper_refuses_shares_declared_without_a_count_and_rate(): + with pytest.raises(HTTPException) as exc: + _validate_ptu_model_info({"ptu_shares": {"team-a": 3, "team-b": 2}, "ptu_effective_from": _SHARED_START}) + assert exc.value.status_code == 400 + assert exc.value.detail == "ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set" + + class _TeamLookup: def __init__(self, existing: frozenset[str]) -> None: self.existing: Final = existing From 96a5eee5af58c61e971d504b301957b0a205fe05 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 1 Oct 2026 17:50:31 -0700 Subject: [PATCH 12/18] fix(ui): seed total_ptu_hours in the empty daily activity metadata --- .../src/components/UsagePage/dailyActivityApi.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/ui/litellm-dashboard/src/components/UsagePage/dailyActivityApi.ts b/ui/litellm-dashboard/src/components/UsagePage/dailyActivityApi.ts index 36c2837d9ac..73600501e6c 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/dailyActivityApi.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/dailyActivityApi.ts @@ -32,6 +32,7 @@ export const EMPTY_DAILY_ACTIVITY_METADATA: DailyActivityMetadata = { total_pages: 1, total_spend: 0, total_flat_cost: 0, + total_ptu_hours: 0, total_api_requests: 0, total_successful_requests: 0, total_failed_requests: 0, From 4a5828ed7a8238617e0c44289d76814fb68beaca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:29:19 -0700 Subject: [PATCH 13/18] fix(ptu): size usage rows by the name the request used --- .../management_endpoints/ptu_consumption.py | 32 +++++++++++---- litellm/router_utils/ptu_shares.py | 4 +- .../test_ptu_consumption.py | 41 ++++++++++++++++++- 3 files changed, 65 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/management_endpoints/ptu_consumption.py b/litellm/proxy/management_endpoints/ptu_consumption.py index 2c8d7d1930e..7d13c293d70 100644 --- a/litellm/proxy/management_endpoints/ptu_consumption.py +++ b/litellm/proxy/management_endpoints/ptu_consumption.py @@ -12,7 +12,8 @@ from typing import Final from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabled from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours from litellm.router import Router -from litellm.router_utils.ptu_shares import model_group_deployments, model_group_ptu_capacity +from litellm.router_utils.common_utils import resolve_model_group_alias +from litellm.router_utils.ptu_shares import model_group_ptu_capacity, routed_deployments from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, MetricWithMetadata, @@ -98,16 +99,31 @@ def attach_ptu_hours( ) +def capacity_by_requested_name(llm_router: Router) -> Callable[[str], PTUCapacity | None]: + """The sizing row behind the name a usage row is keyed by, resolved the way the ceiling + resolves a request: an alias to its group, then a group, a routing group, a deployment id, + or a provider model to the deployments it is served from.""" + listed_rows: Final = llm_router.get_model_list() or () + aliases: Final = llm_router.model_group_alias + + def capacity_for(requested_model: str) -> PTUCapacity | None: + model_group: Final = resolve_model_group_alias(aliases, requested_model) or requested_model + return model_group_ptu_capacity( + routed_deployments( + listed_rows, + llm_router.model_list, # pyright: ignore[reportUnknownArgumentType] # Router.model_list is a bare list + model_group, + ) + ) + + return capacity_for + + def with_ptu_consumption( activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None ) -> SpendAnalyticsPaginatedResponse: - """``activity`` with PTU-hours attached from the router's sized model groups, untouched while + """``activity`` with PTU-hours attached from the router's sized deployments, untouched while PTU cost attribution is off or no router is loaded.""" if llm_router is None or not is_ptu_cost_attribution_enabled(): return activity - return attach_ptu_hours( - activity, - lambda model_group: model_group_ptu_capacity( - model_group_deployments(llm_router.get_model_list() or (), model_group) - ), - ) + return attach_ptu_hours(activity, capacity_by_requested_name(llm_router)) diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index d4f5ffc6e18..ed59a3a40d1 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -130,7 +130,7 @@ def _deployment_model_group(deployment: Mapping[str, object]) -> str | None: return model_name if isinstance(model_name, str) else None -def _routed_deployments( +def routed_deployments( listed_rows: Sequence[Mapping[str, object]], deployments: Sequence[Mapping[str, object]], requested_model: str ) -> tuple[Mapping[str, object], ...]: """The router's own deployments behind ``requested_model``: those of the rows listed under it @@ -150,7 +150,7 @@ def _model_group_of( ) -> str: """The group of the deployment behind ``requested_model`` this team can be served from, one holding its share first.""" - routed: Final = _routed_deployments(listed_rows, deployments, requested_model) + routed: Final = routed_deployments(listed_rows, deployments, requested_model) servable: Final = filter_ptu_shared_deployments(routed, team_id).deployments shared_first: Final = sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None) return next( diff --git a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py index 432bf858899..0b4b048aa1b 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py @@ -4,8 +4,9 @@ from typing import Final import pytest -from litellm.llms.azure.ptu_capacity import PTUCapacity -from litellm.proxy.management_endpoints.ptu_consumption import attach_ptu_hours +from litellm import Router +from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY, PTUCapacity +from litellm.proxy.management_endpoints.ptu_consumption import attach_ptu_hours, with_ptu_consumption from litellm.types.proxy.management_endpoints.common_daily_activity import ( BreakdownMetrics, DailySpendData, @@ -19,6 +20,7 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( _ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0) _CACHED_ROW: Final = PTUCapacity(input_tpm_per_ptu=1_000, output_to_input_ratio=4.0, cached_input_ratio=0.1) _CAPACITY: Final = {"gpt-4.1-ptu": _ROW, "gpt-6-ptu": _CACHED_ROW} +_ONE_PTU_HOUR_OF_INPUT: Final = AZURE_PTU_CAPACITY["gpt-4.1"].normalized_tokens_per_ptu_hour def _metrics(prompt: int, completion: int, cached: int = 0) -> SpendMetrics: @@ -128,3 +130,38 @@ def test_a_page_with_no_ptu_group_is_returned_unchanged(): assert attached.results[0] is day assert attached.metadata.total_ptu_hours == 0.0 assert response.metadata.total_ptu_hours == 0.0 + + +def _shared_ptu_router() -> Router: + return Router( + model_list=[ + { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "shared-ptu", + "base_model": "azure/gpt-4.1", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 30, "team-b": 20}, + }, + } + ], + model_group_alias={"ptu-alias": {"model": "gpt-4.1-ptu", "hidden": True}}, + ) + + +def test_rows_keyed_by_an_alias_a_deployment_id_or_a_provider_model_are_sized_like_the_group(monkeypatch): + """The ceiling charges a request however it names the shared deployment, so the usage row + that request lands in, keyed by the name it used, reports the same PTU-hours as the group.""" + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + names: Final = ("gpt-4.1-ptu", "ptu-alias", "shared-ptu", "azure/gpt-4.1") + one_hour_each: Final = {name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in names} + + attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _shared_ptu_router()) + + groups: Final = attached.results[0].breakdown.model_groups + assert [groups[name].metrics.ptu_hours for name in names] == [pytest.approx(1.0)] * len(names) + assert attached.results[0].metrics.ptu_hours == pytest.approx(float(len(names))) + assert attached.metadata.total_ptu_hours == pytest.approx(float(len(names))) From 5286e6c49afb25f068da16b038789a4254e31420 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:58:35 -0700 Subject: [PATCH 14/18] fix(ptu): size a team's usage rows by the deployment it is served from --- .../daily_activity_routes.py | 6 +- .../management_endpoints/ptu_consumption.py | 39 ++++++---- .../management_endpoints/team_endpoints.py | 4 +- litellm/router_utils/ptu_shares.py | 12 ++-- .../test_daily_activity_routes.py | 26 ++++++- .../test_ptu_consumption.py | 72 ++++++++++++++++++- .../test_team_endpoints.py | 33 +++++++++ 7 files changed, 166 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 15ca658ee38..e5068b020d1 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/litellm/proxy/management_endpoints/daily_activity_routes.py @@ -34,7 +34,7 @@ from litellm.proxy.management_endpoints.daily_activity_scopes import ( EntityScopeResolver, ResolvedScope, ) -from litellm.proxy.management_endpoints.ptu_consumption import with_ptu_consumption +from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, with_ptu_consumption from litellm.proxy.management_endpoints.team_endpoints import aggregated_date_range_error from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import PrismaClient, get_prisma_client_or_throw @@ -272,7 +272,9 @@ def _register_aggregated_route(router: APIRouter, resolver: EntityScopeResolver, include_entity_breakdown=resolver.include_entity_breakdown, api_key_limit=api_key_limit, ) - return with_ptu_consumption(activity, llm_router) if resolver.entity == "team" else activity + if resolver.entity != "team": + return activity + return with_ptu_consumption(activity, llm_router, single_team_id(resolved.scope.entity_ids)) except HTTPException: raise except Exception as exc: diff --git a/litellm/proxy/management_endpoints/ptu_consumption.py b/litellm/proxy/management_endpoints/ptu_consumption.py index 7d13c293d70..7b3752d184a 100644 --- a/litellm/proxy/management_endpoints/ptu_consumption.py +++ b/litellm/proxy/management_endpoints/ptu_consumption.py @@ -5,7 +5,7 @@ cached, and completion tokens a team sent to a PTU model group convert back to t a PTU-hour it consumed, reported next to the raw token counts. """ -from collections.abc import Callable +from collections.abc import Callable, Sequence from types import MappingProxyType from typing import Final @@ -13,7 +13,11 @@ from litellm.litellm_core_utils.ptu_pricing import is_ptu_cost_attribution_enabl from litellm.llms.azure.ptu_capacity import PTUCapacity, normalized_tokens, ptu_hours from litellm.router import Router from litellm.router_utils.common_utils import resolve_model_group_alias -from litellm.router_utils.ptu_shares import model_group_ptu_capacity, routed_deployments +from litellm.router_utils.ptu_shares import ( + model_group_ptu_capacity, + routed_deployments, + team_servable_deployments, +) from litellm.types.proxy.management_endpoints.common_daily_activity import ( DailySpendData, MetricWithMetadata, @@ -99,31 +103,38 @@ def attach_ptu_hours( ) -def capacity_by_requested_name(llm_router: Router) -> Callable[[str], PTUCapacity | None]: +def capacity_by_requested_name(llm_router: Router, team_id: str | None) -> Callable[[str], PTUCapacity | None]: """The sizing row behind the name a usage row is keyed by, resolved the way the ceiling resolves a request: an alias to its group, then a group, a routing group, a deployment id, - or a provider model to the deployments it is served from.""" + or a provider model to the deployments it is served from. For one team those are narrowed + to the deployments it can be served from, its shared one first, so a team with no share on + a name reads no PTU-hours for it; a page spanning teams keeps every deployment behind the name.""" listed_rows: Final = llm_router.get_model_list() or () aliases: Final = llm_router.model_group_alias def capacity_for(requested_model: str) -> PTUCapacity | None: model_group: Final = resolve_model_group_alias(aliases, requested_model) or requested_model - return model_group_ptu_capacity( - routed_deployments( - listed_rows, - llm_router.model_list, # pyright: ignore[reportUnknownArgumentType] # Router.model_list is a bare list - model_group, - ) + routed: Final = routed_deployments( + listed_rows, + llm_router.model_list, # pyright: ignore[reportUnknownArgumentType] # Router.model_list is a bare list + model_group, ) + return model_group_ptu_capacity(routed if team_id is None else team_servable_deployments(routed, team_id)) return capacity_for +def single_team_id(team_ids: Sequence[str] | None) -> str | None: + """The one team a usage page is scoped to, else None for a page spanning several or all teams.""" + return team_ids[0] if team_ids is not None and len(team_ids) == 1 else None + + def with_ptu_consumption( - activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None + activity: SpendAnalyticsPaginatedResponse, llm_router: Router | None, team_id: str | None ) -> SpendAnalyticsPaginatedResponse: - """``activity`` with PTU-hours attached from the router's sized deployments, untouched while - PTU cost attribution is off or no router is loaded.""" + """``activity`` with PTU-hours attached from the deployments ``team_id`` is served from (every + sized deployment when the page spans teams), untouched while PTU cost attribution is off or no + router is loaded.""" if llm_router is None or not is_ptu_cost_attribution_enabled(): return activity - return attach_ptu_hours(activity, capacity_by_requested_name(llm_router)) + return attach_ptu_hours(activity, capacity_by_requested_name(llm_router, team_id)) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 3d7609401df..e8486925d54 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -143,7 +143,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, ) -from litellm.proxy.management_endpoints.ptu_consumption import with_ptu_consumption +from litellm.proxy.management_endpoints.ptu_consumption import single_team_id, with_ptu_consumption from litellm.proxy.management_endpoints.router_weights import validate_router_settings_weights from litellm.proxy.management_endpoints.tag_management_endpoints import ( get_daily_activity, @@ -6689,7 +6689,7 @@ async def get_team_daily_activity( page=page, page_size=page_size, ) - return with_ptu_consumption(activity, llm_router) + return with_ptu_consumption(activity, llm_router, single_team_id(scope.team_ids)) _MAX_AGGREGATED_RANGE_DAYS: Final = 400 diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index ed59a3a40d1..b763b470ff2 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -150,15 +150,19 @@ def _model_group_of( ) -> str: """The group of the deployment behind ``requested_model`` this team can be served from, one holding its share first.""" - routed: Final = routed_deployments(listed_rows, deployments, requested_model) - servable: Final = filter_ptu_shared_deployments(routed, team_id).deployments - shared_first: Final = sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None) + servable: Final = team_servable_deployments(routed_deployments(listed_rows, deployments, requested_model), team_id) return next( - (group for deployment in shared_first if (group := _deployment_model_group(deployment)) is not None), + (group for deployment in servable if (group := _deployment_model_group(deployment)) is not None), requested_model, ) +def team_servable_deployments(deployments: Sequence[_DeploymentT], team_id: str) -> tuple[_DeploymentT, ...]: + """The deployments ``team_id`` can be served from, the ones it holds a share on first.""" + servable: Final = filter_ptu_shared_deployments(deployments, team_id).deployments + return tuple(sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None)) + + def model_group_ptu_capacity(deployments: Sequence[Mapping[str, object]]) -> PTUCapacity | None: """The sizing row of the group's first reserved deployment, single-team or shared, so a team's tokens on that group convert to PTU-hours.""" diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index 69caaba6005..5cb37b16221 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py @@ -3,7 +3,7 @@ import io from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from dataclasses import dataclass, fields from itertools import chain -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock @@ -578,7 +578,7 @@ def test_aggregated_routes_return_scoped_results( assert len(body["results"]) == 2, response.text -def _ptu_sized_router(model_group: str) -> Router: +def _ptu_sized_router(model_group: str, shares: Mapping[str, int] = MappingProxyType({"team-a": 30, "team-b": 20})) -> Router: return Router( model_list=[ { @@ -590,7 +590,7 @@ def _ptu_sized_router(model_group: str) -> Router: "ptu_count": 50, "cost_per_ptu_per_hour": 1.0, "ptu_effective_from": "2024-01-01T00:00:00Z", - "ptu_shares": {"team-a": 30, "team-b": 20}, + "ptu_shares": dict(shares), }, } ] @@ -637,6 +637,26 @@ def test_team_aggregated_route_reports_ptu_hours_for_the_sized_model_group_only_ assert body["metadata"]["total_tokens"] == 90, response.text +def test_team_aggregated_route_reports_no_ptu_hours_to_a_team_holding_no_share( + daily_activity_client: tuple[TestClient, _FakeRepository], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """rare-group is split between team-b and team-c, so team-a's page alone converts none of its + rare-group tokens even though the group is sized.""" + client, _ = daily_activity_client + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + _inject_llm_router(client, _ptu_sized_router("rare-group", shares=MappingProxyType({"team-b": 30, "team-c": 20}))) + + response: Final = client.get("/team/daily/activity/aggregated", params=_entity_params("team_ids", "team-a")) + + assert response.status_code == 200, response.text + body: Final = response.json() + rare_day: Final = {day["date"]: day for day in body["results"]}["2025-01-01"] + assert body["metadata"]["total_ptu_hours"] == 0, response.text + assert rare_day["breakdown"]["model_groups"]["rare-group"]["metrics"]["ptu_hours"] == 0, response.text + assert body["metadata"]["total_tokens"] == 90, response.text + + def test_ptu_hours_stay_zero_outside_the_team_routes( daily_activity_client: tuple[TestClient, _FakeRepository], monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py index 0b4b048aa1b..d05135d94f3 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py @@ -159,9 +159,79 @@ def test_rows_keyed_by_an_alias_a_deployment_id_or_a_provider_model_are_sized_li names: Final = ("gpt-4.1-ptu", "ptu-alias", "shared-ptu", "azure/gpt-4.1") one_hour_each: Final = {name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in names} - attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _shared_ptu_router()) + attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _shared_ptu_router(), "team-a") groups: Final = attached.results[0].breakdown.model_groups assert [groups[name].metrics.ptu_hours for name in names] == [pytest.approx(1.0)] * len(names) assert attached.results[0].metrics.ptu_hours == pytest.approx(float(len(names))) assert attached.metadata.total_ptu_hours == pytest.approx(float(len(names))) + + +def _mixed_ptu_router() -> Router: + """One group split between team-a on a gpt-4.1 PTU deployment and team-b on a gpt-5.5 one, + beside an open pay-as-you-go deployment of gpt-4.1 in its own group.""" + return Router( + model_list=[ + { + "model_name": "ptu", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "ptu-41", + "base_model": "azure/gpt-4.1", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-a": 50}, + }, + }, + { + "model_name": "ptu", + "litellm_params": {"model": "azure/gpt-5.5", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "ptu-55", + "base_model": "azure/gpt-5.5", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + "ptu_shares": {"team-b": 50}, + }, + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-payg", "api_base": "https://payg.example"}, + "model_info": {"id": "payg"}, + }, + ] + ) + + +_ONE_GPT_41_PTU_HOUR_ON_GPT_55: Final = _ONE_PTU_HOUR_OF_INPUT / AZURE_PTU_CAPACITY["gpt-5.5"].normalized_tokens_per_ptu_hour + + +@pytest.mark.parametrize( + ("team_id", "expected_on_the_group", "expected_on_the_provider_model"), + [ + ("team-a", 1.0, 1.0), + ("team-b", _ONE_GPT_41_PTU_HOUR_ON_GPT_55, 0.0), + ("team-c", 0.0, 0.0), + (None, 1.0, 1.0), + ], +) +def test_a_teams_rows_are_sized_by_the_deployment_it_is_served_from( + monkeypatch, team_id: str | None, expected_on_the_group: float, expected_on_the_provider_model: float +): + """A team's page converts its tokens through the deployment the ceiling served it from: team-b's + share is on the gpt-5.5 deployment so its group row counts at that rate, its provider-model row + reached only the open deployment so it counts nothing, team-c holds no share so it counts nothing, + and a page spanning teams keeps the group's first sized deployment.""" + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + one_hour_each: Final = { + name: _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0)) for name in ("ptu", "azure/gpt-4.1") + } + + attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour_each)), _mixed_ptu_router(), team_id) + + groups: Final = attached.results[0].breakdown.model_groups + assert groups["ptu"].metrics.ptu_hours == pytest.approx(expected_on_the_group) + assert groups["azure/gpt-4.1"].metrics.ptu_hours == pytest.approx(expected_on_the_provider_model) + assert attached.metadata.total_ptu_hours == pytest.approx(expected_on_the_group + expected_on_the_provider_model) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 52ed68f2aa9..6c9ed92808a 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -16964,6 +16964,39 @@ async def test_team_daily_activity_sizes_a_team_scoped_deployment_by_its_public_ assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 1.0 +@pytest.mark.asyncio +async def test_team_daily_activity_reports_no_ptu_hours_to_a_team_holding_no_share( + mock_db_client, mock_admin_auth, monkeypatch +): + """The deployment is split between team-a and team-b, so team-c's page alone converts none of + its tokens on the group even though the group is sized.""" + from litellm.proxy.management_endpoints.team_endpoints import get_team_daily_activity + + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + page = _ptu_activity_page() + + with ( + patch("litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", AsyncMock(return_value=page)), + patch("litellm.proxy.proxy_server.llm_router", _shared_ptu_router()), + ): + result = await get_team_daily_activity( + team_ids="team-c", + start_date="2026-09-23", + end_date="2026-09-24", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids=None, + user_api_key_dict=mock_admin_auth, + ) + + assert result.metadata.total_ptu_hours == 0.0 + assert result.results[0].breakdown.model_groups["gpt-4.1-ptu"].metrics.ptu_hours == 0.0 + assert result.results[0].metrics.total_tokens == _ONE_PTU_HOUR_OF_INPUT + + @pytest.mark.parametrize( ("start_date", "end_date"), ( From e65e12d72833bc09e3edd84674320e59f4094891 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:41:16 -0700 Subject: [PATCH 15/18] fix(ptu): never size a team's usage row by another team's reservation --- litellm/router.py | 4 +- litellm/router_utils/common_utils.py | 5 +++ litellm/router_utils/ptu_shares.py | 14 ++++++- .../test_ptu_consumption.py | 42 +++++++++++++++++++ tests/unit/router_utils/test_ptu_shares.py | 15 ++++++- 5 files changed, 75 insertions(+), 5 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 47e440e8ee2..d4b5d3405d2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -185,6 +185,7 @@ from litellm.router_utils.common_utils import ( get_request_team_id, provider_for_generic_call, resolve_model_group_alias, + team_may_use_deployment, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, ) @@ -10325,8 +10326,7 @@ class Router: callers from that same team; deployments without a team owner are shared. """ model_info: Final = model.get("model_info") if isinstance(model, dict) else model.model_info - owner_team_id: Final = model_info.get("team_id") if model_info is not None else None - return owner_team_id is None or owner_team_id == team_id + return team_may_use_deployment(model_info.get("team_id") if model_info is not None else None, team_id) def _get_model_group_deployment_usable_by_team( self, model_group_name: str, team_id: str | None diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 7d111a80264..2d21cf3230e 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -28,6 +28,11 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" +def team_may_use_deployment(owner_team_id: object, team_id: str | None) -> bool: + """A deployment owned by a team (``model_info.team_id``) serves only that team; one with no owner serves every team.""" + return owner_team_id is None or owner_team_id == team_id + + def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None: """The caller's team id, from whichever metadata bucket this surface writes to.""" if request_kwargs is None: diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index b763b470ff2..64d9cf1dc8a 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -12,6 +12,7 @@ from typing import Final, Generic, TypeVar from litellm.litellm_core_utils.ptu_pricing import is_model_info_mapping, parsed_ptu_shares, ptu_terms from litellm.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity, is_azure_deployment +from litellm.router_utils.common_utils import team_may_use_deployment _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) @@ -157,9 +158,18 @@ def _model_group_of( ) +def _deployment_owner(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return model_info.get("team_id") if isinstance(model_info, Mapping) else None + + def team_servable_deployments(deployments: Sequence[_DeploymentT], team_id: str) -> tuple[_DeploymentT, ...]: - """The deployments ``team_id`` can be served from, the ones it holds a share on first.""" - servable: Final = filter_ptu_shared_deployments(deployments, team_id).deployments + """The deployments ``team_id`` can be served from, the ones it holds a share on first: never one + reserved for another team, nor one split into shares it holds none of.""" + shareable: Final = filter_ptu_shared_deployments(deployments, team_id).deployments + servable: Final = tuple( + deployment for deployment in shareable if team_may_use_deployment(_deployment_owner(deployment), team_id) + ) return tuple(sorted(servable, key=lambda deployment: _deployment_shares(deployment) is None)) diff --git a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py index d05135d94f3..6c4366bd7c3 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py @@ -205,6 +205,48 @@ def _mixed_ptu_router() -> Router: ) +def _reserved_ptu_router() -> Router: + """A gpt-4.1 PTU deployment reserved for team-x alone the single-team way, beside an open + pay-as-you-go deployment of gpt-4.1 that serves every other team.""" + return Router( + model_list=[ + { + "model_name": "reserved", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-ptu", "api_base": "https://ptu.example"}, + "model_info": { + "id": "reserved-41", + "team_id": "team-x", + "ptu_count": 50, + "cost_per_ptu_per_hour": 1.0, + "ptu_effective_from": "2026-01-01T00:00:00Z", + }, + }, + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "azure/gpt-4.1", "api_key": "sk-payg", "api_base": "https://payg.example"}, + "model_info": {"id": "payg"}, + }, + ] + ) + + +@pytest.mark.parametrize(("team_id", "expected_on_the_provider_model"), [("team-x", 1.0), ("team-y", 0.0), (None, 1.0)]) +def test_another_teams_single_team_reservation_sizes_nothing_on_a_teams_page( + monkeypatch, team_id: str | None, expected_on_the_provider_model: float +): + """A provider-model row reaches the reserved deployment only for the team it is reserved for: any + other team was served by the open deployment, so its page counts no PTU-hours on that row, while + a page spanning teams keeps the reserved deployment's sizing.""" + monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") + one_hour: Final = {"azure/gpt-4.1": _bucket(_metrics(prompt=_ONE_PTU_HOUR_OF_INPUT, completion=0))} + + attached: Final = with_ptu_consumption(_response(_day("2026-09-23", one_hour)), _reserved_ptu_router(), team_id) + + row: Final = attached.results[0].breakdown.model_groups["azure/gpt-4.1"] + assert row.metrics.ptu_hours == pytest.approx(expected_on_the_provider_model) + assert attached.metadata.total_ptu_hours == pytest.approx(expected_on_the_provider_model) + + _ONE_GPT_41_PTU_HOUR_ON_GPT_55: Final = _ONE_PTU_HOUR_OF_INPUT / AZURE_PTU_CAPACITY["gpt-5.5"].normalized_tokens_per_ptu_hour diff --git a/tests/unit/router_utils/test_ptu_shares.py b/tests/unit/router_utils/test_ptu_shares.py index f0cd0cb20bd..0e1fd165dee 100644 --- a/tests/unit/router_utils/test_ptu_shares.py +++ b/tests/unit/router_utils/test_ptu_shares.py @@ -4,12 +4,13 @@ from typing import Final from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.router_utils.ptu_shares import ( - PTUTeamCeiling, filter_ptu_shared_deployments, model_group_deployments, model_group_ptu_capacity, ptu_capacity_warning, + PTUTeamCeiling, team_ptu_ceiling, + team_servable_deployments, ) _GPT41: Final = AZURE_PTU_CAPACITY["gpt-4.1"] @@ -80,6 +81,18 @@ def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_her assert result.withheld is False +def test_a_team_is_served_from_its_shared_deployment_first_and_never_from_another_teams_reservation(): + """team-b holds a share on the shared deployment and the open one serves everyone, while the + deployment reserved for team-a alone is withheld from it the way the router withholds it.""" + servable: Final = team_servable_deployments([_single_team(), _OPEN, _shared()], "team-b") + assert [d["model_info"]["id"] for d in servable] == ["shared", "open"] + + +def test_the_owning_team_keeps_its_reservation_behind_its_share(): + servable: Final = team_servable_deployments([_single_team(), _OPEN, _shared()], "team-a") + assert [d["model_info"]["id"] for d in servable] == ["shared", "single", "open"] + + def test_a_share_converts_to_the_models_input_tpm_per_ptu(): ceiling: Final = _unaliased_ceiling([_shared()], "team-a", "gpt-4.1-ptu") assert ceiling == PTUTeamCeiling( From 495cf74ccdd9a15489b4a8f9b479b2e5748db3e8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 2 Oct 2026 18:30:58 -0700 Subject: [PATCH 16/18] fix(ptu): reserve an uncapped request against the cap the proxy writes, not its prompt length --- .../hooks/parallel_request_limiter_v3.py | 3 ++- .../hooks/test_parallel_request_limiter_v3.py | 22 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a9fefdef3b3..9d97014d95d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4236,9 +4236,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): min_configured_tpm_limit, ) + capped_request: Final = _REQUEST_RATE_LIMIT_DATA.validate_python(data) ptu_estimated_tokens: Final = self._estimate_ptu_tokens_for_request( ceiling=stash.ptu_ceiling, - data=request_data, + data=capped_request, min_configured_tpm_limit=min_configured_tpm_limit, call_type=call_type, configured_output_tokens=configured_output_tokens, diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 192443a0d32..0dccfd72856 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7561,6 +7561,28 @@ async def test_a_one_ptu_share_admits_four_uncapped_requests_a_minute_and_reject assert all(data["max_tokens"] * 4 <= 3000 // 4 for data in admitted) +@pytest.mark.asyncio +async def test_an_uncapped_prompt_longer_than_the_output_floor_reserves_the_cap_the_proxy_writes(): + """Without a cap the output budget defaults to the prompt's own length, so a long prompt would be + reserved twice, weighted 4:1, and refused on an empty window. The ceiling counts the cap the proxy + writes into the request instead, the same output the deployment can produce.""" + cache = DualCache() + resolve, _ = _ptu_ceiling_for("t", "test-model", tpm_limit=3000, ratio=4.0) + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(cache), ptu_team_ceiling_resolver=resolve + ) + key = UserAPIKeyAuth(api_key=hash_token("sk-ptu"), team_id="t") + data = {"model": "test-model", "messages": [{"role": "user", "content": "word " * 1200}]} + + await handler.async_pre_call_hook(user_api_key_dict=key, cache=cache, data=data, call_type="acompletion") + + stash = get_request_stash() + assert stash is not None + assert data["max_tokens"] * 4 <= 3000 // 4 + assert stash.ptu_reserved_tokens == stash.reserved_tokens + 3 * data["max_tokens"] + assert stash.ptu_reserved_tokens <= 3000 + + @pytest.mark.asyncio async def test_the_ptu_counter_holds_the_normalized_reservation_beside_the_raw_one(): cache = DualCache() From 6fb24d1dad342625209f30c46ef462e051cb7e58 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 2 Oct 2026 19:03:41 -0700 Subject: [PATCH 17/18] fix(ptu): read a share map only through the terms registration requires A ptu_shares map stored without ptu_count, cost_per_ptu_per_hour and ptu_effective_from, or one that does not add up to ptu_count, is refused at registration but was enforced once loaded: every other team got the 400 and the share holder a ceiling while no flat cost accrued. _deployment_shares now reads the shares through ptu_terms, so such a row reserves nothing, the same as before the map was read. --- litellm/router_utils/ptu_shares.py | 13 +++++++++--- tests/unit/router_utils/test_ptu_shares.py | 23 ++++++++++++++++++++-- 2 files changed, 31 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 64d9cf1dc8a..858e5260b54 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -10,7 +10,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Final, Generic, TypeVar -from litellm.litellm_core_utils.ptu_pricing import is_model_info_mapping, parsed_ptu_shares, ptu_terms +from litellm.litellm_core_utils.ptu_pricing import is_model_info_mapping, ptu_terms from litellm.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity, is_azure_deployment from litellm.router_utils.common_utils import team_may_use_deployment @@ -36,10 +36,17 @@ class PTUShareFilterResult(Generic[_DeploymentT]): def _deployment_shares(deployment: Mapping[str, object]) -> Mapping[str, int] | None: + """The teams a deployment's capacity is split across, else None. + + Only a map registration would accept counts, read through the same terms flat cost accrues + under, so a row that reached the table without its count, rate and start reserves nothing + instead of refusing every other team while charging nobody. + """ model_info: Final = deployment.get("model_info") - if not is_model_info_mapping(model_info): + if not is_model_info_mapping(model_info) or model_info.get("ptu_shares") is None: return None - return parsed_ptu_shares(model_info.get("ptu_shares")) + terms: Final = ptu_terms(model_info) + return None if terms is None else terms.shares def filter_ptu_shared_deployments( diff --git a/tests/unit/router_utils/test_ptu_shares.py b/tests/unit/router_utils/test_ptu_shares.py index 0e1fd165dee..ae4e26bac74 100644 --- a/tests/unit/router_utils/test_ptu_shares.py +++ b/tests/unit/router_utils/test_ptu_shares.py @@ -19,13 +19,15 @@ _GPT6SOL: Final = AZURE_PTU_CAPACITY["gpt-6-sol"] _SHARES: Final = {"team-a": 30, "team-b": 20} -def _shared(model: str = "azure/gpt-4.1", shares: object = _SHARES, deployment_id: str = "shared") -> dict: +def _shared( + model: str = "azure/gpt-4.1", shares: object = _SHARES, deployment_id: str = "shared", ptu_count: int | None = None +) -> dict: return { "model_name": "gpt-4.1-ptu", "litellm_params": {"model": model}, "model_info": { "id": deployment_id, - "ptu_count": 50, + "ptu_count": ptu_count or (sum(shares.values()) if isinstance(shares, dict) else 50), "cost_per_ptu_per_hour": 1.0, "ptu_effective_from": "2026-01-01T00:00:00Z", "ptu_shares": shares, @@ -81,6 +83,23 @@ def test_a_single_team_deployment_and_a_malformed_share_map_are_not_filtered_her assert result.withheld is False +_TERMLESS_SHARED: Final = { + "model_name": "gpt-4.1-ptu", + "litellm_params": {"model": "azure/gpt-4.1"}, + "model_info": {"id": "shared", "ptu_shares": _SHARES}, +} + + +def test_a_share_map_registration_would_refuse_reserves_nothing(): + """Registration refuses ``ptu_shares`` without the count, rate and start they split, and a map + that does not add up to the count, so a row carrying either shape holds nobody's capacity: + every team is served from it and no ceiling applies, the same as before the map was read.""" + short: Final = _shared(shares={"team-a": 30}, ptu_count=50) + for deployment in (_TERMLESS_SHARED, short): + assert filter_ptu_shared_deployments([deployment], "team-z").withheld is False + assert _unaliased_ceiling([deployment], "team-a", "gpt-4.1-ptu") is None + + def test_a_team_is_served_from_its_shared_deployment_first_and_never_from_another_teams_reservation(): """team-b holds a share on the shared deployment and the open one serves everyone, while the deployment reserved for team-a alone is withheld from it the way the router withholds it.""" From e4b33bc8a9b45835e579decbff740d19a90a14b7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 2 Oct 2026 19:17:46 -0700 Subject: [PATCH 18/18] fix(ptu): keep a declared split as the access rule and check a config split at boot with the flag off 6fb24d1d read ptu_shares through the pricing terms so a termless map reserved nothing, but with the flag off the router skips ptu_config_error for config.yaml entries, so such an entry loaded and served every team, against the rule that a declared split restricts access whether or not cost attribution is on. _deployment_shares reads the declared map again, and the router now raises ptu_config_error for a config-sourced entry declaring ptu_shares regardless of the flag, so a malformed split is refused at boot instead of being honoured or ignored silently. --- litellm/router.py | 3 ++- litellm/router_utils/ptu_shares.py | 14 ++++++-------- tests/unit/router_utils/test_ptu_shares.py | 17 +++++++++-------- tests/unit/test_router/test_router.py | 10 ++++++++++ 4 files changed, 27 insertions(+), 17 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index d4b5d3405d2..e0270122994 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8957,7 +8957,8 @@ class Router: ptu_error: Final = ( (ptu_config_error(_model_info, model_name=_model_name) or identity_error) if config_sourced else None ) - if ptu_error is not None and is_ptu_cost_attribution_enabled(): + declares_split: Final = _model_info.get("ptu_shares") is not None + if ptu_error is not None and (declares_split or is_ptu_cost_attribution_enabled()): raise ValueError(ptu_error) access_windows_error: Final = access_windows_config_error(_model_info, model_name=_model_name) if access_windows_error is not None: diff --git a/litellm/router_utils/ptu_shares.py b/litellm/router_utils/ptu_shares.py index 858e5260b54..938f3483a4e 100644 --- a/litellm/router_utils/ptu_shares.py +++ b/litellm/router_utils/ptu_shares.py @@ -10,7 +10,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import Final, Generic, TypeVar -from litellm.litellm_core_utils.ptu_pricing import is_model_info_mapping, ptu_terms +from litellm.litellm_core_utils.ptu_pricing import is_model_info_mapping, parsed_ptu_shares, ptu_terms from litellm.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity, is_azure_deployment from litellm.router_utils.common_utils import team_may_use_deployment @@ -36,17 +36,15 @@ class PTUShareFilterResult(Generic[_DeploymentT]): def _deployment_shares(deployment: Mapping[str, object]) -> Mapping[str, int] | None: - """The teams a deployment's capacity is split across, else None. + """The teams a deployment is declared split across, else None. - Only a map registration would accept counts, read through the same terms flat cost accrues - under, so a row that reached the table without its count, rate and start reserves nothing - instead of refusing every other team while charging nobody. + The declared map is the access rule on its own; the pricing terms are read separately, so a + split whose terms are missing still serves only the teams it names rather than everyone. """ model_info: Final = deployment.get("model_info") - if not is_model_info_mapping(model_info) or model_info.get("ptu_shares") is None: + if not is_model_info_mapping(model_info): return None - terms: Final = ptu_terms(model_info) - return None if terms is None else terms.shares + return parsed_ptu_shares(model_info.get("ptu_shares")) def filter_ptu_shared_deployments( diff --git a/tests/unit/router_utils/test_ptu_shares.py b/tests/unit/router_utils/test_ptu_shares.py index ae4e26bac74..dfd41ac3050 100644 --- a/tests/unit/router_utils/test_ptu_shares.py +++ b/tests/unit/router_utils/test_ptu_shares.py @@ -90,14 +90,15 @@ _TERMLESS_SHARED: Final = { } -def test_a_share_map_registration_would_refuse_reserves_nothing(): - """Registration refuses ``ptu_shares`` without the count, rate and start they split, and a map - that does not add up to the count, so a row carrying either shape holds nobody's capacity: - every team is served from it and no ceiling applies, the same as before the map was read.""" - short: Final = _shared(shares={"team-a": 30}, ptu_count=50) - for deployment in (_TERMLESS_SHARED, short): - assert filter_ptu_shared_deployments([deployment], "team-z").withheld is False - assert _unaliased_ceiling([deployment], "team-a", "gpt-4.1-ptu") is None +def test_a_declared_split_is_the_access_rule_even_without_pricing_terms(): + """A map that reached the router without the count, rate and start registration now demands + still names who the deployment is for, so it is served to those teams alone and sized from + their shares; only the flat cost, which needs the terms, goes unattributed.""" + assert filter_ptu_shared_deployments([_TERMLESS_SHARED], "team-z").withheld is True + assert filter_ptu_shared_deployments([_TERMLESS_SHARED], "team-a").withheld is False + ceiling: Final = _unaliased_ceiling([_TERMLESS_SHARED], "team-a", "gpt-4.1-ptu") + assert ceiling is not None + assert ceiling.tpm_limit == 30 * _GPT41.input_tpm_per_ptu def test_a_team_is_served_from_its_shared_deployment_first_and_never_from_another_teams_reservation(): diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index b61aea30b67..324c4585188 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -18845,6 +18845,16 @@ def test_a_shared_ptu_deployment_whose_shares_do_not_add_up_is_refused_at_regist Router(model_list=model_list) +def test_a_config_entry_declaring_shares_without_terms_is_refused_even_while_the_feature_is_off(monkeypatch): + """The split is enforced with the flag off, so its shape is checked with the flag off too: a + malformed one is refused at boot instead of being honoured or ignored without a word.""" + monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False) + model_list = _shared_ptu_model_list()[:1] + model_list[0]["model_info"] = {"id": "shared-deployment", "ptu_shares": {"team-a": 30, "team-b": 20}} + with pytest.raises(ValueError, match=r"ptu_count and cost_per_ptu_per_hour are required when ptu_shares is set"): + Router(model_list=model_list) + + @pytest.mark.asyncio @pytest.mark.parametrize( "verdict",