From 9e16f442860ef0b8c4cdee684a322468d44bb286 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 12:54:49 -0800 Subject: [PATCH] fix(router): replay only the last sibling write per set of cost-row fields, and keep ApplyResult.deleted a list Rebuilding a shared row after a delete replayed every sibling whose write inputs differed, each with a fresh model info lookup, so a delete among 2,000 deployments with limits of their own took about 3 s. Writes that set the same fields of the rebuilt rows now replay only the last one, counting strings and dicts by value and prices as set or unset, and a set never spans a change of mode. Router.apply(ChangeSet(delete=...)).deleted is a list again, as on the base branch. --- litellm/router_api.py | 2 +- litellm/router_python_core.py | 59 +++++++++++----- .../unit/test_router_model_cost_isolation.py | 68 ++++++++++++++++--- 3 files changed, 102 insertions(+), 27 deletions(-) diff --git a/litellm/router_api.py b/litellm/router_api.py index 9d5590155c9..1de381a2897 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1328,7 +1328,7 @@ def apply_changes(router: "Router", changes: ChangeSet) -> ApplyResult: if changes.settings: router._core_update_settings(**{name: detached(value) for name, value in changes.settings.items()}) upserted: Final = [router._core_upsert_deployment(deployment=d) for d in changes.upsert] - deleted: Final = router._core_delete_deployments(changes.delete) + deleted: Final = list(router._core_delete_deployments(changes.delete)) if changes.redis_cache is not None: router._core_update_redis_cache(cache=changes.redis_cache) return ApplyResult(version=router._registry_version, upserted=upserted, deleted=deleted) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 320aa764747..bc1512c148b 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -47,7 +47,7 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial -from itertools import chain, count, groupby +from itertools import accumulate, chain, count, groupby from operator import itemgetter from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast @@ -1152,15 +1152,21 @@ class _SharedCostWrite: return self.model_id in writers or not writers.isdisjoint(self.backend_keys) @property - def may_write_mode(self) -> bool: - """Whether its write can carry ``mode``: set directly, or copied with the backend entry an off-peak schedule - inherits.""" - return isinstance(self.model_info.get("mode"), str) or bool(self.model_info.get("off_peak_pricing")) + def mode(self) -> str | None: + mode: Final = self.model_info.get("mode") + return mode if isinstance(mode, str) else None + + @property + def copies_backend_entry(self) -> bool: + """Whether its write copies the backend entry an off-peak schedule inherits, as that row stands.""" + return bool(self.model_info.get("off_peak_pricing")) def replay_key(self, targets: frozenset[str]) -> tuple[str | None, str]: - """Equal for two entries only when replaying both writes the ``targets`` rows exactly as replaying one does. - A shared row never takes a price, only what a price being set pulls in from the backend model, so prices count - only as set or unset, and an id that names a target row or wrote onto one keeps its write apart by its id.""" + """Equal for two entries whose writes set the same fields of the ``targets`` rows, so the later write replaces + every value the earlier one set. A mapping merges into the row's value, and a string can name the row or the + mode the write lands with, so those count by value and every other value only as set. A shared row never takes + a price, only what a price being set pulls in from the backend model, so prices count only as set or unset, + and an id that names a target row or wrote onto one keeps its write apart by its id.""" write_inputs: Final = ( self.model, self.custom_llm_provider, @@ -1175,11 +1181,15 @@ class _SharedCostWrite: @staticmethod def _write_inputs(fields: Mapping[str, object], inputs: frozenset[str]) -> Mapping[str, object]: return { - name: bool(value) if name in _PRICE_FIELDS else value + name: bool(value) if name in _PRICE_FIELDS else _SharedCostWrite._value_written(value) for name, value in fields.items() if name in inputs and value is not None } + @staticmethod + def _value_written(value: object) -> object: + return value if isinstance(value, (str, Mapping)) else True + MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 @@ -12594,14 +12604,31 @@ class PythonCore(DeprecatedRouterMembers): def _replayed_positions( landing: tuple[tuple[int, _SharedCostWrite], ...], targets: frozenset[str] ) -> tuple[int, ...]: - """The position of the last of each set of ``landing`` writes that leave the ``targets`` rows as one of them - does. Each field keeps the last value written to it, so only the last write of each set counts, unless a write - may carry ``mode``: what a ``mode`` write does depends on the mode the row holds, so then only back-to-back runs - collapse.""" - keyed: Final = tuple((position, write.replay_key(targets)) for position, write in landing) - if any(write.may_write_mode for _position, write in landing): + """The position of the last of each set of ``landing`` writes that set the same fields of the ``targets`` rows. + Each field keeps the last value written to it, so the last write of a set leaves every field as replaying the + whole set does. What a ``mode`` write does depends on the mode the row holds, so a set never spans a change of + mode. An off-peak schedule copies its backend entry as that row stands, so with one only back-to-back writes + with the same inputs collapse.""" + if any(write.copies_backend_entry for _position, write in landing): + keyed: Final = tuple((position, write.replay_key(targets)) for position, write in landing) return tuple(tuple(run)[-1][0] for _replay_key, run in groupby(keyed, key=itemgetter(1))) - return tuple(sorted({replay_key: position for position, replay_key in keyed}.values())) + runs: Final = PythonCore._mode_runs(tuple(write.mode for _position, write in landing)) + last_positions: Final = { + (write.replay_key(targets), run): position for (position, write), run in zip(landing, runs, strict=True) + } + return tuple(sorted(last_positions.values())) + + @staticmethod + def _mode_runs(modes: tuple[str | None, ...]) -> tuple[int, ...]: + """For each write, how many times the ``mode`` the writes carry changed up to it.""" + return tuple(run for _mode, run in accumulate(modes, PythonCore._next_mode_run, initial=(None, 0)))[1:] + + @staticmethod + def _next_mode_run(held: tuple[str | None, int], mode: str | None) -> tuple[str | None, int]: + held_mode, run = held + if mode is None or mode == held_mode: + return held + return (mode, run if held_mode is None else run + 1) def _core_delete_deployment(self, id: str) -> Deployment | None: """ diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 70ecb9af2ad..121d4bb8625 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -20,7 +20,7 @@ import sys import threading from pathlib import Path from types import MappingProxyType -from collections.abc import Callable +from collections.abc import Callable, Mapping from typing import Final, cast from unittest.mock import Mock, patch @@ -3154,6 +3154,31 @@ def test_deleting_a_deployment_keeps_the_mode_off_peak_siblings_inherit_in_model assert litellm.model_cost[backend_model].get("mode") == registered_mode +def test_deleting_a_deployment_keeps_the_mode_an_off_peak_sibling_copies_from_the_row_a_sibling_wrote() -> None: + backend_model: Final = _openai_catalog_model_with_mode("chat") + off_peak: Final = {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-7} + router: Final = Router( + model_list=[ + *( + _cache_point_deployment( + model_id, model=f"openai/{backend_model}", base_model=base_model, off_peak_pricing=off_peak + ) + for model_id, base_model in ( + ("embedding", _openai_catalog_model_with_mode("embedding")), + ("own-model", backend_model), + ) + ), + _cache_point_deployment("doomed", model=f"openai/{backend_model}", max_input_tokens=5), + ] + ) + registered_mode: Final = litellm.model_cost[backend_model].get("mode") + assert registered_mode == "embedding" + + assert router.delete_deployment(id="doomed") is not None + + assert litellm.model_cost[backend_model].get("mode") == registered_mode + + class _CostMapWriteRecordingHost(PythonHost): __slots__ = ("cost_map_writes",) @@ -3223,19 +3248,23 @@ def test_deleting_one_of_many_differently_priced_deployments_rewrites_the_shared assert litellm.model_cost["sibling-7"].get("input_cost_per_token") == pytest.approx(8e-7) -def test_deleting_one_of_many_deployments_alternating_two_flags_rewrites_the_shared_row_once_per_flag() -> None: +def test_deleting_one_of_many_deployments_rewrites_the_shared_row_once_per_set_of_fields_its_siblings_set() -> None: row: Final = _shared_cache_point_row() + first_limit: Final = _unseen_max_output_tokens(row) flag: Final = not bedrock_model_accepts_cache_points(_CACHE_POINT_MODEL) + + def sibling_fields(n: int) -> Mapping[str, object]: + if n == 1: + return {"mode": "chat"} + return {_CACHE_POINT_FLAG: flag} if n % 2 == 0 else {"max_output_tokens": first_limit + n} + host: Final = _CostMapWriteRecordingHost() router: Final = Router.from_config( build_router_config( router_arguments( model_list=[ - *( - _cache_point_deployment(f"sibling-{n}", **{_CACHE_POINT_FLAG: (flag, not flag)[n % 2]}) - for n in range(51) - ), - _cache_point_deployment("doomed", max_output_tokens=_unseen_max_output_tokens(row)), + *(_cache_point_deployment(f"sibling-{n}", **sibling_fields(n)) for n in range(51)), + _cache_point_deployment("doomed", max_input_tokens=5), ] ) ), @@ -3245,9 +3274,28 @@ def test_deleting_one_of_many_deployments_alternating_two_flags_rewrites_the_sha assert router.delete_deployment(id="doomed") is not None - assert host.cost_map_writes.call_count == 2 - assert litellm.model_cost[row].get(_CACHE_POINT_FLAG) is flag - assert litellm.model_cost[row].get("max_output_tokens") == _catalog_field(row, "max_output_tokens") + assert host.cost_map_writes.call_count == 3 + assert litellm.model_cost[row].get("max_output_tokens") == first_limit + 49 + assert litellm.model_cost[row].get("max_input_tokens") == _catalog_field(row, "max_input_tokens") + assert bedrock_model_accepts_cache_points(_CACHE_POINT_MODEL) is flag + + +def test_deleting_a_deployment_keeps_every_entry_its_siblings_merged_into_a_dict_field() -> None: + row: Final = _shared_cache_point_row() + router: Final = Router( + model_list=[ + _cache_point_deployment("us-entry", provider_specific_entry={"us": 1.0}), + _cache_point_deployment("eu-entry", provider_specific_entry={"eu": 2.0}), + _cache_point_deployment("doomed", max_input_tokens=5), + ] + ) + merged: Final = litellm.model_cost[row].get("provider_specific_entry") + assert isinstance(merged, dict) + assert {"us": 1.0, "eu": 2.0}.items() <= merged.items() + + assert router.delete_deployment(id="doomed") is not None + + assert litellm.model_cost[row].get("provider_specific_entry") == merged @pytest.mark.parametrize(