diff --git a/AGENTS.md b/AGENTS.md index f4b1de57807..5682f067732 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,3 +1,5 @@ +For Redis-backed access, typed RedisCache extensions, or explicit Rust ports, follow [the Redis access skill](litellm-rust/.agents/skills/redis-access/SKILL.md) + Do not write comments unless they are any of: - absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear) - used as an input for tools to read and act on. For example: diff --git a/litellm-rust/.agents/skills/redis-access/SKILL.md b/litellm-rust/.agents/skills/redis-access/SKILL.md new file mode 100644 index 00000000000..5e6ba2423dd --- /dev/null +++ b/litellm-rust/.agents/skills/redis-access/SKILL.md @@ -0,0 +1,96 @@ +--- +name: redis-access +description: Use and extend RedisCache's provisional typed operations while keeping Redis transport details private and preserving existing behavior +--- + +# Redis access + +## Use the existing interface + +`litellm/caching/redis_cache.py` owns the caller-facing `RedisCache` interface. Internal workflows use its typed operations. Their value types live in `litellm/types/caching.py`, transport and reply adapters in `litellm/caching/_redis_capabilities.py`, and atomic scripts in `litellm/caching/_redis_scripts.py` + +Choose the existing operation that represents the caller's need. Do not add another Redis facade, a protocol for every operation, or a generic command wrapper + +| Need | Existing operations | +| --- | --- | +| Distinguish missing from unavailable | `async_key_presence` | +| Read ordered counters or cached values | `async_read_atomic_counts`, `async_read_values` | +| Maintain request and token windows | `async_increment_window_counters`, `async_reserve_window_capacity`, `async_increment_guarded_counters` | +| Own concurrency slots | `async_acquire_concurrency`, `async_release_concurrency`, `async_count_concurrency` | +| Increment, reserve or refund counters | `async_increment_counters`, `async_reserve_counter`, `async_refund_counter` | +| Store and consume reservation records | `async_save_record`, `async_pop_record`, `async_delete_record` | +| Enqueue serialized updates | `async_enqueue_updates` with `CacheQueueUpdate` | +| Count session iterations or spend | `async_increment_iterations`, `async_increment_spend` | +| Claim affinity | `async_claim_affinity` | +| Read and record login throttling | `async_read_throttle_ttls`, `async_record_throttle_failure` | +| Own a lease | `async_acquire_lease`, `async_extend_lease`, `async_release_lease`, `async_lease_held` | +| Publish or subscribe | `async_publish`, `async_publish_many`, `async_subscribe` | +| Inspect connection capacity | `get_connection_pool_info` | + +Most atomic operations return an awaitable `CacheOperation[T]`. Await it to obtain `T | CacheFailure`. Pub/sub and value-read methods are coroutines with explicit result unions. `async_key_presence` returns `KeyPresence`, the literal union `"present" | "missing" | "unavailable"`. Check the method signature rather than assuming that every operation has the same return shape + +Legacy cache read/write methods remain available. They are not interchangeable with the new operations when their error handling, encoding or missing-value semantics differ + +## Keep policy in the owner + +The owner supplies domain keys, values, thresholds, tokens and TTL policy. It decides whether an unavailable cache permits a request, rejects it, or triggers a local fallback. `RedisCache` applies its namespace once, executes the supplied TTL, owns atomic command sequences, validates replies and translates backend errors + +For example, the gateway's single-use marker check needs the raw presence distinction, including a stored JSON `null` counting as present and a wrong-type key counting as unavailable + +```python +from litellm.caching.redis_cache import RedisCache +from litellm.types.caching import KeyPresence + +async def marker_state(cache: RedisCache, key: str) -> KeyPresence: + return await cache.async_key_presence(key) +``` + +Do not replace this operation with `async_get_cache`. Its default error handling can hide an outage as a miss, and it decodes stored values. `raise_on_error=True` does not supply a backend-neutral result contract + +`CacheFailure.kind` distinguishes unavailable transport from malformed replies. `reason` distinguishes an open circuit, a timeout and other failures. Do not parse error messages or import Redis exceptions in the owner. Preserve the owner's existing malformed-response behavior separately from outage fallback. The optional validation error exists for legacy validation contracts, not as a transport exception + +Use `log_redis_failure` when preserving existing Redis failure diagnostics. An open circuit stays at debug level and timeouts retain throttling. Do not turn every failed operation into a warning or error storm + +## Batching and subscriptions + +Atomic methods that accept `batch=` can share a `RedisBatch`. The cache declares the command internally, and the caller receives the same typed operation result. Declare independent operations before awaiting them to retain batching. The non-transactional batch is not an atomicity boundary. Each Lua operation retains its own atomic boundary, ordered replies and failure isolation + +`CacheOperation.done` and `on_settled` support the existing cancellation and settlement workflows. Settlement callbacks receive typed results, never clients, pipelines or script handles. Preserve these hooks when a request can be cancelled after a write has reached Redis + +`async_subscribe` returns a `CacheSubscription`. Call `receive(timeout_seconds=...)` for a payload, idle `None`, or `CacheFailure`, and call `close` during shutdown and reconnect cleanup. Redis subscription envelopes, pooling, routing and reconnection transport remain private. The owner retains its reconnect timing, event validation and message application policy + +## Provisional does not mean deprecated + +The new typed API is provisional. Public operations emit `FutureWarning` with the message `RedisCache. is experimental and may change in a future release`. Operation lifecycle methods and subscription methods follow the same policy. Keep the warning at the user's call site with the correct stack level and ordinary warning filtering. Do not use a deprecation decorator for these operations and do not globally silence warnings + +Legacy raw accessors remain callable for compatibility. `init_async_client`, `init_pubsub_client`, and `async_register_script` retain PEP 702 metadata through `typing_extensions.deprecated(..., category=None)`. Existing raw attributes and compatibility exports remain available to external consumers. `RedisBatch.script` is also a deprecated compatibility surface, not an internal extension API + +Only the cache implementation may use the private transport bridge. Its deliberate dispatch through legacy factories preserves subclass overrides and loop-local client reuse. Exact static suppressions there and on compatibility-only re-exports document that narrow exception. They are not permission for another internal caller to use a raw accessor + +## Add a missing capability + +Start with the complete operation and its observable contract. Add a typed method to `RedisCache` and a narrow immutable result type only if the existing operations cannot express the requirement. A small domain protocol is appropriate for injection when it belongs to an existing owner, but do not mirror redis-py + +Keep redis-py imports, raw attributes, connection construction, script registration, Lua execution, pipelines, reply decoding and Redis exception translation inside the cache implementation. Do not expose `get_client`, generic `execute_command`, raw-client callbacks, raw Redis replies or Redis exceptions. A typed protocol that returns a pipeline still leaks the transport + +Preserve existing Lua bodies verbatim when moving them. Do not shorten or rewrite scripts during interface extraction. Preserve command order, arguments, namespace and cluster-slot placement, value formats, TTL units and refresh rules, batching, retry behavior, telemetry and cancellation. In particular, window timestamps used by guarded updates must retain their textual identity, including fractional values + +Return a typed failure when transport or reply validation fails. Keep missing data distinct from an outage where the owner needs that distinction. Do not retry an ambiguous write through a second backend, because a timeout can follow a committed increment, pop or lock acquisition + +Add the provisional warning at the public entry point. Keep implementation helpers private, and never route an internal caller through a compatibility-only raw accessor or alias + +## Verify behavior + +Extend the existing mapped unit tests. Use in-process transports or injected domain dependencies and assert observable values, namespaces, TTLs, ordering, concurrency ownership, malformed replies, missing keys and outage decisions. Do not mock the new method merely to echo the result being asserted + +Test warning category, experimental wording and caller location. Test legacy access under warnings-as-errors and retain subclass dispatch, client reuse and compatibility exports. Audit direct calls, aliased imports, raw attributes, `getattr` and indirect helper paths when checking the boundary. Cache implementation and compatibility-only exports are explicit exemptions. Independent semantic and vector-store Redis integrations are separate from RedisCache + +For command-level verification, use the existing Redis integration grouping with isolated keys and cleanup. Exercise atomicity, contention, expiry, pub/sub lifecycle, cluster behavior where supported, and batch partial failures and NoScript retry. Report untested topologies rather than treating standalone Redis as cluster coverage + +Run scoped lint and typing comparison gates without changing budgets or ceilings. Avoid broad auto-fixes in legacy caller or test files during extraction + +## Backend changes require a separate request + +redis-py remains a valid private implementation. An equivalent Rust operation does not authorize switching callers or changing backend selection. Interface extraction, configuration wiring and backend migration are separate work + +If explicitly asked to port an implementation, reuse `crates/redis-client` and its `AGENTS.md`, resolve configuration in `crates/config`, and keep native conversion at the Python bridge. Preserve the same typed operations and domain policies. Verify standalone and supported cluster behavior against real Redis before switching. Rust pipelines are not substitutes for atomic scripts or transactions diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e62a9463cb5..2a025b533cc 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,5 +1,7 @@ # Rust workspace rules +For Redis-backed operations, typed RedisCache extensions, or explicit Rust ports, follow [.agents/skills/redis-access/SKILL.md](.agents/skills/redis-access/SKILL.md) + For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) diff --git a/litellm/caching/_redis_capabilities.py b/litellm/caching/_redis_capabilities.py new file mode 100644 index 00000000000..b2218221282 --- /dev/null +++ b/litellm/caching/_redis_capabilities.py @@ -0,0 +1,496 @@ +from __future__ import annotations + +import asyncio +import warnings +from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence, Sized +from contextlib import suppress +from types import MappingProxyType +from typing import ( + TYPE_CHECKING, + Final, + Generic, + Literal, + Protocol, + TypeAlias, + TypeVar, + cast, # noqa: TID251 # redis transport protocols only; all replies are validated +) + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.types.caching import ( + CacheFailure, + CacheOperation, + CachePoolInfo, + CachePublish, + CacheQueueUpdate, + CacheSubscription, + CapacityExceeded, + CounterUpdate, + LimitedIncrement, + RedisPipelineRpushOperation, + WindowCount, +) + +if TYPE_CHECKING: + from litellm.caching.redis_batch import BatchResult, RedisBatch + from litellm.caching.redis_cache import RedisCache + +_T: Final = TypeVar("_T") +_Arg: TypeAlias = str | bytes | int | float +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) +_INTS: Final = TypeAdapter(tuple[int, ...]) +_FLOATS: Final = TypeAdapter(tuple[float, ...]) +_INT: Final = TypeAdapter(int) +_FLOAT: Final = TypeAdapter(float) +_TEXT: Final = TypeAdapter(str) +_OPTIONAL_TEXT: Final = TypeAdapter[str | None](str | None) +_PAIR: Final = TypeAdapter(tuple[int, int]) +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_PAYLOAD: Final = TypeAdapter[str | bytes](str | bytes) +_OPTIONAL_INT: Final = TypeAdapter[int | None](int | None) +_CACHE_VALUES: Final = TypeAdapter(tuple[str | bytes | None, ...]) + + +def provisional(name: str) -> None: + warnings.warn( + f"RedisCache.{name} is experimental and may change in a future release", + FutureWarning, + stacklevel=3, + ) + + +def failure(error: BaseException) -> CacheFailure: + from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure + + reason: Final = ( + "circuit_open" + if isinstance(error, RedisCircuitBreakerOpenError) + else "timeout" + if is_redis_timeout_failure(error) + else "other" + ) + return CacheFailure( + kind="unavailable", message=f"Cache operation failed ({type(error).__name__}): {error}", reason=reason + ) + + +def invalid_reply() -> CacheFailure: + return CacheFailure(kind="invalid_response", message="Cache returned an invalid operation response") + + +def validated(decode: Callable[[object], _T | CacheFailure], reply: object) -> _T | CacheFailure: + try: + return decode(reply) + except ValidationError as error: + return CacheFailure(kind="invalid_response", message=str(error), validation_error=error) + except (ValueError, TypeError, UnicodeError, SyntaxError): + return invalid_reply() + + +class _Operation(Generic[_T]): + def __init__( + self, + run: Callable[[], Awaitable[object]], + decode: Callable[[object], _T | CacheFailure], + batched: BatchResult[object] | None, + ) -> None: + self._run = run + self._decode = decode + self._batched = batched + self._task: asyncio.Task[_T | CacheFailure] | None = None + + @property + def done(self) -> bool: + provisional("operation.done") + return self._batched.done if self._batched is not None else self._task is not None and self._task.done() + + def __await__(self) -> Generator[object, None, _T | CacheFailure]: + return self._wait().__await__() + + async def _execute(self) -> _T | CacheFailure: + try: + reply: Final = await self._run() if self._batched is None else await self._batched + except Exception as error: # noqa: BLE001 # backend failures are values; cancellation propagates + return failure(error) + return validated(self._decode, reply) + + def _ensure_task(self) -> asyncio.Task[_T | CacheFailure]: + if self._task is None: + self._task = asyncio.create_task(self._execute()) + return self._task + + async def _wait(self) -> _T | CacheFailure: + return await self._ensure_task() + + def on_settled(self, hook: Callable[[asyncio.Future[_T | CacheFailure]], Awaitable[None] | None]) -> None: + provisional("operation.on_settled") + + async def notify(raw: asyncio.Future[object]) -> None: + result: Final[asyncio.Future[_T | CacheFailure]] = asyncio.get_running_loop().create_future() + if raw.cancelled(): + result.cancel() + elif (error := raw.exception()) is not None: + result.set_result(failure(error)) + else: + result.set_result(validated(self._decode, raw.result())) + follow_up: Final = hook(result) + if follow_up is not None: + await follow_up + + if self._batched is not None: + self._batched.on_settled(notify) + return + + def completed(result: asyncio.Future[_T | CacheFailure]) -> None: + follow_up: Final = hook(result) + if follow_up is not None: + asyncio.ensure_future(follow_up) + + self._ensure_task().add_done_callback(completed) + + +def operation( + cache: RedisCache, + source: str, + keys: Sequence[str], + args: Sequence[_Arg], + decode: Callable[[object], _T | CacheFailure], + batch: RedisBatch | None, +) -> CacheOperation[_T]: + from litellm.caching.redis_batch import _Script # pyright: ignore[reportPrivateUsage] # private cache transport + + runner: Final = cache._script_runner(source) # pyright: ignore[reportPrivateUsage] # private cache implementation + batched: Final = ( + None if batch is None else batch._declare(_Script(cache, source, runner, keys, args)) # pyright: ignore[reportPrivateUsage] # private cache transport + ) + return _Operation(lambda: runner(keys=keys, args=args), decode, batched) + + +def delete_record(cache: RedisCache, key: str) -> CacheOperation[bool]: + return _Operation(lambda: cache.async_delete_cache(key), boolean, None) + + +def enqueue_updates(cache: RedisCache, updates: Sequence[CacheQueueUpdate]) -> CacheOperation[tuple[int, ...]]: + commands: Final = tuple(RedisPipelineRpushOperation(key=item.key, values=item.values) for item in updates) + return _Operation(lambda: cache.async_rpush_pipeline(commands), lambda reply: counts(reply, len(commands)), None) + + +def integer(reply: object) -> int: + return _INT.validate_python(reply) + + +def number(reply: object) -> float: + return _FLOAT.validate_python(reply) + + +def acknowledged(reply: object) -> None | CacheFailure: + return None if integer(reply) == 1 else invalid_reply() + + +def boolean(reply: object) -> bool | CacheFailure: + value: Final = integer(reply) + return bool(value) if value in (0, 1) else invalid_reply() + + +def acquired(reply: object) -> Literal["acquired", "held"] | CacheFailure: + if reply is None or reply is False: + return "held" + return "acquired" if reply in (True, "OK", b"OK") else invalid_reply() + + +def text(reply: object) -> str | None: + return _OPTIONAL_TEXT.validate_python(reply) + + +def pin(reply: object) -> JsonValue | CacheFailure: + try: + value: Final = _TEXT.validate_python(reply.decode("utf-8") if isinstance(reply, bytes) else reply) + except UnicodeError as error: + return failure(error) + try: + return _JSON.validate_json(value) + except ValidationError: + return value + + +def block_ttls(reply: object) -> tuple[int, int]: + return _PAIR.validate_python(reply) + + +def limited(reply: object) -> LimitedIncrement | CacheFailure: + accepted, count = _PAIR.validate_python(reply) + return LimitedIncrement(bool(accepted), count) if accepted in (0, 1) else invalid_reply() + + +def counts(reply: object, size: int) -> tuple[int, ...] | CacheFailure: + values: Final = _INTS.validate_python(reply) + return values if len(values) == size else invalid_reply() + + +def totals(reply: object, size: int) -> tuple[float, ...] | CacheFailure: + values: Final = _FLOATS.validate_python(reply) + return values if len(values) == size else invalid_reply() + + +def optional_counts(reply: object, size: int) -> tuple[int | None, ...] | CacheFailure: + values: Final = _OBJECTS.validate_python(reply) + if len(values) != size: + return invalid_reply() + return tuple(None if value is None or value is False else integer(value) for value in values) + + +def _window_start(reply: object) -> str: + number(reply) + return _TEXT.validate_python(reply) if isinstance(reply, (str, bytes)) else str(reply) + + +def window_counts(reply: object, size: int) -> tuple[WindowCount, ...] | CacheFailure: + values: Final = _OBJECTS.validate_python(reply) + if len(values) != size * 2: + return invalid_reply() + return tuple(WindowCount(_window_start(values[i]), integer(values[i + 1])) for i in range(0, len(values), 2)) + + +def _exceeded(values: tuple[int, ...], size: int) -> CapacityExceeded | CacheFailure: + if len(values) != 4 or not 1 <= values[1] <= size: + return invalid_reply() + return CapacityExceeded(index=values[1] - 1, count=values[2], limit=values[3]) + + +def reservation(reply: object, size: int) -> tuple[WindowCount, ...] | CapacityExceeded | CacheFailure: + values: Final = _OBJECTS.validate_python(reply) + if not values: + return invalid_reply() + if integer(values[0]) == 1: + return _exceeded(_INTS.validate_python(values), size) + if integer(values[0]) != 0 or len(values) != 1 + size * 2: + return invalid_reply() + return tuple(WindowCount(_window_start(values[i + 1]), integer(values[i])) for i in range(1, len(values), 2)) + + +def slots(reply: object, size: int) -> tuple[int, ...] | CapacityExceeded | CacheFailure: + values: Final = _INTS.validate_python(reply) + if not values: + return invalid_reply() + if values[0] == 1: + return _exceeded(values, size) + return values[1:] if values[0] == 0 and len(values) == size + 1 else invalid_reply() + + +def guarded(reply: object, size: int) -> tuple[CounterUpdate, ...] | CacheFailure: + values: Final = _INTS.validate_python(reply) + if len(values) != size * 2 or any(values[i] not in (0, 1) for i in range(0, len(values), 2)): + return invalid_reply() + return tuple(CounterUpdate(bool(values[i]), values[i + 1]) for i in range(0, len(values), 2)) + + +class _RegisteredScript(Protocol): + def __call__(self, keys: Sequence[str], args: Sequence[_Arg], client: object = None) -> Awaitable[object]: ... + + +class _ScriptClient(Protocol): + def register_script(self, source: str) -> _RegisteredScript: ... + + +class _ClusterScripts(Protocol): + def script_load(self, source: str) -> str | Awaitable[str]: ... + async def evalsha(self, sha: str, numkeys: int, *args: _Arg) -> object: ... + + +def script_executor(cache: RedisCache, source: str) -> _RegisteredScript: + transport: Final = cache._get_async_client() # pyright: ignore[reportPrivateUsage] # private cache transport + if hasattr(transport, "register_script"): + registered: Final = cast(_ScriptClient, transport).register_script(source) # cast-ok: redis transport only + + async def standalone(keys: Sequence[str], args: Sequence[_Arg], client: object = None) -> object: + names: Final = tuple(cache.check_and_fix_namespace(key) for key in keys) + return await registered(keys=names, args=args, client=client) + + return standalone + if not hasattr(transport, "script_load"): + raise ValueError("Redis client does not support Lua script registration") + cluster: Final = cast(_ClusterScripts, transport) # cast-ok: RedisCluster omits register_script + initial_load: Final = cluster.script_load(source) + load_task: Final = asyncio.ensure_future(initial_load) if isinstance(initial_load, Awaitable) else None + + async def clustered(keys: Sequence[str], args: Sequence[_Arg], client: object = None) -> object: + from redis.exceptions import NoScriptError + + names: Final = tuple(cache.check_and_fix_namespace(key) for key in keys) + sha: Final = _TEXT.validate_python(await load_task if load_task is not None else initial_load) + try: + return await cluster.evalsha(sha, len(names), *names, *args) + except NoScriptError: + reloaded: Final = cluster.script_load(source) + retry_sha: Final = _TEXT.validate_python(await reloaded if isinstance(reloaded, Awaitable) else reloaded) + return await cluster.evalsha(retry_sha, len(names), *names, *args) + + return clustered + + +class _Pubsub(Protocol): + async def subscribe(self, channel: str) -> object: ... + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> object: ... + async def aclose(self) -> None: ... + + +class _PublishPipeline(Protocol): + def publish(self, channel: str, message: str) -> object: ... + async def execute(self) -> object: ... + async def __aenter__(self) -> _PublishPipeline: ... + async def __aexit__(self, *args: object) -> None: ... + + +class _PubsubClient(Protocol): + async def publish(self, channel: str, message: str) -> object: ... + def pipeline(self, *, transaction: bool) -> _PublishPipeline: ... + def pubsub(self) -> _Pubsub: ... + + +class _Subscription: + def __init__(self, pubsub: _Pubsub) -> None: + self._pubsub: Final = pubsub + self._closed = False + + async def receive(self, timeout_seconds: float = 0.0) -> str | bytes | None | CacheFailure: + provisional("subscription.receive") + if self._closed: + return CacheFailure("unavailable", "Subscription is closed") + try: + reply: Final = await self._pubsub.get_message(ignore_subscribe_messages=True, timeout=timeout_seconds) + if reply is None: + return None + message: Final = TypeAdapter(dict[str, object]).validate_python(reply) + if message.get("type") != "message": + return None + return _PAYLOAD.validate_python(message.get("data")) + except (ValueError, TypeError, UnicodeError): + return invalid_reply() + except Exception as error: # noqa: BLE001 # backend exceptions must stay behind the cache facade + return failure(error) + + async def close(self) -> None | CacheFailure: + provisional("subscription.close") + if self._closed: + return None + self._closed = True + try: + await self._pubsub.aclose() + except Exception as error: # noqa: BLE001 # backend exceptions must stay behind the cache facade + return failure(error) + return None + + +def _pubsub_client(cache: RedisCache) -> _PubsubClient: + return cast( # cast-ok: cache-owned redis client; wire responses are validated separately + _PubsubClient, + cache._get_pubsub_client(), # pyright: ignore[reportPrivateUsage] # private cache transport + ) + + +async def publish(cache: RedisCache, channel: str, message: str) -> int | CacheFailure: + try: + reply: Final = await _pubsub_client(cache).publish(cache.check_and_fix_namespace(channel), message) + return validated(integer, reply) + except Exception as error: # noqa: BLE001 # backend exceptions must stay behind the cache facade + return failure(error) + + +async def publish_many(cache: RedisCache, messages: Sequence[CachePublish]) -> tuple[int, ...] | CacheFailure: + try: + async with _pubsub_client(cache).pipeline(transaction=False) as pipe: + for item in messages: + pipe.publish(cache.check_and_fix_namespace(item.channel), item.message) + reply: Final = await pipe.execute() + return validated(lambda value: counts(value, len(messages)), reply) + except Exception as error: # noqa: BLE001 # backend exceptions must stay behind the cache facade + return failure(error) + + +async def subscribe(cache: RedisCache, channel: str) -> CacheSubscription | CacheFailure: + try: + pubsub: Final = _pubsub_client(cache).pubsub() + except Exception as error: # noqa: BLE001 # construction is part of the backend boundary + return failure(error) + try: + await pubsub.subscribe(cache.check_and_fix_namespace(channel)) + except BaseException as error: # noqa: BLE001 # cancellation must close a partially subscribed connection + with suppress(Exception): + await pubsub.aclose() + if isinstance(error, Exception): + return failure(error) + raise + return _Subscription(pubsub) + + +class _SyncClusterValues(Protocol): + def mget_nonatomic(self, *, keys: Sequence[str]) -> object: ... + + +def cluster_values_sync(cache: RedisCache, keys: Sequence[str]) -> tuple[str | bytes | None, ...]: + client: Final = cast(_SyncClusterValues, cache.redis_client) # cast-ok: cache-owned cluster transport + return _CACHE_VALUES.validate_python(client.mget_nonatomic(keys=keys)) + + +class _ClusterValues(Protocol): + async def mget_nonatomic(self, *, keys: Sequence[str]) -> object: ... + + +async def cluster_values(cache: RedisCache, keys: Sequence[str]) -> tuple[str | bytes | None, ...]: + client: Final = cast( # cast-ok: RedisClusterCache configures a cluster client; validate its wire response + _ClusterValues, + cache._get_async_client(), # pyright: ignore[reportPrivateUsage] # cache-owned cluster transport + ) + return _CACHE_VALUES.validate_python(await client.mget_nonatomic(keys=keys)) + + +async def read_values(cache: RedisCache, keys: Sequence[str]) -> Mapping[str, object] | CacheFailure: + try: + namespaced: Final = [cache.check_and_fix_namespace(key) for key in keys] + from litellm.caching.redis_cache import ( + _run_under_circuit_breaker, # pyright: ignore[reportPrivateUsage] # private cache health boundary + ) + + replies: Final = await _run_under_circuit_breaker( + cache._circuit_breaker, # pyright: ignore[reportPrivateUsage] # same breaker as legacy MGET + "async_read_values", + lambda: cache._admit_async_call( # pyright: ignore[reportPrivateUsage] # preserve cluster connection admission + lambda: cache._async_run_redis_mget_operation(keys=namespaced) # pyright: ignore[reportPrivateUsage] # topology-specific transport + ), + ) + values: Final = _CACHE_VALUES.validate_python(replies) + if len(values) != len(keys): + return invalid_reply() + return MappingProxyType( + { + key: cache._get_cache_logic(value) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # legacy cache decoder owns response shapes + for key, value in zip(keys, values) + } + ) + except (ValueError, TypeError, UnicodeError, SyntaxError): + return invalid_reply() + except Exception as error: # noqa: BLE001 # backend exceptions must stay behind the cache facade + return failure(error) + + +def pool_info(cache: RedisCache) -> CachePoolInfo | CacheFailure | None: + try: + client: Final = TypeAdapter(object).validate_python(cache.redis_client) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # legacy transport narrowed to object + pool: Final[object] = getattr(client, "connection_pool", None) + if pool is None: + return None + maximum: Final = _OPTIONAL_INT.validate_python(getattr(pool, "max_connections", None)) + available: Final[object] = getattr(pool, "_available_connections", None) + in_use: Final[object] = getattr(pool, "_in_use_connections", None) + connection_class: Final[object] = getattr(pool, "connection_class", None) + return CachePoolInfo( + available_connections=len(available) if isinstance(available, Sized) else None, + in_use_connections=len(in_use) if isinstance(in_use, Sized) else None, + max_connections=maximum, + connection_class=_OPTIONAL_TEXT.validate_python(getattr(connection_class, "__name__", None)), + ) + except (ValueError, TypeError, UnicodeError): + return invalid_reply() + except Exception as error: # noqa: BLE001 # optional diagnostics cannot leak redis transport exceptions + return failure(error) diff --git a/litellm/caching/_redis_scripts.py b/litellm/caching/_redis_scripts.py new file mode 100644 index 00000000000..2b7af19461c --- /dev/null +++ b/litellm/caching/_redis_scripts.py @@ -0,0 +1,395 @@ +from typing import Final + +WINDOW_INCREMENT: Final = """ +local results = {} +local now = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) + +-- Process each window/counter pair +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local increment_value = 1 + + -- Check if window exists and is valid + local window_start = redis.call('GET', window_key) + if not window_start or (now - tonumber(window_start)) >= window_size then + -- Reset window and counter + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment_value) + redis.call('EXPIRE', window_key, window_size) + redis.call('EXPIRE', counter_key, window_size) + table.insert(results, tostring(now)) -- window_start + table.insert(results, increment_value) -- counter + else + local counter = redis.call('INCR', counter_key) + -- This happens when window_key exists but counter_key doesn't (e.g., tokens key + -- created after requests key when both share the same window_key) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 then + redis.call('EXPIRE', counter_key, window_size) + end + table.insert(results, window_start) -- window_start + table.insert(results, counter) -- counter + end +end + +return results +""" + +COUNTER_READ: Final = "return redis.call('MGET', unpack(KEYS))" + +WINDOW_RESERVE: Final = """ +-- Atomic check-and-increment-by-N across one or more descriptors. +-- All-or-nothing: if any descriptor would exceed its limit, no counter is +-- modified. +-- +-- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied +-- timestamp so that window resets are deterministic across replicas with +-- skewed wall-clocks. This prevents a clock-skew-induced reopening of the +-- TOCTOU window across multi-replica deployments. +-- +-- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. +-- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: +-- ARGV[(i-1)*4 + 1] = limit +-- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) +-- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) +-- +-- Return on success: +-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... } +-- Return on over-limit: { 1, descriptor_index, current_counter, limit } +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local descriptor_count = #KEYS / 2 +local reset_windows = {} + +-- Pass 1: read state, validate. Abort without writing if any over limit. +local descriptor_state = {} +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local limit = tonumber(ARGV[arg_base]) + local increment = tonumber(ARGV[arg_base + 1]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_start = redis.call('GET', window_key) + local window_expired = (not window_start) or + ((now - tonumber(window_start)) >= window_size) + + local current_counter + if window_expired then + current_counter = 0 + else + current_counter = tonumber(redis.call('GET', counter_key) or 0) + end + + local blocked + if increment > 0 then + blocked = current_counter + increment > limit + else + blocked = current_counter >= limit + end + if blocked then + return { 1, i, current_counter, limit } + end + + descriptor_state[i] = { window_expired, current_counter, window_start } +end + +-- Pass 2: all checks passed. Apply increments. +local results = { 0 } +for i = 1, descriptor_count do + local window_key = KEYS[(i - 1) * 2 + 1] + local counter_key = KEYS[(i - 1) * 2 + 2] + local arg_base = (i - 1) * 4 + 1 + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local window_size = tonumber(ARGV[arg_base + 3]) + + local window_expired = descriptor_state[i][1] + local active_window_start + + if window_expired then + active_window_start = now + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end + redis.call('SET', window_key, tostring(now)) + redis.call('SET', counter_key, increment) + redis.call('EXPIRE', window_key, window_size) + if ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, increment) + else + active_window_start = tonumber(descriptor_state[i][3]) + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, new_counter) + end + table.insert(results, active_window_start) +end + +return results +""" + +GUARDED_INCREMENT: Final = """ +local results = {} +for i = 1, #KEYS, 2 do + local window_key = KEYS[i] + local counter_key = KEYS[i + 1] + local arg_base = ((i - 1) / 2) * 3 + 1 + local expected_window_start = ARGV[arg_base] + local increment = tonumber(ARGV[arg_base + 1]) + local ttl = tonumber(ARGV[arg_base + 2]) + local active_window_start = redis.call('GET', window_key) + + if active_window_start and active_window_start == expected_window_start then + local new_counter = redis.call('INCRBY', counter_key, increment) + local current_ttl = redis.call('TTL', counter_key) + if current_ttl == -1 and ttl > 0 then + redis.call('EXPIRE', counter_key, ttl) + end + table.insert(results, 1) + table.insert(results, new_counter) + else + table.insert(results, 0) + table.insert(results, tonumber(redis.call('GET', counter_key) or 0)) + end +end +return results +""" + +SLOT_ACQUIRE: Final = """ +-- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. +-- Each gauge key is a sorted set of per-request slot ids scored by acquire +-- time (Redis server clock). In-flight requests are counted by ZCARD after +-- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM +-- counters the gauge is never reset while requests are in flight, a +-- rejected request never occupies a slot, and a slot leaked by a crashed +-- worker self-heals after the slot TTL even under continuous traffic. +-- +-- KEYS: one gauge zset key per descriptor. +-- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id). +-- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +for i = 1, #KEYS do + local limit = tonumber(ARGV[(i - 1) * 3 + 1]) + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl) + local in_flight = redis.call('ZCARD', KEYS[i]) + if in_flight + 1 > limit then + return { 1, i, in_flight, limit } + end +end +local results = { 0 } +for i = 1, #KEYS do + local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) + local slot_id = ARGV[(i - 1) * 3 + 3] + redis.call('ZADD', KEYS[i], now, slot_id) + redis.call('EXPIRE', KEYS[i], slot_ttl) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +""" + +SLOT_RELEASE: Final = """ +-- Release one slot per gauge key by removing this request's slot id. +-- ZREM of an absent member (or key) is a no-op, so a release without a +-- matching acquire (proxy-side rejection, double-fired callback, slot +-- already expired) can never free a slot owned by another request. +-- KEYS: gauge zset keys. ARGV: per-key slot_id. +-- Returns the remaining in-flight count per key. +local results = {} +for i = 1, #KEYS do + redis.call('ZREM', KEYS[i], ARGV[i]) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +""" + +SLOT_COUNT: Final = """ +-- Read the current in-flight count per gauge key (prunes expired slots +-- first so leaked slots do not inflate the reading). +-- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds. +local time_reply = redis.call('TIME') +local now = tonumber(time_reply[1]) +local results = {} +for i = 1, #KEYS do + redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i])) + table.insert(results, redis.call('ZCARD', KEYS[i])) +end +return results +""" + +COUNTER_INCREMENT: Final = """ +local results = {} + +-- Process each key/increment_value/ttl triplet +for i = 1, #KEYS do + local key = KEYS[i] + local increment_value = tonumber(ARGV[i * 2 - 1]) + local ttl_seconds = tonumber(ARGV[i * 2]) + + -- Increment the value + local new_value = redis.call('INCRBYFLOAT', key, increment_value) + + -- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL + -- ttl_seconds can be 0 (no TTL) or positive (set TTL) + if ttl_seconds and ttl_seconds > 0 then + local current_ttl = redis.call('TTL', key) + if current_ttl == -1 then + redis.call('EXPIRE', key, ttl_seconds) + end + end + + table.insert(results, new_value) +end + +return results +""" + +LIMITED_INCREMENT: Final = """ +local amount = tonumber(ARGV[1]) +local ttl = tonumber(ARGV[2]) +local limit = tonumber(ARGV[3]) +local current = tonumber(redis.call('GET', KEYS[1]) or '0') +if current + amount > limit then + return {0, current} +end +local updated = redis.call('INCRBY', KEYS[1], amount) +redis.call('EXPIRE', KEYS[1], ttl) +return {1, updated} +""" + +REFUND: Final = """ +local updated = redis.call('DECRBY', KEYS[1], tonumber(ARGV[1])) +if updated <= 0 then + redis.call('DEL', KEYS[1]) +end +return 1 +""" + +SAVE: Final = """ +redis.call('SET', KEYS[1], ARGV[1], 'EX', tonumber(ARGV[2])) +return 1 +""" + +POP: Final = """ +local value = redis.call('GET', KEYS[1]) +if value and value ~= '' then + redis.call('SET', KEYS[1], '', 'EX', tonumber(ARGV[1])) +end +return value +""" + +ITERATION_INCREMENT: Final = """ +local key = KEYS[1] +local ttl = tonumber(ARGV[1]) + +local current = redis.call('INCR', key) +if current == 1 then + redis.call('EXPIRE', key, ttl) +end + +return current +""" + +SPEND_INCREMENT: Final = """ +local key = KEYS[1] +local amount = ARGV[1] +local ttl = tonumber(ARGV[2]) + +local existed = redis.call('EXISTS', key) +local new_val = redis.call('INCRBYFLOAT', key, amount) +if existed == 0 then + redis.call('EXPIRE', key, ttl) +end + +return new_val +""" + +CLAIM: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if ARGV[3] then + local decoded, stored = pcall(cjson.decode, current) + if ARGV[4] then + local legacy = current + if decoded and type(stored) == 'string' then legacy = stored end + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + if eligible['model_id'] == legacy then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end + if decoded and type(stored) == 'table' then + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + local matches = true + for key, value in pairs(eligible) do + if stored[key] ~= value then matches = false; break end + end + for key, _ in pairs(stored) do + if eligible[key] == nil then matches = false; break end + end + if matches then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +""" + +BLOCK_TTLS: Final = "return {redis.call('TTL', KEYS[2]), redis.call('TTL', KEYS[4])}" + +RECORD_FAILURE: Final = ( + "local function bump(count_key, block_key, limit) " + "local blocked = redis.call('TTL', block_key) " + "if blocked > 0 then return blocked end " + "local count = redis.call('INCR', count_key) " + "if redis.call('TTL', count_key) < 0 then redis.call('EXPIRE', count_key, ARGV[3]) end " + "if count > limit then redis.call('SET', block_key, '1', 'EX', ARGV[4]) return tonumber(ARGV[4]) end " + "return 0 end " + "local user_block = bump(KEYS[1], KEYS[2], tonumber(ARGV[1])) " + "local source_block = 0 " + "if tonumber(ARGV[2]) > 0 and user_block == 0 then " + "source_block = bump(KEYS[3], KEYS[4], tonumber(ARGV[2])) end " + "return {user_block, source_block}" +) + +LEASE_ACQUIRE: Final = "return redis.call('SET', KEYS[1], ARGV[1], 'NX', 'PX', ARGV[2])" +LEASE_EXTEND: Final = """ +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('PEXPIRE', KEYS[1], ARGV[2]) +end +return 0 +""" +LEASE_RELEASE: Final = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" +LEASE_HELD: Final = "return redis.call('EXISTS', KEYS[1])" diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index d9a40c44d2b..2c13119ba96 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -1,63 +1,14 @@ """Atomic affinity claims shared by deployment and tier-model selection.""" -import json from collections.abc import Mapping -from typing import ( - Final, - cast, # noqa: TID251 # Redis script results are narrowed only to object, then validated -) - -from pydantic import JsonValue, TypeAdapter, ValidationError +from typing import Final from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.types.caching import CacheFailure ROUTER_SESSION_PINS_TARGET: Final = "router_session_pins" -_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue) - -_CLAIM_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if ARGV[3] then - local decoded, stored = pcall(cjson.decode, current) - if ARGV[4] then - local legacy = current - if decoded and type(stored) == 'string' then legacy = stored end - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - if eligible['model_id'] == legacy then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - if decoded and type(stored) == 'table' then - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - local matches = true - for key, value in pairs(eligible) do - if stored[key] ~= value then matches = false; break end - end - for key, _ in pairs(stored) do - if eligible[key] == nil then matches = false; break end - end - if matches then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if current == ARGV[1] then - redis.call('EXPIRE', KEYS[1], ARGV[2]) -end -return current -""" - def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None: """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry.""" @@ -96,13 +47,6 @@ def claim_affinity_pin_in_memory( return winner -def _decode_pin(value: str) -> object: - try: - return _PIN_JSON_ADAPTER.validate_json(value) - except ValidationError: - return value - - async def claim_affinity_pin( cache: DualCache, cache_key: str, @@ -119,29 +63,23 @@ async def claim_affinity_pin( """ redis_cache: Final = cache.redis_cache if redis_cache is not None: - try: - claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) - args: Final = ( - json.dumps(dict(pin_value)), - int(ttl_seconds), - *( - (json.dumps(tuple(dict(value) for value in eligible_values)),) - if eligible_values is not None - else () - ), - *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), - ) - raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here - object, await claim_script(keys=(cache_key,), args=args) - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - if not isinstance(decoded, str): + winner: Final = await redis_cache.async_claim_affinity( + key=cache_key, + value=pin_value, + ttl_seconds=ttl_seconds, + eligible_values=eligible_values, + preserve_legacy_model_ids=preserve_legacy_model_ids, + ) + match winner: + case CacheFailure(kind="invalid_response"): return pin_value - winner: Final = _decode_pin(decoded) - set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) - return winner - except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims - verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) + case CacheFailure(message=message): + verbose_router_logger.debug( + "Affinity cache: Redis claim failed, using pod-local claim. error=%s", message + ) + case _: + set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) + return winner return claim_affinity_pin_in_memory( cache, cache_key, diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..42e53093316 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import hashlib +import itertools import json import logging import time @@ -22,6 +23,9 @@ from datetime import timedelta from types import MappingProxyType, TracebackType from typing import Final, Generic, Protocol, TypeVar +from pydantic import TypeAdapter +from typing_extensions import deprecated + from litellm._internal_context import ( REDIS_FAMILIES_METADATA_KEY, current_service_target, @@ -42,6 +46,7 @@ _T = TypeVar("_T") _ScriptArg = str | bytes | int | float SettledHook = Callable[[asyncio.Future[_T]], Awaitable[None] | None] POST_CALL_FLUSH_DEADLINE_SECONDS: Final = 1.0 +_MGET_VALUES: Final = TypeAdapter(tuple[str | bytes | None, ...]) class RegisteredScript(Protocol): @@ -102,7 +107,7 @@ class _Op(Generic[_T]): except Exception as e: # noqa: BLE001 # a reply this op cannot decode fails this op alone self.future.set_exception(e) return None - if _is_missing_script(failure): + if isinstance(self, _Script) and _is_missing_script(failure): return self._settle_alone() self.future.set_exception(failure) return None @@ -141,11 +146,14 @@ class _MGet(_Op[Mapping[str, object]]): return 1 def resolve(self, replies: Sequence[object]) -> Mapping[str, object]: - values: Final = replies[0] - if not isinstance(values, (list, tuple)): - raise TypeError(f"MGET reply is not a list: {type(values).__name__}") + raw: Final = replies[0] + if not isinstance(raw, (list, tuple)): + raise TypeError(f"MGET reply is not a list: {type(raw).__name__}") + values: Final = _MGET_VALUES.validate_python(replies[0]) + if len(values) != len(self._keys): + raise ValueError("MGET reply count does not match requested keys") return MappingProxyType( - {key: self._redis_cache._get_cache_logic(value) for key, value in zip(self._keys, values)} # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownArgumentType] # shared decode with async_batch_get_cache + {key: self._redis_cache._get_cache_logic(value) for key, value in zip(self._keys, values)} # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # shared decode with async_batch_get_cache ) async def run_alone(self) -> Mapping[str, object]: @@ -211,7 +219,7 @@ class _Increment(_Op[float]): async def run_alone(self) -> float: value: object = await self._redis_cache.async_increment(key=self._key, value=self._value, ttl=self._ttl) # pyright: ignore[reportUnknownMemberType] # untyped cache API - if not isinstance(value, (int, float)): + if not isinstance(value, (int, float)): # pyright: ignore[reportUnnecessaryIsInstance] # legacy subclasses can return malformed values raise TypeError(f"increment did not return a number: {type(value).__name__}") return float(value) @@ -319,6 +327,7 @@ class RedisBatch: so a per-key GET later in the same request can be answered without another round trip.""" return key in self._misses + @deprecated("Use typed RedisCache operations with batch= instead of script callbacks", category=None) def script( self, source: str, run: RegisteredScript, keys: Sequence[str], args: Sequence[_ScriptArg] ) -> BatchResult[object]: @@ -370,16 +379,14 @@ class RedisBatch: async def _flush_pipeline(self, ops: Sequence[_Op[object]]) -> None: start_time: Final = time.time() target, metadata = _pipeline_service_event(ops) - widths: list[int] = [] # mutable-ok: filled while enqueuing - - async def run() -> list[object]: - client: Final = self.redis_cache.init_async_client() - async with client.pipeline(transaction=False) as pipe: - widths.extend(op.enqueue(pipe) for op in ops) - return await pipe.execute(raise_on_error=False) - try: - replies: Final = await _run_under_circuit_breaker(self.redis_cache._circuit_breaker, self.name, run) # pyright: ignore[reportPrivateUsage] # same breaker as the cache's own methods + widths, replies = await _run_under_circuit_breaker( + self.redis_cache._circuit_breaker, # pyright: ignore[reportPrivateUsage] # same breaker as the cache's own methods + self.name, + lambda: self.redis_cache._execute_batch_pipeline(ops), # pyright: ignore[reportPrivateUsage] # cache owns transport; batch owns settlement + ) + if len(replies) != sum(widths): + raise ValueError("Redis pipeline reply count does not match declared operations") except Exception as e: # noqa: BLE001 # each declaring caller applies its own Redis fallback log_redis_failure(verbose_logger, logging.WARNING, f"{self.name}: pipeline of {len(ops)} ops failed", e) with service_target(target): @@ -408,13 +415,12 @@ class RedisBatch: event_metadata=metadata, ) ) - retries: list[Awaitable[None]] = [] # mutable-ok: collected while slicing replies - offset = 0 - for op, width in zip(ops, widths): - retry = op.settle(replies[offset : offset + width]) - offset += width - if retry is not None: - retries.append(retry) + offsets: Final = tuple(itertools.accumulate(widths, initial=0)) + retries: Final = tuple( + retry + for op, start, end in zip(ops, offsets, offsets[1:]) + if (retry := op.settle(replies[start:end])) is not None + ) if retries: await asyncio.gather(*retries) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..a59848a28b0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -25,9 +25,10 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta from types import FrameType, MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter +from typing_extensions import deprecated import litellm from litellm._internal_context import current_service_caller @@ -43,12 +44,33 @@ from litellm.constants import ( from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.types.caching import ( + CacheFailure, + CacheOperation, + CachePoolInfo, + CachePublish, + CacheQueueUpdate, + CacheSubscription, + CapacityExceeded, + ConcurrencyRead, + ConcurrencyRelease, + ConcurrencySlot, + CounterIncrement, + CounterUpdate, + GuardedIncrement, + KeyPresence, + LimitedIncrement, RedisPipelineIncrementOperation, RedisPipelineLpopOperation, RedisPipelineRpushOperation, + ThrottleKeys, + WindowCount, + WindowCounter, + WindowReservation, ) from litellm.types.services import ServiceTypes +from . import _redis_capabilities as _cap +from . import _redis_scripts as _scripts from .base_cache import BaseCache if TYPE_CHECKING: @@ -58,6 +80,11 @@ if TYPE_CHECKING: from redis.asyncio.cluster import ClusterPipeline from litellm._service_logger import ServiceLogging + from litellm.caching.redis_batch import ( + RedisBatch, + _Op, # pyright: ignore[reportPrivateUsage] # cache-owned batch transport + _RedisPipeline, # pyright: ignore[reportPrivateUsage] # cache-owned batch transport + ) pipeline = Pipeline cluster_pipeline = ClusterPipeline @@ -81,6 +108,8 @@ class _AsyncRedisCommands(Protocol): def ping(self) -> Awaitable[bool]: ... + def get(self, name: str) -> Awaitable[bytes | str | None]: ... + def delete(self, *names: str) -> Awaitable[int]: ... def ttl(self, name: str) -> Awaitable[int]: ... @@ -480,6 +509,11 @@ class RedisCircuitBreakerOpenError(Exception): pass +@deprecated("Use RedisCache.async_publish and async_subscribe instead of raw pub/sub clients", category=None) +def pubsub_capable_client(redis_cache: "RedisCache") -> "Redis[bytes]": + return redis_cache._get_pubsub_client() # pyright: ignore[reportPrivateUsage] # compatibility-only cache transport bridge + + class _RedisTimeoutLogThrottle: """Admits one Redis timeout log line per interval and counts the timeouts it suppressed in between.""" @@ -507,27 +541,32 @@ _redis_timeout_log_throttle: Final = _RedisTimeoutLogThrottle(REDIS_TIMEOUT_LOG_ def log_redis_failure( - logger: logging.Logger, level: int, message: str, exc: BaseException, with_traceback: bool = False + logger: logging.Logger, level: int, message: str, exc: BaseException | CacheFailure, with_traceback: bool = False ) -> None: - if isinstance(exc, RedisCircuitBreakerOpenError): - logger.debug("%s: %s", message, exc, stacklevel=2) + detail: Final = exc.message if isinstance(exc, CacheFailure) else exc + circuit_open: Final = ( + exc.reason == "circuit_open" if isinstance(exc, CacheFailure) else isinstance(exc, RedisCircuitBreakerOpenError) + ) + if circuit_open: + logger.debug("%s: %s", message, detail, stacklevel=2) return - exc_info: Final = exc if with_traceback else None - if not is_redis_timeout_failure(exc): - logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) + exc_info: Final = exc if with_traceback and isinstance(exc, BaseException) else None + timeout: Final = exc.reason == "timeout" if isinstance(exc, CacheFailure) else is_redis_timeout_failure(exc) + if not timeout: + logger.log(level, "%s: %s", message, detail, exc_info=exc_info, stacklevel=2) return suppressed: Final = _redis_timeout_log_throttle.admit() if suppressed is None: - logger.debug("%s: %s", message, exc, stacklevel=2) + logger.debug("%s: %s", message, detail, stacklevel=2) return if suppressed == 0: - logger.log(level, "%s: %s", message, exc, exc_info=exc_info, stacklevel=2) + logger.log(level, "%s: %s", message, detail, exc_info=exc_info, stacklevel=2) return logger.log( level, "%s: %s (%d more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", message, - exc, + detail, suppressed, exc_info=exc_info, stacklevel=2, @@ -778,6 +817,7 @@ class RedisCache(BaseCache): kwargs_hash: Final = hashlib.sha256(kwargs_str.encode()).hexdigest()[:16] return f"async-redis-client-{kwargs_hash}" + @deprecated("Use typed RedisCache operations instead of raw client access", category=None) def init_async_client( self, ) -> async_redis_client | async_redis_cluster_client: @@ -804,12 +844,13 @@ class RedisCache(BaseCache): self.redis_async_client = redis_async_client return redis_async_client + @deprecated("Raw pub/sub client access is deprecated, migrate to typed operations as available", category=None) def init_pubsub_client(self) -> async_redis_client: # pyright: ignore[reportUnknownParameterType] # redis generics from redis.asyncio import RedisCluster from litellm import in_memory_llm_clients_cache - client: Final = self.init_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis generics + client: Final = self._get_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis generics if not isinstance(client, RedisCluster): return client # pyright: ignore[reportUnknownVariableType] # redis generics cache_key: Final = f"{self._get_async_client_cache_key()}-pubsub" @@ -826,8 +867,250 @@ class RedisCache(BaseCache): in_memory_llm_clients_cache.set_cache(key=cache_key, value=pubsub_client, litellm_owned_client=True) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # untyped in-memory client cache return pubsub_client # pyright: ignore[reportUnknownVariableType] # redis generics + def _get_async_client(self) -> "Redis[bytes] | RedisCluster[bytes]": + return self.init_async_client() # pyright: ignore[reportDeprecated, reportUnknownMemberType, reportUnknownVariableType] # preserve dispatch to unparameterized legacy factories + + def _get_pubsub_client(self) -> "Redis[bytes]": + return self.init_pubsub_client() # pyright: ignore[reportDeprecated, reportUnknownMemberType, reportUnknownVariableType] # preserve dispatch to unparameterized legacy factories + def _async_commands(self) -> _AsyncRedisCommands: - return self.init_async_client() + return cast(_AsyncRedisCommands, self._get_async_client()) # cast-ok: RedisCluster stubs omit shared commands + + def async_increment_window_counters( + self, *, counters: Sequence[WindowCounter], now: float, window_seconds: int, batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[WindowCount, ...]]: + _cap.provisional("async_increment_window_counters") + keys: Final = tuple(itertools.chain.from_iterable((item.window_key, item.counter_key) for item in counters)) + return _cap.operation( + self, + _scripts.WINDOW_INCREMENT, + keys, + (now, window_seconds), + lambda reply: _cap.window_counts(reply, len(counters)), + batch, + ) + + def async_read_atomic_counts( + self, *, keys: Sequence[str], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[int | None, ...]]: + _cap.provisional("async_read_atomic_counts") + return _cap.operation( + self, _scripts.COUNTER_READ, tuple(keys), (), lambda reply: _cap.optional_counts(reply, len(keys)), batch + ) + + def async_reserve_window_capacity( + self, *, reservations: Sequence[WindowReservation], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[WindowCount, ...] | CapacityExceeded]: + _cap.provisional("async_reserve_window_capacity") + keys: Final = tuple(itertools.chain.from_iterable((item.window_key, item.counter_key) for item in reservations)) + args: Final = tuple( + itertools.chain.from_iterable( + (item.limit, item.increment, item.ttl_seconds, item.window_seconds) for item in reservations + ) + ) + return _cap.operation( + self, _scripts.WINDOW_RESERVE, keys, args, lambda reply: _cap.reservation(reply, len(reservations)), batch + ) + + def async_increment_guarded_counters( + self, *, increments: Sequence[GuardedIncrement], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[CounterUpdate, ...]]: + _cap.provisional("async_increment_guarded_counters") + keys: Final = tuple(itertools.chain.from_iterable((item.window_key, item.counter_key) for item in increments)) + args: Final = tuple( + itertools.chain.from_iterable( + (item.expected_window_start, item.increment, item.ttl_seconds) for item in increments + ) + ) + return _cap.operation( + self, _scripts.GUARDED_INCREMENT, keys, args, lambda reply: _cap.guarded(reply, len(increments)), batch + ) + + def async_acquire_concurrency( + self, *, slots: Sequence[ConcurrencySlot], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[int, ...] | CapacityExceeded]: + _cap.provisional("async_acquire_concurrency") + args: Final = tuple( + itertools.chain.from_iterable((item.limit, item.ttl_seconds, item.slot_id) for item in slots) + ) + return _cap.operation( + self, + _scripts.SLOT_ACQUIRE, + tuple(item.key for item in slots), + args, + lambda reply: _cap.slots(reply, len(slots)), + batch, + ) + + def async_release_concurrency( + self, *, slots: Sequence[ConcurrencyRelease], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[int, ...]]: + _cap.provisional("async_release_concurrency") + return _cap.operation( + self, + _scripts.SLOT_RELEASE, + tuple(item.key for item in slots), + tuple(item.slot_id for item in slots), + lambda reply: _cap.counts(reply, len(slots)), + batch, + ) + + def async_count_concurrency( + self, *, counters: Sequence[ConcurrencyRead], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[int, ...]]: + _cap.provisional("async_count_concurrency") + return _cap.operation( + self, + _scripts.SLOT_COUNT, + tuple(item.key for item in counters), + tuple(item.ttl_seconds for item in counters), + lambda reply: _cap.counts(reply, len(counters)), + batch, + ) + + def async_increment_counters( + self, *, increments: Sequence[CounterIncrement], batch: "RedisBatch | None" = None + ) -> CacheOperation[tuple[float, ...]]: + _cap.provisional("async_increment_counters") + args: Final = tuple(itertools.chain.from_iterable((item.increment, item.ttl_seconds) for item in increments)) + return _cap.operation( + self, + _scripts.COUNTER_INCREMENT, + tuple(item.key for item in increments), + args, + lambda reply: _cap.totals(reply, len(increments)), + batch, + ) + + def async_reserve_counter( + self, *, key: str, amount: int, limit: int, ttl_seconds: int, batch: "RedisBatch | None" = None + ) -> CacheOperation[LimitedIncrement]: + _cap.provisional("async_reserve_counter") + return _cap.operation( + self, _scripts.LIMITED_INCREMENT, (key,), (amount, ttl_seconds, limit), _cap.limited, batch + ) + + def async_refund_counter(self, *, key: str, amount: int, batch: "RedisBatch | None" = None) -> CacheOperation[None]: + _cap.provisional("async_refund_counter") + return _cap.operation(self, _scripts.REFUND, (key,), (amount,), _cap.acknowledged, batch) + + def async_save_record( + self, *, key: str, value: str, ttl_seconds: int, batch: "RedisBatch | None" = None + ) -> CacheOperation[None]: + _cap.provisional("async_save_record") + return _cap.operation(self, _scripts.SAVE, (key,), (value, ttl_seconds), _cap.acknowledged, batch) + + def async_delete_record(self, *, key: str) -> CacheOperation[bool]: + _cap.provisional("async_delete_record") + return _cap.delete_record(self, key) + + def async_enqueue_updates(self, *, updates: Sequence[CacheQueueUpdate]) -> CacheOperation[tuple[int, ...]]: + _cap.provisional("async_enqueue_updates") + return _cap.enqueue_updates(self, updates) + + def async_pop_record( + self, *, key: str, tombstone_ttl_seconds: int, batch: "RedisBatch | None" = None + ) -> CacheOperation[str | None]: + _cap.provisional("async_pop_record") + return _cap.operation(self, _scripts.POP, (key,), (tombstone_ttl_seconds,), _cap.text, batch) + + def async_increment_iterations( + self, *, key: str, ttl_seconds: int, batch: "RedisBatch | None" = None + ) -> CacheOperation[int]: + _cap.provisional("async_increment_iterations") + return _cap.operation(self, _scripts.ITERATION_INCREMENT, (key,), (ttl_seconds,), _cap.integer, batch) + + def async_increment_spend( + self, *, key: str, amount: float, ttl_seconds: int, batch: "RedisBatch | None" = None + ) -> CacheOperation[float]: + _cap.provisional("async_increment_spend") + return _cap.operation(self, _scripts.SPEND_INCREMENT, (key,), (str(amount), ttl_seconds), _cap.number, batch) + + def async_claim_affinity( + self, + *, + key: str, + value: Mapping[str, str], + ttl_seconds: int, + eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, + ) -> CacheOperation[JsonValue]: + _cap.provisional("async_claim_affinity") + args: Final = ( + json.dumps(dict(value)), + ttl_seconds, + *((json.dumps(tuple(dict(item) for item in eligible_values)),) if eligible_values is not None else ()), + *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), + ) + return _cap.operation(self, _scripts.CLAIM, (key,), args, _cap.pin, None) + + def async_read_throttle_ttls(self, *, keys: ThrottleKeys) -> CacheOperation[tuple[int, int]]: + _cap.provisional("async_read_throttle_ttls") + names: Final = (keys.pair_counter, keys.pair_block, keys.source_counter, keys.source_block) + return _cap.operation(self, _scripts.BLOCK_TTLS, names, (), _cap.block_ttls, None) + + def async_record_throttle_failure( + self, *, keys: ThrottleKeys, pair_limit: int, source_limit: int, window_seconds: int, block_seconds: int + ) -> CacheOperation[tuple[int, int]]: + _cap.provisional("async_record_throttle_failure") + names: Final = (keys.pair_counter, keys.pair_block, keys.source_counter, keys.source_block) + return _cap.operation( + self, + _scripts.RECORD_FAILURE, + names, + (pair_limit, source_limit, window_seconds, block_seconds), + _cap.block_ttls, + None, + ) + + def async_acquire_lease( + self, *, key: str, token: str, ttl_seconds: float + ) -> CacheOperation[Literal["acquired", "held"]]: + _cap.provisional("async_acquire_lease") + return _cap.operation( + self, _scripts.LEASE_ACQUIRE, (key,), (token, int(ttl_seconds * 1000)), _cap.acquired, None + ) + + def async_extend_lease(self, *, key: str, token: str, ttl_seconds: float) -> CacheOperation[bool]: + _cap.provisional("async_extend_lease") + return _cap.operation(self, _scripts.LEASE_EXTEND, (key,), (token, int(ttl_seconds * 1000)), _cap.boolean, None) + + def async_release_lease(self, *, key: str, token: str) -> CacheOperation[bool]: + _cap.provisional("async_release_lease") + return _cap.operation(self, _scripts.LEASE_RELEASE, (key,), (token,), _cap.boolean, None) + + def async_lease_held(self, *, key: str) -> CacheOperation[bool]: + _cap.provisional("async_lease_held") + return _cap.operation(self, _scripts.LEASE_HELD, (key,), (), _cap.boolean, None) + + async def async_publish(self, *, channel: str, message: str) -> int | CacheFailure: + _cap.provisional("async_publish") + return await _cap.publish(self, channel, message) + + async def async_publish_many(self, *, messages: Sequence[CachePublish]) -> tuple[int, ...] | CacheFailure: + _cap.provisional("async_publish_many") + return await _cap.publish_many(self, messages) + + async def async_subscribe(self, *, channel: str) -> CacheSubscription | CacheFailure: + _cap.provisional("async_subscribe") + return await _cap.subscribe(self, channel) + + def get_connection_pool_info(self) -> CachePoolInfo | CacheFailure | None: + _cap.provisional("get_connection_pool_info") + return _cap.pool_info(self) + + async def async_read_values(self, *, keys: Sequence[str]) -> Mapping[str, object] | CacheFailure: + _cap.provisional("async_read_values") + return await _cap.read_values(self, keys) + + async def _execute_batch_pipeline( + self, operations: Sequence["_Op[object]"] + ) -> tuple[tuple[int, ...], Sequence[object]]: + async with self._get_async_client().pipeline(transaction=False) as pipe: + transport: Final = cast("_RedisPipeline", pipe) # cast-ok: private RedisBatch command subset + widths: Final = tuple(operation.enqueue(transport) for operation in operations) + replies: Final = await transport.execute(raise_on_error=False) + return widths, replies def check_and_fix_namespace(self, key: str) -> str: """ @@ -995,7 +1278,7 @@ class RedisCache(BaseCache): start_time: Final = time.time() try: keys: Final = [] - _redis_client: Final = self.init_async_client() + _redis_client: Final = self._get_async_client() if not hasattr(_redis_client, "scan_iter"): verbose_logger.debug( "Redis client does not support scan_iter, potentially using Redis Cluster. Returning empty list." @@ -1040,7 +1323,11 @@ class RedisCache(BaseCache): ) raise e + @deprecated("Use typed RedisCache operations instead of script callbacks", category=None) def async_register_script(self, script: str) -> Callable[..., Awaitable[Any]]: + return self._script_runner(script) + + def _script_runner(self, script: str) -> Callable[..., Awaitable[object]]: """ Register a Lua script with Redis asynchronously. Works with both standalone Redis and Redis Cluster. @@ -1102,38 +1389,10 @@ class RedisCache(BaseCache): Kept separate from async_register_script so each loop caches its own executor; see that method for why the binding must be per loop. """ - _redis_client: Final[Any] = self.init_async_client() - if hasattr(_redis_client, "register_script"): - registered_script: Final = _redis_client.register_script(script) - - async def standalone_executor( - keys: Sequence[str], - args: Sequence[str | bytes | int | float], - client: object = None, - ) -> object: - namespaced_keys: Final = tuple(self.check_and_fix_namespace(key=key) for key in keys) - return await registered_script(keys=namespaced_keys, args=args, client=client) - - return standalone_executor - - if hasattr(_redis_client, "script_load"): - script_sha: Final = _redis_client.script_load(script) - - async def cluster_executor( - keys: Sequence[str], - args: Sequence[str | bytes | int | float], - client: object = None, - ) -> object: - namespaced_keys: Final = tuple(self.check_and_fix_namespace(key=key) for key in keys) - return await _redis_client.evalsha(script_sha, len(namespaced_keys), *namespaced_keys, *args) - - return cluster_executor - - raise ValueError("Redis client does not support Lua script registration") + return _cap.script_executor(self, script) @_redis_circuit_breaker_guard async def async_set_cache(self, key: str | None, value: object, **kwargs): - from redis.asyncio import Redis if key is None: verbose_logger.debug( @@ -1144,7 +1403,9 @@ class RedisCache(BaseCache): start_time: Final = time.time() try: - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final = cast( # cast-ok: RedisCluster stubs omit the shared Redis commands used here + "Redis[bytes]", self._get_async_client() + ) except Exception as e: end_time = time.time() _duration = end_time - start_time @@ -1253,7 +1514,7 @@ class RedisCache(BaseCache): if len(cache_list) == 0: return - _redis_client: Final = self.init_async_client() + _redis_client: Final = self._get_async_client() start_time: Final = time.time() print_verbose( @@ -1319,7 +1580,7 @@ class RedisCache(BaseCache): ) start_time: Final = time.time() try: - async with self.init_async_client().pipeline(transaction=False) as pipe: + async with self._get_async_client().pipeline(transaction=False) as pipe: for cache_key, json_cache_value, ttl in commands: pipe.set(name=cache_key, value=json_cache_value, ex=None if ttl is None else timedelta(seconds=ttl)) await pipe.execute() @@ -1372,11 +1633,12 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_set_cache_sadd(self, key, value: list, ttl: float | None, **kwargs): - from redis.asyncio import Redis start_time: Final = time.time() try: - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final = cast( # cast-ok: RedisCluster stubs omit the shared Redis commands used here + "Redis[bytes]", self._get_async_client() + ) except Exception as e: end_time = time.time() _duration = end_time - start_time @@ -1477,9 +1739,10 @@ class RedisCache(BaseCache): parent_otel_span: Span | None = None, refresh_ttl: bool = False, ) -> float: - from redis.asyncio import Redis - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final = cast( # cast-ok: RedisCluster stubs omit the shared Redis commands used here + "Redis[bytes]", self._get_async_client() + ) start_time: Final = time.time() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) @@ -1544,7 +1807,7 @@ class RedisCache(BaseCache): GET/compare/SET runs in a single Lua call, so it is also atomic across racing callers and pods. Returns the resulting value. """ - _redis_client: Final = self.init_async_client() + _redis_client: Final = self._get_async_client() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) lua: Final = ( @@ -1670,7 +1933,7 @@ class RedisCache(BaseCache): We use a wrapper so RedisCluster can override this method """ - async_redis_client: Final = self.init_async_client() + async_redis_client: Final = self._get_async_client() return await async_redis_client.mget(keys=keys) def batch_get_cache( @@ -1737,13 +2000,23 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) return key_value_dict + async def async_key_presence(self, key: str) -> KeyPresence: + _cap.provisional("async_key_presence") + try: + value: Final = await self._async_commands().get(self.check_and_fix_namespace(key=key)) + except Exception as e: # noqa: BLE001 # Any failed marker read must remain unavailable, never missing + log_redis_failure(verbose_logger, logging.WARNING, "Redis key presence check failed", e) + return "unavailable" + return "missing" if value is None else "present" + @_redis_circuit_breaker_guard async def async_get_cache( self, key, parent_otel_span: Span | None = None, *, raise_on_error: bool = False, **kwargs ): - from redis.asyncio import Redis - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final = cast( # cast-ok: RedisCluster stubs omit the shared Redis commands used here + "Redis[bytes]", self._get_async_client() + ) key = self.check_and_fix_namespace(key=key) start_time: Final = time.time() @@ -2061,9 +2334,9 @@ class RedisCache(BaseCache): if len(increment_list) == 0: return None - from redis.asyncio import Redis - - _redis_client: Final[Redis] = self.init_async_client() + _redis_client: Final = cast( # cast-ok: RedisCluster stubs omit the shared Redis commands used here + "Redis[bytes]", self._get_async_client() + ) start_time: Final = time.time() print_verbose("Increment Async Redis Cache Pipeline: increment list: %s", increment_list) diff --git a/litellm/caching/redis_cluster_cache.py b/litellm/caching/redis_cluster_cache.py index 6fb84df331e..f7438c9b7bf 100644 --- a/litellm/caching/redis_cluster_cache.py +++ b/litellm/caching/redis_cluster_cache.py @@ -14,6 +14,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypeVar from litellm._logging import verbose_logger +from litellm.caching import _redis_capabilities as _cap from litellm.caching.redis_cache import RedisCache from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT @@ -130,14 +131,13 @@ class RedisClusterCache(RedisCache): """ Overrides `_run_redis_mget_operation` in redis_cache.py """ - return self.redis_client.mget_nonatomic(keys=keys) + return _cap.cluster_values_sync(self, keys) async def _async_run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]: """ Overrides `_async_run_redis_mget_operation` in redis_cache.py """ - async_redis_cluster_client: Final = self.init_async_client() - return await async_redis_cluster_client.mget_nonatomic(keys=keys) + return await _cap.cluster_values(self, keys) async def test_connection(self) -> dict: """ diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 70e4bac01b5..2816ed832a3 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -1393,21 +1393,20 @@ class _SingleUseGuard: async def peek(self, key: str) -> Literal["unclaimed", "claimed", "unavailable"]: """Read-only view of a single-use marker, resolved against the same shared authority as :meth:`claim` so introspection observes exactly the record redemption and revocation wrote. - A backend fault is ``"unavailable"`` (fail closed) rather than a guess either way, which is why - the read goes to the Redis client itself: the cache wrapper's ``async_get_cache`` turns a fault - into ``None``, and ``None`` here would pass as unclaimed. The key is namespaced exactly as the - wrapper namespaces every write, or a ``namespace`` configured under ``coordination_redis`` or - ``cache_params`` would hide the marker.""" + A backend fault is ``"unavailable"`` (fail closed), never an unclaimed marker.""" from litellm.proxy.proxy_server import redis_usage_cache # noqa: PLC0415 # circular import at module load redis_cache: Final = redis_usage_cache or getattr(self._cache, "redis_cache", None) if redis_cache is not None: - try: - value = await redis_cache.init_async_client().get(redis_cache.check_and_fix_namespace(key=key)) - except Exception as e: # noqa: BLE001 # ANY Redis fault fails the read closed - verbose_logger.warning("mcp gateway single-use peek: shared cache backend unavailable: %s", e) - return "unavailable" - return "unclaimed" if value is None else "claimed" + match await redis_cache.async_key_presence(key): + case "present": + return "claimed" + case "missing": + return "unclaimed" + case "unavailable": + return "unavailable" + case other: + assert_never(other) local: Final = await self._cache.async_get_cache(key, local_only=True) return "unclaimed" if local is None else "claimed" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py index c7eefed75ed..78aa9b4ae86 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/redis_distributed_lock.py @@ -1,94 +1,52 @@ -"""Concrete ``DistributedLock`` over a Redis client: ``SET NX PX`` / owner-only renew / delete. - -The cross-replica lock the ``RedisRefreshCoordinator`` elects refreshers with. ``acquire`` is an -atomic ``SET key token NX PX ttl`` (only the first caller wins; the entry self-expires so a crashed -holder can't wedge refresh). ``extend`` renews the lease only when the token still matches, and -``release`` deletes the key only when it still holds this caller's token, so a holder whose lock already -PX-expired and was re-acquired by another worker cannot delete the new holder's lock. ``is_held`` is -``EXISTS``. Every key is run through the injected ``namespace_key`` before it reaches Redis, so lock -keys carry the same namespace as cache keys and cannot collide with another deployment sharing Redis. - -The Redis client is injected (in production the async client from LiteLLM's ``RedisCache``), so the -lock is unit-testable with a fake. A transport error on ``acquire`` returns ``LockAcquisition.ERROR`` - -distinct from ``HELD`` - so the coordinator refreshes anyway instead of mistaking a dead backend for a -busy holder; a Redis blip degrades to an extra refresh, never a stale bearer. -""" +"""Map cache lease outcomes to the refresh coordinator's best-effort lock policy.""" from __future__ import annotations -from collections.abc import Callable -from dataclasses import KW_ONLY, dataclass -from typing import Final, Protocol +from dataclasses import dataclass +from typing import Final + +from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.caching.redis_cache import RedisCache from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_coordinator import ( LockAcquisition, ) - -# Delete the key only if it still holds this caller's token, so a holder whose lock already expired -# (PX) and was re-acquired by another worker cannot delete the new holder's lock. -_RELEASE_IF_OWNER = "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('del', KEYS[1]) else return 0 end" -_EXTEND_IF_OWNER: Final = ( - "if redis.call('get', KEYS[1]) == ARGV[1] then return redis.call('pexpire', KEYS[1], ARGV[2]) else return 0 end" -) - - -class RedisCommands(Protocol): - """The slice of the async Redis client this lock needs.""" - - async def set(self, name: str, value: str, *, nx: bool = False, px: int | None = None) -> object | None: ... - - async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: ... - - async def exists(self, *names: str) -> int: ... +from litellm.types.caching import CacheFailure @dataclass(frozen=True, slots=True) class RedisDistributedLock: - client: RedisCommands - _: KW_ONLY - namespace_key: Callable[[str], str] = lambda key: key + cache: RedisCache async def acquire(self, key: str, token: str, ttl_seconds: float) -> LockAcquisition: - try: - result: Final = await self.client.set(self.namespace_key(key), token, nx=True, px=int(ttl_seconds * 1000)) - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - verbose_logger.warning("RedisDistributedLock.acquire failed: %s", exc) - return LockAcquisition.ERROR - return LockAcquisition.ACQUIRED if result is not None else LockAcquisition.HELD + result: Final = await self.cache.async_acquire_lease(key=key, token=token, ttl_seconds=ttl_seconds) + match result: + case "acquired": + return LockAcquisition.ACQUIRED + case "held": + return LockAcquisition.HELD + case CacheFailure(): + verbose_logger.warning("RedisDistributedLock.acquire failed: %s", result.message) + return LockAcquisition.ERROR + case _: + assert_never(result) async def extend(self, key: str, token: str, ttl_seconds: float) -> bool: - try: - result: Final = await self.client.eval( - _EXTEND_IF_OWNER, - 1, - self.namespace_key(key), - token, - str(int(ttl_seconds * 1000)), - ) - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - verbose_logger.warning("RedisDistributedLock.extend failed: %s", exc) + result: Final = await self.cache.async_extend_lease(key=key, token=token, ttl_seconds=ttl_seconds) + if isinstance(result, CacheFailure): + verbose_logger.warning("RedisDistributedLock.extend failed: %s", result.message) return False - return result == 1 + return result async def release(self, key: str, token: str) -> None: - try: - await self.client.eval(_RELEASE_IF_OWNER, 1, self.namespace_key(key), token) - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - verbose_logger.warning("RedisDistributedLock.release failed: %s", exc) + result: Final = await self.cache.async_release_lease(key=key, token=token) + if isinstance(result, CacheFailure): + verbose_logger.warning("RedisDistributedLock.release failed: %s", result.message) async def is_held(self, key: str) -> bool: - try: - return await self.client.exists(self.namespace_key(key)) > 0 - # Degrade on any Redis client error: redis.exceptions narrows only via an import that - # is Unknown under basedpyright, and the lock must never crash the resolve path. - except Exception as exc: # noqa: BLE001 - # On error, report "not held" so a waiter stops waiting and re-reads rather than blocking. - verbose_logger.warning("RedisDistributedLock.is_held failed: %s", exc) + result: Final = await self.cache.async_lease_held(key=key) + if isinstance(result, CacheFailure): + verbose_logger.warning("RedisDistributedLock.is_held failed: %s", result.message) return False + return result diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py index e799838b5b7..1c2a7230a44 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/runtime_refresh_coordinator.py @@ -31,11 +31,4 @@ def runtime_refresh_coordinator() -> RefreshCoordinator | None: redis_cache: Final = user_api_key_cache.redis_cache if redis_cache is None: return None - # The Redis client from init_async_client() is only partially typed; the lock validates every - # reply it depends on, so the untyped boundary is contained here. - redis_client: Final = redis_cache.init_async_client() # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType] # litellm redis wrapper is untyped - lock: Final = RedisDistributedLock( - redis_client, # pyright: ignore[reportArgumentType,reportUnknownArgumentType] # litellm redis wrapper is untyped - namespace_key=redis_cache.check_and_fix_namespace, - ) - return RedisRefreshCoordinator(lock) + return RedisRefreshCoordinator(RedisDistributedLock(redis_cache)) diff --git a/litellm/proxy/auth/login_throttle.py b/litellm/proxy/auth/login_throttle.py index 0536e293f13..b4b65e0e946 100644 --- a/litellm/proxy/auth/login_throttle.py +++ b/litellm/proxy/auth/login_throttle.py @@ -9,7 +9,6 @@ Recovery is the master key over the API, which never passes through here, or wai from __future__ import annotations -import asyncio import hashlib import ipaddress import math @@ -17,16 +16,15 @@ import time from collections.abc import Mapping from dataclasses import dataclass from functools import cache -from typing import Final, Literal, NamedTuple, Protocol, TypeAlias +from typing import Final, Literal, Protocol, TypeAlias from fastapi import Request, status from pydantic import TypeAdapter, ValidationError -from redis.exceptions import RedisError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache -from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCache from litellm.constants import ( EMPTY_MAPPING, LOGIN_THROTTLE_CACHE_KEY_PREFIX, @@ -38,6 +36,7 @@ from litellm.constants import ( from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.network import TrustedProxyConfig, resolve_client_ip from litellm.secret_managers.main import get_secret_bool +from litellm.types.caching import CacheFailure, ThrottleKeys _LOGIN_THROTTLE_TARGET: Final = "login_throttle" @@ -54,7 +53,6 @@ WINDOW_KEY: Final = "failed_login_window_seconds" BLOCK_KEY: Final = "failed_login_block_seconds" TRUSTED_PROXY_RANGES_KEY: Final = "trusted_proxy_ranges" -_REDIS_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, asyncio.TimeoutError) _LOCAL_BLOCK_EXPIRY: Final = TypeAdapter[float | None](float | None) _SOURCE_LIMIT_OVERRIDES: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]) _RANGE_ENTRIES: Final = TypeAdapter[tuple[object, ...]](tuple[object, ...]) @@ -62,7 +60,6 @@ _RANGE_ENTRIES: Final = TypeAdapter[tuple[object, ...]](tuple[object, ...]) Scope: TypeAlias = Literal["user", "source"] _BlockTtls: TypeAlias = tuple[int, int] -_LUA_BLOCK_TTLS: Final = TypeAdapter[_BlockTtls](_BlockTtls) _Network: TypeAlias = ipaddress.IPv4Network | ipaddress.IPv6Network @@ -78,25 +75,6 @@ class LocalStore(Protocol): def delete_cache(self, key: str) -> None: ... -# KEYS: pair counter, pair block, source counter, source block (one cluster slot via the source hash tag) -# ARGV: pair limit, source limit (0 = source scope off), window seconds, block seconds -# Both scripts return {pair block TTL, source block TTL}; 0 or below means not blocked -_BLOCK_TTLS_LUA: Final = "return {redis.call('TTL', KEYS[2]), redis.call('TTL', KEYS[4])}" -_RECORD_FAILURE_LUA: Final = ( - "local function bump(count_key, block_key, limit) " - "local blocked = redis.call('TTL', block_key) " - "if blocked > 0 then return blocked end " - "local count = redis.call('INCR', count_key) " - "if redis.call('TTL', count_key) < 0 then redis.call('EXPIRE', count_key, ARGV[3]) end " - "if count > limit then redis.call('SET', block_key, '1', 'EX', ARGV[4]) return tonumber(ARGV[4]) end " - "return 0 end " - "local user_block = bump(KEYS[1], KEYS[2], tonumber(ARGV[1])) " - "local source_block = 0 " - "if tonumber(ARGV[2]) > 0 and user_block == 0 then " - "source_block = bump(KEYS[3], KEYS[4], tonumber(ARGV[2])) end " - "return {user_block, source_block}" -) - _COUNTERS: Final = InMemoryCache( max_size_in_memory=LOGIN_THROTTLE_MAX_TRACKED_COUNTERS, default_ttl=DEFAULT_FAILED_LOGIN_WINDOW_SECONDS ) @@ -253,13 +231,6 @@ def source_group(client_ip: str) -> str: return str(address) -class _Keys(NamedTuple): - pair_counter: str - pair_block: str - source_counter: str - source_block: str - - @dataclass(frozen=True, slots=True) class Block: scope: Scope @@ -312,10 +283,10 @@ class LoginThrottle: enabled=not exempt and not _rate_limit_disabled(), ) - def _keys(self, username: str) -> _Keys: + def _keys(self, username: str) -> ThrottleKeys: group: Final = source_group(self.client_ip) user: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest() - return _Keys( + return ThrottleKeys( pair_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:user:{user}", pair_block=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:block:user:{user}", source_counter=f"{LOGIN_THROTTLE_CACHE_KEY_PREFIX}:{{{group}}}:source", @@ -338,7 +309,7 @@ class LoginThrottle: ) raise self.refused(block.retry_after) - async def _active_block(self, keys: _Keys) -> Block | None: + async def _active_block(self, keys: ThrottleKeys) -> Block | None: local: Final = self._local_block_ttls(keys) shared: Final = await self._shared_block_ttls(keys) user_ttl: Final = max(local[0], shared[0]) @@ -350,18 +321,13 @@ class LoginThrottle: return None @with_service_target("login_throttle") - async def _shared_block_ttls(self, keys: _Keys) -> _BlockTtls: + async def _shared_block_ttls(self, keys: ThrottleKeys) -> _BlockTtls: if self.redis_cache is None: return LOGIN_THROTTLE_NOT_BLOCKED - try: - return _LUA_BLOCK_TTLS.validate_python( - await self.redis_cache.async_register_script(_BLOCK_TTLS_LUA)(keys, ()) - ) - except _REDIS_FAILURES as err: - self._warn_redis(err) - return LOGIN_THROTTLE_NOT_BLOCKED + result: Final = await self.redis_cache.async_read_throttle_ttls(keys=keys) + return self._shared_ttls_or_fallback(result) or LOGIN_THROTTLE_NOT_BLOCKED - def _local_block_ttls(self, keys: _Keys) -> _BlockTtls: + def _local_block_ttls(self, keys: ThrottleKeys) -> _BlockTtls: return self._local_block_ttl(keys.pair_block), self._local_block_ttl(keys.source_block) def _local_block_ttl(self, block_key: str) -> int: @@ -375,14 +341,16 @@ class LoginThrottle: keys: Final = self._keys(username) source_limit: Final = self.source_limit or 0 if self.redis_cache is not None: - try: - return _LUA_BLOCK_TTLS.validate_python( - await self.redis_cache.async_register_script(_RECORD_FAILURE_LUA)( - keys, (self.user_limit, source_limit, self.window_seconds, self.block_seconds) - ) - ) - except _REDIS_FAILURES as err: - self._warn_redis(err) + result: Final = await self.redis_cache.async_record_throttle_failure( + keys=keys, + pair_limit=self.user_limit, + source_limit=source_limit, + window_seconds=self.window_seconds, + block_seconds=self.block_seconds, + ) + shared: Final = self._shared_ttls_or_fallback(result) + if shared is not None: + return shared user_block: Final = self._local_bump(keys.pair_counter, keys.pair_block, self.user_limit) if source_limit == 0 or user_block > 0: return user_block, 0 @@ -402,13 +370,22 @@ class LoginThrottle: async def clear_pair(self, username: str) -> None: pair_counter: Final = self._keys(username).pair_counter if self.redis_cache is not None: - try: - await self.redis_cache.async_delete_cache(pair_counter) - except _REDIS_FAILURES as err: - self._warn_redis(err) + result: Final = await self.redis_cache.async_delete_record(key=pair_counter) + if isinstance(result, CacheFailure): + self._warn_redis(result.message) self.counters.delete_cache(pair_counter) - def _warn_redis(self, err: Exception) -> None: + def _shared_ttls_or_fallback(self, result: _BlockTtls | CacheFailure) -> _BlockTtls | None: + match result: + case CacheFailure(kind="invalid_response", validation_error=error): + raise error if error is not None else ValueError("Invalid shared login throttle state") + case CacheFailure(message=message): + self._warn_redis(message) + return None + case _: + return result + + def _warn_redis(self, err: str) -> None: verbose_proxy_logger.warning( "Redis failed while counting Admin UI sign-in attempts; using this worker's own counters " "until it recovers: %s", diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index b2640d2a765..ba51de9c439 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -9,15 +9,22 @@ from pydantic import TypeAdapter from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.core_helpers import as_str_mapping -from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports - ConfigSyncPubSub, - _ConfigSyncPubSub, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - _pubsub_capable_client, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export - coordination_redis_cache, - pubsub_capable_client, +from litellm.caching.redis_cache import ( + pubsub_capable_client as _pubsub_capable_client, # noqa: F401 # pyright: ignore[reportDeprecated] # compatibility-only re-export ) +from litellm.caching.redis_cache import ( + pubsub_capable_client as pubsub_capable_client, # noqa: PLC0414 # pyright: ignore[reportDeprecated] # compatibility-only re-export +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.proxy.common_utils.config_sync_pubsub import ( + ConfigSyncPubSub as ConfigSyncPubSub, # noqa: PLC0414 # backwards-compatible package export +) +from litellm.proxy.common_utils.config_sync_pubsub import ( + _ConfigSyncPubSub as _ConfigSyncPubSub, # noqa: PLC0414 # pyright: ignore[reportPrivateUsage] # backwards-compatible package export +) +from litellm.proxy.common_utils.config_sync_pubsub import coordination_redis_cache from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET +from litellm.types.caching import CacheFailure, CachePublish, CacheSubscription if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache @@ -38,9 +45,11 @@ _DECODED_JSON: Final = TypeAdapter(object) def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: - if redis_cache.namespace is None: - return AUTH_CACHE_INVALIDATION_CHANNEL - return f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}" + return ( + AUTH_CACHE_INVALIDATION_CHANNEL + if redis_cache.namespace is None + else f"{redis_cache.namespace}:{AUTH_CACHE_INVALIDATION_CHANNEL}" + ) @dataclass(frozen=True, slots=True) @@ -86,22 +95,26 @@ def _message_from_data(data: object) -> _CacheInvalidationMessage | None: async def _publish_to_redis(redis_cache: "RedisCache", cache_key: str, message: str) -> None: try: - client: Final = pubsub_capable_client(redis_cache) async with _in_flight_publishes: - await client.publish(auth_cache_invalidation_channel(redis_cache), message) + result: Final = await redis_cache.async_publish(channel=AUTH_CACHE_INVALIDATION_CHANNEL, message=message) + if isinstance(result, CacheFailure): + verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, result.message) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) async def _publish_batch_to_redis(redis_cache: "RedisCache", cache_keys: tuple[str, ...]) -> None: try: - client: Final = pubsub_capable_client(redis_cache) - channel: Final = auth_cache_invalidation_channel(redis_cache) - pipeline: Final = client.pipeline(transaction=False) - for cache_key in cache_keys: - _ = pipeline.publish(channel, _cache_invalidation_message_json(cache_key)) + messages: Final = tuple( + CachePublish(channel=AUTH_CACHE_INVALIDATION_CHANNEL, message=_cache_invalidation_message_json(cache_key)) + for cache_key in cache_keys + ) async with _in_flight_publishes: - _ = await pipeline.execute() + result: Final = await redis_cache.async_publish_many(messages=messages) + if isinstance(result, CacheFailure): + verbose_proxy_logger.warning( + "auth cache invalidation publish for %d keys failed: %s", len(cache_keys), result.message + ) except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors verbose_proxy_logger.warning("auth cache invalidation publish for %d keys failed: %s", len(cache_keys), e) @@ -232,7 +245,7 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache") + __slots__ = ("_additional_in_memory_caches", "_backoff_seconds", "_redis_cache", "_task", "_user_api_key_cache") def __init__( self, @@ -243,6 +256,7 @@ class AuthCacheInvalidationSubscriber: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache self._additional_in_memory_caches = tuple(additional_in_memory_caches) + self._backoff_seconds = _BACKOFF_INITIAL_SECONDS self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -262,39 +276,41 @@ class AuthCacheInvalidationSubscriber: pass async def _run(self) -> None: - backoff_seconds = _BACKOFF_INITIAL_SECONDS # rebind-ok: exponential backoff accumulator across reconnects - while True: - try: - client = pubsub_capable_client(self._redis_cache) - pubsub = client.pubsub() - try: - await pubsub.subscribe(auth_cache_invalidation_channel(self._redis_cache)) - backoff_seconds = _BACKOFF_INITIAL_SECONDS - await self._consume(pubsub) - finally: - await self._close_pubsub(pubsub) - except asyncio.CancelledError: - raise - except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect - verbose_proxy_logger.warning( - "auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs", - e, - backoff_seconds, - ) - await asyncio.sleep(backoff_seconds) - backoff_seconds = min(backoff_seconds * 2, _BACKOFF_MAX_SECONDS) + self._backoff_seconds = _BACKOFF_INITIAL_SECONDS + while failure := await self._subscribe_and_consume(): + verbose_proxy_logger.warning( + "auth cache invalidation subscriber redis error: %s; reconnecting in %.0fs", + failure.message, + self._backoff_seconds, + ) + await asyncio.sleep(self._backoff_seconds) + self._backoff_seconds = min(self._backoff_seconds * 2, _BACKOFF_MAX_SECONDS) - async def _consume(self, pubsub: ConfigSyncPubSub) -> None: - while True: - message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + async def _subscribe_and_consume(self) -> CacheFailure: + try: + subscription: Final = await self._redis_cache.async_subscribe(channel=AUTH_CACHE_INVALIDATION_CHANNEL) + if isinstance(subscription, CacheFailure): + return subscription + try: + self._backoff_seconds = _BACKOFF_INITIAL_SECONDS + return await self._consume(subscription) + finally: + await self._close_pubsub(subscription) + except Exception as error: # noqa: BLE001 # custom cache subclasses and eviction errors must not stop reconnects + return CacheFailure("unavailable", str(error)) + + async def _consume(self, subscription: CacheSubscription) -> CacheFailure: + while not isinstance( + message := await subscription.receive(timeout_seconds=_POLL_TIMEOUT_SECONDS), CacheFailure + ): if message is None: continue self._apply_message(message) + return message @with_service_target(AUTH_OBJECTS_TARGET) - def _apply_message(self, message: object) -> None: - data: Final = message.get("data") if isinstance(message, dict) else None - parsed: Final = _message_from_data(data) + def _apply_message(self, message: str | bytes) -> None: + parsed: Final = _message_from_data(message) if parsed is None: return if parsed.new_value is not None: @@ -306,8 +322,10 @@ class AuthCacheInvalidationSubscriber: additional_cache.delete_cache(parsed.cache_key) @staticmethod - async def _close_pubsub(pubsub: ConfigSyncPubSub) -> None: + async def _close_pubsub(subscription: CacheSubscription) -> None: try: - await pubsub.aclose() - except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + failure: Final = await subscription.close() + if failure is not None: + verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", failure.message) + except Exception as e: # noqa: BLE001 # custom subscriptions must not interrupt reconnect or cancellation verbose_proxy_logger.debug("auth cache invalidation pubsub close failed: %s", e) diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index da74851b2d8..a8dd0d9d42f 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -4,10 +4,17 @@ import random import time from collections.abc import Awaitable, Callable from dataclasses import asdict, dataclass -from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast +from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # untyped prisma boundary needs cast from litellm._logging import verbose_proxy_logger +from litellm.caching.redis_cache import ( + pubsub_capable_client as _pubsub_capable_client, # noqa: F401 # pyright: ignore[reportDeprecated] # compatibility-only re-export +) +from litellm.caching.redis_cache import ( + pubsub_capable_client as pubsub_capable_client, # noqa: PLC0414 # pyright: ignore[reportDeprecated] # compatibility-only re-export +) from litellm.repositories.prisma_protocols import RowT_co, TableActions +from litellm.types.caching import CacheFailure, CacheSubscription if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache @@ -21,21 +28,7 @@ class _ConfigSyncPubSub(Protocol): def aclose(self) -> Awaitable[object]: ... -ConfigSyncPubSub = _ConfigSyncPubSub - - -class _PublishPipeline(Protocol): - def publish(self, channel: str, message: str) -> object: ... - - def execute(self) -> Awaitable[object]: ... - - -class _ConfigSyncPubSubClient(Protocol): - def publish(self, channel: str, message: str) -> Awaitable[int]: ... - - def pubsub(self) -> ConfigSyncPubSub: ... - - def pipeline(self, transaction: bool = True) -> _PublishPipeline: ... +ConfigSyncPubSub: Final = _ConfigSyncPubSub CONFIG_SYNC_CHANNEL: Final = "litellm_proxy.config_change" @@ -88,19 +81,7 @@ def coordination_redis_cache() -> "RedisCache | None": def config_sync_channel(redis_cache: "RedisCache") -> str: - if redis_cache.namespace is None: - return CONFIG_SYNC_CHANNEL - return f"{redis_cache.namespace}:{CONFIG_SYNC_CHANNEL}" - - -def pubsub_capable_client(redis_cache: "RedisCache") -> _ConfigSyncPubSubClient: - return cast( # cast-ok: protocol view of the pub/sub-capable async redis client - _ConfigSyncPubSubClient, - redis_cache.init_pubsub_client(), # pyright: ignore[reportUnknownMemberType] # redis generics - ) - - -_pubsub_capable_client: Final = pubsub_capable_client + return CONFIG_SYNC_CHANNEL if redis_cache.namespace is None else f"{redis_cache.namespace}:{CONFIG_SYNC_CHANNEL}" @dataclass(frozen=True, slots=True) @@ -116,8 +97,11 @@ async def publish_config_change(redis_cache: "RedisCache | None", object_type: s if redis_cache is None: return try: - client: Final = pubsub_capable_client(redis_cache) - await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type)) + result: Final = await redis_cache.async_publish( + channel=CONFIG_SYNC_CHANNEL, message=_config_change_message_json(object_type) + ) + if isinstance(result, CacheFailure): + verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, result.message) except Exception as e: # noqa: BLE001 # best-effort publish; writes must never fail on redis errors verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, e) @@ -178,6 +162,7 @@ class ConfigSyncSubscriber: __slots__ = ( "_backoff_initial_seconds", "_backoff_max_seconds", + "_backoff_seconds", "_debounce_seconds", "_jitter_max_seconds", "_last_resync_at", @@ -211,6 +196,7 @@ class ConfigSyncSubscriber: self._min_resync_interval_seconds = min_resync_interval_seconds self._backoff_initial_seconds = backoff_initial_seconds self._backoff_max_seconds = backoff_max_seconds + self._backoff_seconds = backoff_initial_seconds self._rng = rng if rng is not None else random.Random() self._sleep = sleep self._monotonic = monotonic @@ -235,50 +221,61 @@ class ConfigSyncSubscriber: pass async def _run(self) -> None: - backoff_seconds = self._backoff_initial_seconds - while True: - try: - client = pubsub_capable_client(self._redis_cache) - pubsub = client.pubsub() - try: - await pubsub.subscribe(config_sync_channel(self._redis_cache)) - if self._resync_required: - await self._debounce_and_resync(pubsub) - self._resync_required = False - backoff_seconds = self._backoff_initial_seconds - await self._consume(pubsub) - finally: - await self._close_pubsub(pubsub) - except asyncio.CancelledError: - raise - except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect - self._resync_required = True - verbose_proxy_logger.warning( - "config sync subscriber redis error: %s; reconnecting in %.0fs", - e, - backoff_seconds, - ) - await self._sleep(backoff_seconds) - backoff_seconds = min(backoff_seconds * 2, self._backoff_max_seconds) + self._backoff_seconds = self._backoff_initial_seconds + while failure := await self._subscribe_and_consume(): + self._resync_required = True + verbose_proxy_logger.warning( + "config sync subscriber redis error: %s; reconnecting in %.0fs", + failure.message, + self._backoff_seconds, + ) + await self._sleep(self._backoff_seconds) + self._backoff_seconds = min(self._backoff_seconds * 2, self._backoff_max_seconds) - async def _consume(self, pubsub: ConfigSyncPubSub) -> None: - while True: - message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + async def _subscribe_and_consume(self) -> CacheFailure: + try: + subscription: Final = await self._redis_cache.async_subscribe(channel=CONFIG_SYNC_CHANNEL) + if isinstance(subscription, CacheFailure): + return subscription + try: + if self._resync_required: + failure: Final = await self._debounce_and_resync(subscription) + if failure is not None: + return failure + self._resync_required = False + self._backoff_seconds = self._backoff_initial_seconds + return await self._consume(subscription) + finally: + await self._close_pubsub(subscription) + except Exception as error: # noqa: BLE001 # custom cache subclasses and callbacks must not stop reconnects + return CacheFailure("unavailable", str(error)) + + async def _consume(self, subscription: CacheSubscription) -> CacheFailure: + while not isinstance( + message := await subscription.receive(timeout_seconds=_POLL_TIMEOUT_SECONDS), CacheFailure + ): if message is None: continue - await self._debounce_and_resync(pubsub) + if (failure := await self._debounce_and_resync(subscription)) is not None: + return failure + return message - async def _debounce_and_resync(self, pubsub: _ConfigSyncPubSub) -> None: + async def _debounce_and_resync(self, subscription: CacheSubscription) -> CacheFailure | None: await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)) await self._wait_for_min_resync_interval() - await self._drain_pending(pubsub) + failure: Final = await self._drain_pending(subscription) + if failure is not None: + return failure await self._run_resync_callbacks() self._last_resync_at = self._monotonic() + return None async def _wait_for_min_resync_interval(self) -> None: if self._last_resync_at is None: return - seconds_until_next_resync = self._min_resync_interval_seconds - (self._monotonic() - self._last_resync_at) + seconds_until_next_resync: Final = self._min_resync_interval_seconds - ( + self._monotonic() - self._last_resync_at + ) if seconds_until_next_resync <= 0: return verbose_proxy_logger.debug( @@ -288,9 +285,11 @@ class ConfigSyncSubscriber: await self._sleep(seconds_until_next_resync) @staticmethod - async def _drain_pending(pubsub: ConfigSyncPubSub) -> None: - while await pubsub.get_message(ignore_subscribe_messages=True, timeout=0) is not None: - pass + async def _drain_pending(subscription: CacheSubscription) -> CacheFailure | None: + while (message := await subscription.receive(timeout_seconds=0)) is not None: + if isinstance(message, CacheFailure): + return message + return None async def _run_resync_callbacks(self) -> None: for callback in self._resync_callbacks: @@ -300,8 +299,10 @@ class ConfigSyncSubscriber: verbose_proxy_logger.warning("config sync resync callback failed: %s", e) @staticmethod - async def _close_pubsub(pubsub: ConfigSyncPubSub) -> None: + async def _close_pubsub(subscription: CacheSubscription) -> None: try: - await pubsub.aclose() - except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + failure: Final = await subscription.close() + if failure is not None: + verbose_proxy_logger.debug("config sync pubsub close failed: %s", failure.message) + except Exception as e: # noqa: BLE001 # custom subscriptions must not interrupt reconnect or cancellation verbose_proxy_logger.debug("config sync pubsub close failed: %s", e) diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index ec4baa76514..6f724a2c7b5 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -23,8 +23,10 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.bug_report_config import build_proxy_environment_report from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +from litellm.types.caching import CacheFailure, CachePoolInfo if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache from litellm.router import Router router: Final = APIRouter() @@ -501,8 +503,33 @@ def _get_uncollectable_objects_info() -> Mapping[str, object]: } +def _redis_memory_stats(cache: "RedisCache | None") -> Mapping[str, object]: + if cache is None: + return {"enabled": False} + summary: Final[Mapping[str, object]] = {"enabled": True, "cache_type": type(cache).__name__} + try: + pool: Final = cache.get_connection_pool_info() + except Exception as error: + verbose_proxy_logger.debug("Error getting Redis pool info: %s", error) + return summary + match pool: + case CachePoolInfo(): + return { + **summary, + "connection_pool": { + "max_connections": pool.max_connections, + "connection_class": pool.connection_class, + }, + } + case CacheFailure(): + verbose_proxy_logger.debug("Error getting Redis pool info: %s", pool.message) + return summary + case None: + return summary + + def _get_cache_memory_stats( - user_api_key_cache, llm_router: "Router | None", proxy_logging_obj, redis_usage_cache + user_api_key_cache, llm_router: "Router | None", proxy_logging_obj, redis_usage_cache: "RedisCache | None" ) -> Mapping[str, object]: """Calculate memory usage for all caches.""" cache_stats: Final[dict[str, object]] = {} @@ -553,29 +580,7 @@ def _get_cache_memory_stats( "total_size_mb": round((logging_cache_size + logging_ttl_size) / (1024 * 1024), 2), } - # Redis cache info - if redis_usage_cache is not None: - cache_stats["redis_usage_cache"] = { - "enabled": True, - "cache_type": type(redis_usage_cache).__name__, - } - # Try to get Redis connection pool info if available - try: - if hasattr(redis_usage_cache, "redis_client") and redis_usage_cache.redis_client: - if hasattr(redis_usage_cache.redis_client, "connection_pool"): - pool_info: Final = redis_usage_cache.redis_client.connection_pool - cache_stats["redis_usage_cache"]["connection_pool"] = { - "max_connections": ( - pool_info.max_connections if hasattr(pool_info, "max_connections") else None - ), - "connection_class": ( - pool_info.connection_class.__name__ if hasattr(pool_info, "connection_class") else None - ), - } - except Exception as e: - verbose_proxy_logger.debug("Error getting Redis pool info: %s", e) - else: - cache_stats["redis_usage_cache"] = {"enabled": False} + cache_stats["redis_usage_cache"] = _redis_memory_stats(redis_usage_cache) except Exception as e: verbose_proxy_logger.debug("Error calculating cache stats: %s", e) diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 1116cd604a3..95ba17b5edf 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -9,6 +9,7 @@ from litellm._uuid import uuid from litellm.caching.redis_cache import RedisCache, log_redis_failure 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.types.caching import CacheFailure from litellm.types.services import ServiceTypes POD_LOCK_TARGET: Final = "pod_lock" @@ -28,18 +29,9 @@ 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: @@ -127,8 +119,7 @@ end 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. + Falls back to GET + DEL when the atomic release is unavailable. """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping release_lock") @@ -162,31 +153,18 @@ end @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, - ) + if self.redis_cache is None: + return 0 + if isinstance(self.redis_cache, RedisCache): + released: Final = await self.redis_cache.async_release_lease(key=lock_key, token=json.dumps(self.pod_id)) + if not isinstance(released, CacheFailure): + return int(released) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + f"Atomic release failed for lock_key={lock_key}, falling back to GET+DEL", + released, + ) current_value = await self.redis_cache.async_get_cache(lock_key) if isinstance(current_value, bytes): diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 3a4375ac54a..98c72fe5dc5 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -13,7 +13,6 @@ from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import TypeAdapter -from redis.exceptions import RedisError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger @@ -63,6 +62,8 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( from litellm.proxy.db.spend_log_batching import SpendLogRow from litellm.secret_managers.main import str_to_bool from litellm.types.caching import ( + CacheFailure, + CacheQueueUpdate, RedisPipelineLpopOperation, RedisPipelineRpushOperation, ) @@ -549,25 +550,25 @@ class RedisUpdateBuffer: ) rpush_list: Final = tuple( - RedisPipelineRpushOperation(key=redis_key, values=(safe_dumps(transactions),)) + CacheQueueUpdate(key=redis_key, values=(safe_dumps(transactions),)) for transactions, redis_key in restore_configs if transactions ) if len(rpush_list) == 0: return - try: - await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list) - verbose_proxy_logger.info( - "Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.", - len(rpush_list), - ) - except RedisError as e: + result: Final = await self.redis_cache.async_enqueue_updates(updates=rpush_list) + if isinstance(result, CacheFailure): verbose_proxy_logger.error( "Spend tracking - failed to restore uncommitted transactions to Redis. " "These spend updates are lost. Error: %s", - str(e), + result.message, ) + return + verbose_proxy_logger.info( + "Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.", + len(rpush_list), + ) @with_service_target(SPEND_QUEUE_TARGET) async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None: @@ -577,17 +578,14 @@ class RedisUpdateBuffer: items: Final = tool_discovery_queue.flush() if not items: return - try: - _ = await self.redis_cache.async_rpush_pipeline( - rpush_list=[ - RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)]) - ] - ) - except (RedisError, OSError, asyncio.TimeoutError) as e: + result: Final = await self.redis_cache.async_enqueue_updates( + updates=(CacheQueueUpdate(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=(safe_dumps(items),)),) + ) + if isinstance(result, CacheFailure): verbose_proxy_logger.error( "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s", len(items), - str(e), + result.message, ) for item in items: tool_discovery_queue.add_update(item) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index da595e101f1..6cb07847434 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -13,17 +13,18 @@ import logging import math import time import uuid -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias +from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger -from litellm.caching.redis_cache import log_redis_failure +from litellm.caching.redis_cache import RedisCache, log_redis_failure from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQUEUED_TOKEN_TTL_SECONDS from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.caching import CacheFailure from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: @@ -40,40 +41,6 @@ BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset( ScopeKey: TypeAlias = Literal["api_key", "team"] -RESERVE_ENQUEUED_TOKENS_SCRIPT: Final = """ -local amount = tonumber(ARGV[1]) -local ttl = tonumber(ARGV[2]) -local limit = tonumber(ARGV[3]) -local current = tonumber(redis.call('GET', KEYS[1]) or '0') -if current + amount > limit then - return {0, current} -end -local updated = redis.call('INCRBY', KEYS[1], amount) -redis.call('EXPIRE', KEYS[1], ttl) -return {1, updated} -""" - -REFUND_ENQUEUED_TOKENS_SCRIPT: Final = """ -local updated = redis.call('DECRBY', KEYS[1], tonumber(ARGV[1])) -if updated <= 0 then - redis.call('DEL', KEYS[1]) -end -return 1 -""" - -SAVE_RESERVATION_SCRIPT: Final = """ -redis.call('SET', KEYS[1], ARGV[1], 'EX', tonumber(ARGV[2])) -return 1 -""" - -POP_RESERVATION_SCRIPT: Final = """ -local value = redis.call('GET', KEYS[1]) -if value and value ~= '' then - redis.call('SET', KEYS[1], '', 'EX', tonumber(ARGV[1])) -end -return value -""" - @dataclass(frozen=True, slots=True) class BatchEnqueuedTokenScope: @@ -103,17 +70,11 @@ class BatchEnqueuedTokenOverLimit: BatchEnqueuedTokenOutcome: TypeAlias = BatchEnqueuedTokenReservation | BatchEnqueuedTokenOverLimit _LIMIT_ADAPTER: Final = TypeAdapter(Annotated[int, Field(gt=0)]) -_RESERVE_RESULT_ADAPTER: Final = TypeAdapter(tuple[int, int]) -_POPPED_VALUE_ADAPTER: Final = TypeAdapter(str | bytes | None) _STORED_COUNTER_ADAPTER: Final = TypeAdapter(int | None) _RESERVATION_ADAPTER: Final = TypeAdapter(BatchEnqueuedTokenReservation) _LOCAL_RECORD_ADAPTER: Final = TypeAdapter(object) -class _ScriptRunner(Protocol): - def __call__(self, keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> Awaitable[object]: ... - - def _read_metadata_limit(metadata: Mapping[str, object] | None) -> int | None: if not metadata: return None @@ -202,19 +163,7 @@ class BatchEnqueuedTokenStore: self._monotonic: Final = monotonic self._lock = asyncio.Lock() self._owner_token = uuid.uuid4().hex - redis_cache = internal_usage_cache.dual_cache.redis_cache - self._reserve_script: _ScriptRunner | None = ( - redis_cache.async_register_script(RESERVE_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None - ) - self._refund_script: _ScriptRunner | None = ( - redis_cache.async_register_script(REFUND_ENQUEUED_TOKENS_SCRIPT) if redis_cache is not None else None - ) - self._save_script: _ScriptRunner | None = ( - redis_cache.async_register_script(SAVE_RESERVATION_SCRIPT) if redis_cache is not None else None - ) - self._pop_script: _ScriptRunner | None = ( - redis_cache.async_register_script(POP_RESERVATION_SCRIPT) if redis_cache is not None else None - ) + self._redis_cache = internal_usage_cache.dual_cache.redis_cache @staticmethod def _counter_key(scope: BatchEnqueuedTokenScope) -> str: @@ -233,83 +182,77 @@ class BatchEnqueuedTokenStore: ) -> BatchEnqueuedTokenOutcome: if tokens <= 0 or not scopes: return BatchEnqueuedTokenReservation(tokens=max(tokens, 0), scopes=scopes) - reserve_script: Final = self._reserve_script - refund_script: Final = self._refund_script - if reserve_script is not None and refund_script is not None: - try: - return await self._reserve_via_redis(reserve_script, refund_script, tokens=tokens, scopes=scopes) - except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory counters - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - "Redis enqueued-token reserve failed, falling back to in-memory", - e, - ) + if self._redis_cache is not None: + result: Final = await self._reserve_via_redis(self._redis_cache, tokens=tokens, scopes=scopes) + if not isinstance(result, CacheFailure): + return result + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Redis enqueued-token reserve failed, falling back to in-memory", + result, + ) return await self._reserve_in_memory(tokens=tokens, scopes=scopes, span=litellm_parent_otel_span) async def _reserve_via_redis( self, - reserve_script: _ScriptRunner, - refund_script: _ScriptRunner, + redis_cache: RedisCache, tokens: int, scopes: tuple[BatchEnqueuedTokenScope, ...], - ) -> BatchEnqueuedTokenOutcome: + ) -> BatchEnqueuedTokenOutcome | CacheFailure: started: Final = self._monotonic() for index, scope in enumerate(scopes): - result = await self._run_reserve_script( - reserve_script, - refund_script, - tokens=tokens, - scope=scope, - already_reserved=scopes[:index], - ) - if result[0] != 1: - await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index]) - return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1]) + if ( + isinstance( + result := await redis_cache.async_reserve_counter( + key=self._counter_key(scope), + amount=tokens, + limit=scope.limit, + ttl_seconds=BATCH_ENQUEUED_TOKEN_TTL_SECONDS, + ), + CacheFailure, + ) + or not result.accepted + ): + await self._rollback_partial_reserve(redis_cache, tokens=tokens, scopes=scopes[:index]) + return ( + result + if isinstance(result, CacheFailure) + else BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result.count) + ) return BatchEnqueuedTokenReservation( tokens=tokens, scopes=scopes, backend="redis", reserved_at_monotonic=started ) - async def _run_reserve_script( - self, - reserve_script: _ScriptRunner, - refund_script: _ScriptRunner, - tokens: int, - scope: BatchEnqueuedTokenScope, - already_reserved: tuple[BatchEnqueuedTokenScope, ...], - ) -> tuple[int, int]: - try: - raw_result: Final = await reserve_script( - (self._counter_key(scope),), - (tokens, BATCH_ENQUEUED_TOKEN_TTL_SECONDS, scope.limit), - ) - return _RESERVE_RESULT_ADAPTER.validate_python(raw_result) - except Exception: - await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=already_reserved) - raise - async def _rollback_partial_reserve( self, - refund_script: _ScriptRunner, + redis_cache: RedisCache, tokens: int, scopes: tuple[BatchEnqueuedTokenScope, ...], ) -> None: - try: - await self._refund_via_redis(refund_script, tokens=tokens, scopes=scopes) - except Exception as e: # noqa: BLE001 # best-effort rollback: the leak is TTL-bounded and only tightens the allowance + result: Final = await self._refund_via_redis(redis_cache, tokens=tokens, scopes=scopes) + if isinstance(result, CacheFailure): verbose_proxy_logger.warning( "Rollback of partially reserved enqueued tokens failed; leaked increments expire with the TTL: %s", - str(e), + result.message, ) async def _refund_via_redis( self, - refund_script: _ScriptRunner, + redis_cache: RedisCache, tokens: int, scopes: tuple[BatchEnqueuedTokenScope, ...], - ) -> None: + ) -> None | CacheFailure: for scope in scopes: - await refund_script((self._counter_key(scope),), (tokens,)) + if isinstance( + result := await redis_cache.async_refund_counter( + key=self._counter_key(scope), + amount=tokens, + ), + CacheFailure, + ): + return result + return None async def _reserve_in_memory( self, @@ -355,17 +298,16 @@ class BatchEnqueuedTokenStore: await self._set_local_counter(scope, remaining, litellm_parent_otel_span) async def _refund_redis_reservation(self, reservation: BatchEnqueuedTokenReservation) -> None: - refund_script: Final = self._refund_script - if refund_script is None: + redis_cache: Final = self._redis_cache + if redis_cache is None: verbose_proxy_logger.warning( "No Redis client for a Redis-granted enqueued-token refund; leaked increments expire with the TTL" ) return - try: - await self._refund_via_redis(refund_script, tokens=reservation.tokens, scopes=reservation.scopes) - except Exception as e: # noqa: BLE001 # best-effort refund: the leak is TTL-bounded and only tightens the allowance + result: Final = await self._refund_via_redis(redis_cache, tokens=reservation.tokens, scopes=reservation.scopes) + if isinstance(result, CacheFailure): verbose_proxy_logger.warning( - "Redis enqueued-token refund failed; leaked increments expire with the TTL: %s", str(e) + "Redis enqueued-token refund failed; leaked increments expire with the TTL: %s", result.message ) @with_service_target("rate_limits") @@ -378,21 +320,18 @@ class BatchEnqueuedTokenStore: serialized: Final = _RESERVATION_ADAPTER.dump_json(reservation).decode("utf-8") elapsed: Final = self._monotonic() - reservation.reserved_at_monotonic ttl: Final = max(1, BATCH_ENQUEUED_TOKEN_TTL_SECONDS - math.ceil(elapsed)) - if self._save_script is not None: - try: - await self._save_script( - (self._record_key(batch_id),), - (serialized, ttl), - ) - except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - "Redis enqueued-token reservation save failed, falling back to in-memory", - e, - ) - else: + if self._redis_cache is not None: + result: Final = await self._redis_cache.async_save_record( + key=self._record_key(batch_id), value=serialized, ttl_seconds=ttl + ) + if not isinstance(result, CacheFailure): return + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Redis enqueued-token reservation save failed, falling back to in-memory", + result, + ) await self.internal_usage_cache.async_set_cache( key=self._record_key(batch_id), value=serialized, @@ -427,22 +366,21 @@ class BatchEnqueuedTokenStore: verbose_proxy_logger.warning("Discarding malformed enqueued-token reservation record for %s", batch_id) return None - async def _pop_redis_record(self, batch_id: str) -> str | bytes | None: - pop_script: Final = self._pop_script - if pop_script is None: - return None - try: - return _POPPED_VALUE_ADAPTER.validate_python( - await pop_script((self._record_key(batch_id),), (BATCH_ENQUEUED_TOKEN_TTL_SECONDS,)) - ) - except Exception as e: # noqa: BLE001 # any Redis failure must fall back to the in-memory record - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - "Redis enqueued-token reservation pop failed, falling back to in-memory", - e, - ) + async def _pop_redis_record(self, batch_id: str) -> str | None: + if self._redis_cache is None: return None + result: Final = await self._redis_cache.async_pop_record( + key=self._record_key(batch_id), tombstone_ttl_seconds=BATCH_ENQUEUED_TOKEN_TTL_SECONDS + ) + if not isinstance(result, CacheFailure): + return result + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Redis enqueued-token reservation pop failed, falling back to in-memory", + result, + ) + return None async def _pop_local_record(self, batch_id: str, span: "Span | None") -> object: async with self._lock: diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index cf3025d8c7b..ec3f2438e35 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -27,6 +27,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.types.caching import CacheFailure if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -36,23 +37,6 @@ else: InternalUsageCache = Any -# Redis Lua script for atomic float increment with TTL. -# INCRBYFLOAT returns the new value as a string. -# Only sets EXPIRE on first call (when prior value was nil). -MAX_BUDGET_SESSION_INCREMENT_SCRIPT: Final = """ -local key = KEYS[1] -local amount = ARGV[1] -local ttl = tonumber(ARGV[2]) - -local existed = redis.call('EXISTS', key) -local new_val = redis.call('INCRBYFLOAT', key, amount) -if existed == 0 then - redis.call('EXPIRE', key, ttl) -end - -return new_val -""" - # Default TTL for session budget counters (1 hour) DEFAULT_MAX_BUDGET_PER_SESSION_TTL: Final = 3600 @@ -77,12 +61,7 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): ) ) - if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - MAX_BUDGET_SESSION_INCREMENT_SCRIPT - ) - else: - self.increment_script = None + self._redis_cache = self.internal_usage_cache.dual_cache.redis_cache @with_service_target("session_budgets") async def async_pre_call_hook( @@ -239,20 +218,16 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): async def _increment_spend(self, cache_key: str, amount: float) -> float: """Atomically increment the session spend and return the new value.""" - if self.increment_script is not None: - try: - result: Final = await self.increment_script( - keys=[cache_key], - args=[str(amount), self.ttl], - ) - return float(result) - except Exception as e: - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, falling back to in-memory", - e, - ) + if self._redis_cache is not None: + result: Final = await self._redis_cache.async_increment_spend( + key=cache_key, amount=amount, ttl_seconds=self.ttl + ) + if not isinstance(result, CacheFailure): + return result + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: Redis spend increment failed, falling back to in-memory: %s", + result.message, + ) return await self._in_memory_increment_spend(cache_key, amount) diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 5b7d7d04e24..b45e4b532dc 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.types.caching import CacheFailure if TYPE_CHECKING: from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache @@ -30,21 +31,6 @@ else: InternalUsageCache = Any -# Redis Lua script for atomic increment with TTL. -# Returns the new count after increment. -# Only sets EXPIRE on first increment (when count becomes 1). -MAX_ITERATIONS_INCREMENT_SCRIPT: Final = """ -local key = KEYS[1] -local ttl = tonumber(ARGV[1]) - -local current = redis.call('INCR', key) -if current == 1 then - redis.call('EXPIRE', key, ttl) -end - -return current -""" - # Default TTL for session iteration counters (1 hour) DEFAULT_MAX_ITERATIONS_TTL: Final = 3600 @@ -64,22 +50,15 @@ class _PROXY_MaxIterationsHandler(CustomLogger): {session_iterations:}:count Multi-instance support: - Uses Redis Lua script for atomic increment (same pattern as - parallel_request_limiter_v3). Falls back to in-memory cache - when Redis is unavailable. + Uses the cache's atomic iteration increment and falls back to + in-memory counters when Redis is unavailable. """ def __init__(self, internal_usage_cache: InternalUsageCache): self.internal_usage_cache = internal_usage_cache self.ttl = int(os.getenv("LITELLM_MAX_ITERATIONS_TTL", DEFAULT_MAX_ITERATIONS_TTL)) - # Register Lua script with Redis if available (same pattern as v3 limiter) - if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - MAX_ITERATIONS_INCREMENT_SCRIPT - ) - else: - self.increment_script = None + self._redis_cache = self.internal_usage_cache.dual_cache.redis_cache @with_service_target("session_iterations") async def async_pre_call_hook( @@ -186,21 +165,16 @@ class _PROXY_MaxIterationsHandler(CustomLogger): """ Atomically increment the session counter and return the new value. - Tries Redis first (via registered Lua script for atomicity across - instances), falls back to in-memory cache. + Tries the shared atomic counter first, then falls back to in-memory cache. """ - if self.increment_script is not None: - try: - result: Final = await self.increment_script( - keys=[cache_key], - args=[self.ttl], - ) - return int(result) - except Exception as e: - verbose_proxy_logger.warning( - "MaxIterationsHandler: Redis failed, falling back to in-memory: %s", - str(e), - ) + if self._redis_cache is not None: + result: Final = await self._redis_cache.async_increment_iterations(key=cache_key, ttl_seconds=self.ttl) + if not isinstance(result, CacheFailure): + return result + verbose_proxy_logger.warning( + "MaxIterationsHandler: Redis failed, falling back to in-memory: %s", + result.message, + ) # Fallback: in-memory cache return await self._in_memory_increment(cache_key) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 567b8452ae2..d1ef1163187 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -36,8 +36,6 @@ from litellm import DualCache from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.redis_batch import ( - BatchResult, - RegisteredScript, active_post_call_redis_batch, active_request_redis_batch, ) @@ -80,7 +78,20 @@ from litellm.router_utils.add_retry_fallback_headers import ( ) from litellm.router_utils.common_utils import resolve_model_group_alias from litellm.router_utils.ptu_shares import PTUTeamCeiling -from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.caching import ( + CacheFailure, + CacheOperation, + CapacityExceeded, + ConcurrencyRead, + ConcurrencyRelease, + ConcurrencySlot, + CounterIncrement, + GuardedIncrement, + RedisPipelineIncrementOperation, + WindowCount, + WindowCounter, + WindowReservation, +) from litellm.types.extraction import ExtractionResponse from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -213,265 +224,6 @@ def _sibling_counter_keys(window_key: str) -> tuple[str, str]: return f"{prefix}:requests", f"{prefix}:tokens" -BATCH_RATE_LIMITER_SCRIPT: Final = """ -local results = {} -local now = tonumber(ARGV[1]) -local window_size = tonumber(ARGV[2]) - --- Process each window/counter pair -for i = 1, #KEYS, 2 do - local window_key = KEYS[i] - local counter_key = KEYS[i + 1] - local increment_value = 1 - - -- Check if window exists and is valid - local window_start = redis.call('GET', window_key) - if not window_start or (now - tonumber(window_start)) >= window_size then - -- Reset window and counter - local prefix = string.sub(window_key, 1, -(#':window') - 1) - redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') - redis.call('SET', window_key, tostring(now)) - redis.call('SET', counter_key, increment_value) - redis.call('EXPIRE', window_key, window_size) - redis.call('EXPIRE', counter_key, window_size) - table.insert(results, tostring(now)) -- window_start - table.insert(results, increment_value) -- counter - else - local counter = redis.call('INCR', counter_key) - -- This happens when window_key exists but counter_key doesn't (e.g., tokens key - -- created after requests key when both share the same window_key) - local current_ttl = redis.call('TTL', counter_key) - if current_ttl == -1 then - redis.call('EXPIRE', counter_key, window_size) - end - table.insert(results, window_start) -- window_start - table.insert(results, counter) -- counter - end -end - -return results -""" - -BATCH_COUNTER_READ_SCRIPT: Final = "return redis.call('MGET', unpack(KEYS))" - -CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """ --- Atomic check-and-increment-by-N across one or more descriptors. --- All-or-nothing: if any descriptor would exceed its limit, no counter is --- modified. --- --- Uses Redis server time (`redis.call('TIME')`) instead of a client-supplied --- timestamp so that window resets are deterministic across replicas with --- skewed wall-clocks. This prevents a clock-skew-induced reopening of the --- TOCTOU window across multi-replica deployments. --- --- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. --- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: --- ARGV[(i-1)*4 + 1] = limit --- ARGV[(i-1)*4 + 2] = increment --- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) --- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) --- --- Return on success: --- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... } --- Return on over-limit: { 1, descriptor_index, current_counter, limit } -local time_reply = redis.call('TIME') -local now = tonumber(time_reply[1]) -local descriptor_count = #KEYS / 2 -local reset_windows = {} - --- Pass 1: read state, validate. Abort without writing if any over limit. -local descriptor_state = {} -for i = 1, descriptor_count do - local window_key = KEYS[(i - 1) * 2 + 1] - local counter_key = KEYS[(i - 1) * 2 + 2] - local arg_base = (i - 1) * 4 + 1 - local limit = tonumber(ARGV[arg_base]) - local increment = tonumber(ARGV[arg_base + 1]) - local window_size = tonumber(ARGV[arg_base + 3]) - - local window_start = redis.call('GET', window_key) - local window_expired = (not window_start) or - ((now - tonumber(window_start)) >= window_size) - - local current_counter - if window_expired then - current_counter = 0 - else - current_counter = tonumber(redis.call('GET', counter_key) or 0) - end - - local blocked - if increment > 0 then - blocked = current_counter + increment > limit - else - blocked = current_counter >= limit - end - if blocked then - return { 1, i, current_counter, limit } - end - - descriptor_state[i] = { window_expired, current_counter, window_start } -end - --- Pass 2: all checks passed. Apply increments. -local results = { 0 } -for i = 1, descriptor_count do - local window_key = KEYS[(i - 1) * 2 + 1] - local counter_key = KEYS[(i - 1) * 2 + 2] - local arg_base = (i - 1) * 4 + 1 - local increment = tonumber(ARGV[arg_base + 1]) - local ttl = tonumber(ARGV[arg_base + 2]) - local window_size = tonumber(ARGV[arg_base + 3]) - - local window_expired = descriptor_state[i][1] - local active_window_start - - if window_expired then - active_window_start = now - if not reset_windows[window_key] then - local prefix = string.sub(window_key, 1, -(#':window') - 1) - redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') - reset_windows[window_key] = true - end - redis.call('SET', window_key, tostring(now)) - redis.call('SET', counter_key, increment) - redis.call('EXPIRE', window_key, window_size) - if ttl > 0 then - redis.call('EXPIRE', counter_key, ttl) - end - table.insert(results, increment) - else - active_window_start = tonumber(descriptor_state[i][3]) - local new_counter = redis.call('INCRBY', counter_key, increment) - local current_ttl = redis.call('TTL', counter_key) - if current_ttl == -1 and ttl > 0 then - redis.call('EXPIRE', counter_key, ttl) - end - table.insert(results, new_counter) - end - table.insert(results, active_window_start) -end - -return results -""" - -WINDOW_GUARDED_TOKEN_INCREMENT_SCRIPT: Final = """ -local results = {} -for i = 1, #KEYS, 2 do - local window_key = KEYS[i] - local counter_key = KEYS[i + 1] - local arg_base = ((i - 1) / 2) * 3 + 1 - local expected_window_start = ARGV[arg_base] - local increment = tonumber(ARGV[arg_base + 1]) - local ttl = tonumber(ARGV[arg_base + 2]) - local active_window_start = redis.call('GET', window_key) - - if active_window_start and active_window_start == expected_window_start then - local new_counter = redis.call('INCRBY', counter_key, increment) - local current_ttl = redis.call('TTL', counter_key) - if current_ttl == -1 and ttl > 0 then - redis.call('EXPIRE', counter_key, ttl) - end - table.insert(results, 1) - table.insert(results, new_counter) - else - table.insert(results, 0) - table.insert(results, tonumber(redis.call('GET', counter_key) or 0)) - end -end -return results -""" - -PARALLEL_ACQUIRE_SCRIPT: Final = """ --- Atomic check-and-acquire for the max_parallel_requests concurrency gauge. --- Each gauge key is a sorted set of per-request slot ids scored by acquire --- time (Redis server clock). In-flight requests are counted by ZCARD after --- pruning slots older than the slot TTL, so unlike the windowed RPM/TPM --- counters the gauge is never reset while requests are in flight, a --- rejected request never occupies a slot, and a slot leaked by a crashed --- worker self-heals after the slot TTL even under continuous traffic. --- --- KEYS: one gauge zset key per descriptor. --- ARGV: per-key triples (limit, slot_ttl_seconds, slot_id). --- Success: { 0, in_flight_1, ... }. Over-limit: { 1, key_index, in_flight, limit }. -local time_reply = redis.call('TIME') -local now = tonumber(time_reply[1]) -for i = 1, #KEYS do - local limit = tonumber(ARGV[(i - 1) * 3 + 1]) - local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) - redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - slot_ttl) - local in_flight = redis.call('ZCARD', KEYS[i]) - if in_flight + 1 > limit then - return { 1, i, in_flight, limit } - end -end -local results = { 0 } -for i = 1, #KEYS do - local slot_ttl = tonumber(ARGV[(i - 1) * 3 + 2]) - local slot_id = ARGV[(i - 1) * 3 + 3] - redis.call('ZADD', KEYS[i], now, slot_id) - redis.call('EXPIRE', KEYS[i], slot_ttl) - table.insert(results, redis.call('ZCARD', KEYS[i])) -end -return results -""" - -PARALLEL_RELEASE_SCRIPT: Final = """ --- Release one slot per gauge key by removing this request's slot id. --- ZREM of an absent member (or key) is a no-op, so a release without a --- matching acquire (proxy-side rejection, double-fired callback, slot --- already expired) can never free a slot owned by another request. --- KEYS: gauge zset keys. ARGV: per-key slot_id. --- Returns the remaining in-flight count per key. -local results = {} -for i = 1, #KEYS do - redis.call('ZREM', KEYS[i], ARGV[i]) - table.insert(results, redis.call('ZCARD', KEYS[i])) -end -return results -""" - -PARALLEL_COUNT_SCRIPT: Final = """ --- Read the current in-flight count per gauge key (prunes expired slots --- first so leaked slots do not inflate the reading). --- KEYS: gauge zset keys. ARGV: per-key slot_ttl_seconds. -local time_reply = redis.call('TIME') -local now = tonumber(time_reply[1]) -local results = {} -for i = 1, #KEYS do - redis.call('ZREMRANGEBYSCORE', KEYS[i], '-inf', now - tonumber(ARGV[i])) - table.insert(results, redis.call('ZCARD', KEYS[i])) -end -return results -""" - -TOKEN_INCREMENT_SCRIPT: Final = """ -local results = {} - --- Process each key/increment_value/ttl triplet -for i = 1, #KEYS do - local key = KEYS[i] - local increment_value = tonumber(ARGV[i * 2 - 1]) - local ttl_seconds = tonumber(ARGV[i * 2]) - - -- Increment the value - local new_value = redis.call('INCRBYFLOAT', key, increment_value) - - -- Handle TTL: only set expire if ttl_seconds > 0 and key has no current TTL - -- ttl_seconds can be 0 (no TTL) or positive (set TTL) - if ttl_seconds and ttl_seconds > 0 then - local current_ttl = redis.call('TTL', key) - if current_ttl == -1 then - redis.call('EXPIRE', key, ttl_seconds) - end - end - - table.insert(results, new_value) -end - -return results -""" - # Redis cluster slot count REDIS_CLUSTER_SLOTS: Final = 16384 REDIS_NODE_HASHTAG_NAME: Final = "all_keys" @@ -548,16 +300,8 @@ def _values_in_caller_order( return [value_by_position.get(position) for position in range(len(keys))] -def _as_counter_values(reply: object) -> list[CacheCounterValue]: - """A Lua reply read back off the pipeline is the same array the script returns when called directly.""" - if not isinstance(reply, (list, tuple)): - raise TypeError(f"rate limiter script reply is not a list: {type(reply).__name__}") - values: Final[list[CacheCounterValue]] = [] # mutable-ok: each element is narrowed before it is kept - for value in reply: # pyright: ignore[reportUnknownVariableType] # raw Redis reply - if not isinstance(value, (int, float, str, bytes)): - raise TypeError(f"rate limiter script reply holds {type(value).__name__}") # pyright: ignore[reportUnknownArgumentType] # raw Redis reply - values.append(value) - return values +def _window_values(counts: Sequence[WindowCount]) -> tuple[str | int, ...]: + return tuple(itertools.chain.from_iterable((item.window_start, item.count) for item in counts)) ReservationWindowIdentity: TypeAlias = tuple[str, str, Literal["redis", "local"]] @@ -565,12 +309,6 @@ ReservationWindowIdentity: TypeAlias = tuple[str, str, Literal["redis", "local"] ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes -class _AsyncLuaScript(Protocol): - """A Lua script registered against the async Redis client, called with KEYS and ARGV.""" - - def __call__(self, *, keys: Sequence[str], args: Sequence[object]) -> Awaitable[list[CacheCounterValue]]: ... - - class RateLimitDescriptorRateLimitObject(TypedDict, total=False): requests_per_unit: int | None tokens_per_unit: int | None @@ -657,7 +395,7 @@ class AtomicCounterState(TypedDict): window_start: ReadOnly[str] -DescriptorAtomicGroup: TypeAlias = tuple[list[str], list[int], list[AtomicCounterMeta]] +DescriptorAtomicGroup: TypeAlias = list[AtomicCounterMeta] class CallTypeRateLimiter(Protocol): @@ -847,15 +585,6 @@ def _parse_output_cap_value(raw_value: object) -> int | None: class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): - batch_rate_limiter_script: _AsyncLuaScript | None - batch_counter_read_script: _AsyncLuaScript | None - token_increment_script: _AsyncLuaScript | None - check_and_increment_by_n_script: _AsyncLuaScript | None - window_guarded_token_increment_script: _AsyncLuaScript | None - parallel_acquire_script: _AsyncLuaScript | None - parallel_release_script: _AsyncLuaScript | None - parallel_count_script: _AsyncLuaScript | None - def __init__( self, internal_usage_cache: InternalUsageCache, @@ -875,42 +604,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._ptu_team_ceiling_resolver = ptu_team_ceiling_resolver self._fail_closed_resolver = fail_closed_resolver self._force_hash_tag_grouping_resolver = force_hash_tag_grouping_resolver - if self.internal_usage_cache.dual_cache.redis_cache is not None: - self.batch_rate_limiter_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - BATCH_RATE_LIMITER_SCRIPT - ) - self.batch_counter_read_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - BATCH_COUNTER_READ_SCRIPT - ) - self.token_increment_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - TOKEN_INCREMENT_SCRIPT - ) - self.check_and_increment_by_n_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script(CHECK_AND_INCREMENT_BY_N_SCRIPT) - ) - self.window_guarded_token_increment_script = ( - self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - WINDOW_GUARDED_TOKEN_INCREMENT_SCRIPT - ) - ) - self.parallel_acquire_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - PARALLEL_ACQUIRE_SCRIPT - ) - self.parallel_release_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - PARALLEL_RELEASE_SCRIPT - ) - self.parallel_count_script = self.internal_usage_cache.dual_cache.redis_cache.async_register_script( - PARALLEL_COUNT_SCRIPT - ) - else: - self.batch_rate_limiter_script = None - self.batch_counter_read_script = None - self.token_increment_script = None - self.check_and_increment_by_n_script = None - self.window_guarded_token_increment_script = None - self.parallel_acquire_script = None - self.parallel_release_script = None - self.parallel_count_script = None + self._redis_cache = self.internal_usage_cache.dual_cache.redis_cache self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60)) @@ -1315,7 +1009,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) -> CacheCounterValues: """ Implement sliding window rate limiting logic using in-memory cache operations. - This follows the same logic as the Redis Lua script but uses async cache operations. + Apply the same sliding-window policy as the atomic cache operation. """ async with self._check_and_increment_lock: return await self._in_memory_cache_sliding_window(keys=keys, now_int=now_int, window_size=window_size) @@ -1483,21 +1177,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): crc: Final = binascii.crc_hqx(key.encode("utf-8"), 0) return crc % REDIS_CLUSTER_SLOTS - def _pipeline_scripts( - self, - source: str, - run: RegisteredScript, - calls: Sequence[tuple[Sequence[str], Sequence[int]]], - ) -> tuple[BatchResult[object] | None, ...]: - """Declare one Lua call per group on the request's Redis batch, so all groups share one round trip - with whatever else the request declared (the routing read). Returns ``None`` per call when no batch - is open, and the caller runs the script directly as before.""" - redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache - batch: Final = None if redis_cache is None else active_request_redis_batch(redis_cache) - if batch is None: - return (None,) * len(calls) - return tuple(batch.script(source, run, keys, args) for keys, args in calls) - def _group_keys_by_hash_tag(self, keys: list[str]) -> dict[str, list[str]]: """ Group keys by their Redis hash tag to ensure cluster compatibility. @@ -1546,18 +1225,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): local_only=True, ) - async def _read_counter_values_from_redis(self, keys: list[str]) -> CacheCounterValues: - read_script: Final = self.batch_counter_read_script - if read_script is None: - return [] + async def _read_counter_values_from_redis(self, keys: list[str]) -> CacheCounterValues | CacheFailure: + if self._redis_cache is None: + return () key_groups: Final = self._group_keys_by_hash_tag(keys) - group_values: Final[Sequence[CacheCounterValues]] = [ - await read_script(keys=group_keys, args=[]) for group_keys in key_groups.values() - ] + group_values: Final = tuple( + [await self._redis_cache.async_read_atomic_counts(keys=group_keys) for group_keys in key_groups.values()] + ) + failure: Final = next((value for value in group_values if isinstance(value, CacheFailure)), None) + if failure is not None: + return failure values_by_key: Final = dict( zip( itertools.chain.from_iterable(key_groups.values()), - itertools.chain.from_iterable(group_values), + itertools.chain.from_iterable(value for value in group_values if not isinstance(value, CacheFailure)), ) ) return [values_by_key.get(key) for key in keys] @@ -1567,30 +1248,28 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): keys: list[str], parent_otel_span: Span | None, ) -> CacheCounterValues | None: - if self.batch_counter_read_script is None: + if self._redis_cache is None: return await self._batch_get_counter_values(keys=keys, parent_otel_span=parent_otel_span, local_only=False) - try: - return await self._read_counter_values_from_redis(keys) - except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the local mirror unless fail-closed rejects - self._reject_if_rate_limit_unverifiable("batch_counter_read_script", e) - log_redis_failure( - verbose_proxy_logger, logging.WARNING, "batch_counter_read_script failed, using local mirror", e - ) - return await self._batch_get_counter_values(keys=keys, parent_otel_span=parent_otel_span, local_only=True) + result: Final = await self._read_counter_values_from_redis(keys) + if not isinstance(result, CacheFailure): + return result + self._reject_if_rate_limit_unverifiable("atomic counter read", result) + log_redis_failure( + verbose_proxy_logger, logging.WARNING, "Atomic counter read failed, using local mirror", result + ) + return await self._batch_get_counter_values(keys=keys, parent_otel_span=parent_otel_span, local_only=True) - def _reject_if_rate_limit_unverifiable(self, failed_operation: str, error: BaseException) -> None: + def _reject_if_rate_limit_unverifiable(self, failed_operation: str, error: CacheFailure) -> None: if not self._fail_closed_resolver(): return - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - f"fail_closed_rate_limit_enforcement: rejecting request, {failed_operation} could not verify the " - "counters against Redis", - error, + verbose_proxy_logger.warning( + "fail_closed_rate_limit_enforcement: rejecting request, %s could not verify the counters against Redis: %s", + failed_operation, + error.message, ) raise RateLimitUnverifiableError() - async def _execute_redis_batch_rate_limiter_script( + async def _increment_window_counters_by_group( self, keys_to_fetch: list[str], now_int: int, @@ -1605,62 +1284,64 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns: List of cache values """ - if self.batch_rate_limiter_script is None: - return [] + if self._redis_cache is None: + return () - key_groups: Final = list(self._group_keys_by_hash_tag(keys_to_fetch).items()) - all_cache_values: Final[list[CacheCounterValue | None]] = [] - args: Final = (now_int, self.window_size) - pipelined: Final = self._pipeline_scripts( - BATCH_RATE_LIMITER_SCRIPT, - self.batch_rate_limiter_script, - tuple((group_keys, args) for _tag, group_keys in key_groups), + key_groups: Final = tuple(self._group_keys_by_hash_tag(keys_to_fetch).items()) + batch: Final = active_request_redis_batch(self._redis_cache) + operations: Final = tuple( + self._redis_cache.async_increment_window_counters( + counters=tuple(WindowCounter(group_keys[i], group_keys[i + 1]) for i in range(0, len(group_keys), 2)), + now=now_int, + window_seconds=self.window_size, + batch=batch, + ) + for _tag, group_keys in key_groups ) - - for index, ((hash_tag, group_keys), group_result) in enumerate(zip(key_groups, pipelined)): - try: - group_cache_values: CacheCounterValues = ( - await self.batch_rate_limiter_script(keys=group_keys, args=args) - if group_result is None - else _as_counter_values(await group_result) - ) - all_cache_values.extend(group_cache_values) - except Exception as e: - if self._fail_closed_resolver(): - applied_keys = tuple(itertools.chain.from_iterable(keys for _tag, keys in key_groups[:index])) - await self._refund_counter_increments( - self._counter_refunds_from_batch_values(applied_keys, all_cache_values) + all_cache_values: Final[list[CacheCounterValue | None]] = [] + for index, ((hash_tag, group_keys), operation) in enumerate(zip(key_groups, operations)): + if not isinstance(result := await operation, CacheFailure): + all_cache_values.extend(_window_values(result)) + continue + if self._fail_closed_resolver(): + await self._refund_counter_increments( + self._counter_refunds_from_batch_values( + tuple(itertools.chain.from_iterable(keys for _tag, keys in key_groups[:index])), + all_cache_values, ) - await self._refund_later_pipelined_groups(key_groups[index + 1 :], pipelined[index + 1 :]) - self._reject_if_rate_limit_unverifiable("batch_rate_limiter_script", e) - log_redis_failure( - verbose_proxy_logger, logging.WARNING, f"Redis Lua script failed for hash tag {hash_tag}", e ) - # Fallback to in-memory cache for this group - group_cache_values = await self.in_memory_cache_sliding_window( + if batch is not None: + await self._refund_later_pipelined_groups(key_groups[index + 1 :], operations[index + 1 :]) + self._reject_if_rate_limit_unverifiable("window counter increment", result) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + f"Window counter increment failed for hash tag {hash_tag}", + result, + ) + all_cache_values.extend( + await self.in_memory_cache_sliding_window( keys=group_keys, now_int=now_int, window_size=self.window_size, ) - all_cache_values.extend(group_cache_values) + ) return _values_in_caller_order(keys_to_fetch, key_groups, all_cache_values) async def _refund_later_pipelined_groups( self, key_groups: Sequence[tuple[str, list[str]]], - pipelined: Sequence[BatchResult[object] | None], + pipelined: Sequence[CacheOperation[tuple[WindowCount, ...]]], ) -> None: """Groups declared on the request batch ran in the same round trip as the one that failed, so their increments landed even though the loop never read them.""" for (_tag, group_keys), group_result in zip(key_groups, pipelined): - if group_result is None: + if isinstance(result := await group_result, CacheFailure): continue - try: - group_values = _as_counter_values(await group_result) - except Exception: # noqa: BLE001 # a group that failed in Redis incremented nothing to refund - continue - await self._refund_counter_increments(self._counter_refunds_from_batch_values(group_keys, group_values)) + await self._refund_counter_increments( + self._counter_refunds_from_batch_values(group_keys, _window_values(result)) + ) @with_service_target("rate_limits") async def should_rate_limit( @@ -1704,7 +1385,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): current_time: Final = self._get_current_time() now: Final = current_time.timestamp() - now_int: Final = int(now) # Convert to integer for Redis Lua script + now_int: Final = int(now) keys_to_fetch, key_metadata, gauges = self._collect_windowed_keys_and_gauges( descriptors=descriptors, @@ -1740,10 +1421,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): cache_values = [ # rebind-ok: missing keys default to a zeroed window snapshot str(now_int) if key.endswith(":window") else 0 for key in keys_to_fetch ] - elif self.batch_rate_limiter_script is not None: + elif self._redis_cache is not None: # NORMAL MODE: Increment counters in Redis # Group keys by hash tag for Redis cluster compatibility - cache_values = await self._execute_redis_batch_rate_limiter_script( + cache_values = await self._increment_window_counters_by_group( keys_to_fetch=keys_to_fetch, now_int=now_int, ) @@ -1868,7 +1549,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ In-flight count from a cached gauge value: a dict of slot_id -> acquire timestamp when the in-memory registry is authoritative, or - the mirrored integer count from the last Redis script result. + the mirrored integer count from the last concurrency acquisition. """ if raw_value is None: return 0 @@ -1899,21 +1580,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): gauge_keys: Final = [gauge["counter_key"] for gauge in gauges] if read_only: - if self.parallel_count_script is not None: - try: - raw_counts: Final[list[CacheCounterValue]] = await self.parallel_count_script( - keys=gauge_keys, - args=[PARALLEL_REQUEST_SLOT_TTL_SECONDS for _ in gauges], - ) - counts = [max(0, int(value)) for value in raw_counts] - except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the local mirror unless fail-closed rejects - self._reject_if_rate_limit_unverifiable("parallel_count_script", e) - log_redis_failure( - verbose_proxy_logger, logging.WARNING, "parallel_count_script failed, using local mirror", e - ) - counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) - else: - counts = await self._read_local_gauge_counts(gauge_keys, parent_otel_span) + counts: Final = await self._read_gauge_counts(gauge_keys, parent_otel_span) statuses = [] overall_code = "OK" for gauge, in_flight in zip(gauges, counts): @@ -1931,32 +1598,35 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses=[self._gauge_status(gauge, in_flight, "OVER_LIMIT")], ) - if self.parallel_acquire_script is not None: - try: - raw: Final[list[CacheCounterValue]] = await self.parallel_acquire_script( - keys=gauge_keys, - args=[ - arg for gauge in gauges for arg in (gauge["limit"], PARALLEL_REQUEST_SLOT_TTL_SECONDS, slot_id) - ], - ) - except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to in-memory enforcement, never a 500 - self._reject_if_rate_limit_unverifiable("parallel_acquire_script", e) + if self._redis_cache is not None: + result: Final = await self._redis_cache.async_acquire_concurrency( + slots=tuple( + ConcurrencySlot( + key=gauge["counter_key"], + limit=gauge["limit"], + ttl_seconds=PARALLEL_REQUEST_SLOT_TTL_SECONDS, + slot_id=slot_id, + ) + for gauge in gauges + ), + ) + if isinstance(result, CacheFailure): + self._reject_if_rate_limit_unverifiable("concurrency acquire", result) log_redis_failure( verbose_proxy_logger, logging.WARNING, - "parallel_acquire_script failed, falling back to in-memory gauge", - e, + "Concurrency acquire failed, falling back to in-memory gauge", + result, ) async with self._check_and_increment_lock: return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) - if int(raw[0]) == 1: - gauge = gauges[int(raw[1]) - 1] + if isinstance(result, CapacityExceeded): return RateLimitResponse( overall_code="OVER_LIMIT", - statuses=[self._gauge_status(gauge, int(raw[2]), "OVER_LIMIT")], + statuses=[self._gauge_status(gauges[result.index], result.count, "OVER_LIMIT")], ) statuses = [] - for gauge, in_flight in zip(gauges, raw[1:]): + for gauge, in_flight in zip(gauges, result): await self.internal_usage_cache.async_set_cache( key=gauge["counter_key"], value=int(in_flight), @@ -1970,6 +1640,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async with self._check_and_increment_lock: return await self._acquire_parallel_slots_in_memory(gauges, slot_id, parent_otel_span) + async def _read_gauge_counts(self, gauge_keys: list[str], parent_otel_span: Span | None) -> Sequence[int]: + if self._redis_cache is not None: + result: Final = await self._redis_cache.async_count_concurrency( + counters=tuple(ConcurrencyRead(key, PARALLEL_REQUEST_SLOT_TTL_SECONDS) for key in gauge_keys) + ) + if not isinstance(result, CacheFailure): + return tuple(max(0, count) for count in result) + self._reject_if_rate_limit_unverifiable("concurrency count", result) + log_redis_failure( + verbose_proxy_logger, logging.WARNING, "Concurrency count failed, using local mirror", result + ) + return await self._read_local_gauge_counts(gauge_keys, parent_otel_span) + async def _read_local_gauge_counts( self, gauge_keys: list[str], @@ -1993,7 +1676,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): All-or-nothing in-memory slot-registry acquire. Caller holds the lock. A cached dict is the authoritative in-memory registry. A cached - integer is the count mirrored from the last successful Redis script + integer is the count mirrored from the last successful Redis operation call: when Redis fails over to this path, that mirror still counts the slots in flight on the Redis side, so it is carried forward as an integer counter (not discarded as an empty registry, which would @@ -2077,21 +1760,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): slot_id: Final = acquisition["slot_id"] if not counter_keys or not slot_id: return - if self.parallel_release_script is not None: - try: - raw: Final[list[CacheCounterValue]] = await self.parallel_release_script( - keys=counter_keys, - args=[slot_id for _ in counter_keys], - ) - await self._mirror_released_parallel_slots(counter_keys, raw, parent_otel_span) + if self._redis_cache is not None: + result: Final = await self._redis_cache.async_release_concurrency( + slots=tuple(ConcurrencyRelease(key, slot_id) for key in counter_keys) + ) + if not isinstance(result, CacheFailure): + await self._mirror_released_parallel_slots(counter_keys, result, parent_otel_span) return - except Exception as e: # noqa: BLE001 - any Redis/Lua failure degrades to the in-memory release, never a 500 - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - "parallel_release_script failed, falling back to in-memory release", - e, - ) + log_redis_failure( + verbose_proxy_logger, + logging.WARNING, + "Concurrency release failed, falling back to in-memory release", + result, + ) await self._release_parallel_request_slots_in_memory(counter_keys, slot_id, parent_otel_span) async def _defer_parallel_slot_release( @@ -2106,30 +1787,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): counter_keys: Final = acquisition["counter_keys"] slot_id: Final = acquisition["slot_id"] redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache - script: Final = self.parallel_release_script batch: Final = None if redis_cache is None else active_post_call_redis_batch(redis_cache) - if batch is None or script is None or not counter_keys or not slot_id: + if batch is None or redis_cache is None or not counter_keys or not slot_id: return False await self._release_parallel_request_slots_in_memory(counter_keys, slot_id, parent_otel_span) - async def settle(future: asyncio.Future[object]) -> None: - if future.cancelled() or future.exception() is not None: + async def settle(future: asyncio.Future[tuple[int, ...] | CacheFailure]) -> None: + if future.cancelled(): + verbose_proxy_logger.warning("Concurrency release cancelled, the slot stays released in memory only") + return + result: Final = future.result() + if isinstance(result, CacheFailure): log_redis_failure( verbose_proxy_logger, logging.WARNING, - "parallel_release_script failed, the slot stays released in memory only", - future.exception() if not future.cancelled() else asyncio.CancelledError(), + "Concurrency release failed, the slot stays released in memory only", + result, ) - batch.script(PARALLEL_RELEASE_SCRIPT, script, counter_keys, (slot_id,) * len(counter_keys)).on_settled(settle) + redis_cache.async_release_concurrency( + slots=tuple(ConcurrencyRelease(key, slot_id) for key in counter_keys), batch=batch + ).on_settled(settle) return True async def _mirror_released_parallel_slots( - self, counter_keys: list[str], remaining_by_key: Sequence[object], parent_otel_span: Span | None + self, counter_keys: list[str], remaining_by_key: Sequence[int], parent_otel_span: Span | None ) -> None: for counter_key, remaining in zip(counter_keys, remaining_by_key): - if not isinstance(remaining, (int, float, str, bytes)): - continue await self.internal_usage_cache.async_set_cache( key=counter_key, value=max(0, int(remaining)), @@ -2181,11 +1865,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): deployments. Cluster-safety: each descriptor's keys all share a `{key:value}` hash - tag, so the Redis Lua path issues one Lua call per descriptor — every - call's keys co-locate on a single Redis Cluster slot, avoiding - CROSSSLOT errors. Cross-descriptor atomicity is preserved via - refund-on-rollback: if descriptor i is OVER_LIMIT, descriptors 0..i-1 - get a direct INCRBY refund (refunds need no atomicity guarantee). + tag, so each typed capacity reservation uses one Redis Cluster slot. + On rejection, earlier reservations are refunded only in their original + windows. Args: descriptors: rate-limit descriptors to check @@ -2200,35 +1882,22 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if len(descriptors) != len(increments): raise ValueError("atomic_check_and_increment_by_n: descriptors and increments must have the same length") - # Build per-descriptor (keys, args, meta) groups. All keys within a - # group share the descriptor's {key:value} hash tag, so a single Lua - # call per group never triggers CROSSSLOT on Redis Cluster. - descriptor_groups: Final[list[DescriptorAtomicGroup]] = [] - for descriptor, increment_amounts in zip(descriptors, increments): - keys, args, meta = self._build_descriptor_atomic_payload( - descriptor=descriptor, - increment_amounts=increment_amounts, - ) - if keys: - descriptor_groups.append((keys, args, meta)) + groups: Final = tuple( + self._build_descriptor_atomic_payload(descriptor, increment_amounts) + for descriptor, increment_amounts in zip(descriptors, increments) + ) + descriptor_groups: Final = tuple(group for group in groups if group) if not descriptor_groups: return RateLimitResponse(overall_code="OK", statuses=[]) - # Multi-process atomicity via Redis Lua, per descriptor for slot - # co-location. Single-process atomicity falls back to the - # asyncio.Lock + in-memory sliding window below — there are no - # cluster slot concerns locally, so we keep the batched 2-phase - # critical section for true cross-descriptor atomicity. - if self.check_and_increment_by_n_script is not None: - return await self._atomic_lua_per_descriptor( + if self._redis_cache is not None: + return await self._reserve_window_capacity_by_descriptor( descriptor_groups=descriptor_groups, parent_otel_span=parent_otel_span, ) - flat_meta: Final[list[AtomicCounterMeta]] = [ - m for _keys, _args, group_meta in descriptor_groups for m in group_meta - ] + flat_meta: Final = tuple(itertools.chain.from_iterable(descriptor_groups)) async with self._check_and_increment_lock: return await self._atomic_check_and_increment_in_memory( per_counter_meta=flat_meta, @@ -2241,8 +1910,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): increment_amounts: dict[Literal["requests", "tokens"], int], ) -> DescriptorAtomicGroup: """ - Build (KEYS, ARGV, per-counter meta) for a single descriptor's Lua - call. All keys returned share the descriptor's {key:value} hash tag. + Build counter policy for one descriptor with a common cluster hash tag. """ descriptor_key: Final = descriptor["key"] descriptor_value: Final = descriptor["value"] @@ -2252,8 +1920,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_size: Final = rate_limit.get("window_size") or self.window_size window_key: Final = f"{{{descriptor_key}:{descriptor_value}}}:window" - keys: Final[list[str]] = [] - args: Final[list[int]] = [] meta: Final[list[AtomicCounterMeta]] = [] rate_limit_types: Final[tuple[Literal["requests", "tokens"], ...]] = ("requests", "tokens") @@ -2273,10 +1939,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # custom-TTL descriptor doesn't reintroduce a silent expiry bug. ttl_seconds = int(window_size) window_size_seconds = int(window_size) - keys.extend([window_key, counter_key]) - # 4-tuple matches the Lua ARGV layout: - # [limit, increment, ttl_seconds, window_size_seconds]. - args.extend([int(limit_value), inc_amount, ttl_seconds, window_size_seconds]) meta.append( { "descriptor_key": descriptor_key, @@ -2290,18 +1952,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "window_size": window_size_seconds, } ) - return keys, args, meta + return meta - async def _atomic_lua_per_descriptor( + async def _reserve_window_capacity_by_descriptor( self, - descriptor_groups: list[DescriptorAtomicGroup], + descriptor_groups: Sequence[DescriptorAtomicGroup], parent_otel_span: Span | None = None, ) -> RateLimitResponse: """ - Run Lua check-and-increment one descriptor at a time so each call's - keys co-locate on a single Redis Cluster slot. On OVER_LIMIT for - descriptor i, refund descriptors 0..i-1's increments. On Lua failure - mid-loop, refund applied increments and fall back to in-memory. + Reserve each descriptor's capacity on one cluster slot. Refund earlier + reservations on rejection or failure before applying fallback policy. """ if not descriptor_groups: return RateLimitResponse( @@ -2311,52 +1971,36 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): applied: Final[list[tuple[CounterRefund, ...]]] = [] statuses: Final[list[RateLimitStatus]] = [] reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop - raw: list[CacheCounterValue] - - pipelined: Final = self._pipeline_scripts( - CHECK_AND_INCREMENT_BY_N_SCRIPT, - self.check_and_increment_by_n_script, # pyright: ignore[reportArgumentType] # sole caller guards it is not None - tuple((keys, args) for keys, args, _meta in descriptor_groups), + redis_cache: Final = self._redis_cache + if redis_cache is None: + return await self._fallback_atomic_groups(descriptor_groups, parent_otel_span) + batch: Final = active_request_redis_batch(redis_cache) + operations: Final = tuple( + redis_cache.async_reserve_window_capacity(reservations=self._window_reservations(meta), batch=batch) + for meta in descriptor_groups ) - batched: Final = tuple(result for result in pipelined if result is not None) - if len(batched) == len(descriptor_groups): - return await self._settle_pipelined_descriptor_groups(descriptor_groups, batched, parent_otel_span) + if batch is not None: + return await self._settle_pipelined_descriptor_groups(descriptor_groups, operations, parent_otel_span) - for keys, args, meta in descriptor_groups: - try: - raw = await self.check_and_increment_by_n_script( # pyright: ignore[reportOptionalCall] # sole caller guards it is not None - keys=keys, - args=args, - ) - except Exception as e: - # Lua failure (timeout, OOM, network partition) leaves Redis - # state ambiguous. Refund any prior groups so Redis returns - # to its pre-call state, then fall back to in-memory for the - # whole call (counters there are independent of Redis). + for meta, operation in zip(descriptor_groups, operations): + if isinstance(result := await operation, CacheFailure): await self._refund_applied_descriptor_groups(applied) - self._reject_if_rate_limit_unverifiable("check_and_increment_by_n_script", e) + self._reject_if_rate_limit_unverifiable("window capacity reservation", result) log_redis_failure( verbose_proxy_logger, logging.ERROR, - f"atomic_check_and_increment_by_n: Redis Lua execution failed ({type(e).__name__}). Refunded " - f"{len(applied)} prior descriptors and falling back to in-memory enforcement, counters will " - f"diverge from Redis until window expires (window_size={self.window_size}s)", - e, + f"Window capacity reservation failed, refunded {len(applied)} prior descriptors and falling back to memory", + result, ) - flat_meta: list[AtomicCounterMeta] = [m for _k, _a, group_meta in descriptor_groups for m in group_meta] - async with self._check_and_increment_lock: - return await self._atomic_check_and_increment_in_memory( - per_counter_meta=flat_meta, - parent_otel_span=parent_otel_span, - ) + return await self._fallback_atomic_groups(descriptor_groups, parent_otel_span) - response = self._build_atomic_response(raw, meta) - if response["overall_code"] == "OVER_LIMIT": + if (response := self._build_atomic_response(result, meta))["overall_code"] == "OVER_LIMIT": await self._refund_applied_descriptor_groups(applied) return response if len(descriptor_groups) == 1: return response - applied.append(self._counter_refunds_from_atomic_response(raw, meta)) + if not isinstance(result, CapacityExceeded): + applied.append(self._counter_refunds_from_atomic_response(result, meta)) statuses.extend(response["statuses"]) reservation_windows.update(response.get("reservation_windows", frozenset())) @@ -2368,8 +2012,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def _settle_pipelined_descriptor_groups( self, - descriptor_groups: list[DescriptorAtomicGroup], - results: Sequence[BatchResult[object]], + descriptor_groups: Sequence[DescriptorAtomicGroup], + results: Sequence[CacheOperation[tuple[WindowCount, ...] | CapacityExceeded]], parent_otel_span: Span | None, ) -> RateLimitResponse: """Every group's Lua call left in one pipeline, so each group has already checked and incremented on @@ -2377,48 +2021,38 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): incremented, after it as well as before it, where the one-at-a-time loop only unwinds the groups it ran. A Redis denial stands even when another group failed: the in-memory fallback only replaces a verdict Redis never gave.""" - replies: Final = await asyncio.gather(*results, return_exceptions=True) + replies: Final = await asyncio.gather(*results) responses: Final = tuple( - self._pipelined_group_response(reply, meta) - for reply, (_keys, _args, meta) in zip(replies, descriptor_groups) + self._pipelined_group_response(reply, meta) for reply, meta in zip(replies, descriptor_groups) ) applied: Final[list[tuple[CounterRefund, ...]]] = [] # mutable-ok: filled by the group loop statuses: Final[list[RateLimitStatus]] = [] # mutable-ok: filled by the group loop reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop - for reply, response, (_keys, _args, meta) in zip(replies, responses, descriptor_groups): - if isinstance(response, BaseException) or response["overall_code"] != "OK": + for reply, response, meta in zip(replies, responses, descriptor_groups): + if isinstance(reply, (CacheFailure, CapacityExceeded)) or isinstance(response, CacheFailure): continue - applied.append(self._counter_refunds_from_atomic_response(_as_counter_values(reply), meta)) + applied.append(self._counter_refunds_from_atomic_response(reply, meta)) statuses.extend(response["statuses"]) reservation_windows.update(response.get("reservation_windows", frozenset())) over_limit: Final = next( - (r for r in responses if not isinstance(r, BaseException) and r["overall_code"] == "OVER_LIMIT"), None + (r for r in responses if not isinstance(r, CacheFailure) and r["overall_code"] == "OVER_LIMIT"), None ) if over_limit is not None: await self._refund_applied_descriptor_groups(applied) return over_limit - failure: Final = next((r for r in responses if isinstance(r, BaseException)), None) + failure: Final = next((r for r in responses if isinstance(r, CacheFailure)), None) if failure is not None: await self._refund_applied_descriptor_groups(applied) - self._reject_if_rate_limit_unverifiable("check_and_increment_by_n_script", failure) + self._reject_if_rate_limit_unverifiable("window capacity reservation", failure) log_redis_failure( verbose_proxy_logger, logging.ERROR, - f"atomic_check_and_increment_by_n: Redis Lua execution failed ({type(failure).__name__}). Refunding " - f"{len(applied)} pipelined descriptors and falling back to in-memory enforcement, counters will " - f"diverge from Redis until window expires (window_size={self.window_size}s)", + f"Window capacity reservation failed, refunded {len(applied)} batched descriptors and falling back to memory", failure, ) - flat_meta: Final = tuple( - itertools.chain.from_iterable(group_meta for _k, _a, group_meta in descriptor_groups) - ) - async with self._check_and_increment_lock: - return await self._atomic_check_and_increment_in_memory( - per_counter_meta=flat_meta, - parent_otel_span=parent_otel_span, - ) - if len(responses) == 1 and not isinstance(responses[0], BaseException): + return await self._fallback_atomic_groups(descriptor_groups, parent_otel_span) + if len(responses) == 1 and not isinstance(responses[0], CacheFailure): return responses[0] return RateLimitResponse( overall_code="OK", @@ -2427,21 +2061,39 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) def _pipelined_group_response( - self, reply: object, per_counter_meta: list[AtomicCounterMeta] - ) -> RateLimitResponse | BaseException: - if isinstance(reply, BaseException): - return reply - try: - return self._build_atomic_response(_as_counter_values(reply), per_counter_meta) - except Exception as e: # noqa: BLE001 # a reply this group cannot read is that group's Lua failure - return e + self, + reply: tuple[WindowCount, ...] | CapacityExceeded | CacheFailure, + per_counter_meta: list[AtomicCounterMeta], + ) -> RateLimitResponse | CacheFailure: + return reply if isinstance(reply, CacheFailure) else self._build_atomic_response(reply, per_counter_meta) + + @staticmethod + def _window_reservations(meta: Sequence[AtomicCounterMeta]) -> tuple[WindowReservation, ...]: + return tuple( + WindowReservation( + window_key=item["window_key"], + counter_key=item["counter_key"], + limit=item["current_limit"], + increment=item["increment"], + ttl_seconds=item["ttl"], + window_seconds=item["window_size"], + ) + for item in meta + ) + + async def _fallback_atomic_groups( + self, descriptor_groups: Sequence[DescriptorAtomicGroup], parent_otel_span: Span | None + ) -> RateLimitResponse: + flat_meta: Final = tuple(itertools.chain.from_iterable(descriptor_groups)) + async with self._check_and_increment_lock: + return await self._atomic_check_and_increment_in_memory(flat_meta, parent_otel_span) async def _refund_applied_descriptor_groups( self, applied: Sequence[Sequence[CounterRefund]], ) -> None: """ - Decrement counters for descriptor groups already applied via Lua. + Refund descriptor groups already reserved through the cache. Best-effort: refund failures are logged but not raised — the original OVER_LIMIT / fallback decision is what matters to the caller. """ @@ -2449,17 +2101,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): @staticmethod def _counter_refunds_from_atomic_response( - raw: Sequence[CacheCounterValue], + counts: Sequence[WindowCount], per_counter_meta: Sequence[AtomicCounterMeta], ) -> tuple[CounterRefund, ...]: return tuple( CounterRefund( window_key=meta["window_key"], counter_key=meta["counter_key"], - window_start=str(int(raw[2 + index * 2])), + window_start=count.window_start, increment=meta["increment"], ) - for index, meta in enumerate(per_counter_meta) + for count, meta in zip(counts, per_counter_meta) ) @staticmethod @@ -2480,48 +2132,36 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) async def _refund_counter_increments(self, refunds: Sequence[CounterRefund]) -> None: - if self.window_guarded_token_increment_script is None: + if self._redis_cache is None: return for refund in refunds: - try: - await self.window_guarded_token_increment_script( - keys=[refund.window_key, refund.counter_key], - args=[refund.window_start, -refund.increment, 0], - ) - except Exception as e: # noqa: BLE001 # best-effort rollback, the rejection already decided the request + if isinstance( + result := await self._redis_cache.async_increment_guarded_counters( + increments=( + GuardedIncrement( + refund.window_key, + refund.counter_key, + refund.window_start, + -refund.increment, + 0, + ), + ) + ), + CacheFailure, + ): log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - f"Failed to refund {refund.counter_key} on rollback", - e, + verbose_proxy_logger, logging.WARNING, f"Failed to refund {refund.counter_key} on rollback", result ) def _build_atomic_response( self, - raw: list[CacheCounterValue], + result: tuple[WindowCount, ...] | CapacityExceeded, per_counter_meta: list[AtomicCounterMeta], ) -> RateLimitResponse: - """Convert Lua script return value to RateLimitResponse. - - Indexing invariant: `per_counter_meta` and `KEYS` are parallel-indexed - at the COUNTER level, not the descriptor level. A descriptor with both - RPM and TPM limits emits two `(window_key, counter_key)` pairs and - two meta entries — one per counter. The Lua script's loop variable - `i` therefore enumerates counters, and the over-limit return tuple - `{1, i, ...}` carries a counter index that maps directly to - `per_counter_meta[i - 1]`. Keep these arrays parallel at the counter - level when modifying this code. - """ - if not raw: - return RateLimitResponse(overall_code="OK", statuses=[]) - - status_code: Final = int(raw[0]) - if status_code == 1: - # Over limit: { 1, counter_index (1-based), current_counter, limit } - descriptor_index: Final = int(raw[1]) - 1 - current_counter: Final = int(raw[2]) - limit: Final = int(raw[3]) - meta = per_counter_meta[descriptor_index] + if isinstance(result, CapacityExceeded): + current_counter: Final = result.count + limit: Final = result.limit + meta = per_counter_meta[result.index] return RateLimitResponse( overall_code="OVER_LIMIT", statuses=[ @@ -2537,13 +2177,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) statuses: Final[list[RateLimitStatus]] = [] - for index, meta in enumerate(per_counter_meta): - new_counter = raw[1 + index * 2] + for count, meta in zip(result, per_counter_meta): statuses.append( RateLimitStatus( code="OK", current_limit=meta["current_limit"], - limit_remaining=max(0, meta["current_limit"] - int(new_counter)), + limit_remaining=max(0, meta["current_limit"] - count.count), rate_limit_type=meta["rate_limit_type"], descriptor_key=meta["descriptor_key"], descriptor_value=meta["descriptor_value"], @@ -2555,10 +2194,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): reservation_windows=frozenset( ( meta["counter_key"], - str(int(raw[2 + index * 2])), + count.window_start, "redis", ) - for index, meta in enumerate(per_counter_meta) + for count, meta in zip(result, per_counter_meta) ), ) @@ -4582,75 +4221,56 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return response_obj return None - async def _execute_token_increment_script( + async def _execute_token_increments( self, - pipeline_operations: list["RedisPipelineIncrementOperation"], - ) -> None: - """ - Execute token increment script grouped by hash tag for cluster compatibility. - """ - if self.token_increment_script is None: - return - - # Group operations by hash tag for Redis cluster compatibility - operation_keys: Final = [op["key"] for op in pipeline_operations] - key_groups: Final = self._group_keys_by_hash_tag(operation_keys) - - for _hash_tag, group_keys in key_groups.items(): - # Get operations for this hash tag group - group_operations = [op for op in pipeline_operations if op["key"] in group_keys] - - keys = [] - args = [] - - for op in group_operations: - # Convert None TTL to 0 for Lua script - ttl_value = op["ttl"] if op["ttl"] is not None else 0 - - verbose_proxy_logger.debug( - "Executing TTL-preserving increment for key=%s, increment=%s, ttl=%s", - op["key"], - op["increment_value"], - ttl_value, - ) - keys.append(op["key"]) - args.extend([op["increment_value"], ttl_value]) - - if self._defer_token_increment_script(keys, args, group_operations): + pipeline_operations: list[RedisPipelineIncrementOperation], + ) -> CacheFailure | None: + if self._redis_cache is None: + return None + key_groups: Final = self._group_keys_by_hash_tag([op["key"] for op in pipeline_operations]) + grouped: Final = tuple( + [op for op in pipeline_operations if op["key"] in group_keys] for group_keys in key_groups.values() + ) + for group_operations in grouped: + if self._defer_token_increments(group_operations): continue - await self.token_increment_script( - keys=keys, - args=args, - ) + if isinstance( + result := await self._redis_cache.async_increment_counters( + increments=self._counter_increments(group_operations) + ), + CacheFailure, + ): + return result + return None - def _defer_token_increment_script( - self, - keys: list[str], - args: list[int], - group_operations: list["RedisPipelineIncrementOperation"], - ) -> bool: - """Declared into the request's post-call pipeline instead of its own EVALSHA round trip; a failed - script falls back to the plain increment pipeline for its own group, as the direct path does.""" - redis_cache: Final = self.internal_usage_cache.dual_cache.redis_cache - script: Final = self.token_increment_script + @staticmethod + def _counter_increments(operations: Sequence[RedisPipelineIncrementOperation]) -> tuple[CounterIncrement, ...]: + return tuple(CounterIncrement(op["key"], op["increment_value"], op["ttl"] or 0) for op in operations) + + def _defer_token_increments(self, group_operations: list[RedisPipelineIncrementOperation]) -> bool: + redis_cache: Final = self._redis_cache batch: Final = None if redis_cache is None else active_post_call_redis_batch(redis_cache) - if batch is None or script is None: + if batch is None or redis_cache is None: return False - async def fall_back(future: asyncio.Future[object]) -> None: - if future.cancelled() or future.exception() is None: + async def fall_back(future: asyncio.Future[tuple[float, ...] | CacheFailure]) -> None: + if future.cancelled(): + return + result: Final = future.result() + if not isinstance(result, CacheFailure): return log_redis_failure( verbose_proxy_logger, logging.WARNING, "TTL preservation failed, falling back to regular pipeline", - future.exception(), - ) - await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=group_operations, + result, ) + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(increment_list=group_operations) - batch.script(TOKEN_INCREMENT_SCRIPT, script, keys, args).on_settled(fall_back) + redis_cache.async_increment_counters( + increments=self._counter_increments(group_operations), + batch=batch, + ).on_settled(fall_back) return True @with_service_target("rate_limits") @@ -4660,37 +4280,33 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span: Span | None = None, ) -> None: """ - Increment token counters using Lua script to preserve existing TTL. + Increment token counters through the cache while preserving existing TTL. This prevents TTL reset on every token increment. """ if not pipeline_operations: return - # Check if script is available - if self.token_increment_script is None: - verbose_proxy_logger.debug("TTL preservation script not available, using regular pipeline") + if self._redis_cache is None: + verbose_proxy_logger.debug("TTL preservation unavailable, using regular pipeline") await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( increment_list=pipeline_operations, litellm_parent_otel_span=parent_otel_span, ) return - try: - await self._execute_token_increment_script(pipeline_operations) - + failure: Final = await self._execute_token_increments(pipeline_operations) + if failure is None: verbose_proxy_logger.debug( "Successfully executed TTL-preserving increment for %s keys", len(pipeline_operations) ) - - except Exception as e: - log_redis_failure( - verbose_proxy_logger, logging.WARNING, "TTL preservation failed, falling back to regular pipeline", e - ) - # Fallback to regular pipeline on error - await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - litellm_parent_otel_span=parent_otel_span, - ) + return + log_redis_failure( + verbose_proxy_logger, logging.WARNING, "TTL preservation failed, falling back to regular pipeline", failure + ) + await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline( + increment_list=pipeline_operations, + litellm_parent_otel_span=parent_otel_span, + ) async def _apply_local_window_guarded_token_increments( self, @@ -4736,27 +4352,27 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): expected_window_start = operation.get("expected_window_start") if window_key is None or expected_window_start is None: continue - if self.window_guarded_token_increment_script is not None: - try: - await self.window_guarded_token_increment_script( - keys=[ - window_key, - operation["key"], - ], - args=[ - expected_window_start, - operation["increment_value"], - operation["ttl"] or 0, - ], - ) + if self._redis_cache is not None: + if not isinstance( + result := await self._redis_cache.async_increment_guarded_counters( + increments=( + GuardedIncrement( + window_key, + operation["key"], + expected_window_start, + operation["increment_value"], + operation["ttl"] or 0, + ), + ) + ), + CacheFailure, + ): continue - except Exception as e: # noqa: BLE001 # Redis failures use the plain increment fallback - log_redis_failure( - verbose_proxy_logger, - logging.WARNING, - f"Window-guarded token adjustment failed for {operation['key']}", - e, - ) + verbose_proxy_logger.warning( + "Window-guarded token adjustment failed for %s: %s", + operation["key"], + result.message, + ) if operation["increment_value"] > 0: await self.internal_usage_cache.async_increment_cache( key=operation["key"], diff --git a/litellm/types/caching.py b/litellm/types/caching.py index ced2065a674..7e5aa1b35d4 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -1,7 +1,10 @@ -from collections.abc import Sequence +import asyncio +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field from enum import Enum -from typing import Any, Final, Literal, Optional, Union +from typing import Any, Final, Literal, Protocol, TypeVar +from pydantic import ValidationError from typing_extensions import ReadOnly, TypedDict from litellm.types.llms.base import LiteLLMBaseModel @@ -41,6 +44,140 @@ CachingSupportedCallTypes = Literal[ "aanthropic_messages", ] +KeyPresence = Literal["present", "missing", "unavailable"] + + +@dataclass(frozen=True, slots=True) +class CacheFailure: + kind: Literal["unavailable", "invalid_response"] + message: str + reason: Literal["circuit_open", "timeout", "other"] = "other" + validation_error: ValidationError | None = field(default=None, compare=False, repr=False) + + +CacheValue: Final = TypeVar("CacheValue") + + +class CacheOperation(Awaitable[CacheValue | CacheFailure], Protocol[CacheValue]): + @property + def done(self) -> bool: ... + + def on_settled( + self, hook: Callable[[asyncio.Future[CacheValue | CacheFailure]], Awaitable[None] | None] + ) -> None: ... + + +@dataclass(frozen=True, slots=True) +class WindowCounter: + window_key: str + counter_key: str + + +@dataclass(frozen=True, slots=True) +class WindowCount: + window_start: str + count: int + + +@dataclass(frozen=True, slots=True) +class WindowReservation: + window_key: str + counter_key: str + limit: int + increment: int + ttl_seconds: int + window_seconds: int + + +@dataclass(frozen=True, slots=True) +class CapacityExceeded: + index: int + count: int + limit: int + + +@dataclass(frozen=True, slots=True) +class ConcurrencySlot: + key: str + limit: int + ttl_seconds: int + slot_id: str + + +@dataclass(frozen=True, slots=True) +class ConcurrencyRelease: + key: str + slot_id: str + + +@dataclass(frozen=True, slots=True) +class ConcurrencyRead: + key: str + ttl_seconds: int + + +@dataclass(frozen=True, slots=True) +class CounterIncrement: + key: str + increment: float + ttl_seconds: int + + +@dataclass(frozen=True, slots=True) +class GuardedIncrement: + window_key: str + counter_key: str + expected_window_start: str + increment: int + ttl_seconds: int + + +@dataclass(frozen=True, slots=True) +class CounterUpdate: + applied: bool + count: int + + +@dataclass(frozen=True, slots=True) +class LimitedIncrement: + accepted: bool + count: int + + +@dataclass(frozen=True, slots=True) +class ThrottleKeys: + pair_counter: str + pair_block: str + source_counter: str + source_block: str + + +class CacheSubscription(Protocol): + async def receive(self, *, timeout_seconds: float) -> str | bytes | None | CacheFailure: ... + + async def close(self) -> None | CacheFailure: ... + + +@dataclass(frozen=True, slots=True) +class CachePoolInfo: + available_connections: int | None + in_use_connections: int | None + max_connections: int | None + connection_class: str | None + + +@dataclass(frozen=True, slots=True) +class CachePublish: + channel: str + message: str + + +@dataclass(frozen=True, slots=True) +class CacheQueueUpdate: + key: str + values: tuple[str, ...] + + DEFAULT_CACHING_SUPPORTED_CALL_TYPES: tuple[CachingSupportedCallTypes, ...] = ( "completion", "acompletion", diff --git a/tests/integration/sdk/test_redis_typed_capabilities.py b/tests/integration/sdk/test_redis_typed_capabilities.py new file mode 100644 index 00000000000..e6e7a0c6910 --- /dev/null +++ b/tests/integration/sdk/test_redis_typed_capabilities.py @@ -0,0 +1,256 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +from redis import Redis + +from litellm.caching.redis_batch import RedisBatch, active_post_call_redis_batch, request_redis_batch_scope +from litellm.caching.redis_cache import RedisCache +from litellm.types.caching import ( + CacheFailure, + CachePublish, + CacheQueueUpdate, + CapacityExceeded, + ConcurrencyRead, + ConcurrencyRelease, + ConcurrencySlot, + CounterIncrement, + CounterUpdate, + GuardedIncrement, + LimitedIncrement, + ThrottleKeys, + WindowCount, + WindowCounter, + WindowReservation, +) +from tests.integration._support.redis_process import owned_redis + + +@pytest.fixture +def backend(tmp_path: Path) -> Iterator[tuple[RedisCache, Redis[bytes]]]: + with owned_redis(tmp_path) as server: + cache: Final = RedisCache(host=server.host, port=server.port, namespace="typed") + with Redis(host=server.host, port=server.port) as raw: + yield cache, raw + + +@pytest.mark.asyncio +async def test_window_reservations_are_atomic_and_guarded_updates_keep_window_identity( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + reservations: Final = ( + WindowReservation("{scope}:window", "{scope}:requests", 7, 1, 60, 60), + WindowReservation("{scope}:window", "{scope}:tokens", 70, 10, 60, 60), + ) + results: Final = await asyncio.gather( + *(cache.async_reserve_window_capacity(reservations=reservations) for _ in range(20)) + ) + assert sum(isinstance(result, tuple) for result in results) == 7 + assert sum(result == CapacityExceeded(0, 7, 7) for result in results) == 13 + assert raw.mget("typed:{scope}:requests", "typed:{scope}:tokens") == [b"7", b"70"] + assert raw.exists("{scope}:requests") == 0 + winner: Final = next(result for result in results if isinstance(result, tuple)) + assert isinstance(winner[0].window_start, str) + assert await cache.async_increment_guarded_counters( + increments=(GuardedIncrement("{scope}:window", "{scope}:tokens", winner[0].window_start, -5, 0),) + ) == (CounterUpdate(True, 65),) + assert await cache.async_increment_guarded_counters( + increments=(GuardedIncrement("{scope}:window", "{scope}:tokens", "stale", 100, 0),) + ) == (CounterUpdate(False, 65),) + assert 0 < raw.ttl("typed:{scope}:requests") <= 60 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("now", (100, 100.5)) +async def test_legacy_window_increment_resets_siblings_and_preserves_string_starts( + backend: tuple[RedisCache, Redis[bytes]], now: float, +) -> None: + cache, raw = backend + raw.set("typed:lane:tokens", 99) + counter: Final = WindowCounter("lane:window", "lane:requests") + assert await cache.async_increment_window_counters(counters=(counter,), now=now, window_seconds=60) == ( + WindowCount(str(now), 1), + ) + assert raw.get("typed:lane:tokens") is None + assert await cache.async_increment_window_counters(counters=(counter,), now=now + 1, window_seconds=60) == ( + WindowCount(str(now), 2), + ) + assert await cache.async_read_atomic_counts(keys=("lane:requests", "missing")) == (2, None) + + +@pytest.mark.asyncio +async def test_concurrency_prunes_stale_slots_and_rejects_without_partial_acquisition( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + raw.zadd("typed:gauge", {"stale": 0}) + assert await cache.async_acquire_concurrency(slots=(ConcurrencySlot("gauge", 1, 60, "owner"),)) == (1,) + assert await cache.async_acquire_concurrency( + slots=( + ConcurrencySlot("other", 2, 60, "candidate"), + ConcurrencySlot("gauge", 1, 60, "candidate"), + ) + ) == CapacityExceeded(1, 1, 1) + assert raw.exists("typed:other") == 0 + assert await cache.async_count_concurrency(counters=(ConcurrencyRead("gauge", 60),)) == (1,) + assert await cache.async_release_concurrency(slots=(ConcurrencyRelease("gauge", "not-owner"),)) == (1,) + assert await cache.async_release_concurrency(slots=(ConcurrencyRelease("gauge", "owner"),)) == (0,) + + +@pytest.mark.asyncio +async def test_counter_and_record_capabilities_preserve_ttls_and_tombstones( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + assert await cache.async_increment_iterations(key="iterations", ttl_seconds=120) == 1 + raw.expire("typed:iterations", 30) + assert await cache.async_increment_iterations(key="iterations", ttl_seconds=120) == 2 + assert 0 < raw.ttl("typed:iterations") <= 30 + assert await cache.async_increment_spend(key="spend", amount=0.25, ttl_seconds=120) == 0.25 + raw.expire("typed:spend", 30) + assert await cache.async_increment_spend(key="spend", amount=0.5, ttl_seconds=120) == 0.75 + assert 0 < raw.ttl("typed:spend") <= 30 + assert await cache.async_increment_counters(increments=(CounterIncrement("tokens", 1.5, 120),)) == (1.5,) + raw.expire("typed:tokens", 30) + assert await cache.async_increment_counters(increments=(CounterIncrement("tokens", 2.5, 120),)) == (4.0,) + assert 0 < raw.ttl("typed:tokens") <= 30 + assert await cache.async_reserve_counter(key="enqueued", amount=5, limit=8, ttl_seconds=60) == LimitedIncrement( + True, 5 + ) + assert await cache.async_reserve_counter(key="enqueued", amount=4, limit=8, ttl_seconds=60) == LimitedIncrement( + False, 5 + ) + assert await cache.async_refund_counter(key="enqueued", amount=7) is None + assert raw.get("typed:enqueued") is None + assert await cache.async_save_record(key="record", value='{"tokens":5}', ttl_seconds=90) is None + assert await cache.async_pop_record(key="record", tombstone_ttl_seconds=25) == '{"tokens":5}' + assert await cache.async_pop_record(key="record", tombstone_ttl_seconds=25) == "" + assert await cache.async_pop_record(key="missing", tombstone_ttl_seconds=25) is None + assert 0 < raw.ttl("typed:record") <= 25 + + +@pytest.mark.asyncio +async def test_affinity_keeps_eligible_winners_and_refreshes_legacy_model_pins( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + assert await cache.async_claim_affinity(key="pin", value={"model_id": "a"}, ttl_seconds=60) == {"model_id": "a"} + assert await cache.async_claim_affinity(key="pin", value={"model_id": "b"}, ttl_seconds=60) == {"model_id": "a"} + assert await cache.async_claim_affinity( + key="pin", + value={"model_id": "b"}, + ttl_seconds=60, + eligible_values=({"model_id": "b"},), + ) == {"model_id": "b"} + raw.set("typed:legacy", "old") + assert ( + await cache.async_claim_affinity( + key="legacy", + value={"model_id": "new"}, + ttl_seconds=60, + eligible_values=({"model_id": "old"},), + preserve_legacy_model_ids=True, + ) + == "old" + ) + assert 0 < raw.ttl("typed:legacy") <= 60 + + +@pytest.mark.asyncio +async def test_throttle_source_count_stops_after_pair_is_blocked( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + keys: Final = ThrottleKeys("pair", "pair:block", "source", "source:block") + assert await cache.async_read_throttle_ttls(keys=keys) == (-2, -2) + assert await cache.async_record_throttle_failure( + keys=keys, + pair_limit=1, + source_limit=5, + window_seconds=60, + block_seconds=90, + ) == (0, 0) + assert await cache.async_record_throttle_failure( + keys=keys, + pair_limit=1, + source_limit=5, + window_seconds=60, + block_seconds=90, + ) == (90, 0) + assert raw.get("typed:source") == b"1" + pair_ttl, source_ttl = await cache.async_read_throttle_ttls(keys=keys) + assert 0 < pair_ttl <= 90 + assert source_ttl == -2 + + +@pytest.mark.asyncio +async def test_refresh_lease_only_renews_or_deletes_for_its_owner( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + assert await cache.async_acquire_lease(key="lock", token="owner", ttl_seconds=60) == "acquired" + assert await cache.async_acquire_lease(key="lock", token="other", ttl_seconds=60) == "held" + assert await cache.async_extend_lease(key="lock", token="other", ttl_seconds=120) is False + assert await cache.async_release_lease(key="lock", token="other") is False + assert await cache.async_lease_held(key="lock") is True + assert await cache.async_extend_lease(key="lock", token="owner", ttl_seconds=120) is True + assert 60000 < raw.pttl("typed:lock") <= 120000 + assert await cache.async_release_lease(key="lock", token="owner") is True + assert await cache.async_lease_held(key="lock") is False + + +@pytest.mark.asyncio +async def test_typed_batches_retry_noscript_isolate_item_failures_and_settle_post_call_hooks( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + assert await cache.async_increment_iterations(key="warmup", ttl_seconds=60) == 1 + raw.script_flush() + raw.rpush("typed:wrong-type", "value") + batch: Final = RedisBatch(cache) + good: Final = cache.async_increment_iterations(key="good", ttl_seconds=60, batch=batch) + broken: Final = cache.async_increment_iterations(key="wrong-type", ttl_seconds=60, batch=batch) + assert await good == 1 + failure: Final = await broken + assert isinstance(failure, CacheFailure) + assert failure.kind == "unavailable" + settled: Final[asyncio.Future[int | CacheFailure]] = asyncio.get_running_loop().create_future() + with request_redis_batch_scope(post_call_deadline=60) as request: + post_call: Final = active_post_call_redis_batch(cache) + operation: Final = cache.async_increment_iterations(key="post-call", ttl_seconds=60, batch=post_call) + operation.on_settled(lambda result: settled.set_result(result.result())) + assert raw.get("typed:post-call") is None + await request.flush_post_call() + assert await settled == 1 + assert await operation == 1 + assert raw.get("typed:post-call") == b"1" + + +@pytest.mark.asyncio +async def test_pubsub_and_update_queues_preserve_namespace_payloads_and_order( + backend: tuple[RedisCache, Redis[bytes]], +) -> None: + cache, raw = backend + subscription: Final = await cache.async_subscribe(channel="updates") + assert not isinstance(subscription, CacheFailure) + try: + assert await subscription.receive(timeout_seconds=1) is None + assert await cache.async_publish_many(messages=( + CachePublish("updates", "first"), CachePublish("typed:updates", "second"), + )) == (1, 1) + assert await subscription.receive(timeout_seconds=1) == b"first" + assert await subscription.receive(timeout_seconds=1) == b"second" + assert await cache.async_enqueue_updates(updates=( + CacheQueueUpdate("queue", ("first", "second")), CacheQueueUpdate("typed:queue", ("third",)), + )) == (2, 3) + assert raw.lrange("typed:queue", 0, -1) == [b"first", b"second", b"third"] + assert raw.exists("queue") == 0 + assert await cache.async_delete_record(key="queue") is True + assert await cache.async_delete_record(key="queue") is False + finally: + assert await subscription.close() is None diff --git a/tests/unit/caching/test_redis_batch.py b/tests/unit/caching/test_redis_batch.py index cd270035416..532c91796c0 100644 --- a/tests/unit/caching/test_redis_batch.py +++ b/tests/unit/caching/test_redis_batch.py @@ -5,9 +5,10 @@ from __future__ import annotations import asyncio import hashlib import json +import warnings from collections.abc import Awaitable, Callable, Sequence from datetime import timedelta -from typing import Any +from typing import Any, Final import pytest from redis.exceptions import NoScriptError @@ -26,6 +27,7 @@ from litellm.caching.redis_cache import ( _get_call_stack_info, # pyright: ignore[reportPrivateUsage] # the chain the service hook reports ) from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.types.caching import CacheFailure SCRIPT = "return redis.call('GET', KEYS[1])" SHA = hashlib.sha1(SCRIPT.encode()).hexdigest() # noqa: S324 @@ -508,3 +510,137 @@ async def test_failed_pipeline_reports_the_same_family_target_as_a_successful_on assert isinstance(event["error"], ConnectionError) assert (event["call_type"], event["target"]) == ("post_call_redis_batch", MIXED_PIPELINE_TARGET) assert event["event_metadata"] == {"op_count": 2, "families": "auth_objects,spend_counters"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("values", [(), (b"1", b"2", b"3")]) +async def test_mget_wrong_reply_length_fails_only_its_read_and_marks_no_misses(values: tuple[bytes, ...]) -> None: + def reply_for(command: tuple[object, ...]) -> object: + return values if command[0] == "MGET" else b"4.5" + + cache: Final = FakeRedisCache(FakeClient(reply_for)) + batch: Final = RedisBatch(cache) + read: Final = batch.mget(("one", "two")) + increment: Final = batch.increment("count", 1) + + with pytest.raises(ValueError, match="MGET reply count"): + await read + + assert await increment == 4.5 + assert not batch.read_as_missing("one") + assert not batch.read_as_missing("two") + + +@pytest.mark.asyncio +async def test_expire_failure_does_not_replay_an_increment_that_may_have_succeeded() -> None: + def reply_for(command: tuple[object, ...]) -> object: + return NoScriptError("NOSCRIPT") if command[0] == "EXPIRE" else b"3.5" + + cache: Final = FakeRedisCache(FakeClient(reply_for)) + batch: Final = RedisBatch(cache) + failed: Final = batch.increment("count", 2, ttl=60) + sibling: Final = batch.increment("other", 1) + + with pytest.raises(NoScriptError): + await failed + + assert await sibling == 3.5 + assert cache.store == {} + assert cache.alone == [] + + +class _IncompletePipeline(FakePipeline): + async def execute(self, raise_on_error: bool = True) -> list[object]: + return (await super().execute(raise_on_error=raise_on_error))[:-1] + + +class _IncompleteClient(FakeClient): + def pipeline(self, transaction: bool = True) -> _IncompletePipeline: + assert transaction is False + return _IncompletePipeline(self.reply_for, self.fail) + + +@pytest.mark.asyncio +async def test_incomplete_pipeline_response_never_reports_a_write_as_successful() -> None: + cache: Final = FakeRedisCache(_IncompleteClient(replies)) + batch: Final = RedisBatch(cache) + increment: Final = batch.increment("count", 1, ttl=60) + write: Final = batch.set("value", "stored") + + with pytest.raises(ValueError, match="pipeline reply count"): + await increment + with pytest.raises(ValueError, match="pipeline reply count"): + await write + + assert cache.alone == [] + + +@pytest.mark.asyncio +async def test_legacy_batch_script_stays_callable_under_warnings_as_errors() -> None: + cache, _client = make() + batch: Final = RedisBatch(cache) + with warnings.catch_warnings(): + warnings.simplefilter("error") + result: Final = await batch.script(SCRIPT, run_alone_script, ("key",), ()) + + assert result == [1, 2] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "reply, kind", [(ConnectionError("ambiguous write"), "unavailable"), (["bad", "shape"], "invalid_response")] +) +async def test_typed_batch_failures_are_values_and_do_not_disturb_sibling_operations(reply: object, kind: str) -> None: + def reply_for(command: tuple[object, ...]) -> object: + return reply if command[0] == "EVALSHA" else b"3.5" + + client: Final = FakeClient(reply_for) + cache: Final = FakeRedisCache(client) + cache.redis_kwargs = {"host": "typed-batch-failures"} + batch: Final = RedisBatch(cache) + with pytest.warns(FutureWarning, match="async_increment_iterations is experimental and may change"): + operation: Final = cache.async_increment_iterations(key="iteration", ttl_seconds=60, batch=batch) + sibling: Final = batch.increment("counter", 1, ttl=60) + settled: Final[asyncio.Future[int | CacheFailure]] = asyncio.get_running_loop().create_future() + with pytest.warns(FutureWarning, match="operation.on_settled is experimental and may change"): + operation.on_settled(lambda result: settled.set_result(result.result())) + + result: Final = await operation + + assert isinstance(result, CacheFailure) + assert result.kind == kind + assert await settled == result + assert await sibling == 3.5 + assert len(client.pipelines) == 1 + assert cache.alone == [] + + +class _CancelledPipeline(FakePipeline): + async def execute(self, raise_on_error: bool = True) -> list[object]: + raise asyncio.CancelledError + + +class _CancelledClient(FakeClient): + def pipeline(self, transaction: bool = True) -> _CancelledPipeline: + assert transaction is False + return _CancelledPipeline(self.reply_for, self.fail) + + +@pytest.mark.asyncio +async def test_cancelled_typed_batch_notifies_settled_hooks_without_translating_cancellation_to_failure() -> None: + cache: Final = FakeRedisCache(_CancelledClient(replies)) + cache.redis_kwargs = {"host": "typed-batch-cancellation"} + batch: Final = RedisBatch(cache) + settled: Final[asyncio.Future[bool]] = asyncio.get_running_loop().create_future() + with pytest.warns(FutureWarning, match="async_increment_iterations is experimental and may change"): + operation: Final = cache.async_increment_iterations(key="iteration", ttl_seconds=60, batch=batch) + with pytest.warns(FutureWarning, match="operation.on_settled is experimental and may change"): + operation.on_settled(lambda result: settled.set_result(result.cancelled())) + + with pytest.raises(asyncio.CancelledError): + await batch.flush() + with pytest.raises(asyncio.CancelledError): + await operation + + assert settled.result() is True + assert cache.alone == [] diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..8c2a62c959c 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -2,16 +2,20 @@ import asyncio import gc import time import types +import warnings from collections.abc import Iterator from datetime import timedelta from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeRedis as FakeAsyncRedis import litellm from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError +from litellm.types.caching import KeyPresence @pytest.fixture @@ -34,6 +38,89 @@ def sync_batch_redis_cache(redis_no_ping): yield cache +@pytest.fixture +def key_presence_backend(redis_no_ping: None) -> Iterator[tuple[RedisCache, FakeServer]]: + server: Final = FakeServer() + with ( + patch( # test-quality-ok: constructor creates the Redis transport with no injection point + "litellm._redis.get_redis_client", return_value=FakeRedis(server=server) + ), + patch( # test-quality-ok: factory supplies a Redis transport, no cache behavior is mocked + "litellm._redis.get_redis_async_client", return_value=FakeAsyncRedis(server=server) + ), + ): + yield RedisCache(host="127.0.0.1", port=6379, namespace="presence"), server + + +@pytest.mark.asyncio +@pytest.mark.parametrize("key", ("marker", "presence:marker")) +@pytest.mark.parametrize("value", (b"1", b"null", b"", b"0", b"\xff")) +async def test_async_key_presence_reads_raw_values_without_decoding( + key_presence_backend: tuple[RedisCache, FakeServer], key: str, value: bytes +) -> None: + cache, _ = key_presence_backend + await cache.init_async_client().set("presence:marker", value) + + assert await cache.async_key_presence(key) == "present" + + +@pytest.mark.asyncio +async def test_async_key_presence_does_not_read_another_namespace( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + cache, _ = key_presence_backend + await cache.init_async_client().set("marker", b"1") + + assert await cache.async_key_presence("marker") == "missing" + + +@pytest.mark.asyncio +async def test_async_key_presence_wrong_type_is_unavailable( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + cache, _ = key_presence_backend + await cache.init_async_client().rpush("presence:marker", b"1") + + assert await cache.async_key_presence("marker") == "unavailable" + + +@pytest.mark.asyncio +async def test_async_key_presence_outage_is_not_a_miss( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + cache, server = key_presence_backend + server.connected = False + + assert await cache.async_key_presence("marker") == "unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value, expected", ((None, "missing"), (1, "present"))) +async def test_async_key_presence_observes_existing_cache_writes( + key_presence_backend: tuple[RedisCache, FakeServer], value: int | None, expected: KeyPresence +) -> None: + cache, _ = key_presence_backend + if value is not None: + await cache.async_set_cache("marker", value, ttl=60) + + assert await cache.async_key_presence("marker") == expected + + +@pytest.mark.asyncio +async def test_deprecated_client_accessors_still_reuse_a_working_client_without_runtime_warnings( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + cache, _ = key_presence_backend + with warnings.catch_warnings(): + warnings.simplefilter("error", DeprecationWarning) + client: Final = cache.init_async_client() + assert cache.init_async_client() is client + assert cache.init_pubsub_client() is client + assert await client.ping() is True + await client.set("presence:marker", b"1") + assert await cache.async_key_presence("marker") == "present" + + @pytest.mark.parametrize( ("namespace", "key", "expected"), [ @@ -1033,8 +1120,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"]) gc.collect() ServiceLogging() - assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == ( - redis_breaker_state_counts()["closed"] + assert ( + REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) + == (redis_breaker_state_counts()["closed"]) ) @@ -1697,3 +1785,149 @@ async def test_async_set_cache_pipeline_does_not_format_cached_values_for_loggin assert _FormatCountingStr.format_calls == 0 assert pipe.sets == [("k1", '"embedding"', timedelta(seconds=60))] + + +@pytest.mark.asyncio +async def test_typed_iteration_operation_warns_at_caller_and_runs_once( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + cache, _ = key_presence_backend + runner: Final = AsyncMock(return_value=7) + with patch("redis.asyncio.client.Redis.register_script", return_value=runner): + with pytest.warns(FutureWarning, match="async_increment_iterations is experimental and may change") as seen: + operation: Final = cache.async_increment_iterations(key="counter", ttl_seconds=60) + assert seen[0].filename == __file__ + assert await operation == 7 + assert await operation == 7 + runner.assert_awaited_once_with(keys=("presence:counter",), args=(60,), client=None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply", (None, "not-numeric", [1, 2], float("nan"))) +async def test_typed_operation_translates_malformed_replies( + key_presence_backend: tuple[RedisCache, FakeServer], + reply: object, +) -> None: + from litellm.types.caching import CacheFailure + + cache, _ = key_presence_backend + with patch("redis.asyncio.client.Redis.register_script", return_value=AsyncMock(return_value=reply)): + with pytest.warns(FutureWarning): + result: Final = await cache.async_increment_iterations(key="counter", ttl_seconds=60) + assert isinstance(result, CacheFailure) + assert result.kind == "invalid_response" + + +@pytest.mark.asyncio +async def test_typed_operations_translate_outages_and_leave_cancellation_visible( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + from litellm.types.caching import CacheFailure + + cache, _ = key_presence_backend + with patch("redis.asyncio.client.Redis.register_script", return_value=AsyncMock(side_effect=OSError("offline"))): + with pytest.warns(FutureWarning): + result: Final = await cache.async_increment_iterations(key="counter", ttl_seconds=60) + assert isinstance(result, CacheFailure) + assert result.kind == "unavailable" + with patch( + "redis.asyncio.client.Redis.register_script", return_value=AsyncMock(side_effect=asyncio.CancelledError) + ): + with pytest.warns(FutureWarning), pytest.raises(asyncio.CancelledError): + await cache.async_increment_spend(key="spend", amount=0.5, ttl_seconds=60) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("timestamp", (b"123.50", "123.50")) +async def test_typed_window_operations_preserve_fractional_window_identity( + key_presence_backend: tuple[RedisCache, FakeServer], timestamp: str | bytes, +) -> None: + from litellm.types.caching import WindowCount, WindowCounter, WindowReservation + + cache, _ = key_presence_backend + runner: Final = AsyncMock(side_effect=([timestamp, 1], [0, 2, timestamp])) + with patch("redis.asyncio.client.Redis.register_script", return_value=runner): + incremented: Final = await cache.async_increment_window_counters( + counters=(WindowCounter("window", "count"),), now=123.5, window_seconds=60, + ) + reserved: Final = await cache.async_reserve_window_capacity( + reservations=(WindowReservation("window", "count", 5, 1, 60, 60),), + ) + assert incremented == (WindowCount("123.50", 1),) + assert reserved == (WindowCount("123.50", 2),) + + +@pytest.mark.asyncio +async def test_typed_delete_preserves_namespace_and_distinguishes_outage( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + from litellm.types.caching import CacheFailure + + cache, server = key_presence_backend + await cache.async_set_cache("pair", 3, ttl=60) + with pytest.warns(FutureWarning, match="async_delete_record is experimental and may change") as seen: + assert await cache.async_delete_record(key="presence:pair") is True + assert seen[0].filename == __file__ + assert await cache.async_get_cache("pair") is None + assert await cache.async_delete_record(key="pair") is False + server.connected = False + failure: Final = await cache.async_delete_record(key="pair") + assert isinstance(failure, CacheFailure) + assert failure.kind == "unavailable" + + +@pytest.mark.asyncio +async def test_typed_enqueue_keeps_order_and_validates_lengths( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + from litellm.types.caching import CacheQueueUpdate + + cache, _ = key_presence_backend + with pytest.warns(FutureWarning, match="async_enqueue_updates is experimental and may change"): + result: Final = await cache.async_enqueue_updates(updates=( + CacheQueueUpdate(key="queue", values=("one", "two")), + CacheQueueUpdate(key="presence:queue", values=("three",)), + )) + assert result == (2, 3) + assert await cache.async_lpop("queue", count=3) == ["one", "two", "three"] + + +@pytest.mark.asyncio +async def test_typed_pubsub_namespaces_messages_and_closes_the_subscription( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + from litellm.types.caching import CacheFailure, CachePublish + + cache, _ = key_presence_backend + with pytest.warns(FutureWarning, match="async_subscribe is experimental and may change"): + subscription: Final = await cache.async_subscribe(channel="updates") + assert not isinstance(subscription, CacheFailure) + with pytest.warns(FutureWarning): + assert await subscription.receive(timeout_seconds=0.0) is None + with pytest.warns(FutureWarning): + assert await cache.async_publish_many(messages=(CachePublish("presence:updates", "changed"),)) == (1,) + with pytest.warns(FutureWarning): + assert await subscription.receive(timeout_seconds=0.0) == b"changed" + with pytest.warns(FutureWarning): + assert await subscription.close() is None + with pytest.warns(FutureWarning): + assert await subscription.receive(timeout_seconds=0.0) == CacheFailure("unavailable", "Subscription is closed") + + +@pytest.mark.asyncio +async def test_typed_read_values_preserves_cache_decoding_and_distinguishes_outage( + key_presence_backend: tuple[RedisCache, FakeServer], +) -> None: + from litellm.types.caching import CacheFailure + + cache, server = key_presence_backend + await cache.init_async_client().set("presence:json", '{"enabled":true}') + await cache.init_async_client().set("presence:null", "null") + with pytest.warns(FutureWarning): + values: Final = await cache.async_read_values(keys=("json", "null", "missing")) + assert values == {"json": {"enabled": True}, "null": None, "missing": None} + server.connected = False + with pytest.warns(FutureWarning): + result: Final = await cache.async_read_values(keys=("json",)) + assert isinstance(result, CacheFailure) + assert result.kind == "unavailable" diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 85422819926..c1dd02d16ed 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -5,7 +5,7 @@ import ssl import weakref from collections.abc import Awaitable, Callable from importlib import import_module -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -17,6 +17,7 @@ from litellm.caching.evicted_client_closer import EvictedClientCloser from litellm.caching.llm_caching_handler import LLMClientCache from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cluster_cache import ClusterConnectionGate, RedisClusterCache +from litellm.types.caching import CacheFailure @patch("litellm._redis.init_redis_cluster") @@ -52,7 +53,7 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): """ # Create a mock Redis client mock_redis = MagicMock() - mock_redis.mget_nonatomic.return_value = [None, None] # Simulate no cache hits + mock_redis.mget_nonatomic = AsyncMock(return_value=[None, None]) # Create RedisClusterCache instance with mock client cache = RedisClusterCache( @@ -65,10 +66,10 @@ async def test_redis_cluster_async_batch_get(mock_init_redis_cluster): # Test async_batch_get_cache keys = ["key1", "key2"] - await cache.async_batch_get_cache(keys) + result: Final = await cache.async_batch_get_cache(keys) - # Verify mget_nonatomic was called instead of mget - mock_redis.mget_nonatomic.assert_called_once() + assert result == {"key1": None, "key2": None} + mock_redis.mget_nonatomic.assert_awaited_once_with(keys=keys) assert not mock_redis.mget.called @@ -597,3 +598,79 @@ def test_closed_event_loop_is_collected_after_contended_gate_use() -> None: gc.collect() assert loop_ref() is None + + +@pytest.mark.asyncio +async def test_typed_cluster_read_decodes_values_and_namespaces_each_key_once() -> None: + with ( + patch("redis.RedisCluster", return_value=MagicMock()), + patch( + "redis.asyncio.RedisCluster.mget_nonatomic", + new=AsyncMock(return_value=[b'{"id": 7}', None, b"null", b'""']), + ) as mget, + ): + cache: Final = RedisClusterCache( + startup_nodes=[{"host": "typed-cluster-read", "port": 7000}], namespace="ns", max_connections=1 + ) + with pytest.warns(FutureWarning, match="async_read_values is experimental and may change"): + result: Final = await cache.async_read_values(keys=("hit", "ns:miss", "null", "empty")) + + assert result == {"hit": {"id": 7}, "ns:miss": None, "null": None, "empty": ""} + mget.assert_awaited_once_with(keys=["ns:hit", "ns:miss", "ns:null", "ns:empty"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("reply", "failure", "expected_kind"), + [ + pytest.param([], None, "invalid_response", id="missing-key-reply"), + pytest.param([object()], None, "invalid_response", id="unexpected-value-type"), + pytest.param(None, ConnectionError("offline"), "unavailable", id="unavailable-cluster"), + ], +) +async def test_typed_cluster_read_never_turns_failure_into_a_miss( + reply: object, failure: Exception | None, expected_kind: str +) -> None: + with ( + patch("redis.RedisCluster", return_value=MagicMock()), + patch("redis.asyncio.RedisCluster.mget_nonatomic", new=AsyncMock(return_value=reply, side_effect=failure)), + ): + cache: Final = RedisClusterCache(startup_nodes=[{"host": "typed-cluster-failures", "port": 7000}]) + with pytest.warns(FutureWarning, match="async_read_values is experimental and may change"): + result: Final = await cache.async_read_values(keys=("key",)) + + assert isinstance(result, CacheFailure) + assert result.kind == expected_kind + + +@pytest.mark.asyncio +async def test_cancelled_typed_cluster_read_releases_its_connection_admission() -> None: + entered: Final[asyncio.Queue[str]] = asyncio.Queue() + release: Final = asyncio.Event() + + async def read(*, keys: list[str]) -> list[bytes | None]: + await entered.put(keys[0]) + await release.wait() + return [b"1"] + + with ( + patch("redis.RedisCluster", return_value=MagicMock()), + patch("redis.asyncio.RedisCluster.mget_nonatomic", new=AsyncMock(side_effect=read)), + ): + cache: Final = RedisClusterCache( + startup_nodes=[{"host": "typed-cluster-cancellation", "port": 7000}], + connection_gate=ClusterConnectionGate(limit=1, wait_timeout=60), + ) + first: Final = asyncio.create_task(cache.async_read_values(keys=("cancelled",))) + with pytest.warns(FutureWarning, match="async_read_values is experimental and may change"): + entered_key: Final = await asyncio.wait_for(entered.get(), timeout=1) + assert entered_key == "cancelled" + first.cancel() + with pytest.raises(asyncio.CancelledError): + await first + release.set() + with pytest.warns(FutureWarning, match="async_read_values is experimental and may change"): + result: Final = await asyncio.wait_for(cache.async_read_values(keys=("next",)), timeout=1) + + assert result == {"next": 1} + assert entered.get_nowait() == "next" diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py index e06f3bcd5cf..f0e2f738be5 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_distributed_lock.py @@ -1,7 +1,13 @@ -"""Tests for the Redis lock: acquire NX/PX with a token, compare-and-delete release, namespacing.""" +"""Refresh lock policy over the real cache facade with Redis I/O isolated.""" + +import asyncio +from collections.abc import Iterator +from typing import Final, Literal +from unittest.mock import ANY, AsyncMock, call, patch import pytest +from litellm.caching.redis_cache import RedisCache from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_distributed_lock import ( RedisDistributedLock, ) @@ -10,117 +16,146 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.redis_refresh_c ) -class _FakeRedis: - """Models just enough Redis to exercise SET NX (token store) and the compare-and-delete EVAL.""" - - def __init__(self, set_returns=True, exists_returns=1, raise_on=()): - self._set_returns = set_returns - self._exists_returns = exists_returns - self._raise_on = set(raise_on) - self.set_calls = [] - self.eval_calls = [] - self.expire_calls = [] - self.deleted = [] - self.store: dict = {} - - async def set(self, name, value, *, nx=False, px=None): - if "set" in self._raise_on: - raise RuntimeError("redis down") - self.set_calls.append((name, value, nx, px)) - if self._set_returns: - self.store[name] = value - return self._set_returns - - async def eval(self, script, numkeys, *keys_and_args): - if "eval" in self._raise_on: - raise RuntimeError("redis down") - self.eval_calls.append((numkeys, keys_and_args)) - key, token = keys_and_args[0], keys_and_args[1] - if len(keys_and_args) == 3: - ttl_ms = keys_and_args[2] - if self.store.get(key) == token: - self.expire_calls.append((key, ttl_ms)) - return 1 - return 0 - if self.store.get(key) == token: # compare-and-delete: only the owner deletes - del self.store[key] - self.deleted.append(key) - return 1 - return 0 - - async def exists(self, *names): - if "exists" in self._raise_on: - raise RuntimeError("redis down") - return self._exists_returns +@pytest.fixture +def backend() -> Iterator[tuple[RedisDistributedLock, AsyncMock]]: + with ( + patch("redis.client.Redis.execute_command", return_value=True), + patch("redis.asyncio.client.Redis.execute_command", new_callable=AsyncMock) as transport, + ): + yield RedisDistributedLock(RedisCache(host="127.0.0.1", port=6379, namespace="ns")), transport @pytest.mark.asyncio -async def test_acquire_sets_token_with_nx_px_and_reports_acquired(): - redis = _FakeRedis(set_returns=True) - assert await RedisDistributedLock(redis).acquire("k", "tok-1", 10.0) is (LockAcquisition.ACQUIRED) - assert redis.set_calls == [("k", "tok-1", True, 10000)] # token value, NX, px in ms +@pytest.mark.parametrize( + ("reply", "expected"), + ( + (b"OK", LockAcquisition.ACQUIRED), + (None, LockAcquisition.HELD), + (False, LockAcquisition.HELD), + ("invalid", LockAcquisition.ERROR), + ), +) +async def test_acquire_distinguishes_ownership_contention_and_invalid_replies( + backend: tuple[RedisDistributedLock, AsyncMock], reply: object, expected: LockAcquisition +) -> None: + lock, transport = backend + transport.return_value = reply + assert await lock.acquire("k", "tok-1", 10.125) is expected + transport.assert_awaited_once_with("EVALSHA", ANY, 1, "ns:k", "tok-1", 10125) @pytest.mark.asyncio -async def test_acquire_reports_held_when_key_already_held(): - # redis SET NX returns None when the key exists -> another worker holds it. - assert await RedisDistributedLock(_FakeRedis(set_returns=None)).acquire("k", "tok", 10.0) is LockAcquisition.HELD +async def test_acquire_reports_error_on_redis_error_distinct_from_held( + backend: tuple[RedisDistributedLock, AsyncMock], +) -> None: + lock, transport = backend + transport.side_effect = OSError("redis down") + assert await lock.acquire("k", "tok", 10.0) is LockAcquisition.ERROR @pytest.mark.asyncio -async def test_acquire_reports_error_on_redis_error_distinct_from_held(): - # A dead backend must be distinguishable from a busy holder so the coordinator refreshes anyway. - assert await RedisDistributedLock(_FakeRedis(raise_on=["set"])).acquire("k", "tok", 10.0) is LockAcquisition.ERROR +async def test_release_passes_each_holders_token_to_the_atomic_operation( + backend: tuple[RedisDistributedLock, AsyncMock], +) -> None: + lock, transport = backend + transport.side_effect = (0, 1) + assert await lock.release("k", "stale-owner") is None + assert await lock.release("k", "current-owner") is None + assert transport.await_args_list == [ + call("EVALSHA", ANY, 1, "ns:k", "stale-owner"), + call("EVALSHA", ANY, 1, "ns:k", "current-owner"), + ] @pytest.mark.asyncio -async def test_release_deletes_only_when_the_token_matches(): - # Regression: a holder whose lock PX-expired and was re-acquired by another worker must not be - # able to delete the new holder's lock. release with a stale token is a no-op. - redis = _FakeRedis() - lock = RedisDistributedLock(redis) - await lock.acquire("k", "owner-B", 10.0) # B currently holds the lock - await lock.release("k", "owner-A") # A's stale token - assert redis.deleted == [] and redis.store.get("k") == "owner-B" # B's lock survives - await lock.release("k", "owner-B") # the real owner releases - assert redis.deleted == ["k"] and "k" not in redis.store +@pytest.mark.parametrize("key", ("k", "ns:k")) +async def test_all_lease_operations_apply_the_namespace_once( + backend: tuple[RedisDistributedLock, AsyncMock], key: str +) -> None: + lock, transport = backend + transport.side_effect = (b"OK", 1, 1, 0) + assert await lock.acquire(key, "tok", 10.0) is LockAcquisition.ACQUIRED + assert await lock.extend(key, "tok", 10.0) is True + assert await lock.release(key, "tok") is None + assert await lock.is_held(key) is False + assert transport.await_args_list == [ + call("EVALSHA", ANY, 1, "ns:k", "tok", 10000), + call("EVALSHA", ANY, 1, "ns:k", "tok", 10000), + call("EVALSHA", ANY, 1, "ns:k", "tok"), + call("EVALSHA", ANY, 1, "ns:k"), + ] @pytest.mark.asyncio -async def test_keys_are_namespaced_before_reaching_redis(): - redis = _FakeRedis() - lock = RedisDistributedLock(redis, namespace_key=lambda key: f"ns:{key}") - await lock.acquire("k", "tok", 10.0) - await lock.extend("k", "tok", 10.0) - await lock.release("k", "tok") - await lock.is_held("k") - assert redis.set_calls[0][0] == "ns:k" # acquire namespaced - assert redis.eval_calls[0][1][0] == "ns:k" # extend (EVAL KEYS[1]) namespaced - assert redis.eval_calls[1][1][0] == "ns:k" # release (EVAL KEYS[1]) namespaced - assert redis.deleted == ["ns:k"] +async def test_extend_preserves_owner_token_and_millisecond_ttl( + backend: tuple[RedisDistributedLock, AsyncMock], +) -> None: + lock, transport = backend + transport.side_effect = (0, 1) + assert await lock.extend("k", "stale-owner", 10.125) is False + assert await lock.extend("k", "current-owner", 10.125) is True + assert transport.await_args_list == [ + call("EVALSHA", ANY, 1, "ns:k", "stale-owner", 10125), + call("EVALSHA", ANY, 1, "ns:k", "current-owner", 10125), + ] @pytest.mark.asyncio -async def test_extend_refreshes_ttl_only_when_the_token_matches(): - redis = _FakeRedis() - lock = RedisDistributedLock(redis) - await lock.acquire("k", "owner-B", 10.0) - assert await lock.extend("k", "owner-A", 10.0) is False - assert await lock.extend("k", "owner-B", 10.0) is True - assert redis.expire_calls == [("k", "10000")] +@pytest.mark.parametrize(("reply", "held"), ((1, True), (0, False), ("invalid", False))) +async def test_is_held_maps_existence_and_invalid_replies( + backend: tuple[RedisDistributedLock, AsyncMock], reply: object, held: bool +) -> None: + lock, transport = backend + transport.return_value = reply + assert await lock.is_held("k") is held @pytest.mark.asyncio -async def test_extend_degrades_to_false_on_redis_error(): - assert await RedisDistributedLock(_FakeRedis(raise_on=["eval"])).extend("k", "tok", 10.0) is False +@pytest.mark.parametrize("operation", ("extend", "release", "is_held")) +async def test_lease_outages_do_not_block_refresh_or_reread( + backend: tuple[RedisDistributedLock, AsyncMock], operation: Literal["extend", "release", "is_held"] +) -> None: + lock, transport = backend + transport.side_effect = OSError("redis down") + match operation: + case "extend": + assert await lock.extend("k", "tok", 10.0) is False + case "release": + assert await lock.release("k", "tok") is None + case "is_held": + assert await lock.is_held("k") is False @pytest.mark.asyncio -async def test_is_held_reflects_exists(): - assert await RedisDistributedLock(_FakeRedis(exists_returns=1)).is_held("k") is True - assert await RedisDistributedLock(_FakeRedis(exists_returns=0)).is_held("k") is False +@pytest.mark.parametrize("operation", ("acquire", "extend", "release", "is_held")) +async def test_lease_cancellation_is_not_translated_to_an_outage( + backend: tuple[RedisDistributedLock, AsyncMock], operation: Literal["acquire", "extend", "release", "is_held"] +) -> None: + lock, transport = backend + transport.side_effect = asyncio.CancelledError + + async def invoke() -> None: + match operation: + case "acquire": + await lock.acquire("k", "tok", 10.0) + case "extend": + await lock.extend("k", "tok", 10.0) + case "release": + await lock.release("k", "tok") + case "is_held": + await lock.is_held("k") + + with pytest.raises(asyncio.CancelledError): + await invoke() @pytest.mark.asyncio -async def test_is_held_degrades_to_false_on_redis_error(): - assert await RedisDistributedLock(_FakeRedis(raise_on=["exists"])).is_held("k") is False +async def test_lease_apis_warn_as_experimental_at_the_lock_callsite( + backend: tuple[RedisDistributedLock, AsyncMock], +) -> None: + lock, transport = backend + transport.return_value = b"OK" + with pytest.warns(FutureWarning, match="async_acquire_lease is experimental and may change") as seen: + assert await lock.acquire("k", "tok", 10.0) is LockAcquisition.ACQUIRED + warning: Final = seen[0] + assert warning.filename.endswith("/outbound_credentials/redis_distributed_lock.py") diff --git a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py index a6831562018..93a010cf217 100644 --- a/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py +++ b/tests/unit/proxy/_experimental/mcp_server/outbound_credentials/test_redis_refresh_coordinator.py @@ -1,6 +1,7 @@ """Tests for the cross-replica refresh coordinator: winner refreshes, losers wait then re-read.""" import asyncio +from typing import Final import pytest @@ -268,3 +269,67 @@ async def test_winner_stops_renewing_after_the_refresh_budget(): assert result is not None and result.access_token == "new" assert extends_during_budget == 4 # renewed at clock 5, 10, 15, 20 then stopped assert len(lock.extend_calls) == 4 # no further renewals after the budget + + +@pytest.mark.asyncio +async def test_cancelling_a_winner_cancels_refresh_and_renewal_then_releases_its_lease() -> None: + lock: Final = _FakeLock(acquired=LockAcquisition.ACQUIRED) + refresh_started: Final = asyncio.Event() + renewal_started: Final = asyncio.Event() + refresh_cancelled: Final = asyncio.Event() + renewal_cancelled: Final = asyncio.Event() + + async def refresh() -> OAuthToken | None: + refresh_started.set() + try: + await asyncio.Event().wait() + finally: + refresh_cancelled.set() + return None + + async def sleep(_seconds: float) -> None: + renewal_started.set() + try: + await asyncio.Event().wait() + finally: + renewal_cancelled.set() + + async def reread() -> OAuthToken | None: + pytest.fail("A cancelled winner must not reread") + + coord: Final = RedisRefreshCoordinator(lock, new_token=lambda: "owner", sleep=sleep, clock=lambda: 0.0) + task: Final = asyncio.create_task(coord.run("u", "s", refresh, reread)) + await refresh_started.wait() + await renewal_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert refresh_cancelled.is_set() + assert renewal_cancelled.is_set() + assert lock.extend_calls == [] + assert lock.released == [(_KEY, "owner")] + + +@pytest.mark.asyncio +async def test_cancelling_a_waiter_does_not_release_another_workers_lease_or_reread() -> None: + lock: Final = _FakeLock(acquired=LockAcquisition.HELD, held_sequence=(True,)) + waiting: Final = asyncio.Event() + + async def sleep(_seconds: float) -> None: + waiting.set() + await asyncio.Event().wait() + + async def refresh() -> OAuthToken | None: + pytest.fail("A waiting worker must not refresh") + + async def reread() -> OAuthToken | None: + pytest.fail("A cancelled waiter must not reread") + + coord: Final = RedisRefreshCoordinator(lock, sleep=sleep, clock=lambda: 0.0) + task: Final = asyncio.create_task(coord.run("u", "s", refresh, reread)) + await waiting.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert lock.released == [] + assert lock.extend_calls == [] diff --git a/tests/unit/proxy/_experimental/mcp_server/test_byok_credential_cache.py b/tests/unit/proxy/_experimental/mcp_server/test_byok_credential_cache.py index 0ec4b431276..28bcd4b675e 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_byok_credential_cache.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_byok_credential_cache.py @@ -47,10 +47,7 @@ def test_peer_worker_invalidation_message_evicts_the_cached_credential(): ) subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler - { - "type": "message", - "data": json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode(), - } + json.dumps({"cache_key": byok_credential_cache_key("mallory", "srv-byok")}).encode() ) assert get_cached_byok_credential("mallory", "srv-byok") is None diff --git a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 530667fdfdd..d74d6efc779 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -63,6 +63,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i mint_session_token, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper +from litellm.types.caching import KeyPresence MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" @@ -729,9 +730,7 @@ async def test_single_use_guard_fails_closed_when_redis_errors(): @pytest.mark.asyncio async def test_single_use_guard_peek_fails_closed_when_redis_errors(): - """The peek that gates a refresh on its chain's revocation marker must read the Redis client - itself: the cache wrapper's get swallows a fault into ``None``, which would make a revoked chain - look unclaimed for exactly as long as Redis is down.""" + """An unavailable shared marker must not look unclaimed, even when the cache getter misses.""" from unittest.mock import AsyncMock, MagicMock from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard @@ -739,47 +738,39 @@ async def test_single_use_guard_peek_fails_closed_when_redis_errors(): cache = DualCache() cache.redis_cache = MagicMock() cache.redis_cache.async_get_cache = AsyncMock(return_value=None) - cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) - cache.redis_cache.init_async_client.return_value.get = AsyncMock(side_effect=ConnectionError("redis down")) + cache.redis_cache.async_key_presence = AsyncMock(return_value="unavailable") guard = _SingleUseGuard(cache) assert await guard.peek("family-fault") == "unavailable" - cache.redis_cache.init_async_client.return_value.get = AsyncMock(return_value=b"1") + cache.redis_cache.async_key_presence = AsyncMock(return_value="present") assert await guard.peek("family-fault") == "claimed" @pytest.mark.asyncio -async def test_single_use_guard_peek_reads_the_key_under_the_namespace_claim_wrote(): - """A configured ``redis_namespace`` prefixes every key the cache wrapper writes, so the raw client - read behind peek must ask for the same prefixed key: a revocation written under ``ns:...`` and - peeked at the bare key would never be seen, and the revoked chain would keep renewing.""" +async def test_single_use_guard_peek_reads_the_same_marker_claim_wrote(): from unittest.mock import AsyncMock, MagicMock from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard stored: dict[str, int] = {} # mutable-ok: the fake Redis store the test inspects - def namespaced(key: str) -> str: - return f"ns:{key}" - async def increment(key: str, value: int, ttl: int) -> int: - stored[namespaced(key)] = stored.get(namespaced(key), 0) + value - return stored[namespaced(key)] + stored[key] = stored.get(key, 0) + value + return stored[key] - async def raw_get(key: str) -> bytes | None: - return None if key not in stored else str(stored[key]).encode() + async def key_presence(key: str) -> KeyPresence: + return "present" if key in stored else "missing" cache = DualCache() cache.redis_cache = MagicMock() - cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=namespaced) cache.redis_cache.async_increment = AsyncMock(side_effect=increment) - cache.redis_cache.init_async_client.return_value.get = AsyncMock(side_effect=raw_get) + cache.redis_cache.async_key_presence = AsyncMock(side_effect=key_presence) guard = _SingleUseGuard(cache) assert await guard.peek("family-ns") == "unclaimed" assert await guard.claim("family-ns", 60) == "first" - assert set(stored) == {"ns:family-ns"} + assert set(stored) == {"family-ns"} assert await guard.peek("family-ns") == "claimed" @@ -2247,16 +2238,15 @@ async def test_revoke_burns_the_refresh_token_and_answers_200_for_dead_or_unknow assert garbage.status_code == 200 -def _redis_that(async_increment, get=None): +def _redis_that(async_increment, presence: KeyPresence = "missing"): from unittest.mock import AsyncMock, MagicMock cache = DualCache() cache.redis_cache = MagicMock() cache.redis_cache.async_increment = async_increment - cache.redis_cache.check_and_fix_namespace = MagicMock(side_effect=lambda key: key) - cache.redis_cache.init_async_client.return_value.get = get or AsyncMock(return_value=None) + cache.redis_cache.async_key_presence = AsyncMock(return_value=presence) cache.redis_cache.async_get_cache = AsyncMock( - side_effect=AssertionError("peek must read the client, not the wrapper") + side_effect=AssertionError("peek must distinguish missing from unavailable") ) cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory")) return cache @@ -2307,18 +2297,12 @@ async def test_revoke_answers_503_while_the_shared_record_cannot_be_written_then class _RedisThatFaultsOnTheSecondWrite: """A shared Redis double whose first increment lands and whose second raises, the fault between - revoke's two writes; its raw client reads back exactly the keys that landed, and later writes land.""" + revoke's two writes; it reads back exactly the keys that landed, and later writes land.""" def __init__(self) -> None: self.landed: tuple[str, ...] = () self.writes = 0 - def check_and_fix_namespace(self, key: str) -> str: - return key - - def init_async_client(self) -> "_RedisThatFaultsOnTheSecondWrite": - return self - async def async_increment(self, key: str, value: int, **kwargs: object) -> int: self.writes += 1 if self.writes == 2: @@ -2326,8 +2310,8 @@ class _RedisThatFaultsOnTheSecondWrite: self.landed = (*self.landed, key) return self.landed.count(key) - async def get(self, key: str) -> int | None: - return self.landed.count(key) or None + async def async_key_presence(self, key: str) -> KeyPresence: + return "present" if key in self.landed else "missing" @pytest.mark.asyncio @@ -2421,7 +2405,7 @@ async def test_refresh_of_a_rotated_token_answers_503_before_minting_while_redis minter = _Minter() claim = AsyncMock(return_value=1) unreadable = await _refresh_native( - rotated, client_id, minter, _redis_that(claim, get=AsyncMock(side_effect=ConnectionError("redis down"))) + rotated, client_id, minter, _redis_that(claim, presence="unavailable") ) assert unreadable.status_code == 503 assert json.loads(unreadable.body)["error"] == "temporarily_unavailable" @@ -2728,7 +2712,7 @@ async def test_introspect_503s_while_the_single_use_record_cannot_be_read(): keys, now, principal = _introspection_fixtures() minted = mint_session_refresh_token(principal, keys, now) - unreadable = _redis_that(AsyncMock(return_value=1), get=AsyncMock(side_effect=ConnectionError("redis down"))) + unreadable = _redis_that(AsyncMock(return_value=1), presence="unavailable") status, body = await _introspect(minted.token.get_secret_value(), cache=unreadable) assert (status, body["error"]) == (503, "temporarily_unavailable") diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..df7bae07493 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -12,13 +12,15 @@ load_dotenv() from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from copy import deepcopy from types import SimpleNamespace -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest, litellm import httpx from fastapi import Request, status from prisma import Prisma +from redis.asyncio import Redis +from litellm.caching.redis_cache import RedisCache from litellm._service_logger import ServiceTypes from litellm.proxy._types import LiteLLM_OrganizationTable, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import get_org_object, get_user_object @@ -2213,8 +2215,15 @@ class _Pipeline: self._messages.append(message) return self - async def execute(self) -> None: + async def execute(self) -> tuple[int, ...]: await self._send(tuple(self._messages)) + return tuple(1 for _ in self._messages) + + async def __aenter__(self) -> "_Pipeline": + return self + + async def __aexit__(self, *args: object) -> None: + pass class _SlowPublishRedisClient: @@ -2248,14 +2257,14 @@ class _WedgedPublishRedisClient: return _Pipeline(self._send) -class _PubSubRedisCache: +class _PubSubRedisCache(RedisCache): namespace = None def __init__(self, client: _SlowPublishRedisClient | _WedgedPublishRedisClient) -> None: self._client = client - def init_pubsub_client(self) -> _SlowPublishRedisClient | _WedgedPublishRedisClient: - return self._client + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", self._client) async def _drain_publishes() -> None: diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index effdd9725c2..d02adb59ae3 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10517,11 +10517,11 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t published.append((channel, message)) return 1 - class _FakeRedisCache: + class _FakeRedisCache(RedisCache): def __init__(self) -> None: self.namespace = None - def init_pubsub_client(self) -> object: + def init_pubsub_client(self) -> "Redis[bytes]": return _RecordingRedisClient() local_spend_counter_cache = DualCache() @@ -10591,11 +10591,11 @@ async def test_invalidate_team_member_spend_state_self_delivered_broadcast_does_ published.append((channel, message)) return 1 - class _FakeRedisCache: + class _FakeRedisCache(RedisCache): def __init__(self) -> None: self.namespace = None - def init_pubsub_client(self) -> object: + def init_pubsub_client(self) -> "Redis[bytes]": return _RecordingRedisClient() local_spend_counter_cache = DualCache() diff --git a/tests/unit/proxy/auth/test_login_utils.py b/tests/unit/proxy/auth/test_login_utils.py index ec9b9a04f72..70bdc76ce6c 100644 --- a/tests/unit/proxy/auth/test_login_utils.py +++ b/tests/unit/proxy/auth/test_login_utils.py @@ -7,7 +7,7 @@ to login_utils.py for better reusability. import hashlib import os -from collections.abc import Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence from contextlib import ExitStack from typing import TYPE_CHECKING, Final from datetime import datetime, timedelta, timezone @@ -16,6 +16,11 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection +from redis.asyncio import BlockingConnectionPool + +from litellm.caching.redis_cache import RedisCache if TYPE_CHECKING: from litellm.proxy.auth.login_throttle import LoginThrottle @@ -1354,35 +1359,40 @@ async def test_disabling_the_control_lets_every_attempt_through(monkeypatch): assert _local_count(throttle, throttle._keys("admin").pair_counter) == 0 +@pytest.fixture +async def login_redis_backend() -> AsyncIterator[tuple[RedisCache, FakeServer]]: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + yield RedisCache(url="redis://login-throttle.example:6379/0", namespace="login-tests"), server + finally: + await pool.aclose() + sync_client.close() + + class _FakeRedis: - """Redis whose only writes are the throttle's two scripts, run atomically as one call each. + def __init__(self) -> None: + self.values: dict[str, int] = {} + self.ttls: dict[str, int] = {} - Mirrors the Lua: a blocked key returns its remaining block time and is not counted; a counter - is expired on first write; one over the limit installs the block; a blocked pair stops the - source from being counted. The real scripts are exercised against a live Redis in the PR's - proof, this fake only has to be faithful enough for the worker-sharing tests. - """ - - def __init__(self): - self.values: dict = {} - self.ttls: dict = {} - self.scripts: list[str] = [] - - def async_register_script(self, script: str): - from litellm.proxy.auth import login_throttle as lt - - async def _run(keys, args): - self.scripts.append(script) - if script == lt._BLOCK_TTLS_LUA: - return [self._ttl(keys[1]), self._ttl(keys[3])] - assert script == lt._RECORD_FAILURE_LUA + def register_script(self, script: str) -> Callable[..., Awaitable[tuple[int, int]]]: + async def run( + keys: Sequence[str], args: Sequence[str | bytes | int | float], client: object = None + ) -> tuple[int, int]: + if not args: + return self._ttl(keys[1]), self._ttl(keys[3]) user_limit, source_limit, window, block = (int(a) for a in args) - user_block = self._bump(keys[0], keys[1], user_limit, window, block) + user_block: Final = self._bump(keys[0], keys[1], user_limit, window, block) if source_limit > 0 and user_block == 0: - return [user_block, self._bump(keys[2], keys[3], source_limit, window, block)] - return [user_block, 0] + return user_block, self._bump(keys[2], keys[3], source_limit, window, block) + return user_block, 0 - return _run + return run def _ttl(self, key: str) -> int: return self.ttls.get(key, -2) if key in self.values else -2 @@ -1398,57 +1408,50 @@ class _FakeRedis: return block return 0 - async def async_delete_cache(self, key): - self.values.pop(key, None) + async def delete(self, key: str) -> int: + removed: Final = self.values.pop(key, None) self.ttls.pop(key, None) - - -class _DownRedis(_FakeRedis): - """Redis whose every call fails, as during an outage or an open circuit breaker.""" - - def async_register_script(self, script: str): - async def _run(keys, args): - raise ConnectionError("redis is down") - - return _run - - async def async_delete_cache(self, key): - raise ConnectionError("redis is down") + return int(removed is not None) @pytest.mark.asyncio -async def test_redis_is_the_only_counter_while_it_answers(monkeypatch): - """Every worker must spend the same budget, see the same block, and a success must clear the pair for all.""" +async def test_redis_is_the_only_counter_while_it_answers( + monkeypatch: pytest.MonkeyPatch, login_redis_backend: tuple[RedisCache, FakeServer] +) -> None: from litellm.constants import LOGIN_THROTTLE_CACHE_KEY_PREFIX monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") monkeypatch.setenv("DATABASE_URL", "postgresql://stub") - redis = _FakeRedis() - first_worker = _throttle(user_limit=2, stores=_stores(), redis_cache=redis) - second_worker = _throttle(user_limit=2, stores=_stores(), redis_cache=redis) + redis, _ = login_redis_backend + transport: Final = _FakeRedis() + first_worker: Final = _throttle(user_limit=2, stores=_stores(), redis_cache=redis) + second_worker: Final = _throttle(user_limit=2, stores=_stores(), redis_cache=redis) + with ( + patch("redis.asyncio.client.Redis.register_script", side_effect=transport.register_script), + patch("redis.asyncio.client.Redis.delete", side_effect=transport.delete), + ): + assert [await _fail(first_worker, username="user@corp.com") for _ in range(3)] == ["401"] * 3 + assert not [k for k in first_worker.counters.cache_dict if str(k).startswith(LOGIN_THROTTLE_CACHE_KEY_PREFIX)], ( + "with Redis answering, no worker may keep a counter of its own" + ) + assert not first_worker.blocks.cache_dict + assert await _fail(second_worker, username="user@corp.com") == "429", "the second worker sees the block" + block_keys: Final = tuple(k for k in transport.values if ":block:user:" in k) + assert block_keys, "the block lives in Redis, where every worker reads it" + for key in block_keys: + await redis.async_delete_cache(key) + await _db_login(second_worker, "user@corp.com", "right", correct=True) - assert [await _fail(first_worker, username="user@corp.com") for _ in range(3)] == ["401"] * 3 - assert not [k for k in first_worker.counters.cache_dict if str(k).startswith(LOGIN_THROTTLE_CACHE_KEY_PREFIX)], ( - "with Redis answering, no worker may keep a counter of its own" - ) - assert not first_worker.blocks.cache_dict - - assert await _fail(second_worker, username="user@corp.com") == "429", "the second worker sees the block" - - block_keys = [k for k in redis.values if ":block:user:" in k] - assert block_keys, "the block lives in Redis, where every worker reads it" - for key in block_keys: - await redis.async_delete_cache(key) - await _db_login(second_worker, "user@corp.com", "right", correct=True) - - assert not [k for k in redis.values if ":user:" in k and ":block:" not in k], ( + assert not [k for k in transport.values if ":user:" in k and ":block:" not in k], ( "success clears the shared pair counter" ) @pytest.mark.asyncio -async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch): +async def test_a_redis_outage_falls_back_to_this_workers_own_counter( + monkeypatch: pytest.MonkeyPatch, login_redis_backend: tuple[RedisCache, FakeServer] +) -> None: """With Redis raising, guesses are still counted and blocked per worker, with a warning, instead of unbounded.""" import logging @@ -1457,7 +1460,9 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") - throttle = _throttle(user_limit=2, block_seconds=300, redis_cache=_DownRedis()) + redis, server = login_redis_backend + server.connected = False + throttle: Final = _throttle(user_limit=2, block_seconds=300, redis_cache=redis) records: list[logging.LogRecord] = [] handler = logging.Handler() @@ -1476,12 +1481,16 @@ async def test_a_redis_outage_falls_back_to_this_workers_own_counter(monkeypatch @pytest.mark.asyncio -async def test_a_failed_redis_delete_still_clears_this_workers_counter(monkeypatch): +async def test_a_failed_redis_delete_still_clears_this_workers_counter( + monkeypatch: pytest.MonkeyPatch, login_redis_backend: tuple[RedisCache, FakeServer] +) -> None: """The fail-open tradeoff: when Redis cannot clear the pair, the worker clears what it holds and moves on.""" monkeypatch.setenv("UI_USERNAME", "admin") monkeypatch.setenv("UI_PASSWORD", "right") monkeypatch.setenv("DATABASE_URL", "postgresql://stub") - throttle = _throttle(user_limit=5, redis_cache=_DownRedis()) + redis, server = login_redis_backend + server.connected = False + throttle: Final = _throttle(user_limit=5, redis_cache=redis) key = throttle._keys("user@corp.com").pair_counter assert [await _fail(throttle, username="user@corp.com") for _ in range(2)] == ["401", "401"] @@ -1491,6 +1500,49 @@ async def test_a_failed_redis_delete_still_clears_this_workers_counter(monkeypat assert _local_count(throttle, key) == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ("attempt", "failure")) +@pytest.mark.parametrize("reply", (None, (0,), (0, 0, 0), ("broken", 0))) +async def test_invalid_shared_throttle_replies_do_not_allow_unchecked_logins( + login_redis_backend: tuple[RedisCache, FakeServer], phase: str, reply: object +) -> None: + redis, _ = login_redis_backend + throttle: Final = _throttle(user_limit=1, redis_cache=redis) + with patch("redis.asyncio.client.Redis.register_script", return_value=AsyncMock(return_value=reply)): + operation: Final = throttle.attempt if phase == "attempt" else throttle.record_failure + with pytest.raises(ValueError, match="validation error"): + await operation("admin") + assert not throttle.counters.cache_dict + assert not throttle.blocks.cache_dict + + +@pytest.mark.asyncio +async def test_shared_throttle_preserves_keys_thresholds_and_ttl_units( + login_redis_backend: tuple[RedisCache, FakeServer], +) -> None: + redis, _ = login_redis_backend + throttle: Final = _throttle(source_limit=9, user_limit=4, window_seconds=23, block_seconds=71, redis_cache=redis) + username: Final = "Admin" + digest: Final = hashlib.sha256(username.casefold().encode("utf-8")).hexdigest() + expected_keys: Final = ( + f"login-tests:login_fail:{{1.2.3.4}}:user:{digest}", + f"login-tests:login_fail:{{1.2.3.4}}:block:user:{digest}", + "login-tests:login_fail:{1.2.3.4}:source", + "login-tests:login_fail:{1.2.3.4}:block:source", + ) + runner: Final = AsyncMock(return_value=(0, 71)) + with patch("redis.asyncio.client.Redis.register_script", return_value=runner): + with pytest.warns(FutureWarning, match="async_record_throttle_failure is experimental and may change"): + assert await throttle.record_failure(username) == (0, 71) + with pytest.warns(FutureWarning, match="async_read_throttle_ttls is experimental and may change"): + with pytest.raises(ProxyException) as blocked: + await throttle.attempt(username) + assert blocked.value.code == "429" + assert blocked.value.headers == {"Retry-After": "71"} + assert runner.await_args_list[0].kwargs == {"keys": expected_keys, "args": (4, 9, 23, 71), "client": None} + assert runner.await_args_list[1].kwargs == {"keys": expected_keys, "args": (), "client": None} + + @pytest.mark.asyncio async def test_counters_do_not_share_the_key_authentication_cache(monkeypatch): """Regression: throttle entries must not evict cached credentials from user_api_key_cache.""" diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 694b817ab86..71005ed18e5 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -3,6 +3,7 @@ import hashlib import json import time from collections.abc import Iterable +from typing import Final, cast from unittest.mock import patch import pytest @@ -11,6 +12,7 @@ from redis.asyncio import Redis import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module from litellm._internal_context import current_service_target from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AUTH_CACHE_INVALIDATION_CHANNEL, @@ -66,14 +68,23 @@ class _Pipeline: self._messages.append(message) return self - async def execute(self) -> None: + async def execute(self) -> tuple[int, ...]: await self._client.send(tuple(self._messages)) + return tuple(1 for _ in self._messages) + + async def __aenter__(self) -> "_Pipeline": + return self + + async def __aexit__(self, *args: object) -> None: + self._client.closed = True class _PipelinedRedisClient: def __init__(self, wedged: bool = False) -> None: self.published: list[str] = [] self.round_trips = 0 + self.closed = False + self.transaction: bool | None = None self.release = asyncio.Event() if not wedged: self.release.set() @@ -84,6 +95,7 @@ class _PipelinedRedisClient: self.published.extend(messages) def pipeline(self, transaction: bool = True) -> _Pipeline: + self.transaction = transaction return _Pipeline(self) @@ -136,13 +148,13 @@ class _ScriptedPubSubRedisClient(Redis): return next(self._scripted_pubsubs) -class _FakeRedisCache: +class _FakeRedisCache(RedisCache): def __init__(self, client: object, namespace: str | None = None) -> None: self._client = client self.namespace = namespace - def init_pubsub_client(self) -> object: - return self._client + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", self._client) def _invalidation_message(cache_key: str) -> dict: @@ -272,7 +284,8 @@ async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() @pytest.mark.asyncio -async def test_subscriber_ignores_malformed_messages() -> None: +@pytest.mark.parametrize("message", (b"not json", b'{"other": "x"}', "raw string", "null", "[]")) +async def test_subscriber_ignores_malformed_messages(message: str | bytes) -> None: cache = UserApiKeyCache() cache.in_memory_cache.set_cache("project_id:p-1", {"models": []}) @@ -280,10 +293,7 @@ async def test_subscriber_ignores_malformed_messages() -> None: redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])), user_api_key_cache=cache, ) - subscriber._apply_message({"type": "message", "data": b"not json"}) - subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()}) - subscriber._apply_message("raw string") - subscriber._apply_message(None) + subscriber._apply_message(message) assert cache.in_memory_cache.get_cache("project_id:p-1") is not None @@ -351,6 +361,8 @@ async def test_broadcast_evictions_tells_redis_every_key_past_the_pending_publis assert [json.loads(message)["cache_key"] for message in client.published] == list(cache_keys) assert client.round_trips == 1 + assert client.transaction is False + assert client.closed is True @pytest.mark.asyncio @@ -377,6 +389,29 @@ async def test_broadcast_evictions_returns_within_its_bound_while_redis_never_an assert client.published == [json.dumps({"cache_key": "user-1"})] +class _FailingPipeline(_Pipeline): + async def execute(self) -> tuple[int, ...]: + raise ConnectionError("publish batch lost") + + +class _FailingBatchClient(_PipelinedRedisClient): + def pipeline(self, transaction: bool = True) -> _Pipeline: + return _FailingPipeline(self) + + +async def test_batch_publish_failure_closes_resources_without_failing_the_write(caplog: pytest.LogCaptureFixture) -> None: + client: Final = _FailingBatchClient() + cache: Final = _FakeRedisCache(client) + + with pytest.warns(FutureWarning, match="experimental.*may change"): + await pubsub_module._publish_batch_to_redis(cache, ("user-a", "user-b")) + + assert client.closed is True + assert client.published == [] + assert "auth cache invalidation publish for 2 keys failed:" in caplog.text + assert "publish batch lost" in caplog.text + + @pytest.mark.asyncio async def test_evict_and_broadcast_publishes_under_the_auth_objects_service_target() -> None: client = _TargetRecordingRedisClient() diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index 92cfb6a71e5..6d928df5aac 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -1,22 +1,24 @@ import asyncio import json import random -from typing import Callable, Coroutine, Final, Iterable, List, Optional, Tuple +from collections.abc import Callable, Coroutine, Iterable +from typing import Final, List, Tuple, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from redis.asyncio import Redis import litellm +from litellm.caching.redis_cache import RedisCache from litellm.proxy.common_utils.config_sync_pubsub import ( + _CONFIG_SYNCED_TABLE_NAMES, + _RESYNC_APPLIED_CONFIG_PARAM_NAMES, + _WRITE_ACTION_NAMES, CONFIG_SYNC_CHANNEL, CONFIG_SYNC_JITTER_MAX_SECONDS, CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS, ConfigSyncSubscriber, - _CONFIG_SYNCED_TABLE_NAMES, _PublishOnWriteActions, - _RESYNC_APPLIED_CONFIG_PARAM_NAMES, - _WRITE_ACTION_NAMES, publish_config_change, wrap_table_actions_for_config_sync, ) @@ -103,7 +105,7 @@ class _ScriptedPubSubClient: class _QueuePubSub: def __init__(self, initial_messages: Iterable[str] = ()) -> None: - self.queue: "asyncio.Queue[str]" = asyncio.Queue() + self.queue: asyncio.Queue[str] = asyncio.Queue() for message in initial_messages: self.queue.put_nowait(message) self.subscribed_channels: List[str] = [] @@ -112,14 +114,14 @@ class _QueuePubSub: async def subscribe(self, *channels: str) -> None: self.subscribed_channels.extend(channels) - async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, str] | None: if timeout == 0: try: - return self.queue.get_nowait() + return {"type": "message", "data": self.queue.get_nowait()} except asyncio.QueueEmpty: return None try: - return await asyncio.wait_for(self.queue.get(), timeout) + return {"type": "message", "data": await asyncio.wait_for(self.queue.get(), timeout)} except asyncio.TimeoutError: return None @@ -128,7 +130,7 @@ class _QueuePubSub: class _BrokenPubSub(_QueuePubSub): - async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, str] | None: raise ConnectionError("connection lost") @@ -142,7 +144,7 @@ class _EmptyPollsThenMessagePubSub(_QueuePubSub): super().__init__(initial_messages=initial_messages) self.remaining_empty_polls = empty_polls - async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, str] | None: if timeout != 0 and self.remaining_empty_polls > 0: self.remaining_empty_polls -= 1 return None @@ -165,19 +167,20 @@ class _ScriptedPubSubRedisClient(Redis): return next(self._scripted_pubsubs) -class _FakeRedisCache: - def __init__(self, client: object, namespace: Optional[str] = None) -> None: +class _FakeRedisCache(RedisCache): + def __init__(self, client: object, namespace: str | None = None) -> None: self._client = client self.namespace = namespace - def init_pubsub_client(self) -> object: - return self._client + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", self._client) -class _ExplodingRedisCache: - namespace: Optional[str] = None +class _ExplodingRedisCache(RedisCache): + def __init__(self) -> None: + self.namespace = None - def init_pubsub_client(self) -> object: + def init_pubsub_client(self) -> "Redis[bytes]": raise ConnectionError("cannot connect") @@ -211,20 +214,27 @@ async def test_publish_uses_namespaced_channel() -> None: client = _RecordingRedisClient() cache = _FakeRedisCache(client, namespace="prod-eu") - await publish_config_change(redis_cache=cache, object_type="litellm_credentialstable") + with pytest.warns(FutureWarning, match="experimental.*may change"): + await publish_config_change(redis_cache=cache, object_type="litellm_credentialstable") assert client.published[0][0] == "prod-eu:litellm_proxy.config_change" -async def test_publish_swallows_redis_publish_errors() -> None: +async def test_publish_swallows_redis_publish_errors(caplog: pytest.LogCaptureFixture) -> None: cache = _FakeRedisCache(_FailingPublishRedisClient()) await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + assert "config sync publish for litellm_proxymodeltable failed:" in caplog.text + assert "redis down" in caplog.text -async def test_publish_swallows_client_init_errors() -> None: + +async def test_publish_swallows_client_init_errors(caplog: pytest.LogCaptureFixture) -> None: await publish_config_change(redis_cache=_ExplodingRedisCache(), object_type="litellm_proxymodeltable") + assert "config sync publish for litellm_proxymodeltable failed:" in caplog.text + assert "cannot connect" in caplog.text + async def test_publish_reaches_cluster_derived_pubsub_clients() -> None: """LIT-8543: a cluster-backed cache returns a node-level client from @@ -621,6 +631,126 @@ async def test_subscriber_subscribes_on_cluster_derived_pubsub_client() -> None: assert resyncs == ["resync"] +class _PollFailingPubSub(_QueuePubSub): + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, str] | None: + if timeout != 0: + raise ConnectionError("connection lost while polling") + return None + + +class _DrainFailingPubSub(_QueuePubSub): + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> dict[str, str] | None: + if timeout == 0: + raise ConnectionError("connection lost while draining") + return await super().get_message(ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout) + + +class _SubscribeFailingPubSub(_QueuePubSub): + async def subscribe(self, *channels: str) -> None: + raise ConnectionError("subscribe failed") + + +class _BlockingSubscribePubSub(_QueuePubSub): + def __init__(self) -> None: + super().__init__() + self.started = asyncio.Event() + self.release = asyncio.Event() + + async def subscribe(self, *channels: str) -> None: + self.started.set() + await self.release.wait() + + +class _SleepRecorder: + def __init__(self) -> None: + self.delays: tuple[float, ...] = () + + async def __call__(self, seconds: float) -> None: + self.delays += (seconds,) + await asyncio.sleep(0) + + +async def test_drain_failure_reconnects_before_reloading_missed_changes() -> None: + broken: Final = _DrainFailingPubSub(initial_messages=("change",)) + healthy: Final = _QueuePubSub() + recovered: Final = asyncio.Event() + closed_before_reload: Final[asyncio.Queue[bool]] = asyncio.Queue() + + async def reload_config() -> None: + closed_before_reload.put_nowait(broken.closed) + recovered.set() + + subscriber: Final = ConfigSyncSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubClient((broken, healthy))), + resync_callbacks=(reload_config,), + debounce_seconds=0, + jitter_max_seconds=0, + backoff_initial_seconds=0, + sleep=_SleepRecorder(), + ) + subscriber.start() + try: + await asyncio.wait_for(recovered.wait(), timeout=1) + finally: + await subscriber.stop() + + assert closed_before_reload.qsize() == 1 + assert closed_before_reload.get_nowait() is True + assert healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert healthy.closed is True + + +async def test_subscription_backoff_caps_and_resets_after_recovery() -> None: + failed_subscriptions: Final = tuple(_SubscribeFailingPubSub() for _ in range(3)) + broken: Final = _PollFailingPubSub() + healthy: Final = _QueuePubSub() + sleeps: Final = _SleepRecorder() + recovered: Final = asyncio.Event() + reloads: Final[asyncio.Queue[None]] = asyncio.Queue() + + async def reload_config() -> None: + reloads.put_nowait(None) + if reloads.qsize() == 2: + recovered.set() + + subscriber: Final = ConfigSyncSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubClient((*failed_subscriptions, broken, healthy))), + resync_callbacks=(reload_config,), + debounce_seconds=0, + jitter_max_seconds=0, + min_resync_interval_seconds=0, + backoff_initial_seconds=2, + backoff_max_seconds=5, + sleep=sleeps, + ) + subscriber.start() + try: + await asyncio.wait_for(recovered.wait(), timeout=1) + finally: + await subscriber.stop() + + assert sleeps.delays == (2, 4, 5, 0, 2, 0) + assert tuple(subscription.closed for subscription in failed_subscriptions) == (True, True, True) + assert broken.closed is True + assert healthy.closed is True + assert reloads.qsize() == 2 + + +async def test_stop_closes_a_subscription_cancelled_during_subscribe() -> None: + pubsub: Final = _BlockingSubscribePubSub() + subscriber: Final = ConfigSyncSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubClient((pubsub,))), resync_callbacks=() + ) + subscriber.start() + try: + await asyncio.wait_for(pubsub.started.wait(), timeout=1) + finally: + await subscriber.stop() + + assert pubsub.closed is True + assert pubsub.release.is_set() is False + + class _FakeTableActions: def __init__(self, calls: List[Tuple[str, str]]) -> None: self._calls = calls diff --git a/tests/unit/proxy/common_utils/test_debug_utils.py b/tests/unit/proxy/common_utils/test_debug_utils.py index 456be4e8c00..a60a32b3ad1 100644 --- a/tests/unit/proxy/common_utils/test_debug_utils.py +++ b/tests/unit/proxy/common_utils/test_debug_utils.py @@ -4,11 +4,17 @@ import socket from collections.abc import Callable, Iterator, Mapping from dataclasses import asdict from pathlib import Path +from types import SimpleNamespace +from typing import Final +from unittest.mock import MagicMock, patch import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from redis import Redis +from redis.connection import Connection +from litellm.caching.redis_cache import RedisCache from litellm.proxy import proxy_server from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -16,6 +22,7 @@ from litellm.proxy.bug_report_config import build_proxy_bug_report from litellm.proxy.common_utils.debug_utils import ( PSUTIL_MISSING_ERROR, _ProcFilesystemProcess, + _redis_memory_stats, _summary_process_memory, get_memory_summary, ) @@ -198,3 +205,53 @@ def test_memory_routes_count_prompt_caching_pins( assert response.status_code == 200 assert pin_count(response.json()) == expected + + +@pytest.mark.parametrize( + ("pool", "expected"), + [ + pytest.param( + SimpleNamespace( + max_connections=23, + connection_class=Connection, + _available_connections=(object(),), + _in_use_connections=(object(), object()), + ), + { + "enabled": True, + "cache_type": "RedisCache", + "connection_pool": {"max_connections": 23, "connection_class": Connection.__name__}, + }, + id="pool-metadata-preserves-diagnostic-shape", + ), + pytest.param( + SimpleNamespace(), + { + "enabled": True, + "cache_type": "RedisCache", + "connection_pool": {"max_connections": None, "connection_class": None}, + }, + id="unavailable-pool-fields", + ), + pytest.param(None, {"enabled": True, "cache_type": "RedisCache"}, id="no-pool"), + pytest.param( + SimpleNamespace(max_connections=object()), + {"enabled": True, "cache_type": "RedisCache"}, + id="invalid-metadata-remains-best-effort", + ), + ], +) +def test_redis_memory_diagnostics_preserve_pool_shape(pool: object, expected: Mapping[str, object]) -> None: + client: Final = MagicMock(spec=Redis) + client.connection_pool = pool + with patch("redis.Redis", return_value=client): + cache: Final = RedisCache(host="diagnostic-pool", port=6379) + + with pytest.warns(FutureWarning, match="get_connection_pool_info is experimental and may change"): + result: Final = _redis_memory_stats(cache) + + assert result == expected + + +def test_redis_memory_diagnostics_without_cache_remain_disabled() -> None: + assert _redis_memory_stats(None) == {"enabled": False} diff --git a/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py index 4684c3213d6..dbc98d984c9 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -1,13 +1,15 @@ import json import logging -from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock, patch +from collections.abc import AsyncIterator, Sequence +from typing import Final +from unittest.mock import AsyncMock, patch import pytest -from fastapi.testclient import TestClient +import pytest_asyncio +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeRedis as FakeAsyncRedis - -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager @@ -322,144 +324,64 @@ async def test_lock_takeover_race_condition(mock_redis): 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) +@pytest_asyncio.fixture(loop_scope="function") +async def atomic_lock_backend() -> AsyncIterator[tuple[RedisCache, FakeAsyncRedis, AsyncMock]]: + server: Final = FakeServer() + transport: Final = FakeAsyncRedis(server=server) - await pod_lock_manager.release_lock(cronjob_id="test_job") + async def compare_delete(keys: Sequence[str], args: Sequence[str], client: object = None) -> int: + if await transport.get(keys[0]) == args[0].encode(): + return await transport.delete(keys[0]) + return 0 - 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() + release: Final = AsyncMock(side_effect=compare_delete) + with ( + patch("litellm._redis.get_redis_client", return_value=FakeRedis(server=server)), + patch("litellm._redis.get_redis_async_client", return_value=transport), + patch.object(transport, "register_script", return_value=release), + ): + cache: Final = RedisCache(host="127.0.0.1", port=6379, namespace="pod-test") + yield cache, transport, release + await transport.aclose() @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) +async def test_release_lock_deletes_json_encoded_lock_and_allows_the_next_pod(atomic_lock_backend) -> None: + cache, transport, release = atomic_lock_backend + holder: Final = PodLockManager(redis_cache=cache) + other: Final = PodLockManager(redis_cache=cache) + assert await holder.acquire_lock(cronjob_id="spend") is True + assert await transport.get("pod-test:cronjob_lock:spend") == json.dumps(holder.pod_id).encode() - await pod_lock_manager.release_lock(cronjob_id="test_job") - await pod_lock_manager.release_lock(cronjob_id="test_job") + await holder.release_lock(cronjob_id="spend") - assert mock_redis.async_register_script.call_count == 1 + assert await transport.get("pod-test:cronjob_lock:spend") is None + assert await other.acquire_lock(cronjob_id="spend") is True + release.assert_awaited_once_with(keys=("pod-test:cronjob_lock:spend",), args=(json.dumps(holder.pod_id),), client=None) @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) +async def test_release_lock_preserves_lock_held_by_other_pod(atomic_lock_backend) -> None: + cache, transport, _ = atomic_lock_backend + holder: Final = PodLockManager(redis_cache=cache) + other: Final = PodLockManager(redis_cache=cache) + assert await holder.acquire_lock(cronjob_id="spend") is True - with patch.object(pod_lock_manager, "_emit_released_lock_event") as mock_emit: - await pod_lock_manager.release_lock(cronjob_id="test_job") + await other.release_lock(cronjob_id="spend") - 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 + assert await transport.get("pod-test:cronjob_lock:spend") == json.dumps(holder.pod_id).encode() @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") +async def test_release_lock_falls_back_to_get_del_when_atomic_execution_fails(atomic_lock_backend) -> None: + cache, transport, release = atomic_lock_backend + holder: Final = PodLockManager(redis_cache=cache) + assert await holder.acquire_lock(cronjob_id="spend") is True + release.side_effect = OSError("scripting unavailable") - acquired = await pod.acquire_lock(cronjob_id="db_spend_update_job") - assert acquired is True - assert lock_key in redis.store + await holder.release_lock(cronjob_id="spend") - 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 + assert await transport.get("pod-test:cronjob_lock:spend") is None @pytest.mark.asyncio diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 24c12e06c9c..e743fec1537 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,10 +1,14 @@ import json +from collections.abc import Iterator from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeRedis as FakeAsyncRedis +from litellm.caching.redis_cache import RedisCache from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyTagSpendTransaction @@ -388,7 +392,9 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): @pytest.mark.asyncio -async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache): +async def test_restore_transactions_to_redis_pushes_only_provided( + update_queue_cache: RedisCache, update_queue_server: FakeServer, +) -> None: """ restore_transactions_to_redis re-pushes only the transaction sets it was given, to their matching buffer keys, so uncommitted spend can be retried. @@ -398,37 +404,27 @@ async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_b REDIS_UPDATE_BUFFER_KEY, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1, 1]) + db_spend: Final = {"key_list_transactions": {"key1": 1.0}} + daily_user: Final = {"user_key1": {"spend": 1.0}} - db_spend = {"key_list_transactions": {"key1": 1.0}} - daily_user = {"user_key1": {"spend": 1.0}} - - await redis_update_buffer.restore_transactions_to_redis( + await RedisUpdateBuffer(redis_cache=update_queue_cache).restore_transactions_to_redis( db_spend_update_transactions=db_spend, daily_spend_update_transactions=daily_user, ) - mock_redis_cache.async_rpush_pipeline.assert_called_once() - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - pushed_keys = {op["key"] for op in rpush_list} - assert pushed_keys == { - REDIS_UPDATE_BUFFER_KEY, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - } - # Payloads round-trip through the same JSON encoding used on the store path - payloads = {op["key"]: json.loads(op["values"][0]) for op in rpush_list} - assert payloads[REDIS_UPDATE_BUFFER_KEY] == db_spend - assert payloads[REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY] == daily_user + with FakeRedis(server=update_queue_server) as transport: + assert set(transport.keys()) == { + f"updates:{REDIS_UPDATE_BUFFER_KEY}".encode(), f"updates:{REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY}".encode(), + } + assert json.loads(await update_queue_cache.async_lpop(REDIS_UPDATE_BUFFER_KEY)) == db_spend + assert json.loads(await update_queue_cache.async_lpop(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY)) == daily_user @pytest.mark.asyncio -async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): +async def test_restored_window_spend_transactions_drain_back_unchanged(update_queue_cache: RedisCache) -> None: """A window commit that fails after the destructive lpop must be re-pushed in the store path's encoding, so the next drain returns the same increments.""" - from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY - from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( - build_window_spend_transaction, - ) + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import build_window_spend_transaction window_transactions = ( build_window_spend_transaction( @@ -440,17 +436,9 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), ), ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) - - await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) - - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY] - - mock_redis_cache.async_lpop_pipeline = AsyncMock( - return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])] - ) - drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + buffer: Final = RedisUpdateBuffer(redis_cache=update_queue_cache) + await buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) + drained: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline() assert drained[6] == window_transactions @@ -458,23 +446,21 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd @pytest.mark.asyncio async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache): """Nothing to restore -> no Redis call.""" - mock_redis_cache.async_rpush_pipeline = AsyncMock() + mock_redis_cache.async_enqueue_updates = AsyncMock() await redis_update_buffer.restore_transactions_to_redis() - mock_redis_cache.async_rpush_pipeline.assert_not_called() + mock_redis_cache.async_enqueue_updates.assert_not_called() @pytest.mark.asyncio -async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): +async def test_restore_transactions_to_redis_swallows_redis_error( + update_queue_cache: RedisCache, update_queue_server: FakeServer, caplog: pytest.LogCaptureFixture, +) -> None: """A Redis failure during restore must not propagate to the caller's finally block.""" - from redis.exceptions import RedisError - - mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) - - await redis_update_buffer.restore_transactions_to_redis( + update_queue_server.connected = False + assert await RedisUpdateBuffer(redis_cache=update_queue_cache).restore_transactions_to_redis( db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, - ) - - mock_redis_cache.async_rpush_pipeline.assert_called_once() + ) is None + assert "failed to restore uncommitted transactions to Redis" in caplog.text def test_validate_redis_transaction_buffer_raises_without_redis(): @@ -641,20 +627,39 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up ] +@pytest.fixture +def update_queue_server() -> FakeServer: + return FakeServer() + + +@pytest.fixture +def update_queue_cache(update_queue_server: FakeServer) -> Iterator[RedisCache]: + with ( + patch( # test-quality-ok: Redis transport has no constructor injection point + "litellm._redis.get_redis_client", return_value=FakeRedis(server=update_queue_server) + ), + patch( # test-quality-ok: only transport is replaced, cache serialization remains real + "litellm._redis.get_redis_async_client", return_value=FakeAsyncRedis(server=update_queue_server) + ), + ): + yield RedisCache(host="127.0.0.1", port=6379, namespace="updates") + + @pytest.mark.asyncio -async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_update_buffer, mock_redis_cache): +async def test_budget_window_payloads_keep_request_ids_for_older_workers(update_queue_cache: RedisCache) -> None: """A leader from before the field was dropped indexes request_ids while merging what it popped, and the pop is destructive, so a payload without the key would cost a rolling deploy those increments.""" from datetime import datetime, timezone + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendUpdateQueue, build_window_spend_transaction, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) - window_queue = WindowSpendUpdateQueue() + buffer: Final = RedisUpdateBuffer(redis_cache=update_queue_cache) + window_queue: Final = WindowSpendUpdateQueue() await window_queue.add_update( build_window_spend_transaction( entity_type="key", @@ -665,12 +670,13 @@ async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_u ) ) - await redis_update_buffer.restore_transactions_to_redis( + await buffer.restore_transactions_to_redis( window_spend_update_transactions=await window_queue.flush_and_get_aggregated_window_spend_transactions(), ) - rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] - restored = json.loads(rpush_list[0]["values"][0]) + queued: Final = await update_queue_cache.async_lpop(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, count=1) + assert len(queued) == 1 + restored: Final = json.loads(queued[0]) assert [payload["request_ids"] for payload in restored] == [[]] @@ -771,15 +777,16 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ @pytest.mark.asyncio -async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails(): +async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails( + update_queue_cache: RedisCache, update_queue_server: FakeServer, +) -> None: from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue - redis_cache = MagicMock() - redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down")) - queue = ToolDiscoveryQueue() + update_queue_server.connected = False + queue: Final = ToolDiscoveryQueue() queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) - await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue) + await RedisUpdateBuffer(redis_cache=update_queue_cache).store_tool_discoveries_in_redis(queue) assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}] diff --git a/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py b/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py index 40d49f5ab95..e540fbf92dd 100644 --- a/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py @@ -12,11 +12,12 @@ import uuid from collections.abc import Mapping, Sequence from types import MappingProxyType, SimpleNamespace from typing import Final +from unittest.mock import MagicMock, patch import pytest from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.batch_enqueued_tokens import ( @@ -127,7 +128,7 @@ async def test_zero_token_reserve_charges_nothing(): assert isinstance(full, BatchEnqueuedTokenReservation) -class _SingleKeyRedisFake: +class _SingleKeyRedisFake(RedisCache): """Emulates the Redis script path one single-key call at a time, recording every call.""" def __init__( @@ -137,6 +138,11 @@ class _SingleKeyRedisFake: fail_save_keys: frozenset[str] = frozenset(), raise_after_landing_save_keys: frozenset[str] = frozenset(), ) -> None: + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")), + ): + super().__init__(host=f"batch-test-{uuid.uuid4().hex}", port=6379) self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = () self.save_ttls: tuple[int, ...] = () self.counters: Mapping[str, int] = MappingProxyType({}) @@ -146,14 +152,17 @@ class _SingleKeyRedisFake: self.fail_save_keys = fail_save_keys self.raise_after_landing_save_keys = raise_after_landing_save_keys - def async_register_script(self, script: str): + def _get_async_client(self): + return self + + def register_script(self, script: str): kind: Final = ( "reserve" if "INCRBY" in script else "refund" if "DECRBY" in script else "pop" if "GET" in script else "save" ) - async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object: + async def run(keys: Sequence[str], args: Sequence[str | bytes | int | float], client=None) -> object: self.script_calls = (*self.script_calls, (kind, tuple(keys))) return self._run(kind, tuple(keys), tuple(args)) @@ -401,9 +410,9 @@ def test_batch_response_view_accepts_batch_objects_only(): assert batch_response_view("batch_1") is None -class _OpenBreakerRedis: - def async_register_script(self, script: str): - async def refused(keys: Sequence[str], args: Sequence[str | bytes | int | float]) -> object: +class _OpenBreakerRedis(_SingleKeyRedisFake): + def register_script(self, script: str): + async def refused(keys: Sequence[str], args: Sequence[str | bytes | int | float], client=None) -> object: raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") return refused diff --git a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py index 45cc7eb8922..af13f5eadfd 100644 --- a/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py +++ b/tests/unit/proxy/hooks/test_max_budget_per_session_limiter.py @@ -8,9 +8,9 @@ Tests that session-scoped budget tracking works correctly: - Requests without agent_id pass through """ +import logging from unittest.mock import patch -import logging import pytest from fastapi import HTTPException @@ -199,3 +199,33 @@ async def test_an_open_circuit_breaker_reads_session_spend_locally_without_a_war assert spend == 0.0 assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply,expected", [(b"3.75", 3.75), ("malformed", 0.5), (ConnectionError("offline"), 0.5)]) +async def test_session_spend_uses_typed_result_or_local_fallback(reply: object, expected: float) -> None: + from typing import Final + from unittest.mock import AsyncMock, MagicMock + from uuid import uuid4 + + from litellm.caching.redis_cache import RedisCache + + transport: Final = MagicMock() + execute: Final = AsyncMock(side_effect=[reply]) + transport.register_script.return_value = execute + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("litellm._redis.get_redis_async_client", return_value=transport), + ): + with patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")): + redis: Final = RedisCache(host=f"session-spend-{uuid4().hex}", namespace="test") + cache: Final = DualCache(redis_cache=redis) + handler: Final = PROXY_MaxBudgetPerSessionHandler(InternalUsageCache(cache)) + key: Final = handler._make_cache_key("budget-session") + with pytest.warns(FutureWarning, match="experimental.*may change"): + result: Final = await handler._increment_spend(key, 0.5) + + assert result == expected + assert tuple(execute.call_args.kwargs["keys"]) == ("test:{session_budget:budget-session}:spend",) + assert execute.call_args.kwargs["args"] == ("0.5", handler.ttl) + assert cache.in_memory_cache.get_cache(key) == (None if expected == 3.75 else 0.5) diff --git a/tests/unit/proxy/hooks/test_max_iterations_limiter.py b/tests/unit/proxy/hooks/test_max_iterations_limiter.py index 5eb499042b5..5c41e9dbfea 100644 --- a/tests/unit/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/unit/proxy/hooks/test_max_iterations_limiter.py @@ -150,3 +150,33 @@ async def test_max_iterations_no_agent_id_passes(): call_type="", ) assert result is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("reply,expected", [(b"7", 7), ("malformed", 1), (ConnectionError("offline"), 1)]) +async def test_iteration_count_uses_typed_result_or_local_fallback(reply: object, expected: int) -> None: + from typing import Final + from unittest.mock import AsyncMock + from uuid import uuid4 + + from litellm.caching.redis_cache import RedisCache + + transport: Final = MagicMock() + execute: Final = AsyncMock(side_effect=[reply]) + transport.register_script.return_value = execute + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("litellm._redis.get_redis_async_client", return_value=transport), + ): + with patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")): + redis: Final = RedisCache(host=f"iterations-{uuid4().hex}", namespace="test") + cache: Final = DualCache(redis_cache=redis) + handler: Final = PROXY_MaxIterationsHandler(InternalUsageCache(cache)) + key: Final = handler._make_cache_key("iteration-session") + with pytest.warns(FutureWarning, match="experimental.*may change"): + result: Final = await handler._increment_and_get(key) + + assert result == expected + assert tuple(execute.call_args.kwargs["keys"]) == ("test:{session_iterations:iteration-session}:count",) + assert execute.call_args.kwargs["args"] == (handler.ttl,) + assert cache.in_memory_cache.get_cache(key) == (None if expected == 7 else 1) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 345f37f1db4..9cc7003bb7c 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -10,8 +10,8 @@ import time from collections.abc import AsyncGenerator, Iterator, Sequence from contextlib import contextmanager from datetime import datetime, timedelta, timezone -from typing import Any, Dict, Final, List, Optional -from unittest.mock import patch +from typing import Any, Dict, Final, List +from unittest.mock import MagicMock, patch import httpx import pytest @@ -20,8 +20,10 @@ from pydantic import TypeAdapter import litellm from litellm import Router +from litellm.caching import _redis_scripts as redis_scripts from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.litellm_core_utils import ptu_pricing @@ -331,7 +333,7 @@ async def test_sliding_window_rate_limit_v3(monkeypatch, time_controller): results.append(new_counter) return results - parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + parallel_request_handler._redis_cache = _script_cache(mock_batch_rate_limiter) # First request should succeed await parallel_request_handler.async_pre_call_hook( @@ -424,7 +426,7 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller results.append(new_counter) return results - parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + parallel_request_handler._redis_cache = _script_cache(mock_batch_rate_limiter) # Make first request await parallel_request_handler.async_pre_call_hook( @@ -578,7 +580,7 @@ async def test_normal_router_call_tpm_v3( results.append(new_counter) return results - parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + parallel_request_handler._redis_cache = _script_cache(mock_batch_rate_limiter) monkeypatch.setattr(litellm, "callbacks", [parallel_request_handler]) # Helper to get the correct value for key construction @@ -1049,7 +1051,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): captured_calls.append((list(keys), list(args))) return [0, 3] - handler.parallel_acquire_script = fake_acquire + handler._redis_cache = _script_cache(fake_acquire, redis_scripts.SLOT_ACQUIRE) data: Dict[str, Any] = {"model": "gpt-3.5-turbo"} await handler.async_pre_call_hook( @@ -1090,7 +1092,7 @@ async def test_parallel_gauge_uses_atomic_redis_script_v3(): async def fake_acquire_over_limit(keys, args): return [1, 1, 5, 5] - handler.parallel_acquire_script = fake_acquire_over_limit + handler._redis_cache = _script_cache(fake_acquire_over_limit, redis_scripts.SLOT_ACQUIRE) with pytest.raises(HTTPException) as exc_info: await handler.async_pre_call_hook( @@ -1743,7 +1745,7 @@ async def test_async_increment_tokens_fallback_behavior(): ) # Mock the token_increment_script to None to simulate unavailable script - parallel_request_handler.token_increment_script = None + parallel_request_handler._redis_cache = None # Mock the fallback method fallback_called = False @@ -1882,7 +1884,7 @@ def test_keyslot_for_redis_cluster(): @pytest.mark.asyncio -async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): +async def test_increment_window_counters_by_group_cluster_compatibility(): """ Test that the Redis batch rate limiter script execution handles cluster compatibility by grouping keys and falling back gracefully on errors. @@ -1906,7 +1908,7 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): ), # First group fails [1234, 2], # Second group succeeds ] - handler.batch_rate_limiter_script = mock_script + handler._redis_cache = _script_cache(mock_script) # Mock in-memory fallback (returns 2 values for 2 keys: window_start, counter) handler.in_memory_cache_sliding_window = AsyncMock(return_value=[1234, 1]) @@ -1920,11 +1922,11 @@ async def test_execute_redis_batch_rate_limiter_script_cluster_compatibility(): ] # Execute the method - results = await handler._execute_redis_batch_rate_limiter_script( + results = await handler._increment_window_counters_by_group( keys_to_fetch=test_keys, now_int=1234 ) - assert results == [1234, 1, 1234, 2] + assert results == [1234, 1, "1234", 2] # Verify script was called twice (once per slot group) assert mock_script.call_count == 2 @@ -2105,7 +2107,7 @@ async def test_get_rate_limit_type_default_is_total(monkeypatch): ) # Mock general_settings to return empty dict (no token_rate_limit_type set) - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server original_settings = getattr(proxy_server, "general_settings", {}) monkeypatch.setattr(proxy_server, "general_settings", {}) @@ -2130,7 +2132,7 @@ async def test_get_rate_limit_type_invalid_falls_back_to_total(monkeypatch): ) # Mock general_settings to return an invalid token_rate_limit_type - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server original_settings = getattr(proxy_server, "general_settings", {}) monkeypatch.setattr( @@ -2345,8 +2347,8 @@ async def test_execute_token_increment_script_cluster_compatibility(): # Mock _is_redis_cluster to return True for this test with patch.object(handler, "_is_redis_cluster", return_value=True): # Mock script - mock_script = AsyncMock() - handler.token_increment_script = mock_script + mock_script = AsyncMock(side_effect=[[100.0, -1.0], [50.0]]) + handler._redis_cache = _script_cache(mock_script, redis_scripts.COUNTER_INCREMENT) # Create pipeline operations with different hash tags pipeline_operations: List[RedisPipelineIncrementOperation] = [ @@ -2360,7 +2362,7 @@ async def test_execute_token_increment_script_cluster_compatibility(): ] # Execute the method - await handler._execute_token_increment_script(pipeline_operations) + await handler._execute_token_increments(pipeline_operations) # Verify script was called (at least once, possibly more depending on slot grouping) assert mock_script.call_count >= 1, "Script should be called at least once" @@ -2883,7 +2885,7 @@ async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): results.append(new_counter) return results - parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + parallel_request_handler._redis_cache = _script_cache(mock_batch_rate_limiter) with patch( "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", @@ -4002,8 +4004,8 @@ async def test_billed_personal_key_auth_produces_team_rpm_descriptors(monkeypatc from litellm.proxy import proxy_server from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager - from litellm.proxy.auth import auth_checks from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + from litellm.proxy.auth import auth_checks server: Final = MCPServer( server_id="math_a", @@ -4400,7 +4402,7 @@ async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redi raise ConnectionError("redis unavailable") release_script = AsyncMock(side_effect=failing_release) - handler.parallel_release_script = release_script + handler._redis_cache = _script_cache(release_script, redis_scripts.SLOT_RELEASE) await local_cache.async_set_cache(key=parallel_key, value=2, local_only=True) stash = get_or_create_request_stash() stash.owner_litellm_call_id = call_id @@ -4806,7 +4808,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): captured_calls.append((list(keys), list(args))) return [3] - handler.parallel_count_script = fake_count + handler._redis_cache = _script_cache(fake_count, redis_scripts.SLOT_COUNT) response = await handler.should_rate_limit(descriptors=descriptors, read_only=True) assert captured_calls == [ @@ -4827,7 +4829,7 @@ async def test_read_only_gauge_check_counts_without_acquiring_v3(): async def failing_count(keys, args): raise ConnectionError("redis unavailable") - handler.parallel_count_script = failing_count + handler._redis_cache = _script_cache(failing_count, redis_scripts.SLOT_COUNT) await _seed_max_parallel_requests_slots( local_cache, counter_key, ["s1", "s2", "s3", "s4", "s5"] ) @@ -4856,7 +4858,7 @@ async def test_redis_release_script_updates_local_mirror_v3(): captured_calls.append((list(keys), list(args))) return [2] - handler.parallel_release_script = fake_release + handler._redis_cache = _script_cache(fake_release, redis_scripts.SLOT_RELEASE) get_or_create_request_stash().parallel_slot = ParallelSlotAcquisition( slot_id="slot-redis-test", @@ -4943,8 +4945,7 @@ async def test_in_memory_fallback_respects_mirrored_redis_count_v3(): async def failing_script(keys, args): raise ConnectionError("redis unavailable") - handler.parallel_acquire_script = failing_script - handler.parallel_release_script = failing_script + handler._redis_cache = _TransportCache(_UnreachableRedis()) await local_cache.async_set_cache(key=counter_key, value=5, local_only=True) with pytest.raises(HTTPException) as exc_info: @@ -5068,7 +5069,7 @@ async def test_async_data_generator_releases_counter_on_disconnect_v3(disconnect gap that let a disconnect leak the slot in the default limiter-only config. A mid-stream disconnect must still refund the pre-call +1. """ - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server limiter, cache, counter_key, user_api_key_dict = await _build_seeded_limiter() proxy_logging_obj = proxy_server.proxy_logging_obj @@ -5122,8 +5123,8 @@ async def test_async_data_generator_releases_counter_when_wrapped_v3(): from the outer generator: the counter returns to 0 (not -1), proving the nested hook does not also refund and there is no double decrement. """ - import litellm.proxy.proxy_server as proxy_server from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy import proxy_server class _PassthroughIteratorOverride(CustomLogger): async def async_post_call_streaming_iterator_hook( @@ -7067,16 +7068,45 @@ async def test_post_call_success_hook_leaves_raw_provider_dict_untouched(): assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []} +class _TransportCache(RedisCache): + def __init__(self, transport): + self._transport = transport + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")), + ): + super().__init__(host=f"limiter-test-{id(self)}", port=6379) + def _get_async_client(self): + return self._transport + + +class _SingleScriptTransport: + def __init__(self, run, source): + self._run = run + self._source = source + + def register_script(self, source): + async def run(keys, args, client=None): + if source != self._source: + raise ConnectionError("Only the configured operation is available") + return await self._run(keys=list(keys), args=list(args)) + return run + + +def _script_cache(run, source=redis_scripts.WINDOW_INCREMENT): + return _TransportCache(_SingleScriptTransport(run, source)) + + class _OpenBreakerRedis: - def async_register_script(self, script: str): - async def refused(keys, args): + def register_script(self, script: str): + async def refused(keys, args, client=None): from litellm.caching.redis_cache import RedisCircuitBreakerOpenError raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") return refused - async def async_increment_pipeline(self, increment_list, **kwargs): + def pipeline(self, transaction=False): from litellm.caching.redis_cache import RedisCircuitBreakerOpenError raise RedisCircuitBreakerOpenError("Redis circuit breaker is open") @@ -7087,7 +7117,7 @@ async def test_an_open_circuit_breaker_falls_back_to_the_pipeline_without_a_warn from litellm.types.caching import RedisPipelineIncrementOperation handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_OpenBreakerRedis())) # pyright: ignore[reportArgumentType] # duck-typed Redis double + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_TransportCache(_OpenBreakerRedis()))) # pyright: ignore[reportArgumentType] # duck-typed Redis double ) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): @@ -7102,11 +7132,11 @@ async def test_an_open_circuit_breaker_falls_back_to_the_pipeline_without_a_warn @pytest.mark.asyncio async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_a_warning(caplog): handler = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_OpenBreakerRedis())) # pyright: ignore[reportArgumentType] # duck-typed Redis double + internal_usage_cache=InternalUsageCache(DualCache(redis_cache=_TransportCache(_OpenBreakerRedis()))) # pyright: ignore[reportArgumentType] # duck-typed Redis double ) with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): - values = await handler._execute_redis_batch_rate_limiter_script( + values = await handler._increment_window_counters_by_group( ["{quiet}:window", "{quiet}:counter"], now_int=int(time.time()) ) @@ -7116,8 +7146,8 @@ async def test_an_open_circuit_breaker_reads_the_sliding_window_locally_without_ class _UnreachableRedis: - def async_register_script(self, script: str): - async def refused(keys, args): + def register_script(self, script: str): + async def refused(keys, args, client=None): raise ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.") return refused @@ -7144,29 +7174,29 @@ class _ScriptedRedis: self.increments.append((key, value)) return value - def async_register_script(self, script: str): - from litellm.proxy.hooks import parallel_request_limiter_v3 as v3 + def register_script(self, script: str): + from litellm.caching import _redis_scripts as scripts - async def run(keys, args): + async def run(keys, args, client=None): if script == self.failing_script: raise ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.") - if script == v3.WINDOW_GUARDED_TOKEN_INCREMENT_SCRIPT: + if script == scripts.GUARDED_INCREMENT: self.guarded_increments.append((list(keys), list(args))) return [1, 0] * (len(keys) // 2) - if script == v3.BATCH_RATE_LIMITER_SCRIPT: + if script == scripts.WINDOW_INCREMENT: self.batch_calls += 1 self.batch_call_keys.append(list(keys)) self.batch_call_args.append(list(args)) if self.batch_calls == self.failing_batch_call: raise ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.") return [args[0], self.batch_calls] * (len(keys) // 2) - if script == v3.BATCH_COUNTER_READ_SCRIPT: + if script == scripts.COUNTER_READ: return [int(time.time()) if key.endswith(":window") else self.stored_counter_value for key in keys] - if script == v3.PARALLEL_COUNT_SCRIPT: + if script == scripts.SLOT_COUNT: return [0 for _ in keys] - if script == v3.PARALLEL_ACQUIRE_SCRIPT: + if script == scripts.SLOT_ACQUIRE: return [0, *[1 for _ in keys]] - if script == v3.PARALLEL_RELEASE_SCRIPT: + if script == scripts.SLOT_RELEASE: self.released_slots.append((list(keys), list(args))) return [0 for _ in keys] raise AssertionError(f"unexpected script: {script[:60]}") @@ -7179,7 +7209,7 @@ def _handler_with_redis( fail_closed: bool | None = None, force_hash_tag_grouping: bool | None = None, ): - internal_usage_cache = InternalUsageCache(DualCache(redis_cache=redis)) # pyright: ignore[reportArgumentType] # duck-typed Redis double + internal_usage_cache = InternalUsageCache(DualCache(redis_cache=_TransportCache(redis))) if fail_closed is None and force_hash_tag_grouping is None: return PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache) if fail_closed is None: @@ -7213,7 +7243,7 @@ async def test_redis_batch_rate_limiter_respects_hash_tag_grouping_opt_in(force_ "{end_user:u28551}:requests", ] - await handler._execute_redis_batch_rate_limiter_script(keys_to_fetch=keys, now_int=1234) + await handler._increment_window_counters_by_group(keys_to_fetch=keys, now_int=1234) expected_calls: Final = ( [ @@ -7245,11 +7275,11 @@ async def test_redis_batch_rate_limiter_returns_values_in_caller_order(): "{end_user:u28551}:requests", ] - results: Final = await handler._execute_redis_batch_rate_limiter_script( + results: Final = await handler._increment_window_counters_by_group( keys_to_fetch=keys_to_fetch, now_int=now ) - assert results == [now, 1, now, 2, now, 1] + assert results == [str(now), 1, str(now), 2, str(now), 1] def test_default_hash_tag_grouping_resolver_uses_litellm_setting(monkeypatch: pytest.MonkeyPatch): @@ -7345,9 +7375,9 @@ async def test_fail_closed_is_a_no_op_while_redis_answers(): @pytest.mark.asyncio async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired(): - from litellm.proxy.hooks import parallel_request_limiter_v3 as v3 + from litellm.caching import _redis_scripts as scripts - redis = _ScriptedRedis(failing_script=v3.CHECK_AND_INCREMENT_BY_N_SCRIPT) + redis = _ScriptedRedis(failing_script=scripts.WINDOW_RESERVE) handler = _handler_with_redis(redis, fail_closed=True) auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-slot"), max_parallel_requests=1, tpm_limit=1000) data = {"model": "test-model", "messages": [{"role": "user", "content": "hi"}]} @@ -7368,7 +7398,7 @@ async def test_fail_closed_tpm_rejection_releases_the_parallel_slot_it_acquired( @pytest.mark.asyncio async def test_fail_closed_rate_limit_enforcement_is_read_from_general_settings(monkeypatch): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-settings"), rpm_limit=2) @@ -7390,7 +7420,7 @@ async def test_fail_closed_rate_limit_enforcement_is_read_from_general_settings( async def test_fail_closed_rate_limit_enforcement_coerces_the_general_settings_value( monkeypatch, configured_value, rejects ): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server auth = UserAPIKeyAuth(api_key=hash_token("sk-fail-closed-coerced"), rpm_limit=2) monkeypatch.setitem(proxy_server.general_settings, "fail_closed_rate_limit_enforcement", configured_value) @@ -7469,34 +7499,32 @@ async def test_batch_increment_refunds_counters_already_applied_when_a_later_clu @pytest.mark.parametrize("fail_closed", [True, False], ids=["fail_closed", "fail_open"]) @pytest.mark.asyncio async def test_batch_increment_refunds_pipelined_groups_declared_after_the_one_that_failed(fail_closed): - from unittest.mock import patch + from unittest.mock import AsyncMock + + from litellm.caching.redis_batch import request_redis_batch_scope redis = _ScriptedRedis() - handler = _handler_with_redis(redis, fail_closed=fail_closed) - now = int(time.time()) + pipeline = MagicMock() + pipeline.__aenter__.return_value = pipeline + pipeline.execute = AsyncMock(return_value=[ConnectionError("Unavailable slot"), [1234, 1]]) + redis.pipeline = MagicMock(return_value=pipeline) + handler = _handler_with_redis(redis, fail_closed=fail_closed, force_hash_tag_grouping=True) groups = {"a": ["{a}:window", "{a}:requests"], "b": ["{b}:window", "{b}:requests"]} - loop = asyncio.get_running_loop() - failed_group = loop.create_future() - failed_group.set_exception(ConnectionError("Error 61 connecting to 127.0.0.1:6379. Connection refused.")) - landed_group = loop.create_future() - landed_group.set_result([now, 1]) - with ( - patch.object(handler, "_group_keys_by_hash_tag", return_value=groups), - patch.object(handler, "_pipeline_scripts", return_value=[failed_group, landed_group]), - ): + with request_redis_batch_scope(): if fail_closed: with pytest.raises(HTTPException) as exc: - await handler._execute_redis_batch_rate_limiter_script( - keys_to_fetch=[*groups["a"], *groups["b"]], now_int=now + await handler._increment_window_counters_by_group( + keys_to_fetch=[*groups["a"], *groups["b"]], now_int=1234 ) assert exc.value.status_code == 503 else: - await handler._execute_redis_batch_rate_limiter_script( - keys_to_fetch=[*groups["a"], *groups["b"]], now_int=now + await handler._increment_window_counters_by_group( + keys_to_fetch=[*groups["a"], *groups["b"]], now_int=1234 ) - assert redis.guarded_increments == ([(groups["b"], [str(now), -1, 0])] if fail_closed else []) + pipeline.execute.assert_awaited_once_with(raise_on_error=False) + assert redis.guarded_increments == ([(groups["b"], ["1234", -1, 0])] if fail_closed else []) @pytest.mark.parametrize( @@ -7572,7 +7600,7 @@ async def test_request_capacity_enforces_shared_rpm_scopes( @pytest.mark.asyncio async def test_request_capacity_keeps_dynamic_rpm_policy(monkeypatch): - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server router = Router(model_list=[{ "model_name": "test-model", @@ -7945,7 +7973,7 @@ async def _rpm_request(handler: _PROXY_MaxParallelRequestsHandler, cache: DualCa async def test_model_group_alias_shares_deployment_default_rpm_bucket_with_its_target( monkeypatch: pytest.MonkeyPatch, first_name: str, second_name: str ) -> None: - import litellm.proxy.proxy_server as proxy_server + from litellm.proxy import proxy_server router: Final = Router( model_list=[ diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 3d1e494e36c..b735d59ff30 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -16,26 +16,31 @@ Redis. import asyncio import json +from collections.abc import Iterator +from contextlib import contextmanager from datetime import datetime, timedelta from typing import Any, Dict, Final +from unittest.mock import AsyncMock, MagicMock, patch +from uuid import uuid4 import pytest from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _AUDIO_BYTES_PER_TOKEN, PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, RateLimitedModel, - _AUDIO_BYTES_PER_TOKEN, - PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, -) -from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _call_id_from_callback_kwargs, _request_stash, get_or_create_request_stash, get_request_stash, ) +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3 as RateLimitHandler, +) from litellm.proxy.utils import InternalUsageCache, hash_token from litellm.types.llms.openai import ( InputTokensDetails, @@ -701,13 +706,8 @@ async def test_contentless_request_reserves_minimum(rate_limiter): @pytest.mark.asyncio -async def test_atomic_keys_share_hash_tag_per_descriptor(rate_limiter): - """ - Cluster safety: every key in a single descriptor's Lua payload must - share a `{key:value}` hash tag so the call lands on a single Redis - Cluster slot. Otherwise the Lua script raises CROSSSLOT in cluster mode. - """ - handler, _cache = rate_limiter +async def test_atomic_keys_share_hash_tag_per_descriptor(): + """Every atomic reservation groups one descriptor's keys on one cluster slot.""" descriptors = [ { @@ -727,21 +727,17 @@ async def test_atomic_keys_share_hash_tag_per_descriptor(rate_limiter): ] increments = [{"requests": 1, "tokens": 10}, {"tokens": 10}] - for descriptor, inc in zip(descriptors, increments): - keys, _args, _meta = handler._build_descriptor_atomic_payload( - descriptor=descriptor, - increment_amounts=inc, - ) - # All keys in a descriptor's payload must share the same {tag} - # — that's the prefix between the first '{' and '}'. - tags = {k[: k.index("}") + 1] for k in keys} - assert len(tags) == 1, ( - f"Descriptor {descriptor['key']}:{descriptor['value']} produced " - f"keys spanning multiple hash tags: {tags}. Redis Cluster would " - f"reject this Lua call with CROSSSLOT." - ) - expected_tag = f"{{{descriptor['key']}:{descriptor['value']}}}" - assert tags == {expected_tag}, f"Expected hash tag {expected_tag}, got {tags}" + executor: Final = AsyncMock(side_effect=[[0, 1, 10, 1234, 1234], [0, 10, 1234]]) + with _redis_cache_for_executor(executor) as redis_cache: + handler: Final = RateLimitHandler(internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache))) + response: Final = await handler.atomic_check_and_increment_by_n(descriptors, increments) + + assert response["overall_code"] == "OK" + assert executor.await_count == len(descriptors) + for descriptor, call in zip(descriptors, executor.call_args_list): + assert {key[: key.index("}") + 1] for key in call.kwargs["keys"]} == { + f"{{{descriptor['key']}:{descriptor['value']}}}" + } @pytest.mark.asyncio @@ -1671,89 +1667,70 @@ async def test_project_reconciliation_decrements_its_active_window(rate_limiter) assert float(await cache.async_get_cache(key=counter_key, local_only=True) or 0) == 10 +@contextmanager +def _redis_cache_for_executor(executor: AsyncMock) -> Iterator[RedisCache]: + transport: Final = MagicMock() + transport.register_script.return_value = executor + with ( + patch("litellm._redis.get_redis_client", return_value=MagicMock()), + patch("litellm._redis.get_redis_async_client", return_value=transport), + ): + with patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")): + cache: Final = RedisCache(host=f"tpm-concurrent-{uuid4().hex}") + yield cache + + @pytest.mark.asyncio -async def test_redis_window_guard_uses_reservation_identity_and_never_falls_back_negative( - rate_limiter, -): - handler, _cache = rate_limiter - calls = [] - - async def failing_guard(*, keys, args): - calls.append((keys, args)) - raise RuntimeError("redis unavailable") - - unguarded_calls = [] - - async def capture_unguarded(pipeline_operations, **_kwargs): - unguarded_calls.extend(pipeline_operations) - - handler.window_guarded_token_increment_script = failing_guard - handler.async_increment_tokens_with_ttl_preservation = capture_unguarded - await handler.async_increment_reservation_aware_tokens( - pipeline_operations=[ - { - "key": "{model_per_project_itpm:project:model}:tokens", - "increment_value": -90, - "ttl": 60, - "window_key": "{model_per_project_itpm:project:model}:window", - "expected_window_start": "1234", - "reservation_backend": "redis", - } - ] - ) - - assert calls == [ - ( - [ - "{model_per_project_itpm:project:model}:window", - "{model_per_project_itpm:project:model}:tokens", - ], - ["1234", -90, 60], +async def test_redis_window_guard_uses_reservation_identity_and_never_falls_back_negative(): + executor: Final = AsyncMock(side_effect=ConnectionError("redis unavailable")) + with _redis_cache_for_executor(executor) as redis_cache: + cache: Final = DualCache(redis_cache=redis_cache) + cache.in_memory_cache.set_cache("{model_per_project_itpm:project:model}:tokens", 100) + handler: Final = RateLimitHandler(internal_usage_cache=InternalUsageCache(cache)) + await handler.async_increment_reservation_aware_tokens( + pipeline_operations=[ + { + "key": "{model_per_project_itpm:project:model}:tokens", + "increment_value": -90, + "ttl": 60, + "window_key": "{model_per_project_itpm:project:model}:window", + "expected_window_start": "1234", + "reservation_backend": "redis", + } + ] ) - ] - assert unguarded_calls == [] + + assert tuple(executor.call_args.kwargs["keys"]) == ( + "{model_per_project_itpm:project:model}:window", + "{model_per_project_itpm:project:model}:tokens", + ) + assert tuple(executor.call_args.kwargs["args"]) == ("1234", -90, 60) + assert executor.await_count == 1 + assert cache.in_memory_cache.get_cache("{model_per_project_itpm:project:model}:tokens") == 100 @pytest.mark.asyncio -async def test_atomic_lua_response_carries_redis_window_identity(rate_limiter): - handler, _cache = rate_limiter - counter_key = "{model_per_project_itpm:project:model}:tokens" - meta = [ - { - "descriptor_key": PROJECT_ITPM_DESCRIPTOR_KEY, - "descriptor_value": "project:model", - "current_limit": 100, - "rate_limit_type": "tokens", - "counter_key": counter_key, +async def test_atomic_reservation_response_carries_redis_window_identity(): + counter_key: Final = "{model_per_project_itpm:project:model}:tokens" + with _redis_cache_for_executor(AsyncMock(return_value=[0, 25, 1234])) as redis_cache: + handler: Final = RateLimitHandler(internal_usage_cache=InternalUsageCache(DualCache(redis_cache=redis_cache))) + assert await handler.atomic_check_and_increment_by_n([], []) == { + "overall_code": "OK", + "statuses": [], } - ] - - async def successful_reservation(*, keys, args): - return [0, 25, 1234] - - handler.check_and_increment_by_n_script = successful_reservation - assert await handler._atomic_lua_per_descriptor([]) == { - "overall_code": "OK", - "statuses": [], - } - - response = await handler._atomic_lua_per_descriptor( - descriptor_groups=[ - ( - [ - "{model_per_project_itpm:project:model}:window", - counter_key, - ], - [100, 25, 60, 60], - meta, - ) - ] - ) + response: Final = await handler.atomic_check_and_increment_by_n( + descriptors=[ + { + "key": PROJECT_ITPM_DESCRIPTOR_KEY, + "value": "project:model", + "rate_limit": {"tokens_per_unit": 100, "window_size": 60}, + } + ], + increments=[{"tokens": 25}], + ) assert response["statuses"][0]["limit_remaining"] == 75 - assert response["reservation_windows"] == frozenset( - {(counter_key, "1234", "redis")} - ) + assert response["reservation_windows"] == frozenset({(counter_key, "1234", "redis")}) @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index cb8a73d581b..8ab8d7b995f 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -14,6 +14,8 @@ import pytest from fastapi import HTTPException from fastapi.testclient import TestClient from pytest_mock import MockerFixture +from redis.asyncio import Redis +from litellm.caching.redis_cache import RedisCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( @@ -49,7 +51,6 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( client = TestClient(app) if TYPE_CHECKING: - from litellm.caching.redis_cache import RedisCache from litellm.proxy.utils import PrismaClient @@ -2536,11 +2537,12 @@ async def test_user_update_reaches_cached_user_on_every_worker( published.append((channel, message)) return 1 - class _FakeRedisCache: - namespace: str | None = None + class _FakeRedisCache(RedisCache): + def __init__(self) -> None: + self.namespace = None - def init_pubsub_client(self) -> _RecordingRedisClient: - return _RecordingRedisClient() + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", _RecordingRedisClient()) class _UserTableMocks: def __init__( @@ -3200,7 +3202,7 @@ async def test_user_info_v2_redacts_scim_enterprise_metadata(mocker): """ from fastapi import Request - from litellm.proxy._types import UserInfoV2Response + from litellm.proxy._types import LiteLLM_TeamTable, UserInfoV2Response from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 mock_prisma_client = mocker.MagicMock() @@ -3556,7 +3558,7 @@ async def test_user_info_v2_team_admin_can_query_team_member(mocker): """ from fastapi import Request - from litellm.proxy._types import LiteLLM_TeamTable, UserInfoV2Response + from litellm.proxy._types import UserInfoV2Response from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2 mock_prisma_client = mocker.MagicMock() diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index f727448f3ac..999303c0fab 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -14,6 +14,8 @@ import pytest from fastapi import HTTPException, Request from fastapi.testclient import TestClient from pydantic import ValidationError +from redis.asyncio import Redis +from litellm.caching.redis_cache import RedisCache from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger @@ -10253,19 +10255,26 @@ class _BroadcastPipeline: self._messages.append(message) return self - async def execute(self) -> None: + async def execute(self) -> tuple[int, ...]: await self._redis.send(tuple(self._messages)) + return tuple(1 for _ in self._messages) + + async def __aenter__(self) -> "_BroadcastPipeline": + return self + + async def __aexit__(self, *args: object) -> None: + pass -class _BroadcastRecordingRedis: +class _BroadcastRecordingRedis(RedisCache): namespace = None def __init__(self, round_trip_seconds: float = 0.0) -> None: self.broadcast_keys: list[str] = [] self._round_trip_seconds = round_trip_seconds - def init_pubsub_client(self) -> "_BroadcastRecordingRedis": - return self + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", self) async def send(self, messages: Sequence[str]) -> None: if self._round_trip_seconds: diff --git a/tests/unit/proxy/management_helpers/test_team_roster_sync.py b/tests/unit/proxy/management_helpers/test_team_roster_sync.py index 74c3dee9d56..b0babd0fe28 100644 --- a/tests/unit/proxy/management_helpers/test_team_roster_sync.py +++ b/tests/unit/proxy/management_helpers/test_team_roster_sync.py @@ -5,14 +5,16 @@ import json from collections.abc import Awaitable, Callable, Mapping, Sequence from contextlib import asynccontextmanager from datetime import UTC, datetime -from typing import Final, Literal +from typing import Final, Literal, cast from unittest.mock import patch import pytest from pydantic import BaseModel, ConfigDict, TypeAdapter +from redis.asyncio import Redis import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import jwt_key_mapping_cache_key from litellm.proxy.common_utils.user_api_key_cache import ( @@ -553,8 +555,15 @@ class _Pipeline: self._messages.append(message) return self - async def execute(self) -> None: + async def execute(self) -> tuple[int, ...]: await self._send(tuple(self._messages)) + return tuple(1 for _ in self._messages) + + async def __aenter__(self) -> "_Pipeline": + return self + + async def __aexit__(self, *args: object) -> None: + pass class _SlowPublishRedisClient: @@ -576,14 +585,14 @@ class _SlowPublishRedisClient: return _Pipeline(self._send) -class _PubSubRedisCache: +class _PubSubRedisCache(RedisCache): namespace = None def __init__(self, client: "_SlowPublishRedisClient | _WedgedPublishRedisClient") -> None: self._client = client - def init_pubsub_client(self) -> "_SlowPublishRedisClient | _WedgedPublishRedisClient": - return self._client + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", self._client) async def _until_reached_or_done(signal: asyncio.Event, sync: "asyncio.Task[object]") -> None: diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 745488bf93f..5ad9673c6c5 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -11,7 +11,7 @@ import time import types from datetime import datetime, timedelta, timezone from pathlib import Path -from typing import Final +from typing import Final, cast from unittest import mock from unittest.mock import AsyncMock, MagicMock, call, create_autospec, mock_open, patch @@ -15340,14 +15340,14 @@ async def test_auth_cache_invalidation_subscriber_evicts_byok_credentials_cached def pubsub(self) -> _QueuePubSub: return self._scripted_pubsub - class _FakeRedisCache: + class _FakeRedisCache(RedisCache): namespace = None def __init__(self, client: object) -> None: self._client = client - def init_pubsub_client(self) -> object: - return self._client + def init_pubsub_client(self) -> "Redis[bytes]": + return cast("Redis[bytes]", self._client) byok_credential_cache.flush_cache() cache_byok_credential("mallory", "srv-byok", "sk-revoked-elsewhere") diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..b24825f3e82 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -828,78 +828,76 @@ async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployme @pytest.mark.asyncio -async def test_claim_pin_uses_redis_attached_after_construction(): - """The proxy attaches Redis via Router._update_redis_cache after the Router (and - this callback) are built. The claim must resolve the redis tier per call, or pins - silently stay pod-local and cross-pod first-writer-wins is lost.""" - cache = DualCache() - callback = DeploymentAffinityCheck( +async def test_claim_pin_uses_redis_attached_after_construction( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + _, _, redis = shared_session_router + cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=lambda: 100.0)) + callback: Final = DeploymentAffinityCheck( cache=cache, ttl_seconds=3600, enable_user_key_affinity=False, enable_responses_api_affinity=False, ) + cache.attach_redis_cache(redis) + pin_key: Final = _session_pin_key("late-redis-session", "key-1") + cache.in_memory_cache.set_cache(pin_key, {"model_id": "stale-local-winner"}, ttl=10) + runner: Final = AsyncMock(return_value=b'{"model_id": "other-pod-winner"}') - captured = {} - - async def fake_runner(keys, args, client=None): - captured["keys"] = keys - captured["args"] = args - return b'{"model_id": "other-pod-winner"}' - - late_redis = MagicMock() - late_redis.async_register_script = MagicMock(return_value=fake_runner) - cache.redis_cache = late_redis - - import time as time_module - - pin_key = _session_pin_key("late-redis-session", "key-1") - cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) - - claimed = await callback._claim_pin( - cache_key=pin_key, - pin_value={"model_id": "our-deployment"}, - ttl_seconds=777, - ) + with patch("redis.asyncio.client.Redis.register_script", return_value=runner): + with pytest.warns(FutureWarning, match="async_claim_affinity is experimental and may change"): + claimed: Final = await callback._claim_pin( + cache_key=pin_key, + pin_value={"model_id": "our-deployment"}, + ttl_seconds=777, + ) assert claimed == "other-pod-winner" - assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) - assert captured["keys"] == (pin_key,) - assert captured["args"] == ('{"model_id": "our-deployment"}', 777) - assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == { - "model_id": "other-pod-winner" - } + assert cache.in_memory_cache.ttl_dict[pin_key] == 877.0 + runner.assert_awaited_once_with(keys=(pin_key,), args=('{"model_id": "our-deployment"}', 777), client=None) + assert cache.in_memory_cache.get_cache(pin_key) == {"model_id": "other-pod-winner"} @pytest.mark.asyncio -async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): - """A Redis outage must cost cross-pod agreement, never same-pod stickiness. The write - hook only logs this result, so an escaping error would leave the session unpinned and - reshuffle every turn for the whole outage. DualCache's write path, which this claim - replaced, wrote the in-memory tier before ever touching Redis.""" - cache = DualCache() - callback = DeploymentAffinityCheck( +@pytest.mark.parametrize("reply, expected", [(b"\xff", "local"), ([1, 2], "candidate")]) +async def test_claim_pin_distinguishes_invalid_encoding_from_nontext_reply( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], reply: object, expected: str, +) -> None: + _, _, redis = shared_session_router + cache: Final = DualCache(redis_cache=redis, in_memory_cache=InMemoryCache(clock=lambda: 100.0)) + cache.in_memory_cache.set_cache("pin", {"model_id": "local"}, ttl=30) + with patch("redis.asyncio.client.Redis.register_script", return_value=AsyncMock(return_value=reply)): + winner: Final = await claim_affinity_pin(cache, "pin", {"model_id": "candidate"}, 60) + assert winner == {"model_id": expected} + assert cache.in_memory_cache.get_cache("pin") == {"model_id": "local"} + assert cache.in_memory_cache.ttl_dict["pin"] == 130.0 + + +@pytest.mark.asyncio +async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + _, server, redis = shared_session_router + server.connected = False + cache: Final = DualCache(redis_cache=redis, in_memory_cache=InMemoryCache(clock=lambda: 100.0)) + callback: Final = DeploymentAffinityCheck( cache=cache, ttl_seconds=3600, enable_user_key_affinity=False, enable_responses_api_affinity=False, ) - - async def exploding_runner(keys, args, client=None): - raise ConnectionError("redis is down") - - down_redis = MagicMock() - down_redis.async_register_script = MagicMock(return_value=exploding_runner) - cache.redis_cache = down_redis - - key = _session_pin_key("outage-session", "key-1") - claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777) + key: Final = _session_pin_key("outage-session", "key-1") + claimed: Final = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777) assert claimed == "our-deployment" assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"} + assert cache.in_memory_cache.ttl_dict[key] == 877.0 - second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777) + second: Final = await callback._claim_pin( + cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=10 + ) assert second == "our-deployment" + assert cache.in_memory_cache.ttl_dict[key] == 877.0 @pytest.mark.asyncio @@ -1004,10 +1002,12 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( ], ) async def test_redis_deployment_claim_preserves_legacy_result_decoding( - raw: object, expected: str | None, stored: object + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], + raw: object, + expected: str | None, + stored: object, ) -> None: - redis: Final = MagicMock() - redis.async_register_script.return_value = AsyncMock(return_value=raw) + _, _, redis = shared_session_router cache: Final = DualCache(redis_cache=redis) callback: Final = DeploymentAffinityCheck( cache=cache, @@ -1016,7 +1016,8 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding( enable_responses_api_affinity=False, ) - winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) + with patch("redis.asyncio.client.Redis.register_script", return_value=AsyncMock(return_value=raw)): + winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) assert winner == expected assert cache.in_memory_cache.get_cache("deployment-pin") == stored