diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index cadc07a35df..003099f8e3a 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -587,21 +587,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): self, counter: RateLimitCounter, ) -> int | float | str | None: - """ - Get saturation value with configurable local cache TTL. - - Uses DualCache with configurable TTL for local cache storage. - TTL is configurable via litellm.priority_reservation_settings.saturation_check_cache_ttl - - Args: - counter: The saturation counter - - Returns: - Counter value, or None if not found - """ - value: Final = await self.v3_limiter.rate_limits.get( - counter, backfill_ttl=self._get_saturation_check_cache_ttl() - ) + """The model's counter as Redis holds it, so every worker agrees on saturation. A worker's own in-memory copy + carries only the usage it served itself, and steady traffic keeps that copy alive past any TTL. The read is + written back to memory for ``saturation_check_cache_ttl`` seconds; without Redis, or when the read fails, memory is + all there is.""" + rate_limits: Final = self.v3_limiter.rate_limits + try: + value: Final = await rate_limits.get_from_redis( + counter, backfill_ttl=self._get_saturation_check_cache_ttl() + ) + except Exception as e: # noqa: BLE001 # any Redis failure degrades to the local mirror, as the limiter's counter reads do + verbose_proxy_logger.warning("Saturation read failed in Redis, using this worker's copy: %s", e) + local: Final = rate_limits.local((counter,))[0] + return None if isinstance(local, dict) else local return None if isinstance(value, dict) else value def _get_priority_weight(self, priority: str | None, model_info: ModelGroupInfo | None = None) -> float: diff --git a/litellm/redis/stores/rate_limits/__init__.py b/litellm/redis/stores/rate_limits/__init__.py index 6964cf671bc..a49665dd634 100644 --- a/litellm/redis/stores/rate_limits/__init__.py +++ b/litellm/redis/stores/rate_limits/__init__.py @@ -99,6 +99,24 @@ class RateLimitStore: return local return await self._handle.get(entry, backfill_ttl_seconds=backfill_ttl) + @with_service_target(RATE_LIMITS_TARGET) + async def get_from_redis(self, counter: RateLimitCounter, backfill_ttl: float | None = None) -> LocalValue | None: + """Redis, never memory, for a count every worker must agree on: ``DualCache.async_get_cache`` answers from + memory first, and a worker's own copy carries just the usage it served itself. The count read is written + back to a memory that holds nothing for ``backfill_ttl``, as ``DualCache`` backfilled a miss; a copy the limiter + keeps for the whole window is left alone, so a Redis outage still falls back to it. Without Redis memory is + all there is.""" + if not self.has_redis: + return self._handle.get_local(counter) + _window, count = await self.read_counters([counter], batch=Batch.NONE) + if count is None: + return None + number: Final = float(count) + value: Final[LocalValue] = int(number) if number.is_integer() else number + if self._handle.get_local(counter) is None: + self._handle.set_local(counter, value, backfill_ttl) + return value + @with_service_target(RATE_LIMITS_TARGET) def sliding_window( self, counters: Sequence[RateLimitCounter], now: int, window_size: int, *, batch: Batch diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 44af7ddc08f..78ed4cf682a 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -6,7 +6,7 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of import asyncio import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta from types import MappingProxyType from typing import Final @@ -33,9 +33,9 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( claim_request_stash_for_data, ) from litellm.proxy.hooks.priority_class_share import pool_member -from litellm.redis.stores.rate_limits import RateLimitScope +from litellm.redis.stores.rate_limits import LocalValue, RateLimitCounter, RateLimitEntry, RateLimitScope from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse -from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits +from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits, Resolved, counter_pairs class TimeController: @@ -1436,21 +1436,26 @@ async def test_saturation_check_cache_ttl_configuration(monkeypatch): "TTL should be configurable via priority_reservation_settings" ) - # Track read-through (memory then Redis) gets to verify TTL is passed + # The Redis count is written back to an empty memory for the configured TTL get_cache_calls = [] class _RecordingRateLimits(FakeRateLimits): - async def get(self, entry, *, backfill_ttl_seconds=None): - get_cache_calls.append({"entry": entry, "ttl": backfill_ttl_seconds}) + def read_counters(self, counters, send): + return Resolved(lambda: _resolved(counter_pairs(counters, lambda _counter: "1")), send, "read_counters") + def set_local(self, entry, value, ttl_seconds): + get_cache_calls.append({"entry": entry, "ttl": ttl_seconds}) + + async def _resolved(value): + return value + + counter = RateLimitScope("model_saturation_check", model).counter("requests") + fakes = _RecordingRateLimits() handler.v3_limiter = PROXY_MaxParallelRequestsHandler_v3( handler.internal_usage_cache, - rate_limit_store=_RecordingRateLimits().store(), + rate_limit_store=fakes.store(), ) - # Call _get_saturation_value_from_cache - counter = RateLimitScope("model_saturation_check", model).counter("requests") - await handler._get_saturation_value_from_cache(counter) # Verify the read-through get was called with the configured TTL @@ -3797,3 +3802,77 @@ async def test_a_failing_metrics_sink_never_changes_admission(monkeypatch: pytes refusal: Final = await _refused_at_the_hook(handler, _team_key()) assert refusal.headers is not None assert refusal.headers["rate_limited_by"] == "model_per_team_saturated" + + +class _SharedRedisRateLimits(FakeRateLimits): + """One Redis behind several workers, each with its own memory.""" + + def __init__(self, redis: dict[RateLimitEntry, LocalValue]) -> None: + super().__init__() + self.redis = redis + + def read_counters(self, counters: Sequence[RateLimitCounter], send: object) -> Resolved[list[str | None]]: + async def reply() -> list[str | None]: + return counter_pairs(counters, self._redis_text) + + return Resolved(reply, send, "read_counters") + + def _redis_text(self, counter: RateLimitCounter) -> str | None: + value: Final = self.redis.get(counter) + return None if value is None or isinstance(value, dict) else str(value) + + +@pytest.mark.asyncio +async def test_every_worker_reads_the_models_saturation_from_redis_not_from_what_it_served_itself( + monkeypatch: pytest.MonkeyPatch, +): + redis: Final[dict[RateLimitEntry, LocalValue]] = {} + busy, idle = (_borrowing_handler(monkeypatch, borrows=True) for _ in range(2)) + for worker in (busy, idle): + worker.v3_limiter = PROXY_MaxParallelRequestsHandler_v3( + worker.internal_usage_cache, rate_limit_store=_SharedRedisRateLimits(redis).store() + ) + counter: Final = RateLimitScope("model_saturation_check", "shared-model").counter("tokens") + redis[counter] = 700 + busy.v3_limiter.rate_limits.set_local(counter, 700, 60) + idle.v3_limiter.rate_limits.set_local(counter, 30, 60) + group: Final = busy.llm_router.registry().group("shared-model") + assert group is not None + assert await busy._check_model_saturation("shared-model", group) == pytest.approx(0.7) + assert await idle._check_model_saturation("shared-model", group) == pytest.approx(0.7) + + +class _FailingRedisRateLimits(FakeRateLimits): + def read_counters(self, counters: Sequence[RateLimitCounter], send: object) -> Resolved[list[str | None]]: + raise ConnectionError("redis down") + + +@pytest.mark.asyncio +async def test_a_worker_whose_redis_read_fails_reads_the_models_saturation_from_its_own_memory( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + handler.v3_limiter = PROXY_MaxParallelRequestsHandler_v3( + handler.internal_usage_cache, rate_limit_store=_FailingRedisRateLimits().store() + ) + counter: Final = RateLimitScope("model_saturation_check", "shared-model").counter("tokens") + handler.v3_limiter.rate_limits.set_local(counter, 700, 60) + group: Final = handler.llm_router.registry().group("shared-model") + assert group is not None + assert await handler._check_model_saturation("shared-model", group) == pytest.approx(0.7) + + +@pytest.mark.asyncio +async def test_a_worker_without_redis_reads_the_models_saturation_from_its_own_memory( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + fakes: Final = FakeRateLimits(connected=False) + handler.v3_limiter = PROXY_MaxParallelRequestsHandler_v3( + handler.internal_usage_cache, rate_limit_store=fakes.store() + ) + counter: Final = RateLimitScope("model_saturation_check", "shared-model").counter("tokens") + fakes.set_local(counter, 700, 60) + group: Final = handler.llm_router.registry().group("shared-model") + assert group is not None + assert await handler._check_model_saturation("shared-model", group) == pytest.approx(0.7) diff --git a/tests/unit/redis/stores/fake_rate_limits.py b/tests/unit/redis/stores/fake_rate_limits.py index 568a5c8079c..bfc9beb6ee5 100644 --- a/tests/unit/redis/stores/fake_rate_limits.py +++ b/tests/unit/redis/stores/fake_rate_limits.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass, field +from itertools import chain from typing import Final, Generic, TypeVar from pydantic import TypeAdapter @@ -90,6 +91,13 @@ def _over_limit(raw: list[RawValue]) -> tuple[int, int, int] | None: return None +def counter_pairs( + counters: Sequence[RateLimitCounter], text: Callable[[RateLimitCounter], str | None] +) -> list[str | None]: + """Each counter's window start, always ``None`` here, then its count, the shape ``BATCH_COUNTER_READ_SCRIPT`` answers in.""" + return list(chain.from_iterable((None, text(counter)) for counter in counters)) + + @dataclass class FakeRateLimits: """Assign a fake to a ``*_script`` attribute to give the limiter a Redis that runs it; a script the test @@ -152,14 +160,19 @@ class FakeRateLimits: keys: Final = [key for counter in counters for key in (window_key(counter), counter_key(counter))] async def reply() -> list[str | None]: + if self.batch_counter_read_script is None: + return counter_pairs(counters, self._local_text) raw: Final = await self._run(self.batch_counter_read_script, "batch_counter_read_script", keys, []) return [_text(value) for value in raw] return Resolved(reply, send, "read_counters") - def check_and_increment( - self, counters: Sequence[AtomicCounter], send: object - ) -> Resolved[AtomicReply]: + def _local_text(self, counter: RateLimitCounter) -> str | None: + """Memory standing in for Redis when the test gave no read script, as ``get`` does.""" + value: Final = self.memory.get_local(counter) + return None if value is None or isinstance(value, dict) else str(value) + + def check_and_increment(self, counters: Sequence[AtomicCounter], send: object) -> Resolved[AtomicReply]: keys: Final = [key for atomic in counters for key in (window_key(atomic.counter), counter_key(atomic.counter))] args: Final = [ arg for atomic in counters for arg in (atomic.limit, atomic.increment, atomic.ttl, atomic.window_size) @@ -225,7 +238,9 @@ class FakeRateLimits: return Resolved(reply, send, "parallel_release") - def parallel_count(self, gauges: Sequence[RateLimitScope], slot_ttl_seconds: int, send: object) -> Resolved[list[int]]: + def parallel_count( + self, gauges: Sequence[RateLimitScope], slot_ttl_seconds: int, send: object + ) -> Resolved[list[int]]: keys: Final = [scope_key(scope, "max_parallel_requests") for scope in gauges] async def reply() -> list[int]: