fix(proxy): stop leaking periodic tasks on every DB config reload (#42784)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -118,6 +118,7 @@ class SlackAlerting(CustomBatchLogger):
         self.default_webhook_url = default_webhook_url
         self.flush_lock = asyncio.Lock()
         self.periodic_started = False
+        self._periodic_flush_task: asyncio.Task[None] | None = None
         self.hanging_request_check = AlertingHangingRequestCheck(
             slack_alerting_object=self,
         )
@@ -129,6 +130,12 @@ class SlackAlerting(CustomBatchLogger):
         self.digest_lock = asyncio.Lock()
         super().__init__(**kwargs, flush_lock=self.flush_lock)
 
+    def _ensure_periodic_flush_task(self) -> None:
+        if self.periodic_started and (self._periodic_flush_task is None or not self._periodic_flush_task.done()):
+            return
+        self._periodic_flush_task = asyncio.create_task(self.periodic_flush())
+        self.periodic_started = True
+
     def update_values(
         self,
         alerting: list | None = None,
@@ -141,17 +148,14 @@ class SlackAlerting(CustomBatchLogger):
     ):
         if alerting is not None:
             self.alerting = alerting
-            asyncio.create_task(self.periodic_flush())
-            self.periodic_started = True
+            self._ensure_periodic_flush_task()
         if alerting_threshold is not None:
             self.alerting_threshold = alerting_threshold
         if alert_types is not None:
             self.alert_types = alert_types
         if alerting_args is not None:
             self.alerting_args = SlackAlertingArgs(**alerting_args)
-            if not self.periodic_started:
-                asyncio.create_task(self.periodic_flush())
-                self.periodic_started = True
+            self._ensure_periodic_flush_task()
         if alert_type_config is not None:
             for key, val in alert_type_config.items():
                 self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val
@@ -1446,9 +1450,8 @@ Model Info:
             return
 
         # Start periodic flush if not already started
-        if not self.periodic_started and self.alerting is not None and len(self.alerting) > 0:
-            asyncio.create_task(self.periodic_flush())
-            self.periodic_started = True
+        if self.alerting is not None and len(self.alerting) > 0:
+            self._ensure_periodic_flush_task()
 
         if "webhook" in self.alerting and alert_type == "budget_alerts" and user_info is not None:
             await self.send_webhook_alert(webhook_event=user_info)

litellm/router.py

@@ -118,6 +118,7 @@ from litellm.llms.openai_like.model_info import (
     MODEL_INFO_REFRESH_SECONDS,
     get_openai_compatible_model_info,
 )
+from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
 from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
 from litellm.router_strategy.complexity_router.context_compaction import (
     arm_compaction,
@@ -1377,6 +1378,9 @@ class Router:
         `_init_routing_groups`) so repeated `update_settings` calls don't
         accumulate dead selectors that keep receiving callback events.
         """
+        for selector in selectors:
+            if isinstance(selector, BaseRoutingStrategy):
+                selector.retire()
         selector_ids: Final = {id(s) for s in selectors if s is not None}
         if not selector_ids:
             return
@@ -12117,7 +12121,7 @@ class Router:
                                 )
                             rebuild_routing_groups = True
                     elif var == "routing_strategy_args":
-                        routing_args_updated = True
+                        routing_args_updated = value != self.routing_strategy_args
                     setattr(self, var, value)
             else:
                 verbose_router_logger.debug("Setting %s is not allowed", var)

litellm/router_strategy/base_routing_strategy.py

@@ -40,10 +40,24 @@ class BaseRoutingStrategy(ABC):
             self.periodic_sync_in_memory_spend_with_redis(default_sync_interval=default_sync_interval)
         )
 
+    def cancel_sync_task(self) -> None:
+        if self._sync_task is not None:
+            self._sync_task.cancel()
+
+    def retire(self) -> None:
+        self.cancel_sync_task()
+        if not self.redis_increment_operation_queue:
+            return
+        try:
+            loop: Final = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+        loop.create_task(self._push_in_memory_increments_to_redis())
+
     async def cleanup(self):
         """Cleanup method to be called when shutting down"""
         if self._sync_task is not None:
-            self._sync_task.cancel()
+            self.cancel_sync_task()
             try:
                 await self._sync_task
             except asyncio.CancelledError:

tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py

@@ -526,3 +526,29 @@ async def test_async_send_batch_collapses_only_identical_alerts() -> None:
         {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"},
         {"text": CROSSED_ALERT},
     )
+
+
+def _periodic_flush_tasks() -> list[asyncio.Task[object]]:
+    return [
+        t
+        for t in asyncio.all_tasks()
+        if t.get_coro() is not None and t.get_coro().__qualname__ == "SlackAlerting.periodic_flush"
+    ]
+
+
+@pytest.mark.asyncio
+async def test_update_values_repeated_alerting_reload_keeps_single_periodic_flush_task() -> None:
+    slack_alerting: Final = SlackAlerting(alerting=["slack"])
+    try:
+        for _ in range(5):
+            slack_alerting.update_values(alerting=["slack"])
+        await asyncio.sleep(0)
+        flush_tasks: Final = _periodic_flush_tasks()
+        assert len(flush_tasks) == 1, f"expected 1 periodic_flush task, found {len(flush_tasks)}"
+    finally:
+        for t in _periodic_flush_tasks():
+            t.cancel()
+            try:
+                await t
+            except asyncio.CancelledError:
+                pass

tests/test_litellm/router_strategy/test_router_routing_groups.py

@@ -18,6 +18,7 @@ from pydantic import ValidationError
 
 import litellm
 from litellm import Router
+from litellm.caching.redis_cache import RedisPipelineIncrementOperation
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy
 from litellm.utils import Rules, function_setup
@@ -2149,3 +2150,107 @@ async def test_caller_cannot_spoof_a_priority_group_to_bypass_fallback_gates(
             **{metadata_bucket: {"pre_routing_selected_model": "priority-group"}},
         )
     assert checked == ["priority-group"]
+
+
+def _sync_task_count() -> int:
+    return sum(
+        1
+        for t in asyncio.all_tasks()
+        if t.get_coro() is not None
+        and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis"
+    )
+
+
+@pytest.mark.asyncio
+async def test_update_settings_same_routing_strategy_args_does_not_leak_sync_tasks(monkeypatch) -> None:
+    monkeypatch.setattr(litellm, "callbacks", [])
+    monkeypatch.setattr(litellm, "input_callback", [])
+    router: Final = Router(
+        model_list=_model_list(),
+        routing_strategy="usage-based-routing-v2",
+        routing_strategy_args={"ttl": 60},
+    )
+    try:
+        assert _sync_task_count() == 1
+        selector_before: Final = router.lowesttpm_logger_v2
+
+        for _ in range(5):
+            router.update_settings(routing_strategy_args={"ttl": 60})
+        await asyncio.sleep(0)
+        assert _sync_task_count() == 1
+        assert router.lowesttpm_logger_v2 is selector_before, "same routing_strategy_args must not rebuild the selector"
+
+        router.update_settings(routing_strategy_args={"ttl": 120})
+        await asyncio.sleep(0)
+        assert _sync_task_count() == 1
+        assert router.lowesttpm_logger_v2.routing_args.ttl == 120
+    finally:
+        for t in [
+            t
+            for t in asyncio.all_tasks()
+            if t.get_coro() is not None
+            and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis"
+        ]:
+            t.cancel()
+            try:
+                await t
+            except asyncio.CancelledError:
+                pass
+
+
+class _RecordingRedisCache:
+    def __init__(self) -> None:
+        self.increment_lists: list[list[RedisPipelineIncrementOperation]] = []
+
+    async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]:
+        self.increment_lists.append(list(increment_list))
+        return [float(op["increment_value"]) for op in increment_list]
+
+
+@pytest.mark.asyncio
+async def test_update_settings_changed_routing_strategy_args_flushes_replaced_selector_queue(
+    monkeypatch,
+) -> None:
+    monkeypatch.setattr(litellm, "callbacks", [])
+    monkeypatch.setattr(litellm, "input_callback", [])
+    router: Final = Router(
+        model_list=_model_list(),
+        routing_strategy="usage-based-routing-v2",
+        routing_strategy_args={"ttl": 60},
+    )
+    try:
+        redis_cache: Final = _RecordingRedisCache()
+        replaced: Final = router.lowesttpm_logger_v2
+        replaced.dual_cache.redis_cache = redis_cache
+        replaced.redis_increment_operation_queue.append(
+            RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)
+        )
+
+        router.update_settings(routing_strategy_args={"ttl": 120})
+        await asyncio.sleep(0)
+        await asyncio.gather(
+            *(
+                t
+                for t in asyncio.all_tasks()
+                if t.get_coro() is not None
+                and t.get_coro().__qualname__ == "BaseRoutingStrategy._push_in_memory_increments_to_redis"
+            )
+        )
+
+        assert router.lowesttpm_logger_v2 is not replaced
+        assert redis_cache.increment_lists == [
+            [RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)]
+        ]
+        assert replaced.redis_increment_operation_queue == []
+    finally:
+        for t in [
+            t
+            for t in asyncio.all_tasks()
+            if t.get_coro() is not None
+            and t.get_coro().__qualname__ == "BaseRoutingStrategy.periodic_sync_in_memory_spend_with_redis"
+        ]:
+            t.cancel()
+            try:
+                await t
+            except asyncio.CancelledError:
+                pass