fix(router): await budget redis pipeline before sync reads

This commit is contained in:
Emerson Gomes 2026-09-15 12:12:57 -05:00
parent 3ad9a7f336
commit 45a7bf9154
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17
6 changed files with 409 additions and 55 deletions

View file

@ -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": {

View file

@ -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(

View file

@ -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,

View file

@ -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,

View file

@ -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

View file

@ -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;