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