mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Revert "fix(pod-lock): make release lock compare-and-delete atomic (#21226)"
This reverts commit f162371b93.
This commit is contained in:
parent
8e8511a2a3
commit
53dcebc37a
2 changed files with 31 additions and 85 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue