mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-05 08:07:05 +00:00
fix(cache): use sync Redis batch reads (#39358)
* fix(cache): use sync Redis batch reads * fix(cache): type sync circuit breaker decorator * test(cache): isolate sync Redis breaker coverage * fix(cache): keep batch result merge budget compliant * style(cache): format batch read * style(cache): satisfy type-discipline budget * test(cache): mock Redis before sync breaker setup * style(cache): avoid mutable batch placeholder * test(cache): document sync breaker patch target * fix(types): widen batch result params to Sequence * fix(cache): report real callers through breaker guards The sync guard's lambda and runner frames replaced the actual caller in _get_call_stack_info, so Redis service logs attributed every guarded call to the guard machinery. Skip guard-internal frames when walking the stack and ratchet the lint budgets this branch lowered * style(imports): import Sequence from collections.abc * test(cache): cover concurrent sync and async Redis batch reads * refactor: build sync batch_get_cache results as tuples to satisfy the LIT002 gate * chore: ratchet budgets after staging merge * fix: preserve DualCache batch list contract * style: format DualCache batch result * fix: satisfy mutable collection lint gate * fix(caching): keep breaker guard-frame skipping in bytecode-only deploys * chore: preserve staging budget ratchets * test(cache): isolate sync Redis batch reads * fix(cache): isolate service hook failures * fix(cache): preserve sync batch fallback on open breaker
This commit is contained in:
yucheng-berri
GitHub
parent
f39ffbd760
commit
e6e5be0989
7 changed files with 526 additions and 90 deletions
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable, Coroutine
|
||||
from datetime import datetime, timedelta
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
|
|
@ -83,6 +84,47 @@ class ServiceLogging(CustomLogger):
|
|||
return open_telemetry_logger
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _sync_dispatch_loop() -> asyncio.AbstractEventLoop | None:
|
||||
"""The event loop a blocking caller can dispatch on, or ``None`` if it has none."""
|
||||
try:
|
||||
loop: Final = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
return None
|
||||
return None if loop.is_closed() else loop
|
||||
|
||||
@staticmethod
|
||||
async def _emit_guarded(hook: Callable[[], Coroutine[object, object, None]]) -> None:
|
||||
"""Emit one service event, absorbing anything the callbacks raise.
|
||||
|
||||
Monitoring must not break the call it monitors. Sync callers are the ones that
|
||||
swallow their own service failures (a Redis batch read returns an empty dict),
|
||||
so an exception from a misconfigured callback would replace a Redis outage with
|
||||
a callback error and skip the caller's fallback handling.
|
||||
"""
|
||||
try:
|
||||
await hook()
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error emitting service event - %s", e)
|
||||
|
||||
@staticmethod
|
||||
def _dispatch_from_sync(hook: Callable[[], Coroutine[object, object, None]]) -> None:
|
||||
"""Run an async service hook from a blocking caller, whatever event loop it holds.
|
||||
|
||||
Takes a factory rather than a coroutine so the hook is built on the path that
|
||||
runs it, and only ever once.
|
||||
"""
|
||||
loop: Final = ServiceLogging._sync_dispatch_loop()
|
||||
try:
|
||||
if loop is None:
|
||||
asyncio.run(ServiceLogging._emit_guarded(hook))
|
||||
elif loop.is_running():
|
||||
loop.create_task(ServiceLogging._emit_guarded(hook))
|
||||
else:
|
||||
loop.run_until_complete(ServiceLogging._emit_guarded(hook))
|
||||
except Exception as e:
|
||||
verbose_logger.exception("Error dispatching service event - %s", e)
|
||||
|
||||
def service_success_hook(
|
||||
self,
|
||||
service: ServiceTypes,
|
||||
|
|
@ -99,54 +141,45 @@ class ServiceLogging(CustomLogger):
|
|||
if self.mock_testing:
|
||||
self.mock_testing_sync_success_hook += 1
|
||||
|
||||
try:
|
||||
# Try to get the current event loop
|
||||
loop: Final = asyncio.get_event_loop()
|
||||
# Check if the loop is running
|
||||
if loop.is_running():
|
||||
# If we're in a running loop, create a task
|
||||
loop.create_task(
|
||||
self.async_service_success_hook(
|
||||
service=service,
|
||||
duration=duration,
|
||||
call_type=call_type,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# Loop exists but not running, we can use run_until_complete
|
||||
loop.run_until_complete(
|
||||
self.async_service_success_hook(
|
||||
service=service,
|
||||
duration=duration,
|
||||
call_type=call_type,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
)
|
||||
except RuntimeError:
|
||||
# No event loop exists, create a new one and run
|
||||
asyncio.run(
|
||||
self.async_service_success_hook(
|
||||
service=service,
|
||||
duration=duration,
|
||||
call_type=call_type,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
self._dispatch_from_sync(
|
||||
lambda: self.async_service_success_hook(
|
||||
service=service,
|
||||
duration=duration,
|
||||
call_type=call_type,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
)
|
||||
|
||||
def service_failure_hook(self, service: ServiceTypes, duration: float, error: Exception, call_type: str):
|
||||
def service_failure_hook(
|
||||
self,
|
||||
service: ServiceTypes,
|
||||
duration: float,
|
||||
error: Exception,
|
||||
call_type: str,
|
||||
parent_otel_span: Span | None = None,
|
||||
start_time: datetime | float | None = None,
|
||||
end_time: float | datetime | None = None,
|
||||
):
|
||||
"""
|
||||
[TODO] Not implemented for sync calls yet. V0 is focused on async monitoring (used by proxy).
|
||||
Handles both sync and async monitoring by checking for existing event loop.
|
||||
"""
|
||||
if self.mock_testing:
|
||||
self.mock_testing_sync_failure_hook += 1
|
||||
|
||||
self._dispatch_from_sync(
|
||||
lambda: self.async_service_failure_hook(
|
||||
service=service,
|
||||
duration=duration,
|
||||
error=error,
|
||||
call_type=call_type,
|
||||
parent_otel_span=parent_otel_span,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
)
|
||||
|
||||
async def async_service_success_hook(
|
||||
self,
|
||||
service: ServiceTypes,
|
||||
|
|
|
|||
|
|
@ -8,10 +8,9 @@ Has 4 primary methods:
|
|||
- async_get_cache
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from collections.abc import Sequence
|
||||
from threading import Lock
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
|
|
@ -188,31 +187,38 @@ class DualCache(BaseCache):
|
|||
local_only: bool = False,
|
||||
**kwargs,
|
||||
):
|
||||
received_args: Final = locals()
|
||||
received_args.pop("self")
|
||||
|
||||
def run_in_new_loop():
|
||||
"""Run the coroutine in a new event loop within this thread."""
|
||||
new_loop: Final = asyncio.new_event_loop()
|
||||
try:
|
||||
asyncio.set_event_loop(new_loop)
|
||||
return new_loop.run_until_complete(self.async_batch_get_cache(**received_args))
|
||||
finally:
|
||||
new_loop.close()
|
||||
asyncio.set_event_loop(None)
|
||||
|
||||
try:
|
||||
# First, try to get the current event loop
|
||||
_ = asyncio.get_running_loop()
|
||||
# If we're already in an event loop, run in a separate thread
|
||||
# to avoid nested event loop issues
|
||||
with ThreadPoolExecutor(max_workers=1) as executor:
|
||||
future: Final = executor.submit(run_in_new_loop)
|
||||
return future.result()
|
||||
in_memory_result: Final = (
|
||||
self.in_memory_cache.batch_get_cache(keys, **kwargs) if self.in_memory_cache is not None else None
|
||||
)
|
||||
result: Final = in_memory_result if in_memory_result is not None else tuple(None for _ in keys)
|
||||
|
||||
except RuntimeError:
|
||||
# No running event loop, we can safely run in this thread
|
||||
return run_in_new_loop()
|
||||
if None not in result or self.redis_cache is None or local_only:
|
||||
return result
|
||||
|
||||
sublist_keys, previous_access_times = self._reserve_redis_batch_keys(time.time(), keys, result)
|
||||
if len(sublist_keys) == 0:
|
||||
return result
|
||||
|
||||
try:
|
||||
redis_result: Final = self.redis_cache.batch_get_cache(
|
||||
key_list=sublist_keys, parent_otel_span=parent_otel_span
|
||||
)
|
||||
except Exception:
|
||||
# Do not throttle subsequent callers if the Redis read fails.
|
||||
self._rollback_redis_batch_key_reservations(previous_access_times)
|
||||
raise
|
||||
|
||||
if self.in_memory_cache is not None:
|
||||
for key, value in redis_result.items():
|
||||
if value is not None:
|
||||
self.in_memory_cache.set_cache(key, value, **self._backfill_kwargs(kwargs))
|
||||
|
||||
return list( # mutable-ok: public list contract
|
||||
redis_result.get(key) if value is None else value for key, value in zip(keys, result)
|
||||
)
|
||||
except Exception:
|
||||
verbose_logger.error(traceback.format_exc())
|
||||
|
||||
async def async_get_cache(
|
||||
self,
|
||||
|
|
@ -251,7 +257,7 @@ class DualCache(BaseCache):
|
|||
self,
|
||||
current_time: float,
|
||||
keys: list[str],
|
||||
result: list[Any],
|
||||
result: Sequence[Any],
|
||||
) -> tuple[list[str], dict[str, float | None]]:
|
||||
"""
|
||||
Atomically choose keys to fetch from Redis and reserve their access time.
|
||||
|
|
|
|||
|
|
@ -78,10 +78,18 @@ class _AsyncRedisCommands(Protocol):
|
|||
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
|
||||
|
||||
|
||||
_BREAKER_GUARD_FRAME_NAMES: Final = frozenset(
|
||||
{"<lambda>", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"}
|
||||
)
|
||||
|
||||
|
||||
def _get_call_stack_info(num_frames: int = 2) -> str:
|
||||
"""
|
||||
Get the function names from the previous 1-2 functions in the call stack.
|
||||
|
||||
Frames belonging to this module's circuit-breaker guards are skipped so the
|
||||
reported callers stay the real ones even on guarded methods.
|
||||
|
||||
Args:
|
||||
num_frames: Number of previous frames to include (default: 2)
|
||||
|
||||
|
|
@ -102,11 +110,11 @@ def _get_call_stack_info(num_frames: int = 2) -> str:
|
|||
return "unknown"
|
||||
function_names: Final = []
|
||||
|
||||
for _ in range(num_frames):
|
||||
if frame is None:
|
||||
break
|
||||
func_name = frame.f_code.co_name
|
||||
function_names.append(func_name)
|
||||
while frame is not None and len(function_names) < num_frames:
|
||||
if frame.f_code.co_name in _BREAKER_GUARD_FRAME_NAMES and frame.f_globals.get("__name__") == __name__:
|
||||
frame = frame.f_back
|
||||
continue
|
||||
function_names.append(frame.f_code.co_name)
|
||||
frame = frame.f_back
|
||||
|
||||
if not function_names:
|
||||
|
|
@ -241,6 +249,23 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
|
|||
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
|
||||
|
||||
|
||||
def _enter_circuit_breaker(breaker: RedisCircuitBreaker, name: str) -> int:
|
||||
"""Reject the call if the breaker is open, else return the swallowed-failure count to compare against."""
|
||||
if breaker.is_open():
|
||||
raise Exception(f"Redis circuit breaker is open — skipping {name}")
|
||||
return _swallowed_redis_failures.get()
|
||||
|
||||
|
||||
def _exit_circuit_breaker(breaker: RedisCircuitBreaker, swallowed_before: int) -> None:
|
||||
"""Record success only when nothing failed while the call ran.
|
||||
|
||||
Several Redis methods catch their own connection errors and return a default, so a
|
||||
method that returned is not on its own proof of a healthy Redis.
|
||||
"""
|
||||
if _swallowed_redis_failures.get() == swallowed_before:
|
||||
breaker.record_success()
|
||||
|
||||
|
||||
async def _run_under_circuit_breaker(
|
||||
breaker: RedisCircuitBreaker,
|
||||
name: str,
|
||||
|
|
@ -249,20 +274,33 @@ async def _run_under_circuit_breaker(
|
|||
"""Run one Redis coroutine under a circuit breaker.
|
||||
|
||||
Shared by the method decorator and the Lua script executor so both feed the same
|
||||
health signal. Success is recorded only when nothing failed while ``call`` ran,
|
||||
because several Redis methods catch their own connection errors and return a default.
|
||||
health signal.
|
||||
"""
|
||||
if breaker.is_open():
|
||||
raise Exception(f"Redis circuit breaker is open — skipping {name}")
|
||||
swallowed_before: Final = _swallowed_redis_failures.get()
|
||||
swallowed_before: Final = _enter_circuit_breaker(breaker, name)
|
||||
try:
|
||||
result: Final = await call()
|
||||
except Exception as e:
|
||||
if _is_redis_health_failure(e):
|
||||
breaker.record_failure()
|
||||
raise
|
||||
if _swallowed_redis_failures.get() == swallowed_before:
|
||||
breaker.record_success()
|
||||
_exit_circuit_breaker(breaker, swallowed_before)
|
||||
return result
|
||||
|
||||
|
||||
def _run_under_circuit_breaker_sync(
|
||||
breaker: RedisCircuitBreaker,
|
||||
name: str,
|
||||
call: Callable[[], _RedisCallResult],
|
||||
) -> _RedisCallResult:
|
||||
"""Run one blocking Redis call under a circuit breaker, feeding the same health signal as the async path."""
|
||||
swallowed_before: Final = _enter_circuit_breaker(breaker, name)
|
||||
try:
|
||||
result: Final = call()
|
||||
except Exception as e:
|
||||
if _is_redis_health_failure(e):
|
||||
breaker.record_failure()
|
||||
raise
|
||||
_exit_circuit_breaker(breaker, swallowed_before)
|
||||
return result
|
||||
|
||||
|
||||
|
|
@ -288,6 +326,14 @@ def _redis_circuit_breaker_guard(method):
|
|||
return wrapper
|
||||
|
||||
|
||||
def _redis_circuit_breaker_guard_sync(method: Callable[..., _RedisCallResult]) -> Callable[..., _RedisCallResult]:
|
||||
return functools.wraps(method)(
|
||||
lambda self, *args, **kwargs: _run_under_circuit_breaker_sync(
|
||||
self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class RedisCache(BaseCache):
|
||||
# if users don't provider one, use the default litellm cache
|
||||
|
||||
|
|
@ -1146,14 +1192,13 @@ class RedisCache(BaseCache):
|
|||
"""
|
||||
key_value_dict = {}
|
||||
_key_list: Final = [key for key in key_list if key is not None]
|
||||
start_time: Final = time.time()
|
||||
|
||||
try:
|
||||
_keys: Final = []
|
||||
for cache_key in _key_list:
|
||||
cache_key = self.check_and_fix_namespace(key=cache_key or "")
|
||||
_keys.append(cache_key)
|
||||
start_time: Final = time.time()
|
||||
swallowed_before: Final = _enter_circuit_breaker(self._circuit_breaker, "batch_get_cache")
|
||||
_keys: Final = [self.check_and_fix_namespace(key=cache_key or "") for cache_key in _key_list]
|
||||
results: Final = self._run_redis_mget_operation(keys=_keys)
|
||||
_exit_circuit_breaker(self._circuit_breaker, swallowed_before)
|
||||
end_time: Final = time.time()
|
||||
_duration: Final = end_time - start_time
|
||||
self.service_logger_obj.service_success_hook(
|
||||
|
|
@ -1178,7 +1223,18 @@ class RedisCache(BaseCache):
|
|||
|
||||
return decoded_results
|
||||
except Exception as e:
|
||||
failed_at: Final = time.time()
|
||||
self.service_logger_obj.service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=failed_at - start_time,
|
||||
error=e,
|
||||
call_type=f"batch_get_cache <- {_get_call_stack_info()}",
|
||||
start_time=start_time,
|
||||
end_time=failed_at,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
verbose_logger.error("Error occurred in batch get cache - %s", e)
|
||||
_record_swallowed_redis_failure(self._circuit_breaker, e)
|
||||
return key_value_dict
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
#### What this does ####
|
||||
# identifies lowest tpm deployment
|
||||
import random
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
|
|
@ -350,9 +351,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger):
|
|||
model_group: str,
|
||||
healthy_deployments: list,
|
||||
tpm_keys: list,
|
||||
tpm_values: list | None,
|
||||
tpm_values: Sequence | None,
|
||||
rpm_keys: list,
|
||||
rpm_values: list | None,
|
||||
rpm_values: Sequence | None,
|
||||
messages: list[dict[str, str]] | None = None,
|
||||
input: str | list | None = None,
|
||||
) -> dict | None:
|
||||
|
|
|
|||
|
|
@ -240,3 +240,35 @@ async def test_dual_cache_delete(is_async):
|
|||
result = dual_cache.get_cache(test_key)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_concurrent_sync_and_async_redis_reads():
|
||||
"""Sync and async batch reads share one Redis backend in one process, and sync reads never open an async connection"""
|
||||
redis_cache = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT"))
|
||||
dual_cache = DualCache(redis_cache=redis_cache)
|
||||
|
||||
run_id = str(uuid.uuid4())
|
||||
sync_keys = [f"sync_{run_id}_{index}" for index in range(5)]
|
||||
async_keys = [f"async_{run_id}_{index}" for index in range(5)]
|
||||
in_loop_keys = [f"in_loop_{run_id}_{index}" for index in range(3)]
|
||||
survivor_key = f"survivor_{run_id}"
|
||||
expected = {key: {"key": key} for key in [*sync_keys, *async_keys, *in_loop_keys, survivor_key]}
|
||||
for key, value in expected.items():
|
||||
await redis_cache.async_set_cache(key, value, ttl=60)
|
||||
|
||||
concurrent_results = await asyncio.gather(
|
||||
*(asyncio.to_thread(dual_cache.batch_get_cache, keys=[key]) for key in sync_keys),
|
||||
*(dual_cache.async_batch_get_cache(keys=[key]) for key in async_keys),
|
||||
)
|
||||
assert list(concurrent_results) == [[expected[key]] for key in [*sync_keys, *async_keys]]
|
||||
|
||||
with patch.object(
|
||||
redis_cache,
|
||||
"async_batch_get_cache",
|
||||
side_effect=AssertionError("sync batch reads must not call async Redis"),
|
||||
):
|
||||
in_loop_results = [dual_cache.batch_get_cache(keys=[key]) for key in in_loop_keys]
|
||||
|
||||
assert in_loop_results == [[expected[key]] for key in in_loop_keys]
|
||||
assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]]
|
||||
|
|
|
|||
|
|
@ -61,6 +61,137 @@ async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_
|
|||
assert "shared_b" not in dual_cache.last_redis_batch_access_time
|
||||
|
||||
|
||||
def _redis_mock_for_sync_batch(redis_result: dict) -> MagicMock:
|
||||
mock_redis = MagicMock(spec=RedisCache)
|
||||
mock_redis.batch_get_cache.return_value = redis_result
|
||||
return mock_redis
|
||||
|
||||
|
||||
def _assert_sync_batch_used_blocking_client(dual_cache: DualCache, mock_redis: MagicMock) -> None:
|
||||
with patch("asyncio.new_event_loop", side_effect=AssertionError("sync path must not create an event loop")):
|
||||
result = dual_cache.batch_get_cache(keys=["lit6729_key"])
|
||||
|
||||
assert result == ["redis_value"]
|
||||
mock_redis.batch_get_cache.assert_called_once_with(key_list=["lit6729_key"], parent_otel_span=None)
|
||||
mock_redis.async_batch_get_cache.assert_not_called()
|
||||
mock_redis.init_async_client.assert_not_called()
|
||||
assert dual_cache.in_memory_cache.get_cache("lit6729_key") == "redis_value"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_batch_get_cache_uses_sync_redis_client_inside_running_loop():
|
||||
"""
|
||||
Regression test for LIT-6729: sync batch_get_cache ran async_batch_get_cache on a
|
||||
throwaway event loop, reusing an async Redis client created on another loop and
|
||||
corrupting its connection pool. The sync path must use the blocking client, never
|
||||
the async one, and never create an event loop, even when called from a coroutine
|
||||
(e.g. async_raise_no_deployment_exception -> get_min_cooldown).
|
||||
"""
|
||||
mock_redis = _redis_mock_for_sync_batch({"lit6729_key": "redis_value"})
|
||||
dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis)
|
||||
|
||||
_assert_sync_batch_used_blocking_client(dual_cache, mock_redis)
|
||||
|
||||
|
||||
def test_dual_cache_batch_get_cache_uses_sync_redis_client_without_running_loop():
|
||||
mock_redis = _redis_mock_for_sync_batch({"lit6729_key": "redis_value"})
|
||||
dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis)
|
||||
|
||||
_assert_sync_batch_used_blocking_client(dual_cache, mock_redis)
|
||||
|
||||
|
||||
def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis():
|
||||
mock_redis = _redis_mock_for_sync_batch({"miss_key": "from_redis"})
|
||||
dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis)
|
||||
dual_cache.in_memory_cache.set_cache("hit_key", "from_memory")
|
||||
|
||||
result = dual_cache.batch_get_cache(keys=["hit_key", "miss_key"])
|
||||
|
||||
assert result == ["from_memory", "from_redis"]
|
||||
mock_redis.batch_get_cache.assert_called_once_with(key_list=["miss_key"], parent_otel_span=None)
|
||||
|
||||
|
||||
def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads():
|
||||
mock_redis = _redis_mock_for_sync_batch({"absent_key": None})
|
||||
dual_cache = DualCache(
|
||||
in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10
|
||||
)
|
||||
|
||||
first = dual_cache.batch_get_cache(keys=["absent_key"])
|
||||
second = dual_cache.batch_get_cache(keys=["absent_key"])
|
||||
|
||||
assert first == [None]
|
||||
assert second == [None]
|
||||
mock_redis.batch_get_cache.assert_called_once()
|
||||
|
||||
|
||||
def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error():
|
||||
mock_redis = MagicMock(spec=RedisCache)
|
||||
mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable")
|
||||
dual_cache = DualCache(
|
||||
in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10
|
||||
)
|
||||
|
||||
first_result = dual_cache.batch_get_cache(keys=["shared_a"])
|
||||
second_result = dual_cache.batch_get_cache(keys=["shared_a"])
|
||||
|
||||
assert first_result is None
|
||||
assert second_result is None
|
||||
assert mock_redis.batch_get_cache.call_count == 2
|
||||
assert "shared_a" not in dual_cache.last_redis_batch_access_time
|
||||
|
||||
|
||||
def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled():
|
||||
mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"})
|
||||
dual_cache = DualCache(
|
||||
in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10
|
||||
)
|
||||
dual_cache.last_redis_batch_access_time["throttled_key"] = time.time()
|
||||
|
||||
result = dual_cache.batch_get_cache(keys=["throttled_key"])
|
||||
|
||||
assert result == [None]
|
||||
mock_redis.batch_get_cache.assert_not_called()
|
||||
|
||||
|
||||
def test_dual_cache_sync_batch_redis_backfill_injects_default_in_memory_ttl():
|
||||
"""Sync batch_get_cache's Redis-to-memory backfill must honor
|
||||
default_in_memory_ttl, same as the async path."""
|
||||
in_memory_cache = InMemoryCache(default_ttl=600)
|
||||
mock_redis = _redis_mock_for_sync_batch({"batch_backfill_key": "redis_value"})
|
||||
dual_cache = DualCache(
|
||||
in_memory_cache=in_memory_cache,
|
||||
redis_cache=mock_redis,
|
||||
default_in_memory_ttl=60,
|
||||
)
|
||||
|
||||
before = time.time()
|
||||
result = dual_cache.batch_get_cache(keys=["batch_backfill_key"])
|
||||
after = time.time()
|
||||
|
||||
assert result == ["redis_value"]
|
||||
expiry = in_memory_cache.ttl_dict["batch_backfill_key"]
|
||||
assert expiry >= before + 60
|
||||
assert expiry <= after + 60
|
||||
|
||||
|
||||
def test_dual_cache_batch_get_cache_forwards_explicit_ttl_to_backfill():
|
||||
"""An explicit ttl kwarg must reach the in-memory backfill flat, not nested
|
||||
under a 'kwargs' key the way the old locals()-forwarding path sent it."""
|
||||
in_memory_cache = InMemoryCache(default_ttl=600)
|
||||
mock_redis = _redis_mock_for_sync_batch({"explicit_ttl_key": "redis_value"})
|
||||
dual_cache = DualCache(in_memory_cache=in_memory_cache, redis_cache=mock_redis)
|
||||
|
||||
before = time.time()
|
||||
result = dual_cache.batch_get_cache(keys=["explicit_ttl_key"], ttl=5)
|
||||
after = time.time()
|
||||
|
||||
assert result == ["redis_value"]
|
||||
expiry = in_memory_cache.ttl_dict["explicit_ttl_key"]
|
||||
assert expiry >= before + 5
|
||||
assert expiry <= after + 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
import asyncio
|
||||
from unittest.mock import MagicMock, patch
|
||||
from collections.abc import Iterator
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
||||
|
||||
|
|
@ -17,6 +17,17 @@ def redis_no_ping():
|
|||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sync_batch_redis_cache(redis_no_ping):
|
||||
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
"litellm._redis.get_redis_client", return_value=MagicMock()
|
||||
) as get_client:
|
||||
cache = RedisCache(host="127.0.0.1", port=6379)
|
||||
cache.redis_client.mget.side_effect = OSError("redis unavailable")
|
||||
get_client.assert_called_once()
|
||||
yield cache
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("namespace", "key", "expected"),
|
||||
[
|
||||
|
|
@ -504,6 +515,173 @@ async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no
|
|||
await call_method(cache)
|
||||
|
||||
|
||||
def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_redis_cache):
|
||||
"""An open breaker must preserve the sync batch read's dictionary fallback."""
|
||||
from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD
|
||||
|
||||
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD):
|
||||
assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {}
|
||||
|
||||
assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]:
|
||||
service_logger = ServiceLogging(mock_testing=True)
|
||||
failing_client = MagicMock()
|
||||
failing_client.mget.side_effect = OSError("redis unavailable")
|
||||
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
"litellm._redis.get_redis_client", return_value=failing_client
|
||||
):
|
||||
cache = RedisCache(host="127.0.0.1", port=6379, service_logger_obj=service_logger)
|
||||
yield cache, service_logger
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sync_batch_get_cache_reports_a_failed_read_from_a_running_loop(
|
||||
sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging],
|
||||
):
|
||||
"""A swallowed Redis failure must still be reported as a service failure event.
|
||||
|
||||
The routing strategies call this blocking read from inside the request's event loop,
|
||||
and the read hides the Redis error by returning an empty dict. Without an emitted
|
||||
failure event, litellm_redis_failed_requests_total stops moving during a Redis
|
||||
outage while the success path keeps reporting, so the dashboards read healthy.
|
||||
"""
|
||||
cache, service_logger = sync_batch_cache_with_service_logger
|
||||
|
||||
assert cache.batch_get_cache(key_list=["lit6729"]) == {}
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert service_logger.mock_testing_sync_failure_hook == 1
|
||||
assert service_logger.mock_testing_async_failure_hook == 1
|
||||
|
||||
|
||||
def test_sync_batch_get_cache_reports_a_failed_read_from_a_worker_thread(
|
||||
sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging],
|
||||
):
|
||||
"""The same report must reach the async hook when the caller has no event loop at all."""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
cache, service_logger = sync_batch_cache_with_service_logger
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {}
|
||||
|
||||
assert service_logger.mock_testing_async_failure_hook == 1
|
||||
|
||||
|
||||
def test_sync_batch_get_cache_reports_a_failed_read_on_an_idle_event_loop(
|
||||
sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging],
|
||||
):
|
||||
"""The report must also go out when the caller holds an open loop that is not running."""
|
||||
cache, service_logger = sync_batch_cache_with_service_logger
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
asyncio.set_event_loop(loop)
|
||||
assert cache.batch_get_cache(key_list=["lit6729"]) == {}
|
||||
finally:
|
||||
asyncio.set_event_loop(None)
|
||||
loop.close()
|
||||
|
||||
assert service_logger.mock_testing_async_failure_hook == 1
|
||||
|
||||
|
||||
def test_sync_batch_get_cache_survives_a_service_callback_that_raises(
|
||||
sync_batch_cache_with_service_logger: tuple[RedisCache, ServiceLogging],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""A failing service callback must not replace the swallowed Redis failure.
|
||||
|
||||
A misconfigured callback raises while emitting (a datadog callback with no
|
||||
DD_API_KEY raises at construction), and the failure event is emitted from inside
|
||||
the except block that swallows the Redis error. If that exception escapes, a Redis
|
||||
outage surfaces to routing as a callback error and the circuit breaker never
|
||||
records the failed read.
|
||||
"""
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import litellm
|
||||
|
||||
from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD
|
||||
|
||||
cache, service_logger = sync_batch_cache_with_service_logger
|
||||
monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"])
|
||||
monkeypatch.setattr(
|
||||
service_logger,
|
||||
"init_prometheus_services_logger_if_none",
|
||||
AsyncMock(side_effect=Exception("callback is misconfigured")),
|
||||
)
|
||||
|
||||
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD):
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
assert pool.submit(cache.batch_get_cache, key_list=["lit6729"]).result() == {}
|
||||
|
||||
assert cache.batch_get_cache(key_list=["lit6729"]) == {}
|
||||
|
||||
|
||||
def test_call_stack_info_skips_breaker_guard_frames():
|
||||
"""Guarded methods must still report their real callers in service-log call_type.
|
||||
|
||||
The breaker guards put their own frames between a method body and its caller, so
|
||||
without skipping them every guarded method logged the guard machinery instead of
|
||||
who actually issued the Redis call.
|
||||
"""
|
||||
from litellm.caching.redis_cache import (
|
||||
RedisCircuitBreaker,
|
||||
_get_call_stack_info,
|
||||
_redis_circuit_breaker_guard_sync,
|
||||
)
|
||||
|
||||
class Guarded:
|
||||
_circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60)
|
||||
|
||||
@_redis_circuit_breaker_guard_sync
|
||||
def probe(self):
|
||||
return _get_call_stack_info()
|
||||
|
||||
def caller_one():
|
||||
return Guarded().probe()
|
||||
|
||||
def caller_two():
|
||||
return caller_one()
|
||||
|
||||
assert caller_two() == "caller_one <- caller_two"
|
||||
|
||||
|
||||
def test_call_stack_info_skips_guard_frames_when_deployed_without_sources(monkeypatch):
|
||||
"""Guard-frame skipping must survive a bytecode-only deployment.
|
||||
|
||||
Shipping `.pyc` files without their `.py` sources leaves the module's `__file__` pointing
|
||||
at the compiled file while every frame still carries the compile-time source path, so a
|
||||
check comparing those two paths stops skipping and the service log then names the guard
|
||||
machinery instead of the real caller.
|
||||
"""
|
||||
from litellm.caching import redis_cache as redis_cache_module
|
||||
from litellm.caching.redis_cache import (
|
||||
RedisCircuitBreaker,
|
||||
_get_call_stack_info,
|
||||
_redis_circuit_breaker_guard_sync,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(redis_cache_module, "__file__", redis_cache_module.__file__ + "c")
|
||||
|
||||
class Guarded:
|
||||
_circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60)
|
||||
|
||||
@_redis_circuit_breaker_guard_sync
|
||||
def probe(self):
|
||||
return _get_call_stack_info()
|
||||
|
||||
def caller_one():
|
||||
return Guarded().probe()
|
||||
|
||||
def caller_two():
|
||||
return caller_one()
|
||||
|
||||
assert caller_two() == "caller_one <- caller_two"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ping):
|
||||
"""A reachable Redis must keep the breaker closed, however many earlier calls failed.
|
||||
|
|
@ -580,7 +758,6 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure():
|
|||
async def swallows_a_failure():
|
||||
await asyncio.sleep(0.02)
|
||||
_record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable"))
|
||||
return None
|
||||
|
||||
async def succeeds_while_the_other_fails():
|
||||
await asyncio.sleep(0.05)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue