From 1bd2ff7e6331e1fa4167ac7395b617a38e9a7834 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 10:26:05 -0700 Subject: [PATCH] refactor(redis-store): move the pod lock keyspace behind a store Cron job locks now go through the shared Redis transport behind TransportPodLockStore, which owns the lock keys, the owner codec and the compare-and-delete release script. PodLockManager.is_held replaces the three direct lock reads in the budget reset, PTU rollup and spend capture-rate jobs and owns the rule that an unreadable lock counts as unheld. Old and new pods respect each other's locks --- .../proxy/common_utils/reset_budget_job.py | 21 +- .../db_transaction_queue/pod_lock_manager.py | 192 ++---- .../spend_tracking/ptu_flat_cost_rollup.py | 19 +- .../spend_tracking/spend_capture_rate.py | 19 +- litellm/redis/stores/pod_lock.py | 87 +++ .../integrations/focus/test_focus_logger.py | 33 +- .../vantage/test_vantage_logger.py | 33 +- .../common_utils/test_reset_budget_job.py | 11 +- .../test_pod_lock_manager.py | 587 +++++------------- .../proxy/db/test_tool_registry_writer.py | 61 +- .../test_ptu_flat_cost_rollup.py | 20 +- .../spend_tracking/test_spend_capture_rate.py | 5 +- tests/unit/redis/stores/__init__.py | 0 tests/unit/redis/stores/test_pod_lock.py | 91 +++ 14 files changed, 481 insertions(+), 698 deletions(-) create mode 100644 litellm/redis/stores/pod_lock.py create mode 100644 tests/unit/redis/stores/__init__.py create mode 100644 tests/unit/redis/stores/test_pod_lock.py diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..cd6c641ec4d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -49,7 +49,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row from litellm.proxy.db.db_span import db_span, db_spanned -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET, PodLockManager +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -500,23 +500,6 @@ class ResetBudgetJob: new_batch: Final[Callable[[], PrismaBatch]] = self.prisma_client.db.batch_ return new_batch - @with_service_target(POD_LOCK_TARGET) - async def _lease_is_held(self, lock_manager: PodLockManager) -> bool: - """True only when the lease is readable and someone holds it. - - An unreadable lock reports as unheld so the caller sweeps rather than - skipping; being wrong here costs a duplicate sweep, and the alternative - strands every expired budget at its cap. - """ - if lock_manager.redis_cache is None: - return False - try: - lock_key: Final = lock_manager.get_redis_lock_key(RESET_BUDGET_JOB_NAME) - return bool(await lock_manager.redis_cache.async_get_cache(lock_key)) - except Exception as exc: # noqa: BLE001 # an unreadable lease must not strand the sweep - verbose_proxy_logger.warning("Reset budget job: could not read the reset lease: %s", exc) - return False - async def _acquire_lease(self) -> _Lease: """Elect one sweeper per tick. @@ -536,7 +519,7 @@ class ResetBudgetJob: ): return _Lease.LEADER - if await self._lease_is_held(lock_manager): + if await lock_manager.is_held(RESET_BUDGET_JOB_NAME): verbose_proxy_logger.debug("Reset budget job: another pod holds the reset lease, skipping this tick") return _Lease.FOLLOWER 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 1116cd604a3..f04616fb5d6 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -1,14 +1,16 @@ import asyncio -import json -import logging +from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, TypeAlias from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid -from litellm.caching.redis_cache import RedisCache, log_redis_failure +from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj +from litellm.redis.stores import pod_lock +from litellm.redis.stores.pod_lock import Acquired, HeldBy, PodLockStore, TransportPodLockStore, Vacant +from litellm.redis.transport.contract import RedisFailure from litellm.types.services import ServiceTypes POD_LOCK_TARGET: Final = "pod_lock" @@ -28,22 +30,16 @@ class PodLockManager: Ensures that only one pod can run a cron job at a time. """ - _COMPARE_AND_DELETE_LOCK_SCRIPT = """ -if redis.call("get", KEYS[1]) == ARGV[1] then - return redis.call("del", KEYS[1]) -else - return 0 -end -""" - def __init__(self, redis_cache: RedisCache | None = None): self.pod_id = str(uuid.uuid4()) self.redis_cache = redis_cache - self._release_lock_script: Any | None = None @staticmethod def get_redis_lock_key(cronjob_id: str) -> str: - return f"cronjob_lock:{cronjob_id}" + return pod_lock.lock_key(cronjob_id) + + def _store(self) -> PodLockStore | None: + return None if self.redis_cache is None else TransportPodLockStore(self.redis_cache.transport()) @with_service_target(POD_LOCK_TARGET) async def acquire_lock( @@ -54,7 +50,6 @@ end ) -> bool | None: """ Attempt to acquire the lock for a specific cron job using Redis. - Uses the SET command with NX and EX options to ensure atomicity. Args: cronjob_id: The ID of the cron job to lock @@ -66,135 +61,80 @@ end lock marks work as already done for this window, so not even the holder may redo it before the TTL expires. """ - if self.redis_cache is None: + store: Final = self._store() + if store is None: verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") return None - try: - lock_ttl: Final = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS - verbose_proxy_logger.debug( - "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", - self.pod_id, - cronjob_id, - lock_ttl, - ) - # Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX) - # and with an expiration (EX) to avoid deadlocks. - lock_key: Final = PodLockManager.get_redis_lock_key(cronjob_id) - acquired: Final = await self.redis_cache.async_set_cache( - lock_key, - self.pod_id, - nx=True, - ttl=lock_ttl, - ) - if acquired: + lock_ttl: Final = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS + verbose_proxy_logger.debug( + "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", + self.pod_id, + cronjob_id, + lock_ttl, + ) + match await store.acquire(cronjob_id, self.pod_id, ttl=timedelta(seconds=lock_ttl)): + case Acquired(): verbose_proxy_logger.info( - "Pod %s successfully acquired Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, + "Pod %s successfully acquired Redis lock for cronjob_id=%s", self.pod_id, cronjob_id ) - return True - else: - # Check if the current pod already holds the lock - 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 and allow_reentrant: - verbose_proxy_logger.info( - "Pod %s already holds the Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - self._emit_acquired_lock_event(cronjob_id, self.pod_id) - return True - verbose_proxy_logger.info( - "Pod %s could not acquire lock for cronjob_id=%s, held by pod %s.", - self.pod_id, - cronjob_id, - current_value, - ) - return False - except Exception as e: - log_redis_failure(verbose_proxy_logger, logging.ERROR, f"Error acquiring Redis lock for {cronjob_id}", e) - return False + case HeldBy(owner=owner) if owner == self.pod_id and allow_reentrant: + verbose_proxy_logger.info( + "Pod %s already holds the Redis lock for cronjob_id=%s", self.pod_id, cronjob_id + ) + self._emit_acquired_lock_event(cronjob_id, self.pod_id) + return True + case HeldBy(owner=owner): + verbose_proxy_logger.info( + "Pod %s could not acquire lock for cronjob_id=%s, held by pod %s.", self.pod_id, cronjob_id, owner + ) + return False + case Vacant() | RedisFailure(): + return False + @with_service_target(POD_LOCK_TARGET) + async def lock_holder(self, cronjob_id: str) -> HeldBy | None | RedisFailure: + store: Final = self._store() + return None if store is None else await store.holder(cronjob_id) + + @with_service_target(POD_LOCK_TARGET) + async def is_held(self, cronjob_id: str) -> bool: + """An unreadable lock reports as unheld, so callers run a job twice rather than skip it.""" + match await self.lock_holder(cronjob_id): + case HeldBy(): + return True + case None: + return False + case RedisFailure(message=message): + verbose_proxy_logger.warning("Could not read the lock for cronjob_id=%s: %s", cronjob_id, message) + return False + + @with_service_target(POD_LOCK_TARGET) async def release_lock( self, cronjob_id: str, - ): + ) -> None: """ - Release the lock if the current pod holds it. - Uses an atomic Lua compare-and-delete to prevent TOCTOU races where a - stale owner could delete a newly reacquired lock. - Falls back to GET + DEL for cache implementations that don't support - script registration. + Release the lock if the current pod holds it, with an atomic compare-and-delete + so a stale owner can never delete a lock another pod has since taken. """ - if self.redis_cache is None: + store: Final = self._store() + if store is None: verbose_proxy_logger.debug("redis_cache is None, skipping release_lock") return - try: + 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: + 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: verbose_proxy_logger.debug( - "Pod %s attempting to release Redis lock for cronjob_id=%s", + "Pod %s failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", self.pod_id, cronjob_id, ) - lock_key: Final = PodLockManager.get_redis_lock_key(cronjob_id) - result: Final = 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 failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", - self.pod_id, - cronjob_id, - ) - except Exception as e: - log_redis_failure(verbose_proxy_logger, logging.ERROR, f"Error releasing Redis lock for {cronjob_id}", e) - - @with_service_target(POD_LOCK_TARGET) - 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: Final = getattr(self.redis_cache, "async_register_script", None) - if callable(script_register): - try: - if self._release_lock_script is None: - self._release_lock_script = script_register(self._COMPARE_AND_DELETE_LOCK_SCRIPT) - # acquire_lock stores the pod_id via async_set_cache, which - # JSON-encodes the value; compare against the same encoding so - # the Lua equality check matches and the lock is released - result = await self._release_lock_script(keys=[lock_key], args=[json.dumps(self.pod_id)]) - return int(result or 0) - except Exception: - # Lua execution failed (e.g. Redis restart cleared loaded scripts, - # or scripting is disabled). Reset cached script handle and fall - # through to the GET + DEL fallback so the lock is still released. - self._release_lock_script = None - verbose_proxy_logger.warning( - "Lua compare-and-delete failed for lock_key=%s, falling back to GET+DEL", - lock_key, - ) - - current_value = await self.redis_cache.async_get_cache(lock_key) - 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) - return int(result or 0) @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index 64b67ec55dd..b68901ea1a6 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -23,7 +23,6 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter -from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.constants import ( PTU_LAPSED_ALERT_LIMIT, @@ -34,7 +33,6 @@ from litellm.constants import ( PTU_SENTINEL_API_KEY, ) from litellm.litellm_core_utils.ptu_pricing import ptu_terms -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.repositories.model_repository import ModelRepository from litellm.repositories.prisma_protocols import TableActions @@ -705,7 +703,7 @@ async def run_scheduled_ptu_rollup( ) if not await pod_lock_manager.acquire_lock(cronjob_id=PTU_ROLLUP_JOB_ID, ttl=PTU_ROLLUP_LOCK_TTL_SECONDS): - if await _lock_is_held(pod_lock_manager): + if await pod_lock_manager.is_held(PTU_ROLLUP_JOB_ID): verbose_proxy_logger.info("PTU rollup: another pod holds the rollup lock, skipping this run") return None # acquire_lock reports contention and a Redis outage the same way, so an @@ -728,21 +726,6 @@ async def run_scheduled_ptu_rollup( await pod_lock_manager.release_lock(cronjob_id=PTU_ROLLUP_JOB_ID) -@with_service_target(POD_LOCK_TARGET) -async def _lock_is_held(pod_lock_manager: "PodLockManager") -> bool: - """True only when the rollup lock is readable and someone is holding it. - - A Redis that cannot be read is reported as "not held" so the caller runs the day - rather than skipping it; the cost of being wrong here is a duplicate reconcile. - """ - try: - lock_key: Final = pod_lock_manager.get_redis_lock_key(PTU_ROLLUP_JOB_ID) - return bool(await pod_lock_manager.redis_cache.async_get_cache(lock_key)) - except Exception as exc: # noqa: BLE001 # an unreadable lock must not skip the day - verbose_proxy_logger.warning("PTU rollup: could not read the rollup lock: %s", exc) - return False - - async def _run_and_alert( prisma_client: "PrismaClient", *, diff --git a/litellm/proxy/spend_tracking/spend_capture_rate.py b/litellm/proxy/spend_tracking/spend_capture_rate.py index 443010b65da..ff0caf6d79c 100644 --- a/litellm/proxy/spend_tracking/spend_capture_rate.py +++ b/litellm/proxy/spend_tracking/spend_capture_rate.py @@ -14,7 +14,6 @@ from typing import TYPE_CHECKING, Final, TypeAlias from pydantic import ConfigDict, TypeAdapter from typing_extensions import assert_never -from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.constants import ( SPEND_CAPTURE_RATE_CHECK_JOB_ID, @@ -28,7 +27,6 @@ from litellm.llms.openai.organization_costs import ( fetch_openai_daily_costs, provider_billing_get, ) -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import POD_LOCK_TARGET from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.proxy.spend_capture_rate import ( @@ -39,7 +37,6 @@ from litellm.types.proxy.spend_capture_rate import ( ) if TYPE_CHECKING: - from litellm.caching.redis_cache import RedisCache from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.utils import PrismaClient @@ -284,21 +281,9 @@ async def run_scheduled_spend_capture_rate_check( async def _claims_alert_window(pod_lock_manager: "PodLockManager | None") -> bool: """The lock is left to expire, so every replica firing within its TTL of the winner stays quiet.""" - redis_cache: Final = None if pod_lock_manager is None else pod_lock_manager.redis_cache - if pod_lock_manager is None or redis_cache is None: + if pod_lock_manager is None or pod_lock_manager.redis_cache is None: return True acquired: Final = await pod_lock_manager.acquire_lock( cronjob_id=SPEND_CAPTURE_RATE_CHECK_JOB_ID, ttl=SPEND_CAPTURE_RATE_CHECK_LOCK_TTL_SECONDS ) - return acquired or not await _lock_is_held(pod_lock_manager, redis_cache) - - -@with_service_target(POD_LOCK_TARGET) -async def _lock_is_held(pod_lock_manager: "PodLockManager", redis_cache: "RedisCache") -> bool: - try: - return bool( - await redis_cache.async_get_cache(pod_lock_manager.get_redis_lock_key(SPEND_CAPTURE_RATE_CHECK_JOB_ID)) - ) - except Exception as exc: # noqa: BLE001 # an unreadable lock must not silence the alert - verbose_proxy_logger.warning("Spend capture-rate check: could not read the lock: %s", exc) - return False + return acquired or not await pod_lock_manager.is_held(SPEND_CAPTURE_RATE_CHECK_JOB_ID) diff --git a/litellm/redis/stores/pod_lock.py b/litellm/redis/stores/pod_lock.py new file mode 100644 index 00000000000..13ac7c613db --- /dev/null +++ b/litellm/redis/stores/pod_lock.py @@ -0,0 +1,87 @@ +import json +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from pydantic import TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_RELEASE: Final = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" + +_OWNER: Final = TypeAdapter(str) + + +def lock_key(cronjob_id: str) -> str: + return f"cronjob_lock:{cronjob_id}" + + +@dataclass(frozen=True, slots=True) +class Acquired: + pass + + +@dataclass(frozen=True, slots=True) +class HeldBy: + owner: str + + +@dataclass(frozen=True, slots=True) +class Vacant: + pass + + +class PodLockStore(Protocol): + async def acquire( + self, cronjob_id: str, owner: str, *, ttl: timedelta + ) -> Acquired | HeldBy | Vacant | RedisFailure: ... + + async def holder(self, cronjob_id: str) -> HeldBy | None | RedisFailure: ... + + async def release(self, cronjob_id: str, owner: str) -> bool | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportPodLockStore: + transport: RedisTransport + + async def acquire( + self, cronjob_id: str, owner: str, *, ttl: timedelta + ) -> Acquired | HeldBy | Vacant | RedisFailure: + set_reply: Final = await self.transport.set(lock_key(cronjob_id), _encode(owner), ttl=ttl, nx=True) + if isinstance(set_reply, RedisFailure): + return set_reply + if set_reply: + return Acquired() + match await self.holder(cronjob_id): + case None: + return Vacant() + case held: + return held + + async def holder(self, cronjob_id: str) -> HeldBy | None | RedisFailure: + held: Final = await self.transport.get(lock_key(cronjob_id)) + if held is None or isinstance(held, RedisFailure): + return held + return _decode(held) + + async def release(self, cronjob_id: str, owner: str) -> bool | RedisFailure: + reply: Final = await self.transport.eval(_RELEASE, (lock_key(cronjob_id),), (_encode(owner),)) + return reply if isinstance(reply, RedisFailure) else reply == 1 + + +def _encode(owner: str) -> str: + return json.dumps(owner) + + +def _decode(raw: bytes) -> HeldBy | RedisFailure: + try: + return HeldBy(_OWNER.validate_json(raw)) + except ValidationError as e: + return RedisFailure(kind="decode", message=f"pod lock owner is not a JSON string: {e}") diff --git a/tests/unit/integrations/focus/test_focus_logger.py b/tests/unit/integrations/focus/test_focus_logger.py index 5ff5c40df47..2fd76a5f04d 100644 --- a/tests/unit/integrations/focus/test_focus_logger.py +++ b/tests/unit/integrations/focus/test_focus_logger.py @@ -1,29 +1,38 @@ +import json + +import fakeredis import pytest from litellm.integrations.focus.focus_logger import FocusLogger from litellm.proxy import proxy_server +from litellm.redis.transport.contract import RedisTransport +from tests._support.fake_redis import fake_redis_transport FOCUS_LOCK = "cronjob_lock:focus_export_usage_data" VANTAGE_LOCK = "cronjob_lock:vantage_export_usage_data" -class _InMemoryRedis: +class _InMemoryRedis(fakeredis.FakeAsyncRedis): def __init__(self, locks: dict[str, str]) -> None: - self.locks = locks + server = fakeredis.FakeServer() + super().__init__(server=server) + self._sync = fakeredis.FakeRedis(server=server) + for key, owner in locks.items(): + self._sync.set(key, json.dumps(owner)) self.granted: list[tuple[str, str]] = [] - async def async_set_cache(self, key: str, value: str, nx: bool, ttl: int) -> bool: - if nx and key in self.locks: - return False - self.locks[key] = value - self.granted.append((key, value)) - return True + def transport(self) -> RedisTransport: + return fake_redis_transport(self, failure_threshold=1) - async def async_get_cache(self, key: str) -> str | None: - return self.locks.get(key) + @property + def locks(self) -> dict[str, str]: + return {key.decode(): json.loads(self._sync.get(key)) for key in self._sync.scan_iter()} - async def async_delete_cache(self, key: str) -> int: - return 1 if self.locks.pop(key, None) is not None else 0 + async def set(self, name: str, value: str, *, nx: bool, px: int | None) -> bool | None: + stored = await super().set(name, value, nx=nx, px=px) + if stored: + self.granted.append((name, json.loads(value))) + return stored @pytest.mark.asyncio diff --git a/tests/unit/integrations/vantage/test_vantage_logger.py b/tests/unit/integrations/vantage/test_vantage_logger.py index 8ff2c7fa046..0e3e5ea9e8d 100644 --- a/tests/unit/integrations/vantage/test_vantage_logger.py +++ b/tests/unit/integrations/vantage/test_vantage_logger.py @@ -1,29 +1,38 @@ +import json + +import fakeredis import pytest from litellm.integrations.vantage.vantage_logger import VantageLogger from litellm.proxy import proxy_server +from litellm.redis.transport.contract import RedisTransport +from tests._support.fake_redis import fake_redis_transport FOCUS_LOCK = "cronjob_lock:focus_export_usage_data" VANTAGE_LOCK = "cronjob_lock:vantage_export_usage_data" -class _InMemoryRedis: +class _InMemoryRedis(fakeredis.FakeAsyncRedis): def __init__(self, locks: dict[str, str]) -> None: - self.locks = locks + server = fakeredis.FakeServer() + super().__init__(server=server) + self._sync = fakeredis.FakeRedis(server=server) + for key, owner in locks.items(): + self._sync.set(key, json.dumps(owner)) self.granted: list[tuple[str, str]] = [] - async def async_set_cache(self, key: str, value: str, nx: bool, ttl: int) -> bool: - if nx and key in self.locks: - return False - self.locks[key] = value - self.granted.append((key, value)) - return True + def transport(self) -> RedisTransport: + return fake_redis_transport(self, failure_threshold=1) - async def async_get_cache(self, key: str) -> str | None: - return self.locks.get(key) + @property + def locks(self) -> dict[str, str]: + return {key.decode(): json.loads(self._sync.get(key)) for key in self._sync.scan_iter()} - async def async_delete_cache(self, key: str) -> int: - return 1 if self.locks.pop(key, None) is not None else 0 + async def set(self, name: str, value: str, *, nx: bool, px: int | None) -> bool | None: + stored = await super().set(name, value, nx=nx, px=px) + if stored: + self.granted.append((name, json.loads(value))) + return stored @pytest.mark.asyncio diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index 183dd7919c1..1f06f3bdea8 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -12,15 +12,14 @@ import httpx import prisma import pytest - -from litellm.proxy._types import LiteLLM_BudgetTableFull, LiteLLM_VerificationToken -from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MIN_TIME, RESET_BUDGET_JOB_BATCH_SIZE, RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_NAME, ) +from litellm.proxy._types import LiteLLM_BudgetTableFull, LiteLLM_VerificationToken +from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module from litellm.proxy.common_utils.reset_budget_job import ( ResetBudgetJob, _RowReset, @@ -2595,8 +2594,7 @@ class FakePodLockManager: def __init__(self, *, acquired: bool, held_by_other: bool = False, has_redis: bool = True): self.redis_cache = MagicMock() if has_redis else None - if self.redis_cache is not None: - self.redis_cache.async_get_cache = AsyncMock(return_value="another-pod" if held_by_other else None) + self._held = held_by_other self._acquired = acquired self.acquire_calls: List[Dict[str, str | int | None]] = [] self.release_calls: List[str] = [] @@ -2612,6 +2610,9 @@ class FakePodLockManager: async def release_lock(self, cronjob_id: str) -> None: self.release_calls.append(cronjob_id) + async def is_held(self, cronjob_id: str) -> bool: + return self._held + def _make_leader_election_job(monkeypatch, pod_lock_manager): """A ResetBudgetJob wired to one lock manager, with every read observable. 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 4684c3213d6..abe88b913ae 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 @@ -1,473 +1,202 @@ +import asyncio import json import logging -from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import patch +import fakeredis import pytest -from fastapi.testclient import TestClient - -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.redis.stores.pod_lock import HeldBy +from litellm.redis.transport.contract import RedisFailure +from tests._support.fake_redis import fake_redis_cache - -class MockRedisCache: - def __init__(self): - self.async_set_cache = AsyncMock() - self.async_get_cache = AsyncMock() - self.async_delete_cache = AsyncMock() +LOCK_KEY = "cronjob_lock:test_job" @pytest.fixture -def mock_redis(): - return MockRedisCache() +def server() -> fakeredis.FakeServer: + return fakeredis.FakeServer() @pytest.fixture -def pod_lock_manager(mock_redis): - return PodLockManager(redis_cache=mock_redis) +def raw(server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis: + return fakeredis.FakeAsyncRedis(server=server) + + +async def _settle() -> None: + pending = {task for task in asyncio.all_tasks() if task is not asyncio.current_task()} + if pending: + await asyncio.gather(*pending) + await _settle() + + +async def _legacy_acquire(cache: RedisCache, pod_id: str, ttl: int = DEFAULT_CRON_JOB_LOCK_TTL_SECONDS) -> object: + return await cache.async_set_cache(PodLockManager.get_redis_lock_key("test_job"), pod_id, nx=True, ttl=ttl) @pytest.mark.asyncio -async def test_acquire_lock_success(pod_lock_manager, mock_redis): - """ - Test that the lock is acquired successfully when no existing lock exists - """ - # Mock successful acquisition (SET NX returns True) - mock_redis.async_set_cache.return_value = True +async def test_acquire_writes_the_json_encoded_pod_id_with_the_ttl(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server) + pod = PodLockManager(redis_cache=cache) - result = await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - assert result == True + assert await pod.acquire_lock(cronjob_id="test_job", ttl=120) is True - # Verify set_cache was called with correct parameters - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_set_cache.assert_called_once_with( - lock_key, - pod_lock_manager.pod_id, - nx=True, - ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, - ) + assert await raw.get(LOCK_KEY) == json.dumps(pod.pod_id).encode() + assert 119_000 < await raw.pttl(LOCK_KEY) <= 120_000 + assert await cache.async_get_cache(LOCK_KEY) == pod.pod_id @pytest.mark.asyncio -async def test_acquire_lock_existing_active(pod_lock_manager, mock_redis): - """ - Test that the lock is not acquired if there's an active lock by different pod - """ - # Mock failed acquisition (SET NX returns False) - mock_redis.async_set_cache.return_value = False - # Mock get_cache to return a different pod's ID - mock_redis.async_get_cache.return_value = "different_pod_id" +async def test_acquire_defaults_to_the_cron_job_ttl(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server) + assert await PodLockManager(redis_cache=cache).acquire_lock(cronjob_id="test_job") is True - result = await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - assert result == False - - # Verify set_cache was called - mock_redis.async_set_cache.assert_called_once() - # Verify get_cache was called to check existing lock - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - - -@pytest.mark.asyncio -async def test_acquire_lock_expired(pod_lock_manager, mock_redis): - """ - Test that the lock can be acquired if existing lock is expired - """ - # Mock failed acquisition first (SET NX returns False) - mock_redis.async_set_cache.return_value = False - - # Simulate an expired lock by having the TTL return a value - # Since Redis auto-expires keys, an expired lock would be absent - # So we'll simulate a retry after the first check fails - - # First check returns a value (lock exists) - mock_redis.async_get_cache.return_value = "different_pod_id" - - # Then set succeeds on retry (simulating key expiring between checks) - mock_redis.async_set_cache.side_effect = [False, True] - - result = await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - assert result == False # First attempt fails - - # Reset mock for a second attempt - mock_redis.async_set_cache.reset_mock() - mock_redis.async_set_cache.return_value = True - - # Try again (simulating the lock expired) - result = await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - assert result == True - - # Verify set_cache was called again - mock_redis.async_set_cache.assert_called_once() - - -@pytest.mark.asyncio -async def test_release_lock_success(pod_lock_manager, mock_redis): - """ - Test that the release lock works when the current pod holds the lock - """ - # Mock get_cache to return this pod's ID - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id - # Mock successful deletion - mock_redis.async_delete_cache.return_value = 1 - - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) - - # Verify get_cache was called - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - # Verify delete_cache was called - mock_redis.async_delete_cache.assert_called_once_with(lock_key) - - -@pytest.mark.asyncio -async def test_release_lock_different_pod(pod_lock_manager, mock_redis): - """ - Test that the release lock doesn't delete when a different pod holds the lock - """ - # Mock get_cache to return a different pod's ID - mock_redis.async_get_cache.return_value = "different_pod_id" - - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) - - # Verify get_cache was called - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - # Verify delete_cache was NOT called - mock_redis.async_delete_cache.assert_not_called() - - -@pytest.mark.asyncio -async def test_release_lock_no_lock(pod_lock_manager, mock_redis): - """ - Test release lock behavior when no lock exists - """ - # Mock get_cache to return None (no lock) - mock_redis.async_get_cache.return_value = None - - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) - - # Verify get_cache was called - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - # Verify delete_cache was NOT called - mock_redis.async_delete_cache.assert_not_called() - - -@pytest.mark.asyncio -async def test_redis_none(monkeypatch): - """ - Test behavior when redis_cache is None - """ - pod_lock_manager = PodLockManager(redis_cache=None) - - # Test acquire_lock with None redis_cache assert ( - await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - is None - ) - - # Test release_lock with None redis_cache (should not raise exception) - await pod_lock_manager.release_lock( - cronjob_id="test_job", + (DEFAULT_CRON_JOB_LOCK_TTL_SECONDS - 1) * 1000 + < await raw.pttl(LOCK_KEY) + <= DEFAULT_CRON_JOB_LOCK_TTL_SECONDS * 1000 ) @pytest.mark.asyncio -async def test_redis_error_handling(pod_lock_manager, mock_redis): - """ - Test error handling in Redis operations - """ - # Mock exceptions for Redis operations - mock_redis.async_set_cache.side_effect = Exception("Redis error") - mock_redis.async_get_cache.side_effect = Exception("Redis error") - mock_redis.async_delete_cache.side_effect = Exception("Redis error") +async def test_a_lock_taken_by_an_old_pod_is_refused_to_others_and_reentrant_for_it( + server: fakeredis.FakeServer, +) -> None: + cache = fake_redis_cache(server) + holder = PodLockManager(redis_cache=cache) + other = PodLockManager(redis_cache=cache) + await _legacy_acquire(cache, holder.pod_id) - # Test acquire_lock error handling - result = await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - assert result == False - - # Reset side effect for get_cache for the release test - mock_redis.async_get_cache.side_effect = None - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id - - # Test release_lock error handling (should not raise exception) - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) + assert await other.acquire_lock(cronjob_id="test_job") is False + assert await holder.acquire_lock(cronjob_id="test_job") is True + assert await holder.acquire_lock(cronjob_id="test_job", allow_reentrant=False) is False + assert await holder.lock_holder("test_job") == HeldBy(holder.pod_id) + assert await other.is_held("test_job") is True @pytest.mark.asyncio -async def test_lock_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(pod_lock_manager, mock_redis, caplog): - """Every cron job retries its lock on a timer, so an open breaker must not add an error line per cycle.""" - refused = RedisCircuitBreakerOpenError("Redis circuit breaker is open - skipping async_set_cache") - mock_redis.async_set_cache.side_effect = refused - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id - mock_redis.async_delete_cache.side_effect = refused +async def test_only_one_of_many_concurrent_pods_acquires(server: fakeredis.FakeServer) -> None: + cache = fake_redis_cache(server) + pods = [PodLockManager(redis_cache=cache) for _ in range(5)] + + results = await asyncio.gather(*(pod.acquire_lock(cronjob_id="test_job") for pod in pods)) + + assert sorted(results) == [False, False, False, False, True] + + +@pytest.mark.asyncio +async def test_release_deletes_only_the_owners_lock(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server) + 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() + + await holder.release_lock(cronjob_id="test_job") + assert await raw.exists(LOCK_KEY) == 0 + assert await other.acquire_lock(cronjob_id="test_job") is True + + +@pytest.mark.asyncio +async def test_release_frees_a_lock_written_by_an_old_pod(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server) + pod = PodLockManager(redis_cache=cache) + await _legacy_acquire(cache, pod.pod_id) + + await pod.release_lock(cronjob_id="test_job") + + assert await raw.exists(LOCK_KEY) == 0 + + +@pytest.mark.asyncio +async def test_release_survives_a_script_flush(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server) + pod = PodLockManager(redis_cache=cache) + assert await pod.acquire_lock(cronjob_id="test_job") is True + await pod.release_lock(cronjob_id="test_job") + assert await pod.acquire_lock(cronjob_id="test_job") is True + await raw.script_flush() + + await pod.release_lock(cronjob_id="test_job") + + assert await raw.exists(LOCK_KEY) == 0 + + +@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) + pod = PodLockManager(redis_cache=cache) + + with patch.object(pod, "_emit_released_lock_event") as emitted: + await pod.release_lock(cronjob_id="test_job") + assert emitted.call_count == 0 + await pod.acquire_lock(cronjob_id="test_job") + await pod.release_lock(cronjob_id="test_job") + + emitted.assert_called_once_with(cronjob_id="test_job", pod_id=pod.pod_id) + + +@pytest.mark.asyncio +async def test_an_owner_that_does_not_decode_is_a_failure_and_not_acquired(server: fakeredis.FakeServer, raw) -> None: + cache = fake_redis_cache(server) + await raw.set(LOCK_KEY, b"not json") + pod = PodLockManager(redis_cache=cache) + + assert await pod.acquire_lock(cronjob_id="test_job") is False + holder = await pod.lock_holder("test_job") + assert isinstance(holder, RedisFailure) and holder.kind == "decode" + + +@pytest.mark.asyncio +async def test_an_unreadable_lock_is_reported_unheld_with_a_warning_naming_the_job( + server: fakeredis.FakeServer, raw, caplog: pytest.LogCaptureFixture +) -> None: + await raw.set(LOCK_KEY, b"not json") + pod = PodLockManager(redis_cache=fake_redis_cache(server)) + + with caplog.at_level(logging.WARNING): + assert await pod.is_held("test_job") is False + + assert any("test_job" in record.getMessage() for record in caplog.records if record.levelno == logging.WARNING) + + +@pytest.mark.asyncio +async def test_an_unreachable_redis_is_not_acquired_and_logs_no_error( + server: fakeredis.FakeServer, caplog: pytest.LogCaptureFixture +) -> None: + pod = PodLockManager(redis_cache=fake_redis_cache(server)) + await _settle() + server.connected = False + caplog.clear() with caplog.at_level(logging.ERROR): - acquired = await pod_lock_manager.acquire_lock(cronjob_id="test_job") - await pod_lock_manager.release_lock(cronjob_id="test_job") + assert await pod.acquire_lock(cronjob_id="test_job") is False + await pod.release_lock(cronjob_id="test_job") + holder = await pod.lock_holder("test_job") - assert acquired is False - assert caplog.records == [] + assert isinstance(holder, RedisFailure) and holder.kind == "transport" + assert [record for record in caplog.records if record.levelno >= logging.ERROR] == [] @pytest.mark.asyncio -async def test_bytes_handling(pod_lock_manager, mock_redis): - """ - Test handling of bytes values from Redis - """ - # Mock failed acquisition - mock_redis.async_set_cache.return_value = False - # Mock get_cache to return bytes - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id.encode("utf-8") +async def test_without_redis_nothing_is_acquired_or_held() -> None: + pod = PodLockManager(redis_cache=None) - result = await pod_lock_manager.acquire_lock( - cronjob_id="test_job", - ) - assert result == True - - # Reset for release test - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id.encode("utf-8") - mock_redis.async_delete_cache.return_value = 1 - - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) - mock_redis.async_delete_cache.assert_called_once() + assert await pod.acquire_lock(cronjob_id="test_job") is None + await pod.release_lock(cronjob_id="test_job") + assert await pod.lock_holder("test_job") is None + assert await pod.is_held("test_job") is False @pytest.mark.asyncio -async def test_concurrent_lock_acquisition_simulation(): - """ - Simulate multiple pods trying to acquire the lock simultaneously - """ - mock_redis = MockRedisCache() - pod1 = PodLockManager(redis_cache=mock_redis) - pod2 = PodLockManager(redis_cache=mock_redis) - pod3 = PodLockManager(redis_cache=mock_redis) +async def test_a_namespaced_cache_keeps_the_lock_under_its_namespace(server: fakeredis.FakeServer, raw) -> None: + pod = PodLockManager(redis_cache=fake_redis_cache(server, namespace="ns")) - # Simulate first pod getting the lock - mock_redis.async_set_cache.return_value = True + assert await pod.acquire_lock(cronjob_id="test_job") is True - # First pod should get the lock - result1 = await pod1.acquire_lock( - cronjob_id="test_job", - ) - assert result1 == True - - # Simulate other pods failing to get the lock - mock_redis.async_set_cache.return_value = False - mock_redis.async_get_cache.return_value = pod1.pod_id - - # Other pods should fail to acquire - result2 = await pod2.acquire_lock( - cronjob_id="test_job", - ) - result3 = await pod3.acquire_lock( - cronjob_id="test_job", - ) - - # Since other pods don't have the lock, they should get False - assert result2 == False - assert result3 == False - - -@pytest.mark.asyncio -async def test_lock_takeover_race_condition(mock_redis): - """ - Test scenario where multiple pods try to take over an expired lock using Redis - """ - pod1 = PodLockManager(redis_cache=mock_redis) - pod2 = PodLockManager(redis_cache=mock_redis) - - # Simulate first pod's acquisition succeeding - mock_redis.async_set_cache.return_value = True - - # First pod should successfully acquire - result1 = await pod1.acquire_lock( - cronjob_id="test_job", - ) - assert result1 == True - - # Simulate race condition: second pod tries but fails - mock_redis.async_set_cache.return_value = False - mock_redis.async_get_cache.return_value = pod1.pod_id - - # Second pod should fail to acquire - result2 = await pod2.acquire_lock( - 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=[json.dumps(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 - - -@pytest.mark.asyncio -async def test_release_lock_lua_path_emits_released_event(pod_lock_manager, mock_redis): - """ - Test that _emit_released_lock_event is called when the Lua path returns 1 - (successful release). - """ - script_callable = AsyncMock(return_value=1) - mock_redis.async_register_script = MagicMock(return_value=script_callable) - - with patch.object(pod_lock_manager, "_emit_released_lock_event") as mock_emit: - await pod_lock_manager.release_lock(cronjob_id="test_job") - - mock_emit.assert_called_once_with(cronjob_id="test_job", pod_id=pod_lock_manager.pod_id) - - -class FakeRedisLockStore: - """ - Minimal stand-in that mirrors how RedisCache actually stores values: - async_set_cache JSON-encodes the value, and the compare-and-delete Lua - script compares against the raw stored bytes. This is what exposes the - quoted-vs-raw mismatch that a value-agnostic mock cannot catch. - """ - - def __init__(self): - self.store: dict = {} - - async def async_set_cache(self, key, value, nx=False, ttl=None, **kwargs): - if nx and key in self.store: - return None - self.store[key] = json.dumps(value) - return True - - async def async_get_cache(self, key, **kwargs): - raw = self.store.get(key) - return json.loads(raw) if raw is not None else None - - async def async_delete_cache(self, key, **kwargs): - return 1 if self.store.pop(key, None) is not None else 0 - - def async_register_script(self, script): - async def _run(keys, args): - key = keys[0] - if self.store.get(key) == args[0]: - del self.store[key] - return 1 - return 0 - - return _run - - -@pytest.mark.asyncio -async def test_release_lock_deletes_lock_held_by_same_pod(): - """ - Regression: acquire_lock stores the pod_id JSON-encoded, so release_lock's - Lua compare-and-delete must use the same encoding or the comparison never - matches and the lock leaks until its TTL expires (stalling the spend-update - drain and growing the Redis transaction buffers). - """ - redis = FakeRedisLockStore() - pod = PodLockManager(redis_cache=redis) - lock_key = PodLockManager.get_redis_lock_key("db_spend_update_job") - - acquired = await pod.acquire_lock(cronjob_id="db_spend_update_job") - assert acquired is True - assert lock_key in redis.store - - await pod.release_lock(cronjob_id="db_spend_update_job") - assert lock_key not in redis.store - - -@pytest.mark.asyncio -async def test_release_lock_preserves_lock_held_by_other_pod(): - """ - A pod must not release a lock currently held by a different pod, even with - the encoding fix in place. - """ - redis = FakeRedisLockStore() - holder = PodLockManager(redis_cache=redis) - other = PodLockManager(redis_cache=redis) - lock_key = PodLockManager.get_redis_lock_key("db_spend_update_job") - - assert await holder.acquire_lock(cronjob_id="db_spend_update_job") is True - - await other.release_lock(cronjob_id="db_spend_update_job") - assert redis.store.get(lock_key) == json.dumps(holder.pod_id) - - -@pytest.mark.asyncio -async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails(pod_lock_manager, mock_redis): - """ - Test that release_lock falls back to GET+DEL when Lua script execution - raises (e.g. Redis restart cleared loaded scripts). - """ - script_callable = AsyncMock(side_effect=Exception("NOSCRIPT")) - mock_redis.async_register_script = MagicMock(return_value=script_callable) - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id - mock_redis.async_delete_cache.return_value = 1 - - await pod_lock_manager.release_lock(cronjob_id="test_job") - - # Lua failed — should have fallen back to GET+DEL - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - mock_redis.async_delete_cache.assert_called_once_with(lock_key) - # Cached script handle should be reset so next call re-registers - assert pod_lock_manager._release_lock_script is None - - -@pytest.mark.asyncio -async def test_acquire_lock_own_lock_not_reentrant(pod_lock_manager, mock_redis): - """With allow_reentrant=False a live lock means the window's work is done, so even - the holder gets False; the default stays reentrant for leader-election callers.""" - mock_redis.async_set_cache.return_value = False - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id - - assert await pod_lock_manager.acquire_lock(cronjob_id="test_job", allow_reentrant=False) is False - assert await pod_lock_manager.acquire_lock(cronjob_id="test_job") is True + assert await raw.keys("*") == [f"ns:{LOCK_KEY}".encode()] diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index 76898b0eeab..b6f3d3c7046 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -7,6 +7,7 @@ import json from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock +import fakeredis import pytest from prisma.errors import PrismaError, RawQueryError @@ -24,7 +25,7 @@ from litellm.proxy.db.tool_registry_writer import ( list_tools, update_tool_policy, ) -from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation +from tests._support.fake_redis import fake_redis_cache def _mock_row(**kwargs): @@ -69,9 +70,7 @@ def _make_prisma( prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) prisma.db.execute_raw = AsyncMock(return_value=1) - prisma.db.litellm_tooltable.find_many = AsyncMock( - return_value=find_many_rows if find_many_rows is not None else [] - ) + prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else []) prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row) prisma.db.litellm_verificationtoken = MagicMock() prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows)) @@ -227,9 +226,7 @@ async def test_get_tool_found(): result = await get_tool(prisma, "my_tool") assert result is not None assert result.tool_name == "my_tool" - prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio @@ -266,9 +263,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): updated_by="admin", ) prisma = _make_prisma(find_unique_row=row) - result = await update_tool_policy( - prisma, "my_tool", updated_by="admin", input_policy="blocked" - ) + result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked") assert result is not None assert result.input_policy == "blocked" prisma.db.litellm_tooltable.upsert.assert_awaited_once() @@ -431,31 +426,8 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_perms_lookup_failure" -class _InMemoryRedis: - def __init__(self) -> None: - self.store: dict[str, str] = {} - self.lists: dict[str, list[object]] = {} - - async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: - if nx and key in self.store: - return None - self.store[key] = json.dumps(value) - return True - - async def async_get_cache(self, key: str, **kwargs: object) -> str | None: - raw = self.store.get(key) - return json.loads(raw) if raw is not None else None - - async def async_delete_cache(self, key: str) -> int: - return 1 if self.store.pop(key, None) is not None else 0 - - async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: - for op in rpush_list: - self.lists.setdefault(op["key"], []).extend(op["values"]) - return [len(self.lists[op["key"]]) for op in rpush_list] - - async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: - return [self.lists.pop(op["key"], None) for op in lpop_list] +async def _buffered(server: fakeredis.FakeServer) -> list[bytes]: + return await fakeredis.FakeAsyncRedis(server=server).lrange(REDIS_TOOL_DISCOVERY_BUFFER_KEY, 0, -1) def _postgres_rejection(sqlstate: str) -> RawQueryError: @@ -473,7 +445,8 @@ def _postgres_rejection(sqlstate: str) -> RawQueryError: async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_from_the_tool_lock_holder_only( lock_held_by_another_pod: str | None, ): - redis = _InMemoryRedis() + server = fakeredis.FakeServer() + redis = fake_redis_cache(server) if lock_held_by_another_pod is not None: assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=lock_held_by_another_pod) other_worker = ToolDiscoveryQueue() @@ -500,9 +473,9 @@ async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_fro prisma_client.db.litellm_tooltable.upsert.assert_not_awaited() if lock_held_by_another_pod == TOOL_DISCOVERY_COMMIT_JOB_NAME: prisma_client.db.execute_raw.assert_not_awaited() - assert len(redis.lists[REDIS_TOOL_DISCOVERY_BUFFER_KEY]) == 2, "non-holders leave their tools for the holder" + assert len(await _buffered(server)) == 2, "non-holders leave their tools for the holder" return - assert REDIS_TOOL_DISCOVERY_BUFFER_KEY not in redis.lists, "the spend lock must not gate the tool drain" + assert await _buffered(server) == [], "the spend lock must not gate the tool drain" prisma_client.db.execute_raw.assert_awaited_once() sql, *params = prisma_client.db.execute_raw.call_args.args assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') and 'ON CONFLICT ("tool_name")' in sql @@ -512,7 +485,8 @@ async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_fro @pytest.mark.asyncio async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statement_and_never_touches_redis(): - redis = _InMemoryRedis() + server = fakeredis.FakeServer() + redis = fake_redis_cache(server) queue = ToolDiscoveryQueue() queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) @@ -528,7 +502,9 @@ async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statem use_redis_buffer=False, ) - assert redis.lists == {}, "without the buffer flag nothing goes through redis" + assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == [], ( + "without the buffer flag nothing goes through redis" + ) assert queue.flush() == [] prisma_client.db.execute_raw.assert_awaited_once() sql, *params = prisma_client.db.execute_raw.call_args.args @@ -549,7 +525,8 @@ async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statem async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably_did_not_apply( failure: Exception, requeued: bool ): - redis = _InMemoryRedis() + server = fakeredis.FakeServer() + redis = fake_redis_cache(server) queue = ToolDiscoveryQueue() queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) prisma_client = MagicMock() @@ -564,6 +541,6 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably ) prisma_client.db.execute_raw.assert_awaited_once() - left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] + left_in_redis = [json.loads(entry) for entry in await _buffered(server)] expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" diff --git a/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index 766e0c1dfa9..042143d31d8 100644 --- a/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -11,7 +11,6 @@ import pytest import litellm.proxy.spend_tracking.ptu_flat_cost_rollup as ptu_rollup from litellm.constants import PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR -from litellm.types.router import ModelInfo from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( PTUModel, _active_hours_on_day, @@ -21,6 +20,7 @@ from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( run_ptu_flat_cost_rollup, run_scheduled_ptu_rollup, ) +from litellm.types.router import ModelInfo DAY = date(2026, 7, 30) TODAY = date(2026, 7, 31) @@ -652,8 +652,7 @@ def _pod_lock(acquired): lock = MagicMock() lock.pod_id = "this-pod" lock.redis_cache = MagicMock() - lock.redis_cache.async_get_cache = AsyncMock(return_value="this-pod") - lock.get_redis_lock_key = MagicMock(return_value="lock-key") + lock.is_held = AsyncMock(return_value=True) lock.acquire_lock = AsyncMock(return_value=acquired) lock.release_lock = AsyncMock() return lock @@ -784,9 +783,7 @@ class _FakeSentinelTable: """The batch delete of a deployment day's rows. Every clause is honoured and an unexpected one fails, so a dropped scope cannot pass silently.""" assert set(where) == {"date", "api_key", "model"} - doomed = [ - k for k in self.rows if k[1] == where["date"] and k[2] == where["api_key"] and k[3] == where["model"] - ] + doomed = [k for k in self.rows if k[1] == where["date"] and k[2] == where["api_key"] and k[3] == where["model"]] for k in doomed: del self.rows[k] @@ -1032,21 +1029,14 @@ async def test_scheduled_rollup_alerts_from_under_the_pod_lock_too(): @pytest.mark.asyncio -@pytest.mark.parametrize( - "lock_read", - [ - pytest.param(AsyncMock(side_effect=RuntimeError("redis down")), id="redis-unreachable"), - pytest.param(AsyncMock(return_value=None), id="lock-key-missing"), - ], -) -async def test_scheduled_rollup_runs_the_day_when_the_lock_is_unavailable_but_unheld(lock_read): +async def test_scheduled_rollup_runs_the_day_when_the_lock_is_unavailable_but_unheld(): """acquire_lock reports contention and a Redis outage identically. Treating both as "someone else has it" would skip the day on every pod at once, losing every team's charge for that date; the reconcile is safe to run twice, so the day wins.""" rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})] prisma, table = _prisma_with_models(rows) lock = _pod_lock(acquired=False) - lock.redis_cache.async_get_cache = lock_read + lock.is_held = AsyncMock(return_value=False) result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=lock, target_date=DAY) diff --git a/tests/unit/proxy/spend_tracking/test_spend_capture_rate.py b/tests/unit/proxy/spend_tracking/test_spend_capture_rate.py index cb79dd60245..cdb27c02066 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_capture_rate.py +++ b/tests/unit/proxy/spend_tracking/test_spend_capture_rate.py @@ -257,8 +257,7 @@ async def test_check_publishes_no_rate_when_the_provider_billed_nothing(monkeypa def _pod_lock(acquired: bool) -> MagicMock: lock = MagicMock() lock.redis_cache = MagicMock() - lock.redis_cache.async_get_cache = AsyncMock(return_value="other-pod") - lock.get_redis_lock_key = MagicMock(return_value="lock-key") + lock.is_held = AsyncMock(return_value=True) lock.acquire_lock = AsyncMock(return_value=acquired) lock.release_lock = AsyncMock() return lock @@ -335,7 +334,7 @@ async def test_scheduled_check_alerts_and_keeps_the_lock_until_it_expires_when_i @pytest.mark.asyncio async def test_scheduled_check_alerts_when_the_lock_cannot_be_acquired_or_read(monkeypatch): lock = _pod_lock(acquired=False) - lock.redis_cache.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + lock.is_held = AsyncMock(return_value=False) alert, publish = await _scheduled_run_under_threshold(lock, monkeypatch) diff --git a/tests/unit/redis/stores/__init__.py b/tests/unit/redis/stores/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/redis/stores/test_pod_lock.py b/tests/unit/redis/stores/test_pod_lock.py new file mode 100644 index 00000000000..24e9c298996 --- /dev/null +++ b/tests/unit/redis/stores/test_pod_lock.py @@ -0,0 +1,91 @@ +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, NoReturn + +import pytest + +from litellm.redis.stores.pod_lock import TransportPodLockStore, Vacant +from litellm.redis.transport.contract import Arg, EvalReply, RedisFailure + + +@dataclass(frozen=True, slots=True) +class ExpiredLockTransport: + async def get(self, key: str) -> bytes | None | RedisFailure: + return None + + async def set(self, key: str, value: Arg, *, ttl: timedelta | None, nx: bool = False) -> bool | RedisFailure: + return not nx + + async def mget(self, keys: Sequence[str]) -> tuple[bytes | None, ...] | RedisFailure: + raise AssertionError("unexpected mget") + + async def delete(self, *keys: str) -> int | RedisFailure: + raise AssertionError("unexpected delete") + + async def eval(self, script: str, keys: Sequence[str], args: Sequence[Arg]) -> EvalReply | RedisFailure: + raise AssertionError("unexpected eval") + + async def incrby(self, key: str, amount: int) -> NoReturn: + raise AssertionError("unexpected incrby") + + async def incrbyfloat(self, key: str, amount: float) -> NoReturn: + raise AssertionError("unexpected incrbyfloat") + + async def expire(self, key: str, ttl: timedelta) -> NoReturn: + raise AssertionError("unexpected expire") + + async def ttl(self, key: str) -> NoReturn: + raise AssertionError("unexpected ttl") + + async def rpush(self, key: str, *values: Arg) -> NoReturn: + raise AssertionError("unexpected rpush") + + async def lpop(self, key: str, count: int | None = None) -> NoReturn: + raise AssertionError("unexpected lpop") + + async def ltrim(self, key: str, start: int, stop: int) -> NoReturn: + raise AssertionError("unexpected ltrim") + + async def sadd(self, key: str, *members: Arg) -> NoReturn: + raise AssertionError("unexpected sadd") + + async def scan(self, prefix: str, *, limit: int) -> NoReturn: + raise AssertionError("unexpected scan") + + def batch(self, *, transaction: bool = False) -> NoReturn: + raise AssertionError("unexpected batch") + + def request_batch(self) -> NoReturn: + raise AssertionError("unexpected request_batch") + + def post_call_batch(self) -> NoReturn: + raise AssertionError("unexpected post_call_batch") + + async def publish(self, channel: str, message: str | bytes) -> NoReturn: + raise AssertionError("unexpected publish") + + async def publish_all(self, channel: str, messages: Sequence[str | bytes]) -> NoReturn: + raise AssertionError("unexpected publish_all") + + async def subscribe(self, *channels: str) -> NoReturn: + raise AssertionError("unexpected subscribe") + + async def ping(self) -> NoReturn: + raise AssertionError("unexpected ping") + + async def info(self) -> NoReturn: + raise AssertionError("unexpected info") + + async def client_list(self) -> NoReturn: + raise AssertionError("unexpected client_list") + + async def flushall(self) -> NoReturn: + raise AssertionError("unexpected flushall") + + +@pytest.mark.asyncio +async def test_acquire_is_vacant_when_the_lock_expires_after_losing_set_nx() -> None: + store: Final = TransportPodLockStore(ExpiredLockTransport()) + + assert await store.acquire("test_job", "pod", ttl=timedelta(seconds=60)) == Vacant()