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