diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index f04616fb5d6..ac4a4fe68e7 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -123,19 +123,37 @@ class PodLockManager: verbose_proxy_logger.debug("redis_cache is None, skipping release_lock") return verbose_proxy_logger.debug("Pod %s attempting to release Redis lock for cronjob_id=%s", self.pod_id, cronjob_id) - released: Final = await store.release(cronjob_id, self.pod_id) - if released is True: + released: Final = await self._release(store, cronjob_id) + if isinstance(released, RedisFailure): + verbose_proxy_logger.error("Error releasing Redis lock for %s: %s", cronjob_id, released.message) + elif released: verbose_proxy_logger.info( "Pod %s successfully released Redis lock for cronjob_id=%s", self.pod_id, cronjob_id ) self._emit_released_lock_event(cronjob_id=cronjob_id, pod_id=self.pod_id) - elif released is False: + else: verbose_proxy_logger.debug( "Pod %s failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", self.pod_id, cronjob_id, ) + async def _release(self, store: PodLockStore, cronjob_id: str) -> bool | RedisFailure: + released: Final = await store.release(cronjob_id, self.pod_id) + if not isinstance(released, RedisFailure): + return released + verbose_proxy_logger.warning( + "Lua compare-and-delete failed for lock_key=%s, falling back to GET+DEL", + self.get_redis_lock_key(cronjob_id), + ) + match await store.holder(cronjob_id): + case RedisFailure(kind="breaker_open") as refused: + return refused + case HeldBy(owner=owner) if owner == self.pod_id: + return await store.remove(cronjob_id) + case _: + return False + @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): asyncio.create_task( diff --git a/litellm/redis/stores/pod_lock.py b/litellm/redis/stores/pod_lock.py index 13ac7c613db..9e86a5423d2 100644 --- a/litellm/redis/stores/pod_lock.py +++ b/litellm/redis/stores/pod_lock.py @@ -46,6 +46,8 @@ class PodLockStore(Protocol): async def release(self, cronjob_id: str, owner: str) -> bool | RedisFailure: ... + async def remove(self, cronjob_id: str) -> bool | RedisFailure: ... + @dataclass(frozen=True, slots=True) class TransportPodLockStore: @@ -75,6 +77,10 @@ class TransportPodLockStore: reply: Final = await self.transport.eval(_RELEASE, (lock_key(cronjob_id),), (_encode(owner),)) return reply if isinstance(reply, RedisFailure) else reply == 1 + async def remove(self, cronjob_id: str) -> bool | RedisFailure: + deleted: Final = await self.transport.delete(lock_key(cronjob_id)) + return deleted if isinstance(deleted, RedisFailure) else deleted == 1 + def _encode(owner: str) -> str: return json.dumps(owner) diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index 8af6cf88a14..cc5fd18a305 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -17,7 +17,10 @@ def _admit_directly(call: Callable[[], Awaitable[_T]]) -> Awaitable[_T]: def fake_redis_cache( - server: fakeredis.FakeServer, namespace: str | None = None, service_logger: ServiceLogging | None = None + server: fakeredis.FakeServer, + namespace: str | None = None, + service_logger: ServiceLogging | None = None, + client: type[fakeredis.FakeAsyncRedis] = fakeredis.FakeAsyncRedis, ) -> RedisCache: cache = RedisCache( host="127.0.0.1", @@ -28,7 +31,7 @@ def fake_redis_cache( ) litellm.in_memory_llm_clients_cache.set_cache( key=cache._get_async_client_cache_key(), # pyright: ignore[reportPrivateUsage] # seeds the per-loop client slot RedisCache reads - value=fakeredis.FakeAsyncRedis(server=server), + value=client(server=server), ) return cache diff --git a/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py index abe88b913ae..22e88a29ed7 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -5,6 +5,7 @@ from unittest.mock import patch import fakeredis import pytest +from redis.exceptions import ResponseError from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS @@ -127,6 +128,30 @@ async def test_release_survives_a_script_flush(server: fakeredis.FakeServer, raw assert await raw.exists(LOCK_KEY) == 0 +class _ScriptingDisabled(fakeredis.FakeAsyncRedis): + async def eval(self, *args: object) -> object: + raise ResponseError("ERR scripting is disabled") + + async def evalsha(self, *args: object) -> object: + raise ResponseError("ERR scripting is disabled") + + +@pytest.mark.asyncio +async def test_release_falls_back_to_get_and_delete_when_scripts_cannot_run(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server, client=_ScriptingDisabled) + holder = PodLockManager(redis_cache=cache) + other = PodLockManager(redis_cache=cache) + assert await holder.acquire_lock(cronjob_id="test_job") is True + + await other.release_lock(cronjob_id="test_job") + assert await raw.get(LOCK_KEY) == json.dumps(holder.pod_id).encode() + + with patch.object(holder, "_emit_released_lock_event") as emitted: + await holder.release_lock(cronjob_id="test_job") + assert await raw.exists(LOCK_KEY) == 0 + emitted.assert_called_once_with(cronjob_id="test_job", pod_id=holder.pod_id) + + @pytest.mark.asyncio async def test_release_emits_the_released_event_only_when_it_deleted_the_lock(server: fakeredis.FakeServer) -> None: cache = fake_redis_cache(server)