fix(router): finalize Redis batches before propagating cancellation

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)
 

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