From a1afc2f433edf0d4961512713cec94761d112c54 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 19 Aug 2026 13:51:31 -0700 Subject: [PATCH] refactor(ptu): give the rollup a source-agnostic deployment record (#37501) The flat-cost rollup reads deployments only from LiteLLM_ProxyModelTable, so a PTU deployment declared in config.yaml never accrues flat cost. Those deployments live in llm_router.model_list as plain dicts whose id sits in model_info rather than on the entry, so they do not satisfy the shape _parse_ptu_model reads. Adds a frozen record in that shape and a factory that maps a router entry onto it, leaving _parse_ptu_model byte-identical so the existing cases stand as evidence of no behaviour change. Nothing calls the factory yet; the caller lands with the loader union. _decode_model_info also stops handing back valid JSON that is not an object. It decoded a list or a scalar and returned it as a mapping, so the caller read fields off it and raised, losing the whole run rather than the one bad deployment. --- .../spend_tracking/ptu_flat_cost_rollup.py | 46 ++++- .../test_ptu_flat_cost_rollup.py | 177 ++++++++++++++++++ 2 files changed, 220 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index efdbda47fdc..25de9f6d065 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -17,6 +17,7 @@ import json from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from datetime import date, datetime, time, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_proxy_logger @@ -110,17 +111,56 @@ def _public_model_name(row: object, model_info: Mapping[str, object]) -> str: def _decode_model_info(raw: object) -> "Mapping[str, object] | None": - """A deployment's model_info as a dict, decoding a JSON string, else None.""" + """A deployment's model_info as a mapping, decoding a JSON string, else None. + + Valid JSON that is not an object decodes to a list or a scalar, which every caller + would then read fields off, so it is rejected here rather than raised past them. + """ if isinstance(raw, str): try: - return json.loads(raw) + decoded: Final = json.loads(raw) except (TypeError, ValueError): return None - if isinstance(raw, dict): + return decoded if isinstance(decoded, dict) else None + if isinstance(raw, Mapping): return raw return None +@dataclass(frozen=True, slots=True) +class _PTUDeployment: + """A deployment in the shape ``_parse_ptu_model`` 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. + """ + + model_id: str + model_name: str + model_info: Mapping[str, object] + + +def _router_deployment(deployment: Mapping[str, object]) -> _PTUDeployment | None: + """A router ``model_list`` entry in the shape the parser reads, else None. + + An id is required rather than defaulted because it keys the sentinel row: every + deployment without one would collapse onto a single row per team and only the last + would be billed. The mapping is copied because the router rewrites entries in place + while the rollup runs. + """ + model_info: Final = _decode_model_info(deployment.get("model_info")) + if model_info is None: + return None + model_id: Final = model_info.get("id") + if not isinstance(model_id, str) or not model_id: + return None + return _PTUDeployment( + model_id=model_id, + model_name=str(deployment.get("model_name") or ""), + model_info=MappingProxyType(dict(model_info)), + ) + + def _parse_ptu_model(row: object) -> PTUModel | None: """Return a PTUModel when the deployment carries valid manual PTU config, else None. 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 736fc13d137..f41756d2b87 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 @@ -1,5 +1,6 @@ """Tests for the per-model PTU flat-cost daily rollup.""" +import json import types from datetime import date, datetime, timedelta, timezone from unittest.mock import AsyncMock, MagicMock @@ -373,6 +374,182 @@ def test_parse_ptu_model_rejects_an_inverted_window(): assert parsed is None +def _router_entry(model_id="dep-1", model_name="gpt-4o-ptu", model_info=None, with_start=True): + """A deployment as the router stores one: a plain dict whose id lives in model_info.""" + info = dict(model_info or {}) + if ( + with_start + and info.get("ptu_count") is not None + and info.get("cost_per_ptu_per_hour") is not None + and "ptu_effective_from" not in info + ): + info["ptu_effective_from"] = _DEFAULT_PTU_START + if model_id is not None: + info["id"] = model_id + return { + "model_name": model_name, + "litellm_params": {"model": "azure/gpt-4o"}, + "model_info": info, + } + + +# A team-scoped deployment is stored under a synthetic routing key with the operator's name +# in team_public_model_name, while the same deployment in config.yaml carries the operator's +# name directly. Both must resolve to the same PTUModel or the two sources bill differently. +_PARITY_CASES = ( + ( + "an iso string start against the datetime pydantic coerces it to", + {"ptu_effective_from": "2026-07-30T23:00:00Z"}, + {"ptu_effective_from": datetime(2026, 7, 30, 23, 0)}, + datetime(2026, 7, 30, 23, 0, tzinfo=timezone.utc), + None, + ), + ( + "an open window, stored as null and dropped by exclude_none", + {"ptu_effective_from": "2026-07-01T00:00:00Z", "ptu_effective_to": None}, + {"ptu_effective_from": datetime(2026, 7, 1, 0, 0)}, + datetime(2026, 7, 1, tzinfo=timezone.utc), + None, + ), + ( + "a closed window", + {"ptu_effective_from": "2026-07-01T00:00:00Z", "ptu_effective_to": "2026-07-31T00:00:00Z"}, + { + "ptu_effective_from": datetime(2026, 7, 1, 0, 0), + "ptu_effective_to": datetime(2026, 7, 31, 0, 0), + }, + datetime(2026, 7, 1, tzinfo=timezone.utc), + datetime(2026, 7, 31, tzinfo=timezone.utc), + ), +) + + +@pytest.mark.parametrize( + "db_window, router_window, expected_from, expected_to", + [case[1:] for case in _PARITY_CASES], + ids=[case[0] for case in _PARITY_CASES], +) +def test_a_deployment_parses_identically_from_the_db_and_from_config_yaml( + db_window, router_window, expected_from, expected_to +): + # The contract the router union is built on. If these ever diverge, a PTU deployment + # declared in config.yaml is billed differently from the identical one in the database. + from_db = _parse_ptu_model( + _model_row( + model_id="dep-1", + model_name="model_name_t_9f3c2b", + model_info={**_VALID_PTU, **db_window, "team_public_model_name": "gpt-4o-ptu"}, + with_start=False, + ) + ) + from_router = _parse_ptu_model( + ptu_rollup._router_deployment( + _router_entry(model_id="dep-1", model_info={**_VALID_PTU, **router_window}, with_start=False) + ) + ) + expected = PTUModel( + model_id="dep-1", + model_name="gpt-4o-ptu", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + effective_from=expected_from, + effective_to=expected_to, + ) + assert from_db == from_router == expected + + +@pytest.mark.parametrize("model_id", [None, "", 12345], ids=["absent", "blank", "not a string"]) +def test_a_router_deployment_without_a_usable_id_is_dropped(model_id): + # model_id keys the sentinel row, so an unusable one would file every such deployment + # in a team onto one row and bill for a single reservation. + entry = _router_entry(model_id=None, model_info=dict(_VALID_PTU)) + if model_id is not None: + entry["model_info"]["id"] = model_id + assert ptu_rollup._router_deployment(entry) is None + + +def test_a_router_deployment_keeps_the_name_the_operator_wrote(): + priced = _parse_ptu_model( + ptu_rollup._router_deployment(_router_entry(model_name="gpt-4o-ptu", model_info=dict(_VALID_PTU))) + ) + assert priced is not None + assert priced.model_name == "gpt-4o-ptu" + + +def test_a_team_alias_on_a_router_deployment_still_wins_over_the_routing_name(): + # config.yaml can carry team_public_model_name to expose a team-facing alias, and the + # charge has to file under the name that team calls rather than the routing one. + priced = _parse_ptu_model( + ptu_rollup._router_deployment( + _router_entry( + model_name="routing-name", + model_info={**_VALID_PTU, "team_public_model_name": "public-alias"}, + ) + ) + ) + assert priced is not None + assert priced.model_name == "public-alias" + + +def test_a_router_deployment_with_a_stringified_model_info_still_decodes(): + entry = _router_entry(model_info=dict(_VALID_PTU)) + entry["model_info"] = json.dumps(entry["model_info"]) + parsed = _parse_ptu_model(ptu_rollup._router_deployment(entry)) + assert parsed is not None + assert parsed.model_id == "dep-1" + + +@pytest.mark.parametrize( + "model_info", + [None, "not json", 42, "[1, 2, 3]", '"a string"', "42", "true", "null"], + ids=[ + "absent", + "unparseable", + "not a string or dict", + "json array", + "json string", + "json number", + "json bool", + "json null", + ], +) +def test_a_router_deployment_without_usable_model_info_is_dropped(model_info): + # Valid JSON that is not an object decodes to a list or a scalar, and reading fields + # off one raises rather than dropping the single bad deployment the rollup expects + assert ptu_rollup._router_deployment({"model_name": "x", "model_info": model_info}) is None + + +@pytest.mark.parametrize( + "model_info", + ["[1, 2, 3]", '"a string"', "42", "true", "null"], + ids=["json array", "json string", "json number", "json bool", "json null"], +) +def test_a_db_row_holding_non_object_json_is_dropped(model_info): + assert _parse_ptu_model(_model_row(model_info=model_info, with_start=False)) is None + + +def test_a_router_deployment_does_not_alias_the_routers_own_model_info(): + # The rollup runs on a cron while requests are in flight, and the router rewrites + # model_info in place, so a held record must neither observe nor cause those writes. + live = {"id": "dep-1", **_VALID_PTU} + record = ptu_rollup._router_deployment({"model_name": "x", "model_info": live}) + assert record is not None + + live["ptu_count"] = 999 + assert record.model_info["ptu_count"] == _VALID_PTU["ptu_count"] + + with pytest.raises(TypeError): + record.model_info["ptu_count"] = 1 + + +def test_a_raw_router_dict_is_not_a_deployment_record(): + # The parser reads attributes, so a router dict passed straight to it returns None + # instead of raising. Skipping the factory would silently drop every config.yaml + # deployment while leaving the suite green. + assert _parse_ptu_model(_router_entry(model_info=dict(_VALID_PTU))) is None + + @pytest.mark.asyncio async def test_rollup_returns_empty_when_prisma_client_is_none(): result = await run_ptu_flat_cost_rollup(None, target_date=DAY)