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

@@ -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

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."""
         ...

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

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``."""

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

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)