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.
This commit is contained in:
mateo-berri 2026-10-02 19:03:41 -07:00
parent 495cf74ccd
commit 6fb24d1dad
2 changed files with 31 additions and 5 deletions

View file

@ -10,7 +10,7 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass from dataclasses import dataclass
from typing import Final, Generic, TypeVar 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.llms.azure.ptu_capacity import PTUCapacity, deployment_ptu_capacity, is_azure_deployment
from litellm.router_utils.common_utils import team_may_use_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: 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") 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 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( def filter_ptu_shared_deployments(

View file

@ -19,13 +19,15 @@ _GPT6SOL: Final = AZURE_PTU_CAPACITY["gpt-6-sol"]
_SHARES: Final = {"team-a": 30, "team-b": 20} _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 { return {
"model_name": "gpt-4.1-ptu", "model_name": "gpt-4.1-ptu",
"litellm_params": {"model": model}, "litellm_params": {"model": model},
"model_info": { "model_info": {
"id": deployment_id, "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, "cost_per_ptu_per_hour": 1.0,
"ptu_effective_from": "2026-01-01T00:00:00Z", "ptu_effective_from": "2026-01-01T00:00:00Z",
"ptu_shares": shares, "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 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(): 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 """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.""" deployment reserved for team-a alone is withheld from it the way the router withholds it."""