diff --git a/litellm/redis/transport/batch.py b/litellm/redis/transport/batch.py index b62de930cf1..6e283c89ce8 100644 --- a/litellm/redis/transport/batch.py +++ b/litellm/redis/transport/batch.py @@ -147,6 +147,7 @@ class PipelineBatch: backend: PipelineBackend name: str = "redis_batch" transaction: bool = False + raise_on_error: bool = False _pending: list[_Op[object]] = field(default_factory=list) # mutable-ok: drained by flush _flush_hooks: list[Callable[[], None]] = field(default_factory=list) # mutable-ok: append-only registry _lock: asyncio.Lock = field(default_factory=asyncio.Lock) @@ -207,7 +208,7 @@ class PipelineBatch: async def run() -> list[object]: async with self.backend.client().pipeline(transaction=self.transaction) as pipe: widths.extend(op.spec.enqueue(pipe) for op in ops) - return await pipe.execute(raise_on_error=False) + return await pipe.execute(raise_on_error=self.raise_on_error) try: replies: Final = await _run_under_circuit_breaker( diff --git a/litellm/redis/transport/breaker.py b/litellm/redis/transport/breaker.py index c10ef1a732b..daa23cade72 100644 --- a/litellm/redis/transport/breaker.py +++ b/litellm/redis/transport/breaker.py @@ -26,11 +26,11 @@ async def run_under_breaker( return failure -def failure_from(exc: BaseException) -> RedisFailure | None: +def failure_from(exc: Exception) -> RedisFailure | None: if isinstance(exc, RedisCircuitBreakerOpenError): - return RedisFailure(kind="breaker_open", message=str(exc)) + return RedisFailure(kind="breaker_open", message=str(exc), exception=exc) if _is_redis_health_failure(exc): - return RedisFailure(kind="transport", message=str(exc), timeout=is_redis_timeout_failure(exc)) + return RedisFailure(kind="transport", message=str(exc), timeout=is_redis_timeout_failure(exc), exception=exc) if is_redis_error(exc): - return RedisFailure(kind="server", message=str(exc)) + return RedisFailure(kind="server", message=str(exc), exception=exc) return None diff --git a/litellm/redis/transport/contract.py b/litellm/redis/transport/contract.py index 983da44a89f..2552a54f9c2 100644 --- a/litellm/redis/transport/contract.py +++ b/litellm/redis/transport/contract.py @@ -1,8 +1,8 @@ import math from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import timedelta -from typing import Final, Literal, Protocol, TypeAlias, TypeVar +from typing import Final, Literal, NoReturn, Protocol, TypeAlias, TypeVar Arg: TypeAlias = str | bytes | int | float EvalReply: TypeAlias = int | bytes | None | tuple["EvalReply", ...] @@ -23,6 +23,13 @@ class RedisFailure: kind: Literal["transport", "server", "decode", "breaker_open"] message: str timeout: bool = False + exception: Exception | None = field(default=None, compare=False, repr=False) + + +def raise_redis_failure(failure: RedisFailure) -> NoReturn: + if failure.exception is not None: + raise failure.exception + raise RuntimeError(failure.message) @dataclass(frozen=True, slots=True) @@ -121,7 +128,9 @@ class RedisTransport(Protocol): async def eval(self, script: str, keys: Sequence[str], args: Sequence[Arg]) -> EvalReply | RedisFailure: ... - def batch(self, *, transaction: bool = False) -> RedisCommandBatch: ... + def batch( + self, *, transaction: bool = False, raise_on_error: bool = False, name: str = "redis_batch" + ) -> RedisCommandBatch: ... def request_batch(self) -> RedisCommandBatch | None: ... diff --git a/litellm/redis/transport/events.py b/litellm/redis/transport/events.py index e6e3fc35b41..ab24a7796c5 100644 --- a/litellm/redis/transport/events.py +++ b/litellm/redis/transport/events.py @@ -55,7 +55,7 @@ class RedisServiceEvents: lambda: self.service_logging.async_service_failure_hook( service=ServiceTypes.REDIS, duration=end_time - start_time, - error=result.message, + error=result.exception if result.exception is not None else result.message, call_type=call_type, start_time=start_time, end_time=end_time, diff --git a/litellm/redis/transport/redis_py.py b/litellm/redis/transport/redis_py.py index 1729b0f1ca2..af6c0adfa73 100644 --- a/litellm/redis/transport/redis_py.py +++ b/litellm/redis/transport/redis_py.py @@ -115,8 +115,15 @@ class RedisPyTransport: await self._run("eval", lambda: _evalsha_or_eval(self.client(), script, sha, namespaced, args)) ) - def batch(self, *, transaction: bool = False) -> "RedisPyCommandBatch": - return RedisPyCommandBatch(self, PipelineBatch(self._backend(runs_alone=False), transaction=transaction)) + def batch( + self, *, transaction: bool = False, raise_on_error: bool = False, name: str = "redis_batch" + ) -> "RedisPyCommandBatch": + return RedisPyCommandBatch( + self, + PipelineBatch( + self._backend(runs_alone=False), name=name, transaction=transaction, raise_on_error=raise_on_error + ), + ) def request_batch(self) -> "RedisPyCommandBatch | None": engine: Final = self.request_engine() diff --git a/tests/unit/redis/stores/test_pod_lock.py b/tests/unit/redis/stores/test_pod_lock.py index f538d7694fb..e06abc9937a 100644 --- a/tests/unit/redis/stores/test_pod_lock.py +++ b/tests/unit/redis/stores/test_pod_lock.py @@ -59,7 +59,7 @@ class ExpiredLockTransport: async def scan(self, prefix: str, *, limit: int) -> NoReturn: raise AssertionError("unexpected scan") - def batch(self, *, transaction: bool = False) -> NoReturn: + def batch(self, *, transaction: bool = False, raise_on_error: bool = False, name: str = "redis_batch") -> NoReturn: raise AssertionError("unexpected batch") def request_batch(self) -> NoReturn: diff --git a/tests/unit/redis/transport/test_redis_py.py b/tests/unit/redis/transport/test_redis_py.py index 76d77a22cba..075ced876a5 100644 --- a/tests/unit/redis/transport/test_redis_py.py +++ b/tests/unit/redis/transport/test_redis_py.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, patch import fakeredis import pytest +from prometheus_client import REGISTRY from redis.asyncio.client import Pipeline, Redis from redis.exceptions import ConnectionError as RedisConnectionError from redis.exceptions import ResponseError @@ -18,7 +19,7 @@ from litellm._service_logger import ServiceLogging from litellm.caching.redis_batch import active_request_redis_batch, request_redis_batch_scope from litellm.caching.redis_cluster_cache import ClusterConnectionGate, RedisClusterCache from litellm.redis.transport.batch import PipelineBatch -from litellm.redis.transport.contract import PubSubMessage, RedisFailure +from litellm.redis.transport.contract import PubSubMessage, RedisFailure, raise_redis_failure from litellm.redis.transport.redis_py import RedisPyCommandBatch, RedisPyTransport from litellm.types.services import ServiceTypes from tests._support.fake_redis import fake_redis_cache, fake_redis_transport @@ -258,6 +259,7 @@ class _RecordingEvents(ServiceLogging): def __init__(self) -> None: super().__init__() self.seen: tuple[_Event, ...] = () + self.failures: tuple[str | Exception, ...] = () def has_sinks(self) -> bool: return True @@ -290,6 +292,7 @@ class _RecordingEvents(ServiceLogging): ) -> None: assert service == ServiceTypes.REDIS and 0 <= duration self.seen = (*self.seen, _Event("failure", call_type, current_service_target(), str(error))) + self.failures = (*self.failures, error) @pytest.mark.asyncio @@ -313,6 +316,51 @@ async def test_every_command_reports_one_service_event_under_the_callers_target( ) +@pytest.mark.asyncio +async def test_failure_events_keep_original_exception_objects( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + events: Final = _RecordingEvents() + transport: Final = fake_redis_transport(raw, failure_threshold=1, events=events) + await raw.set("wrong", "text") + + failed_command: Final = await transport.rpush("wrong", "item") + server.connected = False + failed_connection: Final = await transport.get("wrong") + refused: Final = await transport.get("wrong") + + assert isinstance(failed_command, RedisFailure) + assert isinstance(failed_connection, RedisFailure) + assert isinstance(refused, RedisFailure) and refused.kind == "breaker_open" + assert events.failures == (failed_command.exception, failed_connection.exception) + assert isinstance(events.failures[0], ResponseError) + assert isinstance(events.failures[1], RedisConnectionError) + + +@pytest.mark.asyncio +async def test_failed_request_metrics_preserve_redis_exception_classes( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) + labels: Final = {"redis": "redis", "error_class": "ResponseError", "function_name": "rpush"} + connection_labels: Final = {"redis": "redis", "error_class": "ConnectionError", "function_name": "get"} + transport: Final = fake_redis_transport(raw, failure_threshold=1, events=ServiceLogging()) + before: Final = REGISTRY.get_sample_value("litellm_redis_failed_requests_total", labels) or 0 + before_connection: Final = REGISTRY.get_sample_value("litellm_redis_failed_requests_total", connection_labels) or 0 + await raw.set("wrong", "text") + + failed_command: Final = await transport.rpush("wrong", "item") + server.connected = False + failed_connection: Final = await transport.get("wrong") + refused: Final = await transport.get("wrong") + + assert isinstance(failed_command, RedisFailure) and failed_command.kind == "server" + assert isinstance(failed_connection, RedisFailure) and failed_connection.kind == "transport" + assert isinstance(refused, RedisFailure) and refused.kind == "breaker_open" + assert REGISTRY.get_sample_value("litellm_redis_failed_requests_total", labels) == before + 1 + assert REGISTRY.get_sample_value("litellm_redis_failed_requests_total", connection_labels) == before_connection + 1 + + @pytest.mark.asyncio async def test_a_call_the_breaker_refuses_reports_no_service_event( server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis @@ -692,6 +740,51 @@ async def test_a_batch_sends_every_declared_command_in_one_round_trip( assert 0 < await raw.pttl("ns:b") <= 30_000 +@pytest.mark.asyncio +@pytest.mark.parametrize("strict", (False, True)) +async def test_explicit_batch_error_policy_preserves_other_writes_and_event_outcome( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis, strict: bool +) -> None: + events: Final = _RecordingEvents() + cache: Final = fake_redis_cache(server, namespace="ns", service_logger=events) + await _settle() + await raw.set("ns:wrong", "text") + before: Final = len(events.seen) + batch: Final = cache.transport().batch(raise_on_error=strict, name="legacy_batch") + bad: Final = batch.rpush("wrong", "item") + good: Final = batch.set("written", "value", ttl=None) + + failure: Final = await bad + assert isinstance(failure, RedisFailure) and failure.kind == "server" + with pytest.raises(ResponseError, match="WRONGTYPE") as caught: + raise_redis_failure(failure) + assert caught.value is failure.exception + assert await good == (failure if strict else True) + assert await raw.get("ns:written") == b"value" + assert await raw.get("ns:wrong") == b"text" + await _settle() + assert tuple((event.outcome, event.call_type) for event in events.seen[before:]) == ( + ("failure" if strict else "success", "legacy_batch"), + ) + + +@pytest.mark.asyncio +async def test_transport_failure_retains_original_exception_for_public_facade( + server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis +) -> None: + transport: Final = fake_redis_transport(raw, failure_threshold=1) + server.connected = False + + failure: Final = await transport.get("key") + assert isinstance(failure, RedisFailure) and failure.kind == "transport" + with pytest.raises(RedisConnectionError) as caught: + raise_redis_failure(failure) + assert caught.value is failure.exception + assert str(caught.value) == failure.message + assert failure == RedisFailure(kind="transport", message=failure.message) + assert "exception=" not in repr(failure) + + @pytest.mark.asyncio async def test_a_failing_command_in_a_batch_fails_only_its_own_reply( server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis