From 60ba1c5da11ba3b1a7cbbd89fb672b42edd30d66 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 03:23:08 -0800 Subject: [PATCH] fix(router): rebuild shared rows oldest router first and keep fields an SDK user added to an unmapped model's row --- .../litellm_core_utils/get_model_cost_map.py | 3 -- litellm/router_python_core.py | 19 +++++++---- litellm/utils.py | 8 +---- .../unit/test_router_model_cost_isolation.py | 32 +++++++++++++++++++ 4 files changed, 46 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 826726d8102..1fc76728e74 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -562,7 +562,6 @@ def get_model_cost_map_loaded_at() -> "datetime | None": class AliasedCostEntry(dict[str, object]): __slots__ = ( - "created_by_registration", "displaced_values", "latest_unsaved_write_fields", "names", @@ -578,7 +577,6 @@ class AliasedCostEntry(dict[str, object]): entry: Mapping[str, object], names: tuple[str, ...], registered_fields: frozenset[str] = frozenset(), - created_by_registration: bool = False, registered_by: frozenset[str] = frozenset(), written_fields: frozenset[str] = frozenset(), displaced_values: tuple[tuple[str, object], ...] = (), @@ -589,7 +587,6 @@ class AliasedCostEntry(dict[str, object]): super().__init__(entry) self.names: tuple[str, ...] = names self.registered_fields: frozenset[str] = registered_fields - self.created_by_registration: bool = created_by_registration self.registered_by: frozenset[str] = registered_by self.written_fields: frozenset[str] = written_fields self.displaced_values: tuple[tuple[str, object], ...] = displaced_values diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index ff41b96c7a1..a1849fc5471 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, groupby +from itertools import 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 @@ -1301,13 +1301,19 @@ class RoutingArgs(enum.Enum): # as the per-request one built from a caller-supplied user_config, drops out on its # own rather than leaving entries behind that nothing can withdraw. _live_routers: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet() +_live_router_creation_ticks: Final["weakref.WeakKeyDictionary[PythonCore, int]"] = weakref.WeakKeyDictionary() +_next_router_creation_tick: Final = count() _routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet() _routers_holding_global_callbacks_lock: Final = threading.Lock() +def _live_routers_oldest_first() -> tuple["PythonCore", ...]: + return tuple(sorted(_live_routers, key=_live_router_creation_ticks.__getitem__)) + + def _replay_live_router_model_cost() -> None: - """Re-assert every live router's deployments after the cost map is refreshed.""" - for router in tuple(_live_routers): + """Re-assert every live router's deployments, oldest router first, after the cost map is refreshed.""" + for router in _live_routers_oldest_first(): router.replay_model_cost_registrations() @@ -1939,6 +1945,7 @@ class PythonCore(DeprecatedRouterMembers): # Track this router so a price data reload can rebuild its deployments' # cost-map entries from the list it is serving at that moment, whether # those deployments came from model_list or were added later from the DB. + _live_router_creation_ticks[self] = next(_next_router_creation_tick) _live_routers.add(self) if allowed_fails is not None: @@ -12511,8 +12518,8 @@ class PythonCore(DeprecatedRouterMembers): def _rebuild_cost_rows_written_by(self, written_keys: frozenset[str]) -> None: """Rebuild the rows ``register_model`` writes of ``written_keys`` landed on: the loaded catalog row, then the - ``register_model`` registrations replayed on every reload, then every live deployment that wrote it in - ``model_list`` order, then again the fields of those registrations no deployment wrote over since, then the + ``register_model`` registrations replayed on every reload, then every live deployment that wrote it, oldest router + first and each router's in ``model_list`` order, then again the fields of those registrations no deployment wrote over since, then the fields set in place after their last write, so each field keeps the order its writes came in, whether the registration came before the router or after it. The replay writes only those rows, so a row it leaves alone keeps what later deployments and registrations put there, and deployments with the same write replay it once, @@ -12522,7 +12529,7 @@ class PythonCore(DeprecatedRouterMembers): return writers: Final = model_cost_row_writers(stale) targets: Final = stale | writers - routers: Final = tuple(_live_routers) + routers: Final = _live_routers_oldest_first() entries: Final = tuple(chain.from_iterable(router._model_list_entries() for router in routers)) writes: Final = (PythonCore._entry_shared_cost_write(entry) for entry in entries) landing: Final = tuple( diff --git a/litellm/utils.py b/litellm/utils.py index 0bfc3e7bb74..7fa9cda1992 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3889,7 +3889,6 @@ def _replace_cost_map_entry( **dict(previous.displaced_values), }.items() ) - created_by_registration: Final = previous.created_by_registration if current is previous else current is None written_now: Final = {name: value for name, value in entry.items() if name in changed_fields or name in set_fields} cost_map.update( dict.fromkeys( @@ -3898,7 +3897,6 @@ def _replace_cost_map_entry( merged, previous.names, registered_fields, - created_by_registration, registered_by=previous.registered_by | {registered_key}, written_fields=previous.written_fields | registered_fields | changed_fields, displaced_values=displaced_values, @@ -4218,7 +4216,7 @@ def restore_model_cost_rows_to_catalog(rows: Iterable[str]) -> None: """Put each of ``rows`` back to its loaded catalog entry, keeping every alias on the one shared entry and every field written into it outside ``register_model``, such as an SDK user's own price, at the value it had before a registration first overwrote it, so registrations replayed onto them start where a fresh boot does. A row the - catalog lacks goes once no such field is left or a registration created it.""" + catalog lacks goes once no such field is left.""" cost_map: Final = _builtin_cost_map() catalog: Final = GetModelCostMap.loaded_model_cost_map() entries: Final = ((row, cost_map.get(row)) for row in rows) @@ -4236,8 +4234,6 @@ def restore_model_cost_rows_to_catalog(rows: Iterable[str]) -> None: def _fields_outside_registrations(entry: AliasedCostEntry) -> Mapping[str, object]: - if entry.created_by_registration: - return {} unwritten: Final = { name: value for name, value in entry.items() if name not in entry.written_fields and value is not None } @@ -4274,7 +4270,6 @@ def _without_fields_left_by_earlier_writes( }, current.names, registered_fields=current.registered_fields - dropped, - created_by_registration=current.created_by_registration, registered_by=current.registered_by, written_fields=current.written_fields - dropped, displaced_values=tuple((name, value) for name, value in current.displaced_values if name not in dropped), @@ -4319,7 +4314,6 @@ def _with_fields_set_in_place(entry: AliasedCostEntry, values: Mapping[str, obje {**entry, **values}, entry.names, registered_fields=entry.registered_fields, - created_by_registration=entry.created_by_registration, registered_by=entry.registered_by, written_fields=entry.written_fields, displaced_values=entry.displaced_values, diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 3b617e0af96..15fbb7afc4a 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -2700,6 +2700,38 @@ def test_an_unmapped_models_shared_row_lasts_exactly_as_long_as_a_deployment_ser litellm.get_model_info(backend_model) +def test_deleting_a_deployment_keeps_a_limit_an_sdk_user_added_to_an_unmapped_models_row() -> None: + backend_model: Final = "hosted_vllm/withdraw-test-unmapped-in-place-model" + sdk_limit: Final = 987_656 + router: Final = Router( + model_list=[ + _cache_point_deployment("kept", model=backend_model), + _cache_point_deployment("doomed", model=backend_model), + ] + ) + litellm.model_cost[backend_model]["max_input_tokens"] = sdk_limit + + assert router.delete_deployment(id="doomed") is not None + + assert litellm.get_model_info(backend_model).get("max_input_tokens") == sdk_limit + + +def test_deleting_a_deployment_keeps_the_limit_the_newest_router_wrote_to_the_shared_row() -> None: + row: Final = _shared_cache_point_row() + oldest_limit: Final = _unseen_max_output_tokens(row) + routers: Final = tuple( + Router(model_list=[_cache_point_deployment(f"router-{index}", max_output_tokens=oldest_limit + index)]) + for index in range(16) + ) + newest_limit: Final = oldest_limit + len(routers) - 1 + assert routers[0].upsert_deployment(Deployment(**_cache_point_deployment("plain"))) is not None + assert litellm.model_cost[row].get("max_output_tokens") == newest_limit + + assert routers[0].delete_deployment(id="plain") is not None + + assert litellm.model_cost[row].get("max_output_tokens") == newest_limit + + def test_deleting_a_flagged_deployment_keeps_the_flag_another_router_sets() -> None: row: Final = _shared_cache_point_row() flag: Final = not bedrock_model_accepts_cache_points(_CACHE_POINT_MODEL)