From 74d76fd758fbba675f1a3c5370dcff49b42dc20a Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Wed, 11 Nov 2026 16:58:52 -0800 Subject: [PATCH] perf(router): list a group's deployments once per registry change on the acompletion path (#50749) * perf(router): list a group's deployments once per registry change on the acompletion path * perf(router): drop the new docstrings * perf(router): keep the counter cache's mutable-ok on its annotation * perf(router): define the counter cache's alias in the registry --- litellm/router_rust/core.py | 9 ++++- litellm/router_rust/invoke/endpoints.py | 2 ++ .../router_rust/invoke/prompt_management.py | 6 ++-- litellm/router_rust/surface/model_info.py | 25 ++++++++----- litellm/router_rust/surface/registry.py | 35 ++++++++++++++++++- ...st_registry_characterization_model_info.py | 15 ++++++++ 6 files changed, 79 insertions(+), 13 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 69724f0b9fe..3d8627cd2b7 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -1351,6 +1351,9 @@ class _EndpointPort: def group_deployments(self, model_name: str) -> Sequence[Mapping[str, object]] | None: return self._core._hosted._core_get_model_list(model_name=model_name) # pyright: ignore[reportPrivateUsage] # see the class docstring + def group_members(self, model_name: str) -> Sequence[tuple[object, object]] | None: + return self._core._hosted._group_members(model_name) # pyright: ignore[reportPrivateUsage] # see the class docstring + def routed_name(self, model: object) -> str: return self._core.routed_name(model) @@ -1683,7 +1686,11 @@ class RustCore( try: selector: Final = self._prefetch_usage_selector(model, request) - deployments: Final = self._hosted._core_get_model_list(model_name=model) or [] # pyright: ignore[reportPrivateUsage] # see __init__ + deployments: Final = ( + self._hosted._core_get_model_list(model_name=model) or [] # pyright: ignore[reportPrivateUsage] # see __init__ + if isinstance(selector, LowestTPMLoggingHandler_v2) + else [] + ) except Exception: # noqa: BLE001 # the batched read is an optimisation: the reads stay apart yield return diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index c170aaab6e6..5c1205c70ca 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -122,6 +122,8 @@ class EndpointCore(Protocol): """``_core_get_model_list(model_name=...)`` on the host registry.""" ... + def group_members(self, model_name: str) -> Sequence[tuple[object, object]] | None: ... + def routed_name(self, model: object) -> str: """``RustCore.routed_name``: what a ``model`` that is no string raises, or the name it routes under.""" ... diff --git a/litellm/router_rust/invoke/prompt_management.py b/litellm/router_rust/invoke/prompt_management.py index 62bd77e1b48..74a7991b06c 100644 --- a/litellm/router_rust/invoke/prompt_management.py +++ b/litellm/router_rust/invoke/prompt_management.py @@ -21,10 +21,10 @@ _PROMPT_MANAGEMENT_PARAMS: Final = frozenset( def is_prompt_management_model(core: EndpointCore, model: str) -> bool: """`_is_prompt_management_model`: one deployment, whose model's provider prefix is a known custom-logger-compatible callback.""" - model_list: Final = core.group_deployments(model_name=model) - if model_list is None or len(model_list) != 1: + members: Final = core.group_members(model) + if members is None or len(members) != 1: return False - litellm_model: Final = as_dict(model_list[0].get("litellm_params")).get("model") + litellm_model: Final = members[0][1] if not isinstance(litellm_model, str) or "/" not in litellm_model: return False known: Final = core.host.dynamic().known_custom_logger_compatible_callbacks diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py index 508d2bb1145..dd67e252370 100644 --- a/litellm/router_rust/surface/model_info.py +++ b/litellm/router_rust/surface/model_info.py @@ -35,6 +35,7 @@ from typing_extensions import overload import litellm from litellm._internal_context import service_target from litellm._logging import verbose_router_logger +from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.add_retry_fallback_headers import ( apply_quality_router_decision_headers, apply_remaining_usage_headers, @@ -74,6 +75,7 @@ from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.utils import get_utc_datetime, is_generalized_model_info if TYPE_CHECKING: + from litellm.router_rust.surface.registry import GroupCounters from litellm.rust_bridge._native import ( RouterCall, RouterCoreHandle, @@ -255,9 +257,12 @@ class HostModelInfo: """Mixed into ``HostRegistry``, which holds the model list, the indexes and the host.""" _core_usage_store: RouterUsage + _group_counters_cache: dict[str, "GroupCounters"] # mutable-ok: by model group def _live_handle(self) -> "RouterCoreHandle | None": ... # the registry's: the handle ``RustCore`` set + def _group_members(self, model_group: str) -> tuple[tuple[object, object], ...] | None: ... # the registry's + def _core_handle(self) -> "RouterCoreHandle": handle: Final = self._live_handle() if handle is None: @@ -798,17 +803,21 @@ class HostModelInfo: """This minute's tpm and rpm counters of every deployment in ``model_group`` (can be a wildcard route), under the model sent to ``litellm.completion()``, as the usage writes name them.""" current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock - model_list: Final = self._core_get_model_list(model_name=model_group) - if model_list is None: # no matching deployments + members: Final = self._group_members(model_group) + if members is None: # no matching deployments return None - named: Final = tuple( - (model.get("model_info", {}).get("id"), model["litellm_params"].get("model")) for model in model_list - ) - deployments: Final = tuple(pair for pair in named if pair[0] is not None and pair[1] is not None) - return tuple( - tuple(UsageCounter(id, litellm_model, current_minute, kind) for id, litellm_model in deployments) + cached: Final = self._group_counters_cache.get(model_group) + if cached is not None and cached[0] is members and cached[1] == current_minute: + return cached[2] + deployments: Final = tuple(pair for pair in members if pair[0] is not None and pair[1] is not None) + counters: Final = tuple( + tuple(UsageCounter(id, litellm_model, current_minute, kind) for id, litellm_model in deployments) # pyright: ignore[reportArgumentType] # what the listing names, as the Python core passed it for kind in ("tpm", "rpm") ) + if len(self._group_counters_cache) >= DEFAULT_MAX_LRU_CACHE_SIZE: + self._group_counters_cache.clear() + self._group_counters_cache[model_group] = (members, current_minute, counters) + return counters async def _core_get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]: """Current (tpm, rpm) usage of ``model_group``.""" diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index 59d3d7389be..18bc493b878 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -96,7 +96,7 @@ from litellm.litellm_core_utils.routing_groups import apply_routing_group_priori from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.redis.stores.connection import connect -from litellm.redis.stores.router_usage import RouterUsage, RouterUsageStore +from litellm.redis.stores.router_usage import RouterUsage, RouterUsageStore, UsageCounter from litellm.responses.bridge_check import keep_registered_mode from litellm.router_api import RouterRegistry, deployment_id from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET @@ -697,6 +697,21 @@ def _index_key(value: object) -> str | None: return value if isinstance(value, str) else None +GroupMembers: TypeAlias = tuple[tuple[object, object], ...] +GroupCounters: TypeAlias = tuple[object, str, tuple[tuple[UsageCounter, ...], ...]] + + +def _members_of(model_list: Sequence[Mapping[str, object]] | None) -> GroupMembers | None: + if model_list is None: + return None + return tuple((_field(entry, "model_info", "id"), _field(entry, "litellm_params", "model")) for entry in model_list) + + +def _field(entry: Mapping[str, object], section: str, key: str) -> object: + held: Final = entry.get(section) + return held.get(key) if isinstance(held, Mapping) else None + + def _round_trips(value: object) -> bool: """Whether JSON reads ``value`` back equal to itself.""" try: @@ -787,6 +802,8 @@ class HostRegistry(HostModelInfo): # `zero_cost()`: the per-deployment zero-cost lookup the auth checks fill, emptied with the # group info caches. self._zero_cost_cache: dict[str, bool] = {} # mutable-ok: the auth checks write into it + self._group_members_cache: dict[str, tuple[object, GroupMembers | None]] = {} # mutable-ok: by model group + self._group_counters_cache = {} # The routing cache usage reads come from; RustCore sets it to the one its Rust core writes. self._core_cache: DualCache = DualCache() # The tpm/rpm store usage reads come from; RustCore sets it to the one its Rust core writes. @@ -939,6 +956,7 @@ class HostRegistry(HostModelInfo): self._core_cached_deployment_model_info.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache self.cached_deployment_tracks_usage.cache_clear() # pyright: ignore[reportFunctionMemberAccess] # an lru_cache self._zero_cost_cache.clear() + self._group_members_cache.clear() self._routing_group_rows = None def set_routing_groups(self, groups: Sequence[RoutingGroup]) -> None: @@ -1887,6 +1905,21 @@ class HostRegistry(HostModelInfo): return sum(len(held(member)) for member in group.models) > 1 + def _group_members(self, model_group: str) -> GroupMembers | None: + # An alias or a wildcard name reads state the registry version does not cover: those list every time. + handle: Final = self._live_handle() + if handle is None or _index_key(model_group) is None or model_group in (self._core_model_group_alias or {}): + return _members_of(self._core_get_model_list(model_name=model_group)) + stamp: Final = (handle, handle.registry_version(), self._registry_version) + cached: Final = self._group_members_cache.get(model_group) + if cached is not None and cached[0] == stamp: + return cached[1] + members: Final = _members_of(self._core_get_model_list(model_name=model_group)) + if len(self._group_members_cache) >= DEFAULT_MAX_LRU_CACHE_SIZE: + self._group_members_cache.clear() + self._group_members_cache[model_group] = (stamp, members) + return members + def _core_get_model_list_from_model_alias( self, model_name: str | None = None ) -> list[UntypedEntry]: # mutable-ok: a fresh list diff --git a/tests/unit/test_router/core/test_registry_characterization_model_info.py b/tests/unit/test_router/core/test_registry_characterization_model_info.py index 578183d7f02..b613aed041d 100644 --- a/tests/unit/test_router/core/test_registry_characterization_model_info.py +++ b/tests/unit/test_router/core/test_registry_characterization_model_info.py @@ -8,6 +8,7 @@ import pytest import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap +from litellm.router_api import ChangeSet from litellm.router_core import RouterCore from litellm.types.router import Deployment, RouterCacheEnum from litellm.utils import get_utc_datetime @@ -481,6 +482,20 @@ async def test_group_usage_sums_each_deployments_counters_for_this_minute(make_c assert await core.health().usage("al") == (34, 3) +@pytest.mark.asyncio +async def test_group_usage_follows_deployments_added_and_removed_after_a_read(make_core: MakeCore) -> None: + core = _usage_core(make_core) + await _serve(core, "a", "openai/gpt-4o", 30) + assert await core.health().usage("gpt") == (30, 1) + + core.apply(ChangeSet(upsert=[typed("gpt", "c", model="openai/gpt-4o-mini", tpm=100, rpm=10)])) + await _serve(core, "c", "openai/gpt-4o-mini", 5) + assert await core.health().usage("gpt") == (35, 2) + + core.apply(ChangeSet(delete=["a"])) + assert await core.health().usage("gpt") == (5, 1) + + @pytest.mark.asyncio async def test_io_token_usage_sums_input_and_output_counters(make_core: MakeCore) -> None: core = _usage_core(make_core)