From 45a7bf91541bf17237ba2783a92654128dee525c Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:12:57 -0500 Subject: [PATCH 1/9] fix(router): await budget redis pipeline before sync reads --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- .../proxy/hooks/model_max_budget_limiter.py | 4 + litellm/router_strategy/budget_limiter.py | 178 ++++++++---- ...test_unit_test_max_model_budget_limiter.py | 18 ++ .../router_strategy/test_budget_limiter.py | 258 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 6 files changed, 409 insertions(+), 55 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter.py diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fb2f014e3d8..81889e4da0c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18974,7 +18974,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index efaaab277a9..11191dc98a3 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,3 +1,4 @@ +import asyncio import json import time from collections.abc import Iterable, Mapping, Sequence @@ -268,6 +269,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): def __init__(self, dual_cache: DualCache): self.dual_cache = dual_cache self.redis_increment_operation_queue = [] + self._redis_increment_queue_lock = asyncio.Lock() + self._redis_increment_flush_lock = asyncio.Lock() + self._detached_increment_operations = None self.deployment_budget_config = None async def is_key_within_model_budget( diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 3e094df7ac8..23586569b87 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -109,6 +109,9 @@ class RouterBudgetLimiting(CustomLogger): ): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] + self._redis_increment_queue_lock = asyncio.Lock() + self._redis_increment_flush_lock = asyncio.Lock() + self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None @@ -402,7 +405,81 @@ class RouterBudgetLimiting(CustomLogger): increment_value=response_cost, ttl=ttl, ) - self.redis_increment_operation_queue.append(increment_op) + async with self._get_redis_increment_queue_lock(): + self.redis_increment_operation_queue.append(increment_op) + + def _get_redis_increment_queue_lock(self) -> asyncio.Lock: + return self._redis_increment_queue_lock + + async def _detach_queued_increment_operations(self) -> tuple[RedisPipelineIncrementOperation, ...]: + async with self._get_redis_increment_queue_lock(): + if self._detached_increment_operations is not None: + return self._detached_increment_operations + increment_operations_to_flush: Final = tuple(self.redis_increment_operation_queue) + if not increment_operations_to_flush: + return increment_operations_to_flush + self.redis_increment_operation_queue = [] # mutable-ok: emptied queue must stay appendable + self._detached_increment_operations = increment_operations_to_flush + return increment_operations_to_flush + + async def _clear_detached_increment_operations(self) -> None: + async with self._get_redis_increment_queue_lock(): + self._detached_increment_operations = None + + async def _requeue_detached_increment_operations(self) -> None: + async with self._get_redis_increment_queue_lock(): + detached_increment_operations: Final = self._detached_increment_operations + if detached_increment_operations is None: + return + self.redis_increment_operation_queue = ( + list( # mutable-ok: restored flush batch must stay appendable + detached_increment_operations + ) + + self.redis_increment_operation_queue + ) + self._detached_increment_operations = None + + async def _finish_increment_pipeline_after_cancellation( + self, + pipeline_task: asyncio.Task[object], + ) -> None: + try: + await pipeline_task + except Exception: + await self._requeue_detached_increment_operations() + verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + return + await self._clear_detached_increment_operations() + + async def _flush_queued_increment_operations(self, redis_cache: RedisCache) -> bool: + increment_operations_to_flush: Final = await self._detach_queued_increment_operations() + if len(increment_operations_to_flush) == 0: + await self._clear_detached_increment_operations() + return True + + verbose_router_logger.debug( + "Pushing Redis Increment Pipeline for queue: %s", + increment_operations_to_flush, + ) + increment_list: Final = list( # mutable-ok: Redis pipeline contract requires a list + increment_operations_to_flush + ) + pipeline_task: Final = asyncio.create_task( + redis_cache.async_increment_pipeline( + increment_list=increment_list, + ) + ) + try: + await asyncio.shield(pipeline_task) + except Exception: + await asyncio.shield(self._requeue_detached_increment_operations()) + verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + return False + except asyncio.CancelledError: + await asyncio.shield(self._finish_increment_pipeline_after_cancellation(pipeline_task)) + raise + await self._clear_detached_increment_operations() + return True async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" @@ -528,29 +605,25 @@ class RouterBudgetLimiting(CustomLogger): DEFAULT_REDIS_SYNC_INTERVAL ) # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying - async def _push_in_memory_increments_to_redis(self): + async def _push_in_memory_increments_to_redis(self) -> bool: """ How this works: - async_log_success_event collects all provider spend increments in `redis_increment_operation_queue` - This function pushes all increments to Redis in a batched pipeline to optimize performance - Only runs if Redis is initialized + Only runs if Redis is initialized. Returns False when the detached batch could not be + written, so callers must not treat Redis as up to date. """ - try: - if not self.dual_cache.redis_cache: - return # Redis is not initialized + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None: + return True - verbose_router_logger.debug( - "Pushing Redis Increment Pipeline for queue: %s", - self.redis_increment_operation_queue, - ) - queued: Final = self.redis_increment_operation_queue - self.redis_increment_operation_queue = [] - if queued: - asyncio.create_task(_push_increments_to_redis(self.dual_cache.redis_cache, queued)) - - except Exception as e: - log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) + async with self._redis_increment_flush_lock: + try: + return await self._flush_queued_increment_operations(redis_cache) + except asyncio.CancelledError: + await asyncio.shield(self._requeue_detached_increment_operations()) + raise async def _sync_in_memory_spend_with_redis(self): """ @@ -569,44 +642,45 @@ class RouterBudgetLimiting(CustomLogger): # No need to sync if Redis cache is not initialized if self.dual_cache.redis_cache is None: return - - # 1. Push all provider spend increments to Redis - await self._push_in_memory_increments_to_redis() - - # 2. Fetch all current provider spend from Redis to update in-memory cache - cache_keys: Final = [] - - if self.provider_budget_config is not None: - for provider, config in self.provider_budget_config.items(): - if config is None: - continue - cache_keys.append(f"provider_spend:{provider}:{config.budget_duration}") - - if self.deployment_budget_config is not None: - for model_id, config in self.deployment_budget_config.items(): - if config is None: - continue - cache_keys.append(f"deployment_spend:{model_id}:{config.budget_duration}") - - if self.tag_budget_config is not None: - for tag, config in self.tag_budget_config.items(): - if config is None: - continue - cache_keys.append(f"tag_spend:{tag}:{config.budget_duration}") - - # Batch fetch current spend values from Redis - redis_values: Final = await self.dual_cache.redis_cache.async_batch_get_cache(key_list=cache_keys) - - # Update in-memory cache with Redis values - if isinstance(redis_values, dict): # Check if redis_values is a dictionary - for key, value in redis_values.items(): - if value is not None: - await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value)) - verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value) - + await self._flush_increments_then_copy_redis_spend() except Exception as e: log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) + async def _flush_increments_then_copy_redis_spend(self) -> None: + if not await self._push_in_memory_increments_to_redis(): + return + + cache_keys: Final = [] + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None: + return + + if self.provider_budget_config is not None: + for provider, config in self.provider_budget_config.items(): + if config is None: + continue + cache_keys.append(f"provider_spend:{provider}:{config.budget_duration}") + + if self.deployment_budget_config is not None: + for model_id, config in self.deployment_budget_config.items(): + if config is None: + continue + cache_keys.append(f"deployment_spend:{model_id}:{config.budget_duration}") + + if self.tag_budget_config is not None: + for tag, config in self.tag_budget_config.items(): + if config is None: + continue + cache_keys.append(f"tag_spend:{tag}:{config.budget_duration}") + + redis_values: Final = await redis_cache.async_batch_get_cache(key_list=cache_keys) + + if isinstance(redis_values, dict): + for key, value in redis_values.items(): + if value is not None: + await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value)) + verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value) + def _get_budget_config_for_deployment( self, model_id: str, diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 096efc33aaf..89ad0db7f93 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -19,6 +19,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( resolve_model_budget, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -487,6 +488,23 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config mock_push.assert_awaited_once() +@pytest.mark.asyncio +async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + spend_key = "virtual_key_spend:test-key:gpt-4:1d" + + await limiter._increment_spend_in_current_window( + spend_key=spend_key, response_cost=0.01, ttl=86400 + ) + + assert limiter.redis_increment_operation_queue == [ + RedisPipelineIncrementOperation( + key=spend_key, increment_value=0.01, ttl=86400 + ) + ] + + @pytest.mark.asyncio async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( budget_limiter, diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..bb0b72f757d --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,258 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.utils import BudgetConfig + +_SPEND_KEY = "provider_spend:openai:1d" + + +def _increment(increment_value: float) -> RedisPipelineIncrementOperation: + return RedisPipelineIncrementOperation(key=_SPEND_KEY, increment_value=increment_value, ttl=86400) + + +class _MockRedisCache: + def __init__( + self, + initial_values: dict[str, float], + pipeline_started: asyncio.Event | None = None, + allow_pipeline_to_complete: asyncio.Event | None = None, + should_fail_pipeline: bool = False, + ) -> None: + self.values = initial_values + self.events: list[str] = [] + self.pipeline_started = pipeline_started + self.allow_pipeline_to_complete = allow_pipeline_to_complete + self.should_fail_pipeline = should_fail_pipeline + + async def async_increment_pipeline( + self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object + ) -> None: + self.events.append("increment_pipeline:start") + if self.pipeline_started is not None: + self.pipeline_started.set() + if self.allow_pipeline_to_complete is not None: + await self.allow_pipeline_to_complete.wait() + if self.should_fail_pipeline: + raise RuntimeError("redis down") + for op in increment_list: + key = op["key"] + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(op["increment_value"]) + self.events.append("increment_pipeline:done") + + async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, float | None]: + self.events.append("batch_get") + return {key: self.values.get(key) for key in key_list} + + +class _MockInMemoryCache: + def __init__(self, initial_values: dict[str, float]) -> None: + self.values = initial_values + + async def async_increment(self, key: str, value: float, ttl: int, **kwargs: object) -> float: + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(value) + return self.values[key] + + async def async_set_cache(self, key: str, value: float, **kwargs: object) -> None: + self.values[key] = float(value) + + +def _new_router_budget_limiter( + *, + redis_cache: object, + in_memory_cache: object | None = None, + redis_increment_operation_queue: list[RedisPipelineIncrementOperation] | None = None, + provider_budget_config: dict[str, BudgetConfig] | None = None, +) -> RouterBudgetLimiting: + budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) + budget_limiter.dual_cache = SimpleNamespace( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache if in_memory_cache is not None else SimpleNamespace(), + ) + budget_limiter.provider_budget_config = provider_budget_config + budget_limiter.deployment_budget_config = None + budget_limiter.tag_budget_config = None + budget_limiter.redis_increment_operation_queue = ( + list(redis_increment_operation_queue) if redis_increment_operation_queue is not None else [] + ) + budget_limiter._redis_increment_queue_lock = asyncio.Lock() + budget_limiter._redis_increment_flush_lock = asyncio.Lock() + budget_limiter._detached_increment_operations = None + return budget_limiter + + +@pytest.mark.asyncio +async def test_should_await_redis_pipeline_before_sync_reads() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 100.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + ) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(60.0)], + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)}, + ) + + sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert "batch_get" not in redis_cache.events + allow_pipeline_to_complete.set() + await sync_task + + assert redis_cache.values[_SPEND_KEY] == 160.0 + assert in_memory_cache.values[_SPEND_KEY] == 160.0 + assert budget_limiter.redis_increment_operation_queue == [] + assert redis_cache.events == [ + "increment_pipeline:start", + "increment_pipeline:done", + "batch_get", + ] + + +@pytest.mark.asyncio +async def test_should_requeue_increments_when_redis_pipeline_fails() -> None: + redis_cache = _MockRedisCache(initial_values={}, should_fail_pipeline=True) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + flush_succeeded = await budget_limiter._push_in_memory_increments_to_redis() + + assert flush_succeeded is False + assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)] + assert budget_limiter._detached_increment_operations is None + + +@pytest.mark.asyncio +async def test_should_keep_new_increments_when_pipeline_flush_fails() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + ) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 0.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + await budget_limiter._increment_spend_in_current_window(spend_key=_SPEND_KEY, response_cost=20.0, ttl=86400) + allow_pipeline_to_complete.set() + await push_task + + assert budget_limiter.redis_increment_operation_queue == [_increment(10.0), _increment(20.0)] + + +@pytest.mark.asyncio +async def test_should_keep_in_memory_spend_when_redis_pipeline_fails() -> None: + redis_cache = _MockRedisCache(initial_values={_SPEND_KEY: 100.0}, should_fail_pipeline=True) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(60.0)], + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)}, + ) + + await budget_limiter._sync_in_memory_spend_with_redis() + + assert in_memory_cache.values[_SPEND_KEY] == 160.0 + assert redis_cache.values[_SPEND_KEY] == 100.0 + assert budget_limiter.redis_increment_operation_queue == [_increment(60.0)] + assert "batch_get" not in redis_cache.events + + +@pytest.mark.asyncio +async def test_should_keep_increments_when_flush_is_cancelled_after_success() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 0.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + ) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + push_task.cancel() + allow_pipeline_to_complete.set() + with pytest.raises(asyncio.CancelledError): + await push_task + + assert redis_cache.values[_SPEND_KEY] == 10.0 + assert budget_limiter.redis_increment_operation_queue == [] + assert budget_limiter._detached_increment_operations is None + + +@pytest.mark.asyncio +async def test_empty_flush_does_not_block_later_increment_sync() -> None: + redis_cache = _MockRedisCache(initial_values={_SPEND_KEY: 100.0}) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 100.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)}, + ) + + empty_flush_succeeded = await budget_limiter._push_in_memory_increments_to_redis() + await budget_limiter._increment_spend_in_current_window(spend_key=_SPEND_KEY, response_cost=20.0, ttl=86400) + await budget_limiter._sync_in_memory_spend_with_redis() + + assert empty_flush_succeeded is True + assert budget_limiter._detached_increment_operations is None + assert budget_limiter.redis_increment_operation_queue == [] + assert redis_cache.values[_SPEND_KEY] == 120.0 + assert in_memory_cache.values[_SPEND_KEY] == 120.0 + assert redis_cache.events == [ + "increment_pipeline:start", + "increment_pipeline:done", + "batch_get", + ] + + +@pytest.mark.asyncio +async def test_should_requeue_increments_when_flush_is_cancelled_and_redis_fails() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 0.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + ) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + push_task.cancel() + allow_pipeline_to_complete.set() + with pytest.raises(asyncio.CancelledError): + await push_task + + assert redis_cache.values[_SPEND_KEY] == 0.0 + assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)] + assert budget_limiter._detached_increment_operations is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b26f5e25b6f..6b7337ea8d1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35343,7 +35343,7 @@ export interface components { default_model?: string | null; /** * Deployment Affinity - * @description When True and a session_id is resolvable on the request, pin the deployment chosen inside each routed model group and reuse it whenever the session returns to that group, without pinning which group the session routes to. Independent of session_affinity, which pins the model group instead (and always carries this deployment pin with it): with session_affinity off, every turn is still classified on its own merits while a session that escalates to a stronger tier and comes back still lands on the deployment it used before, which is what keeps a provider prompt cache warm. Pins are held per model group, so switching tiers does not disturb the pin left behind in the previous group. On by default because re-shuffling a conversation across deployments of the same model discards that cache for no benefit; set False to keep every turn load-balanced across the group, which is what a deployment set with tight per-deployment rate limits wants. Inert when no session_id is resolvable, since there is nothing to key a pin on, and suppressed when plugins are configured, for the same reason session_affinity is. + * @description When True and a client session_id is resolvable, reuse the session's chosen model for each classified tier and its deployment within each model group. With session_affinity off, every turn is still classified: moving to another tier leaves the previous tier's model pin intact for a later return. Pins yield to current candidate, context, modality, and availability constraints. Adaptive selection chooses the initial model from its eligible pool, then reuses that choice per tier. This reduces avoidable provider prompt-cache misses; it does not guarantee cache hits. Set False to select models and load-balance deployments on every turn, unless session_affinity or user_turn classification requires a pin. Inert without a client session_id and suppressed when plugins are configured. * @default true */ deployment_affinity: boolean; @@ -35487,7 +35487,7 @@ export interface components { session_affinity: boolean; /** * Session Affinity Ttl Seconds - * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity deployment pin, so it measures idle time for the session's routing decisions rather than total session length + * @description TTL for the session affinity pin; refreshed on every cache hit. Bounds both the session_affinity model pin and the deployment_affinity per-tier model and deployment pins, so it measures idle time for the session's routing decisions rather than total session length * @default 3600 */ session_affinity_ttl_seconds: number; From 1cf840454eb29b9ef7f2e44082fac888ac230cfa Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:30:11 -0500 Subject: [PATCH 2/9] refactor(router): remove superseded Redis flush helper --- litellm/router_strategy/budget_limiter.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 23586569b87..2838c3c6509 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -93,13 +93,6 @@ class _LiteLLMParamsDictView: return dict(self._params) -async def _push_increments_to_redis(redis_cache: RedisCache, queued: list[RedisPipelineIncrementOperation]) -> None: - try: - await redis_cache.async_increment_pipeline(increment_list=queued) - except Exception as e: - log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) - - class RouterBudgetLimiting(CustomLogger): def __init__( self, From 7a2e70020da6047a44b04a557cac0cd2cc9e778f Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:33:59 -0500 Subject: [PATCH 3/9] fix(proxy): preserve CI-compatible OpenAPI snapshot formatting --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 81889e4da0c..fb2f014e3d8 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18974,7 +18974,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 6cb7a6f6009e2265f303fa5d1a1b5de530d25db2 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:55:35 -0500 Subject: [PATCH 4/9] fix(router): preserve concurrent spend during Redis synchronization --- litellm/router_strategy/budget_limiter.py | 40 ++++++++++------ .../router_strategy/test_budget_limiter.py | 48 ++++++++++++++++++- 2 files changed, 73 insertions(+), 15 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 2838c3c6509..6fd2a2534ce 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -388,17 +388,17 @@ class RouterBudgetLimiting(CustomLogger): - Increments the spend in memory cache (so spend instantly updated in memory) - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) """ - await self.dual_cache.in_memory_cache.async_increment( - key=spend_key, - value=response_cost, - ttl=ttl, - ) increment_op: Final = RedisPipelineIncrementOperation( key=spend_key, increment_value=response_cost, ttl=ttl, ) async with self._get_redis_increment_queue_lock(): + await self.dual_cache.in_memory_cache.async_increment( + key=spend_key, + value=response_cost, + ttl=ttl, + ) self.redis_increment_operation_queue.append(increment_op) def _get_redis_increment_queue_lock(self) -> asyncio.Lock: @@ -635,18 +635,21 @@ class RouterBudgetLimiting(CustomLogger): # No need to sync if Redis cache is not initialized if self.dual_cache.redis_cache is None: return - await self._flush_increments_then_copy_redis_spend() + async with self._redis_increment_flush_lock: + try: + await self._flush_increments_then_copy_redis_spend() + except asyncio.CancelledError: + await asyncio.shield(self._requeue_detached_increment_operations()) + raise except Exception as e: log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) async def _flush_increments_then_copy_redis_spend(self) -> None: - if not await self._push_in_memory_increments_to_redis(): + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None or not await self._flush_queued_increment_operations(redis_cache): return cache_keys: Final = [] - redis_cache: Final = self.dual_cache.redis_cache - if redis_cache is None: - return if self.provider_budget_config is not None: for provider, config in self.provider_budget_config.items(): @@ -668,11 +671,20 @@ class RouterBudgetLimiting(CustomLogger): redis_values: Final = await redis_cache.async_batch_get_cache(key_list=cache_keys) - if isinstance(redis_values, dict): + if not isinstance(redis_values, dict): + return + async with self._get_redis_increment_queue_lock(): for key, value in redis_values.items(): - if value is not None: - await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value)) - verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value) + if value is None: + continue + pending_spend: Final = sum( + operation["increment_value"] + for operation in self.redis_increment_operation_queue + if operation["key"] == key + ) + updated_spend: Final = float(value) + pending_spend + await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=updated_spend) + verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, updated_spend) def _get_budget_config_for_deployment( self, diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index bb0b72f757d..29a646e963f 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -21,12 +21,16 @@ class _MockRedisCache: pipeline_started: asyncio.Event | None = None, allow_pipeline_to_complete: asyncio.Event | None = None, should_fail_pipeline: bool = False, + read_started: asyncio.Event | None = None, + allow_read_to_complete: asyncio.Event | None = None, ) -> None: self.values = initial_values self.events: list[str] = [] self.pipeline_started = pipeline_started self.allow_pipeline_to_complete = allow_pipeline_to_complete self.should_fail_pipeline = should_fail_pipeline + self.read_started = read_started + self.allow_read_to_complete = allow_read_to_complete async def async_increment_pipeline( self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object @@ -46,7 +50,12 @@ class _MockRedisCache: async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, float | None]: self.events.append("batch_get") - return {key: self.values.get(key) for key in key_list} + snapshot = {key: self.values.get(key) for key in key_list} + if self.read_started is not None: + self.read_started.set() + if self.allow_read_to_complete is not None: + await self.allow_read_to_complete.wait() + return snapshot class _MockInMemoryCache: @@ -256,3 +265,40 @@ async def test_should_requeue_increments_when_flush_is_cancelled_and_redis_fails assert redis_cache.values[_SPEND_KEY] == 0.0 assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)] assert budget_limiter._detached_increment_operations is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pause_during", ["write", "read"]) +async def test_sync_preserves_spend_recorded_during_redis_io(pause_during: str) -> None: + io_started = asyncio.Event() + allow_io_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 100.0}, + pipeline_started=io_started if pause_during == "write" else None, + allow_pipeline_to_complete=allow_io_to_complete if pause_during == "write" else None, + read_started=io_started if pause_during == "read" else None, + allow_read_to_complete=allow_io_to_complete if pause_during == "read" else None, + ) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(60.0)], + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=175.0)}, + ) + + sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis()) + await asyncio.wait_for(io_started.wait(), timeout=1) + await budget_limiter._increment_spend_in_current_window(_SPEND_KEY, 20.0, 86400) + allow_io_to_complete.set() + await sync_task + + assert in_memory_cache.values[_SPEND_KEY] == 180.0 + assert redis_cache.values[_SPEND_KEY] == 160.0 + assert budget_limiter.redis_increment_operation_queue == [_increment(20.0)] + + await budget_limiter._sync_in_memory_spend_with_redis() + + assert in_memory_cache.values[_SPEND_KEY] == 180.0 + assert redis_cache.values[_SPEND_KEY] == 180.0 + assert budget_limiter.redis_increment_operation_queue == [] From 16e5642c25aa6aedc068f0f55f0ac41b12ecd732 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:04:45 -0500 Subject: [PATCH 5/9] fix(router): type per-key spend totals without loop Final bindings --- litellm/router_strategy/budget_limiter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 6fd2a2534ce..fe736c56b5e 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -677,12 +677,12 @@ class RouterBudgetLimiting(CustomLogger): for key, value in redis_values.items(): if value is None: continue - pending_spend: Final = sum( + pending_spend = sum( # rebind-ok: each cache key has independent queued spend operation["increment_value"] for operation in self.redis_increment_operation_queue if operation["key"] == key ) - updated_spend: Final = float(value) + pending_spend + updated_spend = float(value) + pending_spend # rebind-ok: each cache key has an independent total await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=updated_spend) verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, updated_spend) From d43ca9b7891590082a4bb72268b13c8bb2d8e96f Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:07:20 -0500 Subject: [PATCH 6/9] fix(router): log Redis failures before cancellable cleanup --- litellm/router_strategy/budget_limiter.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index fe736c56b5e..63c49d03a27 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -439,8 +439,8 @@ class RouterBudgetLimiting(CustomLogger): try: await pipeline_task except Exception: - await self._requeue_detached_increment_operations() verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + await self._requeue_detached_increment_operations() return await self._clear_detached_increment_operations() @@ -465,8 +465,8 @@ class RouterBudgetLimiting(CustomLogger): try: await asyncio.shield(pipeline_task) except Exception: - await asyncio.shield(self._requeue_detached_increment_operations()) verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + await asyncio.shield(self._requeue_detached_increment_operations()) return False except asyncio.CancelledError: await asyncio.shield(self._finish_increment_pipeline_after_cancellation(pipeline_task)) From 8dc7a0c71580e2bddce7fc9ad76aa242c8c13b3f Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:13:19 -0500 Subject: [PATCH 7/9] fix(router): finalize Redis batches before propagating cancellation --- litellm/router_strategy/budget_limiter.py | 49 +++++++------------ .../router_strategy/test_budget_limiter.py | 34 +++++++++++++ 2 files changed, 51 insertions(+), 32 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 63c49d03a27..43ee2658a30 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -432,19 +432,20 @@ class RouterBudgetLimiting(CustomLogger): ) self._detached_increment_operations = None - async def _finish_increment_pipeline_after_cancellation( - self, - pipeline_task: asyncio.Task[object], - ) -> None: - try: - await pipeline_task - except Exception: - verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") - await self._requeue_detached_increment_operations() - return - await self._clear_detached_increment_operations() - async def _flush_queued_increment_operations(self, redis_cache: RedisCache) -> bool: + flush_task: Final = asyncio.create_task(self._write_queued_increment_operations(redis_cache)) + try: + return await asyncio.shield(flush_task) + except asyncio.CancelledError: + while not flush_task.done(): + try: + await asyncio.shield(flush_task) + except asyncio.CancelledError: + continue + flush_task.result() + raise + + async def _write_queued_increment_operations(self, redis_cache: RedisCache) -> bool: increment_operations_to_flush: Final = await self._detach_queued_increment_operations() if len(increment_operations_to_flush) == 0: await self._clear_detached_increment_operations() @@ -457,20 +458,12 @@ class RouterBudgetLimiting(CustomLogger): increment_list: Final = list( # mutable-ok: Redis pipeline contract requires a list increment_operations_to_flush ) - pipeline_task: Final = asyncio.create_task( - redis_cache.async_increment_pipeline( - increment_list=increment_list, - ) - ) try: - await asyncio.shield(pipeline_task) + await redis_cache.async_increment_pipeline(increment_list=increment_list) except Exception: verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") - await asyncio.shield(self._requeue_detached_increment_operations()) + await self._requeue_detached_increment_operations() return False - except asyncio.CancelledError: - await asyncio.shield(self._finish_increment_pipeline_after_cancellation(pipeline_task)) - raise await self._clear_detached_increment_operations() return True @@ -612,11 +605,7 @@ class RouterBudgetLimiting(CustomLogger): return True async with self._redis_increment_flush_lock: - try: - return await self._flush_queued_increment_operations(redis_cache) - except asyncio.CancelledError: - await asyncio.shield(self._requeue_detached_increment_operations()) - raise + return await self._flush_queued_increment_operations(redis_cache) async def _sync_in_memory_spend_with_redis(self): """ @@ -636,11 +625,7 @@ class RouterBudgetLimiting(CustomLogger): if self.dual_cache.redis_cache is None: return async with self._redis_increment_flush_lock: - try: - await self._flush_increments_then_copy_redis_spend() - except asyncio.CancelledError: - await asyncio.shield(self._requeue_detached_increment_operations()) - raise + await self._flush_increments_then_copy_redis_spend() except Exception as e: log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index 29a646e963f..7a0e9bd2b3c 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -21,6 +21,7 @@ class _MockRedisCache: pipeline_started: asyncio.Event | None = None, allow_pipeline_to_complete: asyncio.Event | None = None, should_fail_pipeline: bool = False, + pipeline_completed: asyncio.Event | None = None, read_started: asyncio.Event | None = None, allow_read_to_complete: asyncio.Event | None = None, ) -> None: @@ -29,6 +30,7 @@ class _MockRedisCache: self.pipeline_started = pipeline_started self.allow_pipeline_to_complete = allow_pipeline_to_complete self.should_fail_pipeline = should_fail_pipeline + self.pipeline_completed = pipeline_completed self.read_started = read_started self.allow_read_to_complete = allow_read_to_complete @@ -47,6 +49,8 @@ class _MockRedisCache: current = float(self.values.get(key, 0.0) or 0.0) self.values[key] = current + float(op["increment_value"]) self.events.append("increment_pipeline:done") + if self.pipeline_completed is not None: + self.pipeline_completed.set() async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, float | None]: self.events.append("batch_get") @@ -302,3 +306,33 @@ async def test_sync_preserves_spend_recorded_during_redis_io(pause_during: str) assert in_memory_cache.values[_SPEND_KEY] == 180.0 assert redis_cache.values[_SPEND_KEY] == 180.0 assert budget_limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancellations", [1, 2]) +async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: int) -> None: + pipeline_started = asyncio.Event() + pipeline_completed = asyncio.Event() + allow_pipeline = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 0.0}, + pipeline_started=pipeline_started, + pipeline_completed=pipeline_completed, + allow_pipeline_to_complete=allow_pipeline, + ) + limiter = _new_router_budget_limiter(redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)]) + push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + async with limiter._redis_increment_queue_lock: + allow_pipeline.set() + await asyncio.wait_for(pipeline_completed.wait(), timeout=1) + for _ in range(cancellations): + push_task.cancel() + await asyncio.sleep(0) + assert not push_task.done() + with pytest.raises(asyncio.CancelledError): + await push_task + await limiter._push_in_memory_increments_to_redis() + assert redis_cache.values[_SPEND_KEY] == 10.0 + assert limiter.redis_increment_operation_queue == [] + assert limiter._detached_increment_operations is None From 64c6e242e00a09bb5f888e9f9f6788d41860e2d4 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:16:37 -0500 Subject: [PATCH 8/9] test(router): reproduce cancellation while Redis cleanup is blocked --- .../router_strategy/test_budget_limiter.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index 7a0e9bd2b3c..e8b0f713725 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -14,6 +14,17 @@ def _increment(increment_value: float) -> RedisPipelineIncrementOperation: return RedisPipelineIncrementOperation(key=_SPEND_KEY, increment_value=increment_value, ttl=86400) +class _ObservedLock(asyncio.Lock): + def __init__(self) -> None: + super().__init__() + self.waiter_started = asyncio.Event() + + async def acquire(self) -> bool: + if self.locked(): + self.waiter_started.set() + return await super().acquire() + + class _MockRedisCache: def __init__( self, @@ -78,6 +89,7 @@ class _MockInMemoryCache: def _new_router_budget_limiter( *, redis_cache: object, + queue_lock: asyncio.Lock | None = None, in_memory_cache: object | None = None, redis_increment_operation_queue: list[RedisPipelineIncrementOperation] | None = None, provider_budget_config: dict[str, BudgetConfig] | None = None, @@ -93,7 +105,7 @@ def _new_router_budget_limiter( budget_limiter.redis_increment_operation_queue = ( list(redis_increment_operation_queue) if redis_increment_operation_queue is not None else [] ) - budget_limiter._redis_increment_queue_lock = asyncio.Lock() + budget_limiter._redis_increment_queue_lock = queue_lock if queue_lock is not None else asyncio.Lock() budget_limiter._redis_increment_flush_lock = asyncio.Lock() budget_limiter._detached_increment_operations = None return budget_limiter @@ -320,12 +332,16 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: pipeline_completed=pipeline_completed, allow_pipeline_to_complete=allow_pipeline, ) - limiter = _new_router_budget_limiter(redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)]) + queue_lock = _ObservedLock() + limiter = _new_router_budget_limiter( + redis_cache=redis_cache, queue_lock=queue_lock, redis_increment_operation_queue=[_increment(10.0)] + ) push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) await asyncio.wait_for(pipeline_started.wait(), timeout=1) async with limiter._redis_increment_queue_lock: allow_pipeline.set() await asyncio.wait_for(pipeline_completed.wait(), timeout=1) + await asyncio.wait_for(queue_lock.waiter_started.wait(), timeout=1) for _ in range(cancellations): push_task.cancel() await asyncio.sleep(0) From 53478e86517b8ec4ea5801feb76b4468243ecec1 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 13:39:49 -0500 Subject: [PATCH 9/9] test: align budget hotpath checks with awaited Redis flush --- litellm/router_strategy/budget_limiter.py | 4 ++-- .../test_budget_limiter_hotpath.py | 22 +++++++++++-------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 43ee2658a30..089e95d10e3 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -460,8 +460,8 @@ class RouterBudgetLimiting(CustomLogger): ) try: await redis_cache.async_increment_pipeline(increment_list=increment_list) - except Exception: - verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + except Exception as error: + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", error) await self._requeue_detached_increment_operations() return False await self._clear_detached_increment_operations() diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 4cc8fe78811..272ecd4179a 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -338,7 +338,9 @@ async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_ta assert caplog.records == [] unretrieved.assert_not_called() - assert limiter.redis_increment_operation_queue == [] + assert limiter.redis_increment_operation_queue == [ + {"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60} + ] assert redis_cache.async_increment_pipeline.await_count == 1 @@ -353,25 +355,27 @@ async def _limiter_with_redis(redis_cache: MagicMock) -> RouterBudgetLimiting: @pytest.mark.asyncio -async def test_push_returns_before_redis_answers(disable_budget_sync): - """The push runs inside the request success callback, so it must hand the Redis round trip to a task instead of waiting on it.""" +async def test_push_waits_for_redis_before_completing(disable_budget_sync): + redis_started = asyncio.Event() redis_answered = asyncio.Event() async def wait_for_redis(**_: object) -> None: + redis_started.set() await redis_answered.wait() redis_cache = MagicMock(spec=RedisCache) redis_cache.async_increment_pipeline = AsyncMock(side_effect=wait_for_redis) limiter = await _limiter_with_redis(redis_cache) - await asyncio.wait_for(limiter._push_in_memory_increments_to_redis(), timeout=1) - await asyncio.sleep(0) - - assert not redis_answered.is_set() + push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(redis_started.wait(), timeout=1) + assert not push_task.done() + assert limiter._detached_increment_operations is not None + redis_answered.set() + assert await asyncio.wait_for(push_task, timeout=1) is True assert redis_cache.async_increment_pipeline.await_count == 1 assert limiter.redis_increment_operation_queue == [] - redis_answered.set() - await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + assert limiter._detached_increment_operations is None @pytest.mark.asyncio