From 2e5c5611c09a6e67ef7f83f32ab0ea09e24d8a3b Mon Sep 17 00:00:00 2001 From: Charan Kotra Date: Tue, 27 Oct 2026 10:39:25 -0400 Subject: [PATCH] fix(router): count concurrent successes atomically in usage-based routing Each success read the model group's {deployment: count} dict, added to it and wrote it back. With Redis, concurrent callbacks interleave between the read and the write and overwrite each other, so TPM/RPM were undercounted. Each deployment now has its own counter updated with a single atomic increment; the in-memory dict routing reads is filled from those totals. Fixes #44917 --- litellm/router_strategy/lowest_tpm_rpm.py | 92 ++++++++++++------- .../router_strategy/test_lowest_tpm_rpm.py | 82 +++++++++++++++-- 2 files changed, 134 insertions(+), 40 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 2d373e0c266..5b3dc99a28b 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -1,6 +1,7 @@ #### What this does #### # identifies lowest tpm deployment import traceback +from collections.abc import Mapping, Sequence from datetime import datetime from typing import Final @@ -28,6 +29,40 @@ class LowestTPMLoggingHandler(CustomLogger): self.router_cache = router_cache self.routing_args = RoutingArgs(**routing_args) + def _cached_usage(self, usage_key: str) -> Mapping[str, float]: + """The {deployment_id: count} dict this process has recorded for a model group this minute.""" + cached: Final = self.router_cache.get_cache(key=usage_key, local_only=True) + if not isinstance(cached, Mapping): + return {} + return {k: float(v) for k, v in cached.items() if isinstance(k, str) and isinstance(v, (int, float))} + + def _record_usage(self, usage_key: str, deployment_id: str, total: float | None) -> None: + """Copy a deployment's running total into the in-memory {deployment_id: count} dict for its + model group, which routing and the router's pre-call RPM check read. Nothing awaits between + the read and the write, so concurrent callbacks in this process can't overwrite each other.""" + if total is None: + return + recorded: Final = self._cached_usage(usage_key) + usage: Final = {**recorded, deployment_id: max(recorded.get(deployment_id, 0.0), total)} + self.router_cache.set_cache(key=usage_key, value=usage, local_only=True, ttl=self.routing_args.ttl) + + def _get_usage(self, usage_key: str, deployment_ids: Sequence[str]) -> Mapping[str, float]: + """This process's recorded usage, plus the shared counter for any deployment it hasn't served + this minute (read from Redis, if there is one, so other instances' traffic counts). + + Counters are kept per deployment, f"{usage_key}:{deployment_id}", so each success is one + atomic increment (INCRBYFLOAT in Redis) instead of a read-modify-write of a shared dict, + which loses updates when callbacks run concurrently.""" + recorded: Final = self._cached_usage(usage_key) + unseen: Final = tuple(i for i in deployment_ids if i not in recorded) + if not unseen: + return recorded + counts: Final = self.router_cache.batch_get_cache(keys=[f"{usage_key}:{i}" for i in unseen]) + if not isinstance(counts, Sequence): + return recorded + shared: Final = {i: float(c) for i, c in zip(unseen, counts) if c is not None} + return {**shared, **recorded} + @with_service_target("router_usage") def log_success_event(self, kwargs, response_obj, start_time, end_time): if is_batch_retrieve_call_type(kwargs.get("call_type")): @@ -62,17 +97,20 @@ class LowestTPMLoggingHandler(CustomLogger): # Update usage # ------------ - ## TPM - request_count_dict = self.router_cache.get_cache(key=tpm_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens + if not isinstance(id, str) or not isinstance(total_tokens, int): + return - self.router_cache.set_cache(key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl) + ## TPM + tpm_total: Final = self.router_cache.increment_cache( + key=f"{tpm_key}:{id}", value=total_tokens, ttl=self.routing_args.ttl + ) + self._record_usage(tpm_key, id, tpm_total) ## RPM - request_count_dict = self.router_cache.get_cache(key=rpm_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 - - self.router_cache.set_cache(key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl) + rpm_total: Final = self.router_cache.increment_cache( + key=f"{rpm_key}:{id}", value=1, ttl=self.routing_args.ttl + ) + self._record_usage(rpm_key, id, rpm_total) ### TESTING ### if self.test_flag: @@ -123,21 +161,20 @@ class LowestTPMLoggingHandler(CustomLogger): # ------------ # update cache - ## TPM - request_count_dict = await self.router_cache.async_get_cache(key=tpm_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens + if not isinstance(id, str) or not isinstance(total_tokens, (int, float)): + return - await self.router_cache.async_set_cache( - key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl + ## TPM + tpm_total: Final = await self.router_cache.async_increment_cache( + key=f"{tpm_key}:{id}", value=total_tokens, ttl=self.routing_args.ttl ) + self._record_usage(tpm_key, id, tpm_total) ## RPM - request_count_dict = await self.router_cache.async_get_cache(key=rpm_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 - - await self.router_cache.async_set_cache( - key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl + rpm_total: Final = await self.router_cache.async_increment_cache( + key=f"{rpm_key}:{id}", value=1, ttl=self.routing_args.ttl ) + self._record_usage(rpm_key, id, rpm_total) ### TESTING ### if self.test_flag: @@ -169,8 +206,9 @@ class LowestTPMLoggingHandler(CustomLogger): tpm_key: Final = f"{model_group}:tpm:{current_minute}" rpm_key: Final = f"{model_group}:rpm:{current_minute}" - tpm_dict = self.router_cache.get_cache(key=tpm_key) - rpm_dict: Final = self.router_cache.get_cache(key=rpm_key) + deployment_ids: Final = tuple(str(d["model_info"]["id"]) for d in healthy_deployments) + tpm_dict: Final = dict(self._get_usage(tpm_key, deployment_ids)) + rpm_dict: Final = self._get_usage(rpm_key, deployment_ids) verbose_router_logger.debug("tpm_key=%s, tpm_dict: %s, rpm_dict: %s", tpm_key, tpm_dict, rpm_dict) try: @@ -183,16 +221,8 @@ class LowestTPMLoggingHandler(CustomLogger): # ---------------------- lowest_tpm = float("inf") - if tpm_dict is None: # base case - none of the deployments have been used - # initialize a tpm dict with {model_id: 0} - tpm_dict = {} - for deployment in healthy_deployments: - tpm_dict[deployment["model_info"]["id"]] = 0 - else: - for d in healthy_deployments: - ## if healthy deployment not yet used - if d["model_info"]["id"] not in tpm_dict: - tpm_dict[d["model_info"]["id"]] = 0 + for deployment_id in deployment_ids: ## if healthy deployment not yet used + _ = tpm_dict.setdefault(deployment_id, 0) all_deployments: Final = tpm_dict @@ -229,7 +259,7 @@ class LowestTPMLoggingHandler(CustomLogger): if ( item_tpm + input_tokens > _deployment_tpm - or (rpm_dict is not None and item in rpm_dict) + or (item in rpm_dict) and (rpm_dict[item] + 1 >= _deployment_rpm) ): continue diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index f1c56b9d033..4d4587712dd 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -1,12 +1,16 @@ import asyncio from datetime import datetime, timedelta from typing import Final -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock, patch +import fakeredis import pytest from litellm import Router, token_counter from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache +from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict, RouterErrors @@ -43,14 +47,11 @@ def test_usage_based_routing_v1_selects_the_lowest_recorded_tpm() -> None: } now: Final = datetime.now() cache_keys: Final = tuple( - f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" - for offset in range(60) + f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" for offset in range(60) ) for cache_key in cache_keys: - router.cache.set_cache( - key=cache_key, value=usage_by_deployment, ttl=float("inf") - ) + router.cache.set_cache(key=cache_key, value=usage_by_deployment, ttl=float("inf")) deployment: Final = router.get_available_deployment( model=MODEL_GROUP, @@ -135,9 +136,7 @@ async def test_v2_subclass_overriding_async_get_available_deployments_with_the_o ) router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache, routing_args={}) - response: Final = await router.acompletion( - model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}] - ) + response: Final = await router.acompletion(model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}]) assert response.choices[0].message.content in { f"from {HIGH_USAGE_DEPLOYMENT_ID}", @@ -209,3 +208,68 @@ def test_usage_based_routing_v1_rejects_a_deployment_that_reached_its_recorded_r router: Final = _router_with_recorded_rpm(recorded=RPM_LIMIT, enable_pre_call_checks=enable_pre_call_checks) with pytest.raises(ValueError, match=RouterErrors.no_deployments_available.value): router.completion(model=MODEL_GROUP, messages=[{"role": "user", "content": "hello"}]) + + +# --------------------------------------------------------------------------- concurrent successes, Redis (#44917) +class _RedisWithReplyLatency(fakeredis.FakeAsyncRedis): + """fakeredis replies instantly, so concurrent callers never overlap. A real Redis reads the value + when the command arrives and the reply then spends time on the network while other callers run. + Yielding after the read reproduces that window.""" + + async def get(self, *args, **kwargs): + value = await super().get(*args, **kwargs) + await asyncio.sleep(0) + return value + + +def _limited_deployment(deployment_id: str, **limits: int) -> DeploymentTypedDict: + deployment: Final = _deployment(deployment_id) + return {**deployment, "litellm_params": {**deployment["litellm_params"], **limits}} + + +@pytest.fixture +def redis_backed_handler(): + with patch("litellm._redis.get_redis_client", return_value=MagicMock()): + redis_cache: Final = RedisCache(host="localhost", port=6379) + with patch.object(redis_cache, "init_async_client", return_value=_RedisWithReplyLatency()): + yield LowestTPMLoggingHandler(router_cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)) + + +async def _record_successes( + handler: LowestTPMLoggingHandler, deployment_id: str, total_tokens: int, times: int = 1 +) -> None: + kwargs: Final = {"litellm_params": {"metadata": {"model_group": MODEL_GROUP}, "model_info": {"id": deployment_id}}} + response: Final = {"usage": {"total_tokens": total_tokens}} + await asyncio.gather(*(handler.async_log_success_event(kwargs, response, None, None) for _ in range(times))) + + +@pytest.mark.asyncio +async def test_usage_based_routing_v1_counts_every_concurrent_success_toward_tpm(redis_backed_handler) -> None: + await _record_successes(redis_backed_handler, LOW_USAGE_DEPLOYMENT_ID, total_tokens=3, times=20) + await _record_successes(redis_backed_handler, HIGH_USAGE_DEPLOYMENT_ID, total_tokens=40) + + deployment: Final = redis_backed_handler.get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=[ + _limited_deployment(LOW_USAGE_DEPLOYMENT_ID, tpm=50), + _limited_deployment(HIGH_USAGE_DEPLOYMENT_ID, tpm=1000), + ], + ) + + assert deployment["model_info"]["id"] == HIGH_USAGE_DEPLOYMENT_ID + + +@pytest.mark.asyncio +async def test_usage_based_routing_v1_counts_every_concurrent_success_toward_rpm(redis_backed_handler) -> None: + await _record_successes(redis_backed_handler, LOW_USAGE_DEPLOYMENT_ID, total_tokens=1, times=20) + await _record_successes(redis_backed_handler, HIGH_USAGE_DEPLOYMENT_ID, total_tokens=40) + + deployment: Final = redis_backed_handler.get_available_deployments( + model_group=MODEL_GROUP, + healthy_deployments=[ + _limited_deployment(LOW_USAGE_DEPLOYMENT_ID, rpm=20), + _limited_deployment(HIGH_USAGE_DEPLOYMENT_ID, rpm=1000), + ], + ) + + assert deployment["model_info"]["id"] == HIGH_USAGE_DEPLOYMENT_ID