diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 8d0d044ff93..17ec3ed787d 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/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) diff --git a/litellm/router.py b/litellm/router.py index 8960cd92cd8..da92eef9102 100644 --- a/litellm/router.py +++ b/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) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 686d57e2b77..79d457f2836 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/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: diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 2d5eb78950c..b9e5ff2eeb7 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/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 diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 425f68dda18..534aea47885 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/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