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 bb5424b0e90..5fee1b28e71 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -24,6 +24,15 @@ class PodLockManager: def __init__(self, redis_cache: Optional[RedisCache] = None): self.pod_id = str(uuid.uuid4()) self.redis_cache = redis_cache + self._release_lock_script: Optional[Any] = None + + _COMPARE_AND_DELETE_LOCK_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" @staticmethod def get_redis_lock_key(cronjob_id: str) -> str: @@ -106,39 +115,20 @@ class PodLockManager: cronjob_id, ) lock_key = PodLockManager.get_redis_lock_key(cronjob_id) - - current_value = await self.redis_cache.async_get_cache(lock_key) - if current_value is not None: - if isinstance(current_value, bytes): - current_value = current_value.decode("utf-8") - if current_value == self.pod_id: - result = await self.redis_cache.async_delete_cache(lock_key) - if result == 1: - 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, - ) - else: - verbose_proxy_logger.debug( - "Pod %s failed to release Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - else: - verbose_proxy_logger.debug( - "Pod %s cannot release Redis lock for cronjob_id=%s because it is held by pod %s", - self.pod_id, - cronjob_id, - current_value, - ) + result = await self._compare_and_delete_lock(lock_key=lock_key) + if result == 1: + 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, + ) else: verbose_proxy_logger.debug( - "Pod %s attempted to release Redis lock for cronjob_id=%s, but no lock was found", + "Pod %s failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", self.pod_id, cronjob_id, ) @@ -147,6 +137,31 @@ class PodLockManager: f"Error releasing Redis lock for {cronjob_id}: {e}" ) + async def _compare_and_delete_lock(self, lock_key: str) -> int: + """ + Atomically delete lock key only if current pod owns it. + + Falls back to get/delete for non-RedisCache implementations that do not + expose Lua script registration. + """ + script_register = getattr(self.redis_cache, "async_register_script", None) + if callable(script_register): + if self._release_lock_script is None: + self._release_lock_script = script_register( + self._COMPARE_AND_DELETE_LOCK_SCRIPT + ) + script_callable = self._release_lock_script + result = await script_callable(keys=[lock_key], args=[self.pod_id]) + return int(result or 0) + + current_value = await self.redis_cache.async_get_cache(lock_key) # type: ignore + if isinstance(current_value, bytes): + current_value = current_value.decode("utf-8") + if current_value != self.pod_id: + return 0 + result = await self.redis_cache.async_delete_cache(lock_key) # type: ignore + return int(result or 0) + @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): asyncio.create_task( diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index e83fd75c3a0..7790961eb16 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -307,3 +307,42 @@ async def test_lock_takeover_race_condition(mock_redis): cronjob_id="test_job", ) assert result2 == False + + +@pytest.mark.asyncio +async def test_release_lock_uses_atomic_compare_delete_script_when_available( + pod_lock_manager, mock_redis +): + """ + Test that release_lock prefers atomic compare-and-delete Lua script when + redis cache exposes script registration. + """ + script_callable = AsyncMock(return_value=1) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + await pod_lock_manager.release_lock(cronjob_id="test_job") + + lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") + mock_redis.async_register_script.assert_called_once_with( + PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT + ) + script_callable.assert_called_once_with( + keys=[lock_key], args=[pod_lock_manager.pod_id] + ) + mock_redis.async_get_cache.assert_not_called() + mock_redis.async_delete_cache.assert_not_called() + + +@pytest.mark.asyncio +async def test_release_lock_reuses_registered_script(pod_lock_manager, mock_redis): + """ + Test script registration is cached on manager instance and reused. + """ + script_callable = AsyncMock(return_value=0) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + await pod_lock_manager.release_lock(cronjob_id="test_job") + await pod_lock_manager.release_lock(cronjob_id="test_job") + + assert mock_redis.async_register_script.call_count == 1 + assert script_callable.call_count == 2