mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
Merge f69ff24678 into f285229b51
This commit is contained in:
commit
91920b990b
2 changed files with 41 additions and 6 deletions
|
|
@ -171,6 +171,7 @@ class ConfigSyncSubscriber:
|
|||
"_monotonic",
|
||||
"_redis_cache",
|
||||
"_resync_callbacks",
|
||||
"_resync_required",
|
||||
"_rng",
|
||||
"_sleep",
|
||||
"_task",
|
||||
|
|
@ -201,6 +202,7 @@ class ConfigSyncSubscriber:
|
|||
self._monotonic = monotonic
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
self._last_resync_at: float | None = None
|
||||
self._resync_required = False
|
||||
|
||||
def start(self) -> None:
|
||||
if self._task is not None:
|
||||
|
|
@ -226,6 +228,9 @@ class ConfigSyncSubscriber:
|
|||
pubsub = client.pubsub()
|
||||
try:
|
||||
await pubsub.subscribe(config_sync_channel(self._redis_cache))
|
||||
if self._resync_required:
|
||||
await self._debounce_and_resync(pubsub)
|
||||
self._resync_required = False
|
||||
backoff_seconds = self._backoff_initial_seconds
|
||||
await self._consume(pubsub)
|
||||
finally:
|
||||
|
|
@ -233,6 +238,7 @@ class ConfigSyncSubscriber:
|
|||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect
|
||||
self._resync_required = True
|
||||
verbose_proxy_logger.warning(
|
||||
"config sync subscriber redis error: %s; reconnecting in %.0fs",
|
||||
e,
|
||||
|
|
@ -246,11 +252,14 @@ class ConfigSyncSubscriber:
|
|||
message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS)
|
||||
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._drain_pending(pubsub)
|
||||
await self._run_resync_callbacks()
|
||||
self._last_resync_at = self._monotonic()
|
||||
await self._debounce_and_resync(pubsub)
|
||||
|
||||
async def _debounce_and_resync(self, pubsub: _ConfigSyncPubSub) -> None:
|
||||
await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds))
|
||||
await self._wait_for_min_resync_interval()
|
||||
await self._drain_pending(pubsub)
|
||||
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:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from typing import Callable, Coroutine, Iterable, List, Optional, Tuple
|
||||
from typing import Callable, Coroutine, Final, Iterable, List, Optional, Tuple
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -521,6 +521,32 @@ async def test_redis_error_leads_to_backoff_and_resubscribe() -> None:
|
|||
assert resyncs == ["resync"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("broken_type", (_BrokenPubSub, _CloseFailingBrokenPubSub))
|
||||
async def test_reconnect_recovers_missed_config_without_another_message(broken_type: type[_BrokenPubSub]) -> None:
|
||||
broken: Final = broken_type()
|
||||
healthy: Final = _QueuePubSub()
|
||||
cache: Final = _FakeRedisCache(_ScriptedPubSubRedisClient((broken, healthy)))
|
||||
recovered: Final = asyncio.Event()
|
||||
reload_config: Final = AsyncMock(side_effect=recovered.set)
|
||||
subscriber: Final = ConfigSyncSubscriber(
|
||||
redis_cache=cache,
|
||||
resync_callbacks=(reload_config,),
|
||||
debounce_seconds=0.0,
|
||||
jitter_max_seconds=0.0,
|
||||
backoff_initial_seconds=0.0,
|
||||
backoff_max_seconds=0.0,
|
||||
)
|
||||
|
||||
subscriber.start()
|
||||
try:
|
||||
await asyncio.wait_for(recovered.wait(), timeout=1)
|
||||
finally:
|
||||
await subscriber.stop()
|
||||
|
||||
assert healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL]
|
||||
reload_config.assert_awaited_once_with()
|
||||
|
||||
|
||||
async def test_failing_resync_callback_does_not_kill_subscriber() -> None:
|
||||
pubsub = _QueuePubSub()
|
||||
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue