scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(ptu): size a team's usage rows by the deployment it is served from

Browse source
This commit is contained in:
mateo-berri 2026-10-02 15:58:35 -07:00
parent 4a5828ed7a
commit 5286e6c49a
7 changed files with 166 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/management_endpoints/daily_activity_routes.py
View file

@ -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:

39
litellm/proxy/management_endpoints/ptu_consumption.py
View file

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

4
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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

12
litellm/router_utils/ptu_shares.py
View file

@ -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."""

26
tests/unit/proxy/management_endpoints/test_daily_activity_routes.py
View file

@ -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:

72
tests/unit/proxy/management_endpoints/test_ptu_consumption.py
View file

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

33
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -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"),
(