diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 0ec17bbea5d..851f5438dad 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -544,6 +544,65 @@ class SlackAlerting(CustomBatchLogger): ttl=self.alerting_args.budget_alert_ttl, ) + async def _try_claim_budget_alert(self, cache: DualCache, cache_key: str) -> bool: + """ + Claim the alert key before sending so concurrent requests do not all alert. + + Redis-backed caches use SET NX EX for an atomic cross-process claim. The + fallback path preserves the existing best-effort behavior for local-only + caches. + """ + ttl = self.alerting_args.budget_alert_ttl + + if cache.redis_cache is not None: + try: + redis_set_result = await cache.redis_cache.async_set_cache( + key=cache_key, + value="SENT", + ttl=ttl, + nx=True, + ) + if redis_set_result: + if cache.in_memory_cache is not None: + await cache.async_set_cache( + key=cache_key, + value="SENT", + ttl=ttl, + local_only=True, + ) + return True + + cached_result = await cache.async_get_cache(key=cache_key) + if cached_result is not None: + return False + except Exception as e: + verbose_proxy_logger.exception( + "Failed atomic budget alert claim for cache key %s: %s", + cache_key, + str(e), + ) + + result = await cache.async_get_cache(key=cache_key) + if result is not None: + return False + + await cache.async_set_cache( + key=cache_key, + value="SENT", + ttl=ttl, + ) + return True + + async def _clear_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None: + try: + await cache.async_delete_cache(cache_key) + except Exception as e: + verbose_proxy_logger.exception( + "Failed clearing budget alert claim for cache key %s: %s", + cache_key, + str(e), + ) + async def budget_alerts( self, type: Literal[ @@ -613,25 +672,30 @@ class SlackAlerting(CustomBatchLogger): # send alert if event is not None and user_info.event_group is not None: _cache_key = "budget_alerts:{}:{}".format(event, _id) - result = await _cache.async_get_cache(key=_cache_key) - if result is None: + should_send_alert = await self._try_claim_budget_alert( + cache=_cache, + cache_key=_cache_key, + ) + if should_send_alert: webhook_event = WebhookEvent( event=event, event_message=event_message, **user_info_json, ) - await self.send_alert( - message=event_message + "\n\n" + user_info_str, - level="High", - alert_type=AlertType.budget_alerts, - user_info=webhook_event, - alerting_metadata={}, - ) - await _cache.async_set_cache( - key=_cache_key, - value="SENT", - ttl=self.alerting_args.budget_alert_ttl, - ) + try: + await self.send_alert( + message=event_message + "\n\n" + user_info_str, + level="High", + alert_type=AlertType.budget_alerts, + user_info=webhook_event, + alerting_metadata={}, + ) + except Exception: + await self._clear_budget_alert_claim( + cache=_cache, + cache_key=_cache_key, + ) + raise return return diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py index 134056de807..1f2f23cf6a0 100644 --- a/tests/logging_callback_tests/test_alerting.py +++ b/tests/logging_callback_tests/test_alerting.py @@ -226,6 +226,72 @@ async def test_budget_alerts_crossed_again(slack_alerting): mock_send_alert.assert_not_awaited() +@pytest.mark.asyncio +async def test_budget_alerts_concurrent_requests_use_atomic_redis_claim(): + class AtomicRedisCache: + def __init__(self): + self.values = {} + self.lock = asyncio.Lock() + self.set_calls = [] + + async def async_set_cache(self, key, value, **kwargs): + self.set_calls.append((key, kwargs)) + await asyncio.sleep(0) + async with self.lock: + if kwargs.get("nx") is True and key in self.values: + return None + self.values[key] = value + return True + + async def async_get_cache(self, key, **kwargs): + await asyncio.sleep(0) + return self.values.get(key) + + async def async_delete_cache(self, key): + self.values.pop(key, None) + + async def slow_send_alert(*args, **kwargs): + await asyncio.sleep(0.01) + + redis_cache = AtomicRedisCache() + slack_alerting = SlackAlerting( + alerting=["slack"], + alerting_args={"budget_alert_ttl": 60}, + internal_usage_cache=DualCache(redis_cache=redis_cache), + ) + user_info = CallInfo( + token="test_token", + spend=101, + max_budget=100, + team_id="team-atomic", + event_group=Litellm_EntityType.TEAM, + ) + + with patch.object( + slack_alerting, + "send_alert", + new=AsyncMock(side_effect=slow_send_alert), + ) as mock_send_alert: + await asyncio.gather( + *( + slack_alerting.budget_alerts( + "team_budget", + user_info=user_info, + ) + for _ in range(20) + ) + ) + + mock_send_alert.assert_awaited_once() + + cache_key = "budget_alerts:budget_crossed:team-atomic" + assert redis_cache.values[cache_key] == "SENT" + assert len(redis_cache.set_calls) == 20 + assert all(call[0] == cache_key for call in redis_cache.set_calls) + assert all(call[1]["nx"] is True for call in redis_cache.set_calls) + assert all(call[1]["ttl"] == 60 for call in redis_cache.set_calls) + + # Test for send_alert - should be called once @pytest.mark.asyncio async def test_send_alert(slack_alerting):