scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-03 07:23:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 166

refactor(redis-store): move the pod lock keyspace behind a store

Browse source 
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
This commit is contained in:
Yujong Lee 2026-10-30 10:26:05 -07:00 • committed by Devin AI
parent 911ddf8ead
commit 1bd2ff7e63
14 changed files with 481 additions and 698 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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

192
litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

@ -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):

19
litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
View file

@ -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",
*,

19
litellm/proxy/spend_tracking/spend_capture_rate.py
View file

@ -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)

87
litellm/redis/stores/pod_lock.py Normal file
View file

@ -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}")

33
tests/unit/integrations/focus/test_focus_logger.py
View file

@ -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

33
tests/unit/integrations/vantage/test_vantage_logger.py
View file

@ -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

11
tests/unit/proxy/common_utils/test_reset_budget_job.py
View file

@ -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.

587
tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py
View file

@ -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()]

61
tests/unit/proxy/db/test_tool_registry_writer.py
View file

@ -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}"

20
tests/unit/proxy/spend_tracking/test_ptu_flat_cost_rollup.py
View file

@ -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)

5
tests/unit/proxy/spend_tracking/test_spend_capture_rate.py
View file

@ -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)

0
tests/unit/redis/stores/__init__.py Normal file
View file

91
tests/unit/redis/stores/test_pod_lock.py Normal file
View file

@ -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()