Fix duplicate budget alerts under concurrent traffic

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

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):