diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index af5441c2e94..9826d05d1d4 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -185,6 +185,7 @@ class ConfigSyncSubscriber: "_monotonic", "_redis_cache", "_resync_callbacks", + "_resync_on_subscribe", "_rng", "_sleep", "_task", @@ -205,6 +206,7 @@ class ConfigSyncSubscriber: ) -> None: self._redis_cache = redis_cache self._resync_callbacks = resync_callbacks + self._resync_on_subscribe = False self._debounce_seconds = debounce_seconds self._jitter_max_seconds = jitter_max_seconds self._min_resync_interval_seconds = min_resync_interval_seconds @@ -241,12 +243,16 @@ class ConfigSyncSubscriber: try: await pubsub.subscribe(config_sync_channel(self._redis_cache)) backoff_seconds = self._backoff_initial_seconds + if self._resync_on_subscribe: + await self._resync(pubsub) + self._resync_on_subscribe = False await self._consume(pubsub) finally: await self._close_pubsub(pubsub) except asyncio.CancelledError: raise except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect + self._resync_on_subscribe = True verbose_proxy_logger.warning( "config sync subscriber redis error: %s; reconnecting in %.0fs", e, @@ -261,10 +267,14 @@ class ConfigSyncSubscriber: if message is None: continue await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)) - await self._wait_for_min_resync_interval() + await self._resync(pubsub) + + async def _resync(self, pubsub: ConfigSyncPubSub | None = None) -> None: + await self._wait_for_min_resync_interval() + if pubsub is not None: await self._drain_pending(pubsub) - await self._run_resync_callbacks() - self._last_resync_at = self._monotonic() + await self._run_resync_callbacks() + self._last_resync_at = self._monotonic() async def _wait_for_min_resync_interval(self) -> None: if self._last_resync_at is None: diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index 83ed3afa293..7c31c1aeb9f 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -1,6 +1,7 @@ import asyncio import json import random +import typing from typing import Callable, Coroutine, Iterable, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch @@ -433,6 +434,52 @@ async def test_writes_during_the_throttle_wait_collapse_into_the_next_resync() - assert pubsub.queue.empty() +async def test_write_arriving_during_throttle_wait_does_not_trigger_an_extra_resync() -> None: + pubsub: typing.Final = _QueuePubSub() + cache: typing.Final = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + clock: typing.Final = _FakeClock() + first_resync: typing.Final = asyncio.Event() + second_resync: typing.Final = asyncio.Event() + inserted_during_wait: typing.Final = asyncio.Event() + resyncs: typing.Final[List[str]] = [] + + async def recording_sleep(seconds: float) -> None: + if seconds > 0 and not inserted_during_wait.is_set(): + inserted_during_wait.set() + pubsub.queue.put_nowait("change-during-wait") + await asyncio.sleep(0) + + async def resync() -> None: + resyncs.append("resync") + if len(resyncs) == 1: + first_resync.set() + if len(resyncs) == 2: + second_resync.set() + + subscriber: typing.Final = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(resync,), + debounce_seconds=0.0, + jitter_max_seconds=0.0, + min_resync_interval_seconds=10.0, + sleep=recording_sleep, + monotonic=clock, + ) + + subscriber.start() + pubsub.queue.put_nowait("first-change") + await asyncio.wait_for(first_resync.wait(), timeout=5) + clock.now += 4.0 + pubsub.queue.put_nowait("second-change") + await asyncio.wait_for(second_resync.wait(), timeout=5) + await asyncio.sleep(0.05) + await subscriber.stop() + + assert inserted_during_wait.is_set() + assert resyncs == ["resync", "resync"] + assert pubsub.queue.empty() + + async def test_polls_without_messages_do_not_trigger_resyncs() -> None: pubsub = _EmptyPollsThenMessagePubSub(empty_polls=3, initial_messages=["change"]) cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) @@ -493,15 +540,21 @@ async def test_second_start_does_not_open_a_second_subscription() -> None: assert pubsub.subscribed_channels == [CONFIG_SYNC_CHANNEL] -async def test_redis_error_leads_to_backoff_and_resubscribe() -> None: +async def test_redis_error_resubscribes_then_resyncs_without_new_message() -> None: broken = _BrokenPubSub() - healthy = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_credentialstable"})]) + healthy = _QueuePubSub() cache = _FakeRedisCache(_ScriptedPubSubRedisClient([broken, healthy])) resyncs: List[str] = [] fired = asyncio.Event() + + async def resync() -> None: + resyncs.append("resync") + if healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL]: + fired.set() + subscriber = ConfigSyncSubscriber( redis_cache=cache, - resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + resync_callbacks=(resync,), debounce_seconds=0.01, jitter_max_seconds=0.0, backoff_initial_seconds=0.02,