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 5fee1b28e71..bb5424b0e90 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -24,15 +24,6 @@ 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: @@ -115,20 +106,39 @@ end cronjob_id, ) lock_key = PodLockManager.get_redis_lock_key(cronjob_id) - 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, - ) + + 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, + ) else: verbose_proxy_logger.debug( - "Pod %s failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", + "Pod %s attempted to release Redis lock for cronjob_id=%s, but no lock was found", self.pod_id, cronjob_id, ) @@ -137,31 +147,6 @@ end 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 7790961eb16..e83fd75c3a0 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,42 +307,3 @@ 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