degrade budget reservation cache failures

litellm/proxy/spend_tracking/budget_reservation.py

@@ -573,33 +573,54 @@ async def _reserve_counter(
     from litellm.proxy.proxy_server import (
         _ensure_spend_counter_initialized,
         _ensure_window_spend_counter_initialized,
+        _invalidate_spend_counter,
         _increment_spend_counter_cache,
     )
 
-    if counter.source_cache_key is not None:
-        await _ensure_spend_counter_initialized(
-            counter_key=counter.counter_key,
-            source_cache_key=counter.source_cache_key,
-        )
-    elif counter.spend_log_entity_id is not None and counter.window_start is not None:
-        initialized = await _ensure_window_spend_counter_initialized(
-            counter_key=counter.counter_key,
-            entity_type=counter.entity_type,
-            entity_id=counter.spend_log_entity_id,
-            window_start=counter.window_start,
-        )
-        if initialized is False:
-            verbose_proxy_logger.warning(
-                "Skipping budget reservation for %s because window spend could not be loaded",
-                counter.counter_key,
+    try:
+        if counter.source_cache_key is not None:
+            await _ensure_spend_counter_initialized(
+                counter_key=counter.counter_key,
+                source_cache_key=counter.source_cache_key,
             )
-            raise _CounterReservationUnavailable
+        elif (
+            counter.spend_log_entity_id is not None and counter.window_start is not None
+        ):
+            initialized = await _ensure_window_spend_counter_initialized(
+                counter_key=counter.counter_key,
+                entity_type=counter.entity_type,
+                entity_id=counter.spend_log_entity_id,
+                window_start=counter.window_start,
+            )
+            if initialized is False:
+                verbose_proxy_logger.warning(
+                    "Skipping budget reservation for %s because window spend could not be loaded",
+                    counter.counter_key,
+                )
+                raise _CounterReservationUnavailable
 
-    reserved_value = await _increment_spend_counter_cache(
-        counter_key=counter.counter_key,
-        increment=reservation_cost,
-    )
-    return float(reserved_value) if reserved_value is not None else None
+        reserved_value = await _increment_spend_counter_cache(
+            counter_key=counter.counter_key,
+            increment=reservation_cost,
+        )
+        return float(reserved_value) if reserved_value is not None else None
+    except _CounterReservationUnavailable:
+        raise
+    except Exception:
+        verbose_proxy_logger.warning(
+            "Skipping budget reservation for %s because spend counter reservation failed",
+            counter.counter_key,
+            exc_info=True,
+        )
+        try:
+            await _invalidate_spend_counter(counter_key=counter.counter_key)
+        except Exception:
+            verbose_proxy_logger.warning(
+                "Failed to invalidate spend counter after budget reservation failure for %s",
+                counter.counter_key,
+                exc_info=True,
+            )
+        raise _CounterReservationUnavailable
 
 
 async def _get_current_counter_value(counter: _BudgetCounter) -> float:

tests/test_litellm/proxy/test_budget_reservation.py

@@ -891,6 +891,102 @@ async def test_should_skip_window_reservation_when_db_baseline_unavailable(
     )
 
 
+@pytest.mark.asyncio
+async def test_should_skip_reservation_when_counter_increment_fails(
+    spend_counter_state,
+    monkeypatch,
+):
+    counter_cache, key_cache = spend_counter_state
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
+    valid_token = UserAPIKeyAuth(
+        token="key-budget-reserve-unavailable",
+        spend=0.0,
+        max_budget=1.0,
+    )
+
+    async def fail_increment_cache(*args, **kwargs):
+        raise RuntimeError("counter unavailable")
+
+    monkeypatch.setattr(counter_cache, "async_increment_cache", fail_increment_cache)
+
+    with (
+        patch(
+            "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
+            return_value=0.5,
+        ),
+        patch(
+            "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.warning"
+        ) as mock_warning,
+    ):
+        reservation = await reserve_budget_for_request(
+            request_body=_request_body(),
+            route="/chat/completions",
+            llm_router=None,
+            valid_token=valid_token,
+            team_object=None,
+            user_object=None,
+            prisma_client=None,
+            user_api_key_cache=key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+    assert reservation is None
+    assert mock_warning.call_count >= 1
+    assert (
+        counter_cache.in_memory_cache.get_cache(
+            key="spend:key:key-budget-reserve-unavailable"
+        )
+        is None
+    )
+
+
+@pytest.mark.asyncio
+async def test_should_skip_reservation_when_counter_initialization_fails(
+    spend_counter_state,
+):
+    counter_cache, key_cache = spend_counter_state
+    proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
+    valid_token = UserAPIKeyAuth(
+        token="key-budget-reserve-init-unavailable",
+        spend=0.0,
+        max_budget=1.0,
+    )
+
+    with (
+        patch(
+            "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
+            return_value=0.5,
+        ),
+        patch(
+            "litellm.proxy.proxy_server._ensure_spend_counter_initialized",
+            side_effect=RuntimeError("redis unavailable"),
+        ),
+        patch(
+            "litellm.proxy.spend_tracking.budget_reservation.verbose_proxy_logger.warning"
+        ) as mock_warning,
+    ):
+        reservation = await reserve_budget_for_request(
+            request_body=_request_body(),
+            route="/chat/completions",
+            llm_router=None,
+            valid_token=valid_token,
+            team_object=None,
+            user_object=None,
+            prisma_client=None,
+            user_api_key_cache=key_cache,
+            proxy_logging_obj=proxy_logging_obj,
+        )
+
+    assert reservation is None
+    assert mock_warning.call_count >= 1
+    assert (
+        counter_cache.in_memory_cache.get_cache(
+            key="spend:key:key-budget-reserve-init-unavailable"
+        )
+        is None
+    )
+
+
 @pytest.mark.asyncio
 async def test_should_not_re_read_uncapped_budget_after_reservation_fallback(
     spend_counter_state,