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] 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(