From caf9bbbd5a535e61da844e1a0225a4f8a6312b8d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:20:05 -0700 Subject: [PATCH] fix(least-busy): keep the shared count readable, counted once, and off the loop A Redis outage read as "every deployment is idle", because batch_get_cache swallows the failure and answers with an empty dict. batch_get_counts and its async twin raise instead, so a worker that cannot reach Redis falls back to its own numbers rather than routing on zeros. The counter's TTL is now set only on a key that has none, so a +1 left behind by a worker that died mid-request ages out an hour after the key was created. It used to be refreshed on every touch, which kept that stuck count alive for as long as the group took traffic. Two least-busy groups counted the same request twice, since the pre-call list kept a selector per group while the success list deduped by class. The selector now goes on through add_litellm_input_callback, which dedupes the same way. A prompt-management model picked its deployment on the synchronous path, so the new Redis read landed on the event loop and configured routing plugins never ran. It awaits the async selector now. --- basedpyright-code-budget.json | 12 ++-- litellm/caching/redis_cache.py | 36 +++++++++- litellm/router.py | 8 ++- litellm/router_strategy/least_busy.py | 10 +-- ruff-strict-budget.json | 10 +-- .../test_litellm/caching/test_redis_cache.py | 44 ++++++++++++ .../router_strategy/test_least_busy.py | 67 +++++++++---------- .../test_router_routing_groups.py | 38 +++++++++++ .../test_router_routing_plugins.py | 30 +++++++++ type-discipline-budget.json | 8 +-- 10 files changed, 199 insertions(+), 64 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c21dc76743a..32d3730aa32 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5551 + "limit": 5532 }, "reportMissingTypeArgument": { - "limit": 15277 + "limit": 15273 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38240 + "limit": 38208 }, "reportUnknownParameterType": { - "limit": 19558 + "limit": 19532 }, "reportUnknownVariableType": { - "limit": 29781 + "limit": 29751 }, "reportUnnecessaryCast": { "limit": 110 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 138 + "limit": 137 }, "reportUnusedImport": { "limit": 542 diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 4eccde5742b..3abc6e6f3c9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -92,11 +92,18 @@ _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( _INCREMENT_WITH_FLOOR_LUA: Final = ( "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " - "redis.call('EXPIRE', KEYS[1], ARGV[2]) " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " "return count" ) _LUA_COUNT: Final = TypeAdapter(int) +_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) + + +def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: + return _OPTIONAL_COUNTS.validate_python( + tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) + ) def _get_call_stack_info(num_frames: int = 2) -> str: @@ -749,13 +756,19 @@ class RedisCache(BaseCache): ) raise e + @_redis_circuit_breaker_guard_sync def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - """Add ``value`` to ``key``, clamp the result at zero, and refresh the TTL, in one Lua call. + """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. A counter whose key expired while a request was still in flight would otherwise be recreated negative by that request's decrement. Clamping inside the same call is what keeps it safe: a separate corrective write could land after another pod's increment and - erase it. Returns the resulting count. + erase it. + + The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was + created rather than ``ttl`` after it was last touched. Refreshing it on every touch + would keep a count a dead worker never decremented alive for as long as the group + takes traffic. Returns the resulting count. """ namespaced_key: Final = self.check_and_fix_namespace(key=key) count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval @@ -763,6 +776,23 @@ class RedisCache(BaseCache): ) return _LUA_COUNT.validate_python(count) + @_redis_circuit_breaker_guard_sync + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + + ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller + cannot tell apart from "every counter is unset". A caller that has to fall back to its + own numbers when Redis is unreachable needs the failure, not a dict of zeros. + """ + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + + @_redis_circuit_breaker_guard + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Async twin of ``batch_get_counts``, raising on failure the same way.""" + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() diff --git a/litellm/router.py b/litellm/router.py index 3f450661946..f5b7924fb58 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1212,7 +1212,7 @@ class Router: selector = LeastBusyLoggingHandler(router_cache=self.cache) if register_callbacks: if isinstance(litellm.input_callback, list): - litellm.input_callback.append(selector) + litellm.logging_callback_manager.add_litellm_input_callback(selector) else: litellm.input_callback = [selector] case RoutingStrategy.USAGE_BASED_ROUTING.value: @@ -4139,10 +4139,12 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - prompt_management_deployment: Final = self.get_available_deployment( + specific_deployment: Final = kwargs.pop("specific_deployment", None) + prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), + specific_deployment=specific_deployment, + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 771d2bb4328..ab6d702bbc2 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -38,7 +38,6 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) -_REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) _MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) @@ -72,11 +71,6 @@ def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: return tuple(0 if value is None else int(value) for value in values) -def _shared_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: - by_key: Final = _REDIS_COUNTS.validate_python(raw) - return _as_counts([by_key.get(key) for key in keys]) - - def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: values: Final = _MEMORY_COUNTS.validate_python(raw) if values is None or len(values) != len(keys): @@ -151,7 +145,7 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _shared_counts(redis_cache.batch_get_cache(key_list=list(keys)), keys) + shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) except Exception as e: _warn_unreadable(model_group, e) else: @@ -166,7 +160,7 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _shared_counts(await redis_cache.async_batch_get_cache(key_list=list(keys)), keys) + shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) except Exception as e: _warn_unreadable(model_group, e) else: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ceff8c6e15c..d63a69de76f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2937 + "limit": 2918 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1972 + "limit": 1965 }, "ANN202": { - "limit": 830 + "limit": 829 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2915 + "limit": 2914 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 212 + "limit": 207 }, "S112": { "limit": 22 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2f412e7382b..6b2df118611 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} +def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): + """A caller that must fall back when Redis is unreachable needs the failure, not zeros. + + The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell + apart from "every counter is unset". Least-busy routing read that as an idle deployment and + kept sending traffic to it instead of falling back to this worker's own in-flight counts. + """ + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + sync_batch_redis_cache.batch_get_counts(["lit7039"]) + + +@pytest.mark.asyncio +async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): + """Async twin: the async batch read hides the same failure behind an empty dict.""" + failing_client = AsyncMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + with patch.object(cache, "init_async_client", return_value=failing_client): + assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_batch_get_counts(["lit7039"]) + + +@pytest.mark.parametrize("stored", [b"3", "3"]) +def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): + """Counters come back positionally, so an unset key has to stay a hole rather than shift the + rest of the row onto the wrong deployments, and a count has to survive whether the client + hands it back as bytes or as text.""" + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.return_value = [stored, None, b"0"] + + assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) + + @pytest.fixture def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: service_logger = ServiceLogging(mock_testing=True) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index eb9591811ac..9efa526fc02 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -1,4 +1,3 @@ -import json from typing import Final import pytest @@ -18,44 +17,34 @@ def _call_kwargs(deployment_id: str) -> dict[str, object]: class SharedRedisCounters: - """Stores JSON strings and hands back a fresh object per read, the way a real Redis client does.""" + """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when + the key is created, and ordered reads that raise rather than invent a value.""" def __init__(self) -> None: - self.encoded: dict[str, str] = {} - self.ttls: dict[str, float] = {} + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} - def count(self, key: str) -> object: - raw: Final = self.encoded.get(key) - return None if raw is None else json.loads(raw) + def count(self, key: str) -> int | None: + return self.counts.get(key) - def get_cache(self, key: str, **kwargs: object) -> object: - return self.count(key) - - def set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.encoded[key] = json.dumps(value) - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - return self.count(key) - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.set_cache(key, value) + def expire(self, key: str) -> None: + self.counts.pop(key, None) + self.ttls.pop(key, None) def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - current: Final = self.count(key) or 0 - assert isinstance(current, int) - incremented: Final = max(0, current + value) - self.encoded[key] = json.dumps(incremented) - self.ttls[key] = ttl + incremented: Final = max(0, self.counts.get(key, 0) + value) + self.counts[key] = incremented + self.ttls.setdefault(key, ttl) return incremented async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: return self.increment_with_floor(key, value, ttl) - def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: - return {key: self.count(key) for key in key_list} + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return tuple(self.counts.get(key) for key in key_list) - async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: - return self.batch_get_cache(key_list) + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return self.batch_get_counts(key_list) def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: @@ -98,17 +87,25 @@ def test_sync_pick_reads_the_shared_counts() -> None: assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A -def test_redis_counts_keep_a_refreshed_ttl() -> None: +def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: + """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that + stuck count an hour after the key was created, which only works while nothing writes the TTL + again: a handler that refreshed it on every touch would keep the count alive for as long as + the group takes traffic, and the deployment would read busier than it is forever.""" shared: Final = SharedRedisCounters() worker: Final = _worker(shared) key: Final = f"{GROUP}_request_count:dep-a" - shared.ttls[key] = 5 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + + shared.ttls[key] = 5 worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(key) == 0 - assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + assert shared.count(key) == 1 + assert shared.ttls == {key: 5} @pytest.mark.asyncio @@ -127,10 +124,10 @@ async def test_counts_stay_in_memory_without_redis() -> None: class UnavailableRedis(SharedRedisCounters): - def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: raise ConnectionError("redis is down") - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: raise ConnectionError("redis is down") @@ -153,7 +150,7 @@ def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: worker: Final = _worker(shared) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - shared.encoded.clear() + shared.expire(f"{GROUP}_request_count:dep-a") worker.log_success_event(_call_kwargs("dep-a"), None, None, None) assert shared.count(f"{GROUP}_request_count:dep-a") == 0 @@ -187,4 +184,4 @@ def test_calls_without_a_deployment_are_ignored() -> None: worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) worker.log_pre_api_call(model="m", messages=[], kwargs={}) - assert shared.encoded == {} + assert shared.counts == {} diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5599c5aad63..0ce0ed5b37e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy @@ -435,6 +436,43 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp assert router._group_selectors == {} +def test_two_least_busy_groups_count_a_request_once(monkeypatch): + """ + Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and + back down from the success hooks on `litellm.callbacks`. The success list drops a second + selector of the same class, so a pre-call list that kept both counted every request twice + and released it once, and the deployment's in-flight count climbed until it looked pinned. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + router = _build_router( + routing_strategy="least-busy", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "least-busy", + } + ], + ) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index 293af36080a..c49b22ea367 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -164,6 +164,36 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) +@pytest.mark.asyncio +async def test_prompt_management_model_still_runs_the_plugin_pipeline(): + """ + A prompt-management model routes through its own factory, which picked the deployment + on the synchronous path. Plugins never run there, so the guard turned every such request + into an error message about the caller's own API choice, on an async call the caller made + correctly. It also read the in-flight counts with a blocking call inside the event loop. + """ + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="lit-7039", + ) + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e413e6db2c6..0c0952289e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22176 + "limit": 22174 }, "LIT002": { - "limit": 26721 + "limit": 26715 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16412 + "limit": 16398 }, "LIT011": { - "limit": 5505 + "limit": 5504 }, "LIT012": { "limit": 4486