mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge a35cf6bcb0 into f82e4c80ac
This commit is contained in:
yucheng-berri
GitHub
commit
2edfc192ff
8 changed files with 2560 additions and 39 deletions
|
|
@ -16,6 +16,8 @@ import inspect
|
|||
import itertools
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import weakref
|
||||
|
|
@ -24,6 +26,7 @@ from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
|||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from decimal import Decimal
|
||||
from types import FrameType, MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast
|
||||
|
||||
|
|
@ -140,10 +143,34 @@ _INCREMENT_IF_PRESENT_LUA: Final = (
|
|||
)
|
||||
|
||||
_LUA_COUNT: Final = TypeAdapter(int)
|
||||
_PLAIN_DECIMAL: Final = re.compile(r"-?[0-9]+(?:\.[0-9]+)?(?:[eE][-+]?[0-9]+)?")
|
||||
_PIPELINE_COMMAND_NUMBER: Final = re.compile(r"Command # ([0-9]+) \(")
|
||||
_LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
|
||||
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
|
||||
|
||||
|
||||
def _as_incrbyfloat_writes(value: float) -> str:
|
||||
"""``value`` as INCRBYFLOAT writes it into a missing key: 17 decimal places with the trailing zeros cut."""
|
||||
text: Final = format(Decimal(repr(value)), ".17f").rstrip("0").rstrip(".")
|
||||
return "0" if text == "-0" else text
|
||||
|
||||
|
||||
def _failed_after_the_first_command(error: Exception) -> bool:
|
||||
"""Whether redis-py names a pipeline command after the first as the one that failed, so the first one applied."""
|
||||
message: Final = error.args[0] if error.args else None
|
||||
number: Final = _PIPELINE_COMMAND_NUMBER.match(message) if isinstance(message, str) else None
|
||||
return number is not None and int(number.group(1)) > 1
|
||||
|
||||
|
||||
def _plain_float(reply: object) -> float | None:
|
||||
"""The finite float a stored value spells in plain decimal or exponent form, or None for anything else."""
|
||||
text: Final = reply.decode(errors="replace") if isinstance(reply, bytes) else reply
|
||||
if not isinstance(text, str) or _PLAIN_DECIMAL.fullmatch(text) is None:
|
||||
return None
|
||||
number: Final = float(text)
|
||||
return number if math.isfinite(number) else None
|
||||
|
||||
|
||||
def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]:
|
||||
return _OPTIONAL_COUNTS.validate_python(
|
||||
tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values)
|
||||
|
|
@ -723,6 +750,11 @@ class RedisCache(BaseCache):
|
|||
def _admit_async_call(self, call: Callable[[], Awaitable[_RedisCallResult]]) -> Awaitable[_RedisCallResult]:
|
||||
return call()
|
||||
|
||||
@property
|
||||
def circuit_breaker_closed(self) -> bool:
|
||||
"""Whether calls reach Redis without a recovery probe. Reading it, unlike making a call, never starts one."""
|
||||
return not self._circuit_breaker.enabled or self._circuit_breaker.state == RedisCircuitBreaker.CLOSED
|
||||
|
||||
def _setup_health_pings(self):
|
||||
"""Setup async and sync health pings for Redis."""
|
||||
# ASYNC HEALTH PING
|
||||
|
|
@ -1482,15 +1514,69 @@ class RedisCache(BaseCache):
|
|||
await _redis_client.expire(key, ttl)
|
||||
return float(result)
|
||||
|
||||
async def _set_if_missing_and_get(self, key: str, value: float, ttl: int | None) -> tuple[bool, float] | None:
|
||||
"""SET NX GET and its TTL command in one round trip: whether the key was created and the value it holds, or
|
||||
None when the server refused SET NX GET rather than failed to answer, or the key holds a value this does not
|
||||
read as INCRBYFLOAT would. SET NX GET is one command, so a refused one wrote nothing. A refused TTL command
|
||||
raises, as it does in main's add."""
|
||||
from redis.exceptions import ClusterError, RedisError
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
|
||||
if not math.isfinite(value):
|
||||
return None
|
||||
stored: Final = _as_incrbyfloat_writes(value)
|
||||
try:
|
||||
async with self._async_commands().pipeline(transaction=False) as pipe:
|
||||
pipe.set(name=key, value=stored, nx=True, get=True)
|
||||
if ttl is not None:
|
||||
pipe.expire(key, ttl)
|
||||
replies: Final[Sequence[object]] = await pipe.execute() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # redis generics
|
||||
except RedisError as e:
|
||||
failed_to_answer: Final = isinstance(e, (RedisConnectionError, RedisTimeoutError, ClusterError))
|
||||
if failed_to_answer or _failed_after_the_first_command(e):
|
||||
raise
|
||||
return None
|
||||
if replies[0] is None:
|
||||
return True, float(stored)
|
||||
held: Final = _plain_float(replies[0]) # pyright: ignore[reportUnknownArgumentType] # redis generics
|
||||
return None if held is None else (False, held)
|
||||
|
||||
async def _set_if_missing_with_ttl(
|
||||
self,
|
||||
_redis_client: "Redis[bytes]",
|
||||
key: str,
|
||||
value: float,
|
||||
ttl: int | None,
|
||||
refresh_ttl: bool,
|
||||
on_added: Callable[[], None] | None,
|
||||
) -> float:
|
||||
"""The value the key holds after setting it if it was missing. A server that refuses the commands this sends,
|
||||
such as a user denied SET or a Redis before 7.0, and a key holding a value INCRBYFLOAT might read differently,
|
||||
get a plain add instead. ``on_added`` runs when this call added the value."""
|
||||
written: Final = await self._set_if_missing_and_get(key=key, value=value, ttl=ttl)
|
||||
current: Final = (
|
||||
await self._incrbyfloat_with_ttl(_redis_client, key=key, value=value, ttl=ttl, refresh_ttl=refresh_ttl)
|
||||
if written is None
|
||||
else written[1]
|
||||
)
|
||||
if on_added is not None and (written is None or written[0]):
|
||||
on_added()
|
||||
return current
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_increment(
|
||||
self,
|
||||
key,
|
||||
key: str,
|
||||
value: float,
|
||||
ttl: int | None = None,
|
||||
parent_otel_span: Span | None = None,
|
||||
refresh_ttl: bool = False,
|
||||
only_if_missing: bool = False,
|
||||
on_added: Callable[[], None] | None = None,
|
||||
) -> float:
|
||||
"""Add ``value`` to the key and return its new value. With ``only_if_missing``, a key that exists keeps its
|
||||
value, which is returned, and only the TTL is refreshed, and ``on_added`` runs if this call added the value."""
|
||||
from redis.asyncio import Redis
|
||||
|
||||
_redis_client: Final[Redis] = self.init_async_client()
|
||||
|
|
@ -1498,8 +1584,19 @@ class RedisCache(BaseCache):
|
|||
_used_ttl: Final = self.get_ttl(ttl=ttl)
|
||||
key = self.check_and_fix_namespace(key=key)
|
||||
try:
|
||||
result: Final = await self._incrbyfloat_with_ttl(
|
||||
_redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl
|
||||
result: Final = await (
|
||||
self._set_if_missing_with_ttl(
|
||||
_redis_client, # pyright: ignore[reportUnknownArgumentType] # redis generics
|
||||
key=key,
|
||||
value=value,
|
||||
ttl=_used_ttl,
|
||||
refresh_ttl=refresh_ttl,
|
||||
on_added=on_added,
|
||||
)
|
||||
if only_if_missing
|
||||
else self._incrbyfloat_with_ttl(
|
||||
_redis_client, key=key, value=value, ttl=_used_ttl, refresh_ttl=refresh_ttl
|
||||
)
|
||||
)
|
||||
|
||||
## LOGGING ##
|
||||
|
|
|
|||
|
|
@ -1942,6 +1942,8 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = get_env_float("SPEND_LOG_QUEUE_POLL_INTER
|
|||
RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS: Final = max(1, get_env_int("RESPONSES_SESSION_LOOKUP_MAX_ATTEMPTS", 3))
|
||||
RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL: Final = get_env_float("RESPONSES_SESSION_LOOKUP_RETRY_INTERVAL", 0.2)
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = get_env_int("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)
|
||||
SPEND_COUNTER_COLD_SEED_JOIN_SECONDS: Final = 0.05
|
||||
SPEND_COUNTER_INCRBYFLOAT_ROUNDING: Final = 1e-16
|
||||
PROXY_DB_LOOKUP_MAX_CONCURRENCY: Final = max(1, get_env_int("PROXY_DB_LOOKUP_MAX_CONCURRENCY", 25))
|
||||
PROXY_DB_LOOKUP_DEADLINE_SECONDS: Final = max(0.1, get_env_float("PROXY_DB_LOOKUP_DEADLINE_SECONDS", 10.0))
|
||||
PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS: Final = max(0.0, get_env_float("PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS", 30.0))
|
||||
|
|
|
|||
|
|
@ -13,22 +13,31 @@ memory in long-lived deployments.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import functools
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Mapping
|
||||
from contextvars import ContextVar, copy_context
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Optional
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
|
||||
from litellm.constants import (
|
||||
SPEND_COUNTER_COLD_SEED_JOIN_SECONDS,
|
||||
SPEND_COUNTER_INCRBYFLOAT_ROUNDING,
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE,
|
||||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
SPEND_COUNTERS_TARGET,
|
||||
SpendCounterRead,
|
||||
read_batched_spend_counter,
|
||||
record_spend_counter_value,
|
||||
spend_counter_tick,
|
||||
)
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.project_repository import ProjectRepository
|
||||
|
|
@ -68,6 +77,93 @@ _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class _CachedSpendSeed:
|
||||
"""This process adding its cached spend to one cold counter: the add in flight, the tick it started at, the tick
|
||||
the last write of its spend landed at, and that spend."""
|
||||
|
||||
started_at: int = 0
|
||||
landed_at: int = 0
|
||||
landed_spend: float = 0.0
|
||||
in_flight: asyncio.Future[float | None] | None = None
|
||||
|
||||
def on_this_loop(self) -> bool:
|
||||
return self.in_flight is None or self.in_flight.get_loop() is asyncio.get_running_loop()
|
||||
|
||||
def run(self, add: Callable[[], Awaitable[float | None]]) -> Coroutine[object, object, float | None]:
|
||||
"""Make ``add`` the add in flight from this call on, and return the coroutine that runs it and publishes its
|
||||
result to the requests waiting on ``in_flight``, which ends cancelled when ``add`` fails or is cancelled."""
|
||||
in_flight: Final[asyncio.Future[float | None]] = asyncio.get_running_loop().create_future()
|
||||
self.started_at = spend_counter_tick()
|
||||
self.in_flight = in_flight
|
||||
|
||||
async def add_and_publish() -> float | None:
|
||||
try:
|
||||
added: Final = await add()
|
||||
in_flight.set_result(added)
|
||||
return added
|
||||
finally:
|
||||
if self.in_flight is in_flight:
|
||||
self.in_flight = None
|
||||
in_flight.cancel()
|
||||
|
||||
return add_and_publish()
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE)
|
||||
def _cached_spend_seed(spend_counter_cache: "DualCache", counter_key: str) -> _CachedSpendSeed:
|
||||
return _CachedSpendSeed()
|
||||
|
||||
|
||||
_cold_seed_waited: Final[ContextVar[float]] = ContextVar("cold_spend_seed_waited", default=0.0)
|
||||
|
||||
|
||||
async def _wait_for_cold_seed(in_flight: asyncio.Future[float | None]) -> None:
|
||||
"""Wait on ``in_flight`` for what this context has left of ``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS``."""
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
waited: Final = _cold_seed_waited.get()
|
||||
started: Final = loop.time()
|
||||
await asyncio.wait((in_flight,), timeout=max(0.0, SPEND_COUNTER_COLD_SEED_JOIN_SECONDS - waited))
|
||||
_cold_seed_waited.set(waited + loop.time() - started)
|
||||
|
||||
|
||||
_late_writes: Final[set[asyncio.Task[float | None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks
|
||||
_late_write_abandoned: Final[ContextVar[asyncio.Event | None]] = ContextVar("late_spend_write_abandoned", default=None)
|
||||
|
||||
|
||||
def late_write_abandoned() -> bool:
|
||||
"""Whether this runs in a late write whose request was cancelled, which main's add would not have outlived."""
|
||||
abandoned: Final = _late_write_abandoned.get()
|
||||
return abandoned is not None and abandoned.is_set()
|
||||
|
||||
|
||||
def _late_write_finished(task: asyncio.Task[float | None]) -> None:
|
||||
_late_writes.discard(task)
|
||||
if not task.cancelled():
|
||||
task.exception()
|
||||
|
||||
|
||||
async def _landing_even_if_cancelled(write: Coroutine[object, object, float | None]) -> float | None:
|
||||
"""Await ``write`` in a task of its own, so cancelling the caller leaves it to land as a command already sent does,
|
||||
and marks it abandoned from then on, so ``late_write_abandoned`` holds inside it."""
|
||||
abandoned: Final = asyncio.Event()
|
||||
context: Final = copy_context()
|
||||
context.run(_late_write_abandoned.set, abandoned)
|
||||
task: Final = context.run(asyncio.ensure_future, write)
|
||||
_late_writes.add(task)
|
||||
task.add_done_callback(_late_write_finished)
|
||||
try:
|
||||
return await asyncio.shield(task)
|
||||
except asyncio.CancelledError:
|
||||
abandoned.set()
|
||||
raise
|
||||
|
||||
|
||||
def _redis_breaker_closed(spend_counter_cache: "DualCache") -> bool:
|
||||
redis_cache: Final = spend_counter_cache.redis_cache
|
||||
return redis_cache is None or redis_cache.circuit_breaker_closed
|
||||
|
||||
|
||||
def _as_utc(value: datetime) -> datetime:
|
||||
return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
|
@ -112,15 +208,97 @@ class SpendCounterReseed:
|
|||
SpendCounterReseed._locks.popitem(last=False)
|
||||
return lock
|
||||
|
||||
@staticmethod
|
||||
async def add_cached_spend_once(
|
||||
spend_counter_cache: "DualCache",
|
||||
counter_key: str,
|
||||
read: SpendCounterRead,
|
||||
spend: float,
|
||||
add: Callable[[], Awaitable[float | None]],
|
||||
set_if_missing: Callable[[Callable[[], None]], Awaitable[float | None]],
|
||||
) -> None:
|
||||
"""Run ``add``, this process adding ``spend`` to a cold counter, once for all of this process's requests that
|
||||
read the counter cold before that add landed. A request joins an add that started after its read came back and
|
||||
gets its result, and starts or joins another add when that one fails, is cancelled or wrote nothing after its
|
||||
read went out. A request whose read went out before an add landed, and one still waiting on an add after
|
||||
``SPEND_COUNTER_COLD_SEED_JOIN_SECONDS`` across its counters, runs ``set_if_missing``, which calls its argument
|
||||
if it added the spend, in place of ``add``, then runs ``add`` too if it did not and the counter holds less than
|
||||
the spend it should already hold, beyond what INCRBYFLOAT's 17 decimal places can round away.
|
||||
Those two steps finish even if the request is cancelled, and stay the add in flight until they do, but once it is
|
||||
cancelled a failure in them leaves the counter as main's cancelled add does, without the DEL. A request on
|
||||
another event loop than the add in flight, one that finds the Redis circuit breaker not closed, and each other
|
||||
process, still run ``add``."""
|
||||
seed: Final = _cached_spend_seed(spend_counter_cache, counter_key)
|
||||
|
||||
def spend_landed() -> None:
|
||||
seed.landed_at = spend_counter_tick()
|
||||
seed.landed_spend = spend
|
||||
|
||||
async def add_spend() -> float | None:
|
||||
added: Final = await add()
|
||||
spend_landed()
|
||||
return added
|
||||
|
||||
async def write_if_short() -> float | None:
|
||||
least: Final = min(spend, seed.landed_spend) if seed.landed_at >= read.sent_at else spend
|
||||
set_added: Final = asyncio.Event()
|
||||
|
||||
def on_added() -> None:
|
||||
set_added.set()
|
||||
spend_landed()
|
||||
|
||||
current: Final = await set_if_missing(on_added)
|
||||
if current is None or set_added.is_set():
|
||||
return current
|
||||
return await add_spend() if least - current > SPEND_COUNTER_INCRBYFLOAT_ROUNDING else current
|
||||
|
||||
async def add_if_short() -> float | None:
|
||||
return await _landing_even_if_cancelled(seed.run(write_if_short))
|
||||
|
||||
async def add_unless_covered() -> float | None:
|
||||
in_flight: Final = seed.in_flight
|
||||
if not seed.on_this_loop() or not _redis_breaker_closed(spend_counter_cache):
|
||||
return await add()
|
||||
if in_flight is None:
|
||||
return await (seed.run(add_spend) if seed.landed_at < read.sent_at else add_if_short())
|
||||
covers: Final = seed.started_at > read.returned_at
|
||||
await _wait_for_cold_seed(in_flight)
|
||||
if not in_flight.done():
|
||||
return await add_if_short()
|
||||
if covers and not in_flight.cancelled() and seed.landed_at >= read.sent_at:
|
||||
return in_flight.result()
|
||||
return await add_unless_covered()
|
||||
|
||||
added: Final = await add_unless_covered()
|
||||
if added is not None:
|
||||
record_spend_counter_value(counter_key, added)
|
||||
|
||||
@staticmethod
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None:
|
||||
"""Apply local deltas after an in-flight reseed establishes the spend balance."""
|
||||
async def increment_in_memory(
|
||||
spend_counter_cache: "DualCache",
|
||||
counter_key: str,
|
||||
increment: float,
|
||||
only_if_missing: bool = False,
|
||||
on_added: Callable[[], None] | None = None,
|
||||
) -> float | None:
|
||||
"""Apply local deltas after an in-flight reseed establishes the spend balance. With ``only_if_missing``, a
|
||||
counter that exists keeps its value, which is returned, and ``on_added`` runs if this call added to it."""
|
||||
lock: Final = await SpendCounterReseed._get_lock(counter_key)
|
||||
async with lock:
|
||||
return await spend_counter_cache.async_increment_cache(
|
||||
current: Final[float | None] = ( # pyright: ignore[reportUnknownVariableType] # untyped in-memory cache
|
||||
spend_counter_cache.in_memory_cache.get_cache(key=counter_key) # pyright: ignore[reportUnknownMemberType] # untyped in-memory cache
|
||||
if only_if_missing
|
||||
else None
|
||||
)
|
||||
if current is not None:
|
||||
return float(current)
|
||||
added: Final = await spend_counter_cache.async_increment_cache(
|
||||
key=counter_key, value=increment, local_only=True, refresh_ttl=True
|
||||
)
|
||||
if on_added is not None:
|
||||
on_added()
|
||||
return added
|
||||
|
||||
@staticmethod
|
||||
async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> float | None:
|
||||
|
|
|
|||
|
|
@ -682,7 +682,7 @@ from litellm.proxy.db.request_error_tracking import (
|
|||
flush_request_errors,
|
||||
request_error_accumulator,
|
||||
)
|
||||
from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed
|
||||
from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed, late_write_abandoned
|
||||
from litellm.proxy.discovery_endpoints import (
|
||||
agent_skills_discovery_router,
|
||||
ui_discovery_endpoints_router,
|
||||
|
|
@ -981,6 +981,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
|
|||
read_batched_spend_counter,
|
||||
record_spend_counter_value,
|
||||
spend_counter_batch_scope,
|
||||
spend_counter_read,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_management_endpoints import (
|
||||
router as spend_management_router,
|
||||
|
|
@ -4299,6 +4300,7 @@ async def ensure_spend_counter_initialized(
|
|||
) -> None:
|
||||
is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key)
|
||||
if is_warm is False:
|
||||
read: Final = spend_counter_read(counter_key)
|
||||
# Shares the per-counter lock with get_current_spend.
|
||||
db_spend: Final = await SpendCounterReseed.coalesced(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -4310,7 +4312,16 @@ async def ensure_spend_counter_initialized(
|
|||
# DB unavailable - fall back to in-process cache (may be stale).
|
||||
base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key)
|
||||
if base_spend > 0:
|
||||
await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend)
|
||||
await SpendCounterReseed.add_cached_spend_once(
|
||||
spend_counter_cache=spend_counter_cache,
|
||||
counter_key=counter_key,
|
||||
read=read,
|
||||
spend=base_spend,
|
||||
add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend),
|
||||
set_if_missing=lambda on_added: increment_spend_counter_cache(
|
||||
counter_key=counter_key, increment=base_spend, only_if_missing=True, on_added=on_added
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
_ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized
|
||||
|
|
@ -4413,16 +4424,32 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool:
|
|||
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def increment_spend_counter_cache(counter_key: str, increment: float) -> float | None:
|
||||
async def increment_spend_counter_cache(
|
||||
counter_key: str,
|
||||
increment: float,
|
||||
only_if_missing: bool = False,
|
||||
on_added: Callable[[], None] | None = None,
|
||||
) -> float | None:
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
current_value: Final = await spend_counter_cache.redis_cache.async_increment(
|
||||
key=counter_key,
|
||||
value=increment,
|
||||
refresh_ttl=True,
|
||||
current_value: Final = await (
|
||||
spend_counter_cache.redis_cache.async_increment(
|
||||
key=counter_key,
|
||||
value=increment,
|
||||
refresh_ttl=True,
|
||||
only_if_missing=True,
|
||||
on_added=on_added,
|
||||
)
|
||||
if only_if_missing
|
||||
else spend_counter_cache.redis_cache.async_increment(
|
||||
key=counter_key,
|
||||
value=increment,
|
||||
refresh_ttl=True,
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
await invalidate_spend_counter(counter_key=counter_key)
|
||||
if not late_write_abandoned():
|
||||
await invalidate_spend_counter(counter_key=counter_key)
|
||||
raise
|
||||
spend_counter_cache.in_memory_cache.set_cache(
|
||||
key=counter_key,
|
||||
|
|
@ -4432,7 +4459,11 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f
|
|||
return current_value
|
||||
|
||||
return await SpendCounterReseed.increment_in_memory(
|
||||
spend_counter_cache=spend_counter_cache, counter_key=counter_key, increment=increment
|
||||
spend_counter_cache=spend_counter_cache,
|
||||
counter_key=counter_key,
|
||||
increment=increment,
|
||||
only_if_missing=only_if_missing,
|
||||
on_added=on_added,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
"""One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each."""
|
||||
|
||||
import asyncio
|
||||
import itertools
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from contextvars import ContextVar, Token
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType, TracebackType
|
||||
from typing import Final
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
|
|
@ -22,6 +23,24 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
_CounterValues: Final = TypeAdapter(dict[str, float | None])
|
||||
_NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
|
||||
SPEND_COUNTERS_TARGET: Final = "spend_counters"
|
||||
_ticks: Final = itertools.count(1)
|
||||
|
||||
|
||||
def spend_counter_tick() -> int:
|
||||
"""The next number in a per-process sequence that orders spend counter reads and cold seeds."""
|
||||
return next(_ticks)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SpendCounterRead:
|
||||
"""When a read of a spend counter went out and when its reply came back, as spend counter ticks."""
|
||||
|
||||
sent_at: int
|
||||
returned_at: int
|
||||
|
||||
|
||||
_NO_READS: Final[Mapping[str, SpendCounterRead]] = MappingProxyType({})
|
||||
_PendingMGet: TypeAlias = tuple[BatchResult[Mapping[str, object]], int]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -39,7 +58,17 @@ class SpendCounterBatch:
|
|||
hook declares whatever is bound but unread, so whoever flushes first (the auth object prefetch, usually)
|
||||
carries the spend counters in the same round trip."""
|
||||
|
||||
__slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_redis_cache", "_request_batch")
|
||||
__slots__ = (
|
||||
"_fetched",
|
||||
"_inflight",
|
||||
"_keys",
|
||||
"_loaded",
|
||||
"_lock",
|
||||
"_open",
|
||||
"_reads",
|
||||
"_redis_cache",
|
||||
"_request_batch",
|
||||
)
|
||||
|
||||
def __init__(self, redis_cache: RedisCache) -> None:
|
||||
self._redis_cache: Final = redis_cache
|
||||
|
|
@ -48,7 +77,8 @@ class SpendCounterBatch:
|
|||
self._keys: frozenset[str] = frozenset()
|
||||
self._fetched: frozenset[str] = frozenset()
|
||||
self._loaded: Mapping[str, float | None] = _NO_VALUES
|
||||
self._inflight: Final[list[BatchResult[Mapping[str, object]]]] = [] # mutable-ok: drained by _load
|
||||
self._reads: Mapping[str, SpendCounterRead] = _NO_READS
|
||||
self._inflight: Final[list[_PendingMGet]] = [] # mutable-ok: drained by _load
|
||||
self._request_batch: Final[RedisBatch | None] = active_request_redis_batch(redis_cache)
|
||||
if self._request_batch is not None:
|
||||
self._request_batch.add_flush_hook(self._declare_pending)
|
||||
|
|
@ -78,6 +108,10 @@ class SpendCounterBatch:
|
|||
return None
|
||||
return loaded[counter_key], True
|
||||
|
||||
def counter_read(self, counter_key: str) -> SpendCounterRead | None:
|
||||
"""When the MGET behind this batch's value for the counter went out and came back, None without that value."""
|
||||
return self._reads.get(counter_key) if self._open and counter_key in self._loaded else None
|
||||
|
||||
def record(self, counter_key: str, value: float) -> None:
|
||||
"""A write returned the counter's new value; later reads in this scope see it instead of the MGET value."""
|
||||
if not self._open:
|
||||
|
|
@ -93,6 +127,7 @@ class SpendCounterBatch:
|
|||
self._keys = self._keys - key
|
||||
self._fetched = self._fetched - key
|
||||
self._loaded = MappingProxyType({k: v for k, v in self._loaded.items() if k != counter_key})
|
||||
self._reads = MappingProxyType({k: v for k, v in self._reads.items() if k != counter_key})
|
||||
|
||||
async def _load(self) -> Mapping[str, float | None]:
|
||||
async with self._lock:
|
||||
|
|
@ -103,10 +138,15 @@ class SpendCounterBatch:
|
|||
pending: Final = self._keys - self._fetched
|
||||
if pending:
|
||||
self._fetched = self._fetched | pending
|
||||
sent_at: Final = spend_counter_tick()
|
||||
fetched: Final = await self._fetch(pending)
|
||||
self._loaded = MappingProxyType({**fetched, **self._loaded})
|
||||
self._merge(fetched, SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick()))
|
||||
return self._loaded
|
||||
|
||||
def _merge(self, fetched: Mapping[str, float | None], read: SpendCounterRead) -> None:
|
||||
self._loaded = MappingProxyType({**fetched, **self._loaded})
|
||||
self._reads = MappingProxyType({**dict.fromkeys(fetched, read), **self._reads})
|
||||
|
||||
def _declare_pending(self) -> None:
|
||||
"""Flush hook: put every bound-but-unread counter on the request pipeline that is about to go out."""
|
||||
if self._request_batch is None or not self._open:
|
||||
|
|
@ -114,19 +154,21 @@ class SpendCounterBatch:
|
|||
pending: Final = self._keys - self._fetched
|
||||
if pending:
|
||||
self._fetched = self._fetched | pending
|
||||
sent_at: Final = spend_counter_tick()
|
||||
with service_target(SPEND_COUNTERS_TARGET):
|
||||
self._inflight.append(self._request_batch.mget(sorted(pending)))
|
||||
read: Final = self._request_batch.mget(sorted(pending))
|
||||
self._inflight.append((read, sent_at))
|
||||
|
||||
async def _collect_inflight(self) -> None:
|
||||
results: Final = tuple(self._inflight)
|
||||
reads: Final = tuple(self._inflight)
|
||||
self._inflight.clear()
|
||||
for result in results:
|
||||
for read, sent_at in reads:
|
||||
try:
|
||||
fetched: Mapping[str, float | None] = _CounterValues.validate_python(await result)
|
||||
fetched: Mapping[str, float | None] = _CounterValues.validate_python(await read)
|
||||
except Exception as e: # noqa: BLE001 # per-key reads take over and apply their own Redis fallback
|
||||
verbose_proxy_logger.debug("spend counter batch read failed, falling back to per-key reads: %s", e)
|
||||
continue
|
||||
self._loaded = MappingProxyType({**fetched, **self._loaded})
|
||||
self._merge(fetched, SpendCounterRead(sent_at=sent_at, returned_at=spend_counter_tick()))
|
||||
|
||||
async def _fetch(self, keys: frozenset[str]) -> Mapping[str, float | None]:
|
||||
try:
|
||||
|
|
@ -145,6 +187,17 @@ def active_spend_counter_batch() -> SpendCounterBatch | None:
|
|||
return _active_batch.get()
|
||||
|
||||
|
||||
def spend_counter_read(counter_key: str) -> SpendCounterRead:
|
||||
"""When this context's read of the counter went out and came back: the open batch's MGET for it, else now, since
|
||||
``SpendCounterReseed.coalesced`` reads the counter again after a read outside the batch."""
|
||||
batch: Final = _active_batch.get()
|
||||
batched: Final = batch.counter_read(counter_key) if batch is not None else None
|
||||
if batched is not None:
|
||||
return batched
|
||||
now: Final = spend_counter_tick()
|
||||
return SpendCounterRead(sent_at=now, returned_at=now)
|
||||
|
||||
|
||||
class spend_counter_batch_scope:
|
||||
"""Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it.
|
||||
Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET."""
|
||||
|
|
|
|||
|
|
@ -1,17 +1,20 @@
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.database import read_rows, write_rows
|
||||
from integration._support.upstream import delete_scenario, register_scenario
|
||||
from integration.cost_calculation.cost_tracking_case import JsonResponse
|
||||
from pydantic import BaseModel, JsonValue
|
||||
from redis import Redis
|
||||
|
||||
PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key"
|
||||
CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint."
|
||||
|
|
@ -649,3 +652,141 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) ->
|
|||
assert "data:" not in denied.text, denied.text
|
||||
_assert_budget_refusal(denied, end_user, 0.05)
|
||||
assert _observed_requests(upstream) == []
|
||||
|
||||
|
||||
_COLD_BURST: Final = 40
|
||||
_COLD_SEED_SPEND: Final = 6.0
|
||||
_COLD_BUDGET: Final = 10.0
|
||||
_CHAT_COST: Final = 20 * 0.001 + 20 * 0.002
|
||||
|
||||
|
||||
def _serves_model(gateway: Gateway, model: str) -> bool:
|
||||
response: Final = gateway.request("GET", "/model/info")
|
||||
return response.status_code == 200 and any(
|
||||
object_value(entry)["model_name"] == model for entry in object_value(response.json())["data"]
|
||||
)
|
||||
|
||||
|
||||
def _counter_cache() -> Redis:
|
||||
return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
|
||||
|
||||
def _counter(cache: Redis, end_user: str) -> float | None:
|
||||
raw: Final = cache.get(f"spend:end_user:{end_user}")
|
||||
return None if raw is None else float(raw)
|
||||
|
||||
|
||||
def _lose_counter(cache: Redis, end_user: str) -> None:
|
||||
cache.delete(f"spend:end_user:{end_user}")
|
||||
|
||||
|
||||
def _set_end_user_row_spend(end_user: str, spend: float) -> None:
|
||||
write_rows('UPDATE "LiteLLM_EndUserTable" SET spend = %s WHERE user_id = %s', (str(spend), end_user))
|
||||
|
||||
|
||||
def _await_row_spend(end_user: str, spend: float) -> None:
|
||||
eventually(
|
||||
lambda: _end_user_row(end_user),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(spend),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _refused_spend(denied: httpx.Response, end_user: str) -> float:
|
||||
error: Final = _denial_error(denied)
|
||||
assert error["type"] == "budget_exceeded", denied.text
|
||||
matched: Final = _REFUSAL_MESSAGE.match(string_value(error["message"]))
|
||||
assert matched is not None, denied.text
|
||||
assert matched.group("user") == end_user
|
||||
return float(matched.group("spend"))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", ("user", "responses_user", "messages_metadata_user_id"))
|
||||
def test_overlapping_cold_starts_count_the_end_user_spend_once(gateway: Gateway, spelling: str) -> None:
|
||||
chosen: Final = SPELLINGS[spelling]
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
_counter_cache() as cache,
|
||||
):
|
||||
deployment: Final = _deployment(scenario, chosen)
|
||||
key: Final = scenario.key(models=[deployment.model])
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
_new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
||||
assert _customer_info(gateway, end_user).spend == _COLD_SEED_SPEND
|
||||
assert _counter(cache, end_user) is None
|
||||
upstream.delete("/__observations").raise_for_status()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=_COLD_BURST) as pool:
|
||||
burst: Final = tuple(
|
||||
pool.map(
|
||||
lambda index: _send(gateway, chosen, deployment, key, end_user, f"cold {index}"),
|
||||
range(_COLD_BURST),
|
||||
)
|
||||
)
|
||||
refused: Final = [response.text for response in burst if response.status_code != 200]
|
||||
assert refused == [], refused
|
||||
follow_up: Final = _send(gateway, chosen, deployment, key, end_user, "cold follow-up")
|
||||
_assert_served(chosen, follow_up)
|
||||
assert len(_observed_requests(upstream)) == _COLD_BURST + 1
|
||||
|
||||
expected_spend: Final = _COLD_SEED_SPEND + (_COLD_BURST + 1) * _CHAT_COST
|
||||
_await_row_spend(end_user, expected_spend)
|
||||
assert _customer_info(gateway, end_user).spend == pytest.approx(expected_spend)
|
||||
counted: Final = eventually(
|
||||
lambda: _counter(cache, end_user),
|
||||
lambda value: value is not None and value == pytest.approx(expected_spend),
|
||||
seconds=30,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
assert counted == pytest.approx(expected_spend)
|
||||
|
||||
|
||||
def test_a_lost_counter_keeps_spend_this_proxy_has_not_written_yet(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario, _counter_cache() as cache:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.25)
|
||||
key: Final = scenario.key(models=[model])
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
call_cost: Final = 20 * 0.001 + 20 * 0.25
|
||||
_new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
||||
served: Final = _chat(gateway, model, key, end_user)
|
||||
assert served.status_code == 200, served.text
|
||||
_await_row_spend(end_user, _COLD_SEED_SPEND + call_cost)
|
||||
_set_end_user_row_spend(end_user, _COLD_SEED_SPEND)
|
||||
_lose_counter(cache, end_user)
|
||||
|
||||
denied: Final = _chat(gateway, model, key, end_user)
|
||||
|
||||
assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + call_cost)
|
||||
|
||||
|
||||
def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: Gateway, peer: Gateway) -> None:
|
||||
with gateway.scenario() as scenario, _counter_cache() as cache:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.149)
|
||||
eventually(lambda: _serves_model(peer, model), bool, seconds=30)
|
||||
key: Final = scenario.key(models=[model])
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
call_cost: Final = 20 * 0.001 + 20 * 0.149
|
||||
_new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
|
||||
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
||||
first: Final = _chat(peer, model, key, end_user)
|
||||
assert first.status_code == 200, first.text
|
||||
_await_row_spend(end_user, _COLD_SEED_SPEND + call_cost)
|
||||
_set_end_user_row_spend(end_user, _COLD_SEED_SPEND)
|
||||
_lose_counter(cache, end_user)
|
||||
second: Final = _chat(peer, model, key, end_user)
|
||||
assert second.status_code == 200, second.text
|
||||
expected_spend: Final = _COLD_SEED_SPEND + 2 * call_cost
|
||||
counted: Final = eventually(
|
||||
lambda: _counter(cache, end_user),
|
||||
lambda value: value is not None and value == pytest.approx(expected_spend),
|
||||
seconds=30,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
|
||||
denied: Final = _chat(gateway, model, key, end_user)
|
||||
|
||||
assert counted == pytest.approx(expected_spend)
|
||||
assert _denial_error(denied)["type"] == "budget_exceeded", denied.text
|
||||
|
|
|
|||
|
|
@ -8,6 +8,10 @@ from typing import Final, cast
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fakeredis import FakeAsyncRedis, FakeServer
|
||||
from redis.exceptions import ClusterDownError, NoPermissionError, RedisError, ResponseError
|
||||
from redis.exceptions import ConnectionError as RedisConnectionError
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
|
||||
import litellm
|
||||
from litellm._service_logger import ServiceLogging
|
||||
|
|
@ -1033,8 +1037,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"])
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1474,6 +1479,411 @@ async def test_async_increment_pipelines_the_ttl_command(
|
|||
assert client.round_trips == expected_round_trips
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("existing", "expected"),
|
||||
[
|
||||
pytest.param(None, 1.5, id="missing key: set to the value"),
|
||||
pytest.param(0.25, 0.25, id="key below the value: kept"),
|
||||
pytest.param(4.0, 4.0, id="key above the value: kept"),
|
||||
],
|
||||
)
|
||||
async def test_async_increment_only_if_missing_sets_a_missing_key_and_refreshes_the_ttl(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, expected: float
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
client: Final = FakeAsyncRedis(server=FakeServer())
|
||||
if existing is not None:
|
||||
await client.set("ns:spend:key:k", existing, ex=5)
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client):
|
||||
current: Final = await redis_cache.async_increment(
|
||||
key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True
|
||||
)
|
||||
after_add: Final = await redis_cache.async_increment(key="spend:key:k", value=2.0, ttl=60, refresh_ttl=True)
|
||||
|
||||
assert current == expected
|
||||
assert after_add == expected + 2.0
|
||||
assert await client.ttl("ns:spend:key:k") == 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("value", "stored"),
|
||||
[(6.0, b"6"), (1e-05, b"0.00001"), (6.000000002, b"6.000000002"), (0.1, b"0.1"), (1e-20, b"0"), (-0.0, b"0")],
|
||||
)
|
||||
async def test_async_increment_only_if_missing_stores_a_missing_key_the_way_an_add_stores_it(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float, stored: bytes
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
client: Final = FakeAsyncRedis(server=FakeServer())
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client):
|
||||
current: Final = await redis_cache.async_increment(key="spend:key:k", value=value, ttl=60, only_if_missing=True)
|
||||
|
||||
assert await client.get("ns:spend:key:k") == stored
|
||||
assert current == float(stored)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("existing", "reported"), [pytest.param(None, ["created"], id="created"), pytest.param(7.0, [], id="kept")]
|
||||
)
|
||||
async def test_async_increment_only_if_missing_reports_only_a_key_it_created(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, existing: float | None, reported: list[str]
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
client: Final = FakeAsyncRedis(server=FakeServer())
|
||||
if existing is not None:
|
||||
await client.set("ns:spend:key:k", existing)
|
||||
created: Final[list[str]] = []
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client):
|
||||
await redis_cache.async_increment(
|
||||
key="spend:key:k", value=1.5, ttl=60, only_if_missing=True, on_added=lambda: created.append("created")
|
||||
)
|
||||
|
||||
assert created == reported
|
||||
|
||||
|
||||
class _RaisingPipeline:
|
||||
"""A pipeline whose execute raises the first error in its replies, as redis-py's does by default."""
|
||||
|
||||
def __init__(self, replies: list[object]) -> None:
|
||||
self.replies = replies
|
||||
self.commands: list[tuple[object, ...]] = []
|
||||
|
||||
async def __aenter__(self) -> "_RaisingPipeline":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
return None
|
||||
|
||||
def set(self, name: str, value: str, nx: bool = False, get: bool = False) -> None:
|
||||
self.commands.append(("SET", name, value, nx, get))
|
||||
|
||||
def incrbyfloat(self, name: str, amount: float) -> None:
|
||||
self.commands.append(("INCRBYFLOAT", name, amount))
|
||||
|
||||
def expire(self, name: str, ttl: int) -> None:
|
||||
self.commands.append(("EXPIRE", name, ttl))
|
||||
|
||||
async def execute(self) -> list[object]:
|
||||
first_error: Final = next((reply for reply in self.replies if isinstance(reply, Exception)), None)
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
return self.replies
|
||||
|
||||
|
||||
class _ClusterPipelineAfterADroppedConnection(_RaisingPipeline):
|
||||
"""The first attempt reads a dropped connection. As in redis-py 5.3.1's ClusterPipeline.execute, raising is what
|
||||
drives the retry, and ``raise_on_error=False`` hands the error back as a reply instead."""
|
||||
|
||||
async def execute(self, raise_on_error: bool = True) -> list[object]:
|
||||
if not raise_on_error:
|
||||
return [RedisConnectionError("Connection closed by server."), self.replies[1]]
|
||||
return self.replies
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_increment_only_if_missing_keeps_the_cluster_retry_of_a_dropped_connection(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.return_value = _ClusterPipelineAfterADroppedConnection([None, True])
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client):
|
||||
current: Final = await redis_cache.async_increment(
|
||||
key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True
|
||||
)
|
||||
|
||||
assert current == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"refusal",
|
||||
[
|
||||
pytest.param(NoPermissionError("no permissions to run the 'set' command"), id="denied"),
|
||||
pytest.param(ResponseError("syntax error"), id="redis-before-7-has-no-set-nx-get"),
|
||||
pytest.param(ResponseError("unknown command 'SET'"), id="renamed"),
|
||||
pytest.param(RedisError("SET command doesn't exist in Redis commands"), id="renamed-on-a-cluster-client"),
|
||||
],
|
||||
)
|
||||
async def test_async_increment_only_if_missing_adds_the_value_when_the_server_refuses_the_set(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, refusal: Exception
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
add: Final = _RaisingPipeline([7.5, True])
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.side_effect = [_RaisingPipeline([refusal, True]), add]
|
||||
failure_hook: Final = AsyncMock()
|
||||
added: Final = MagicMock()
|
||||
|
||||
with (
|
||||
patch.object(redis_cache, "init_async_client", return_value=client),
|
||||
patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook),
|
||||
):
|
||||
current: Final = await redis_cache.async_increment(
|
||||
key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert current == 7.5
|
||||
added.assert_called_once_with()
|
||||
assert add.commands == [("INCRBYFLOAT", "ns:spend:key:k", 1.5), ("EXPIRE", "ns:spend:key:k", 60)]
|
||||
failure_hook.assert_not_awaited()
|
||||
assert redis_cache._circuit_breaker._hard_failure_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_increment_only_if_missing_reports_a_write_redis_refuses_with_the_add_and_its_error_text(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
add_error: Final = ResponseError("Command # 1 (INCRBYFLOAT ns:spend:key:k 1.5) of pipeline caused error")
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.side_effect = [
|
||||
_RaisingPipeline([ResponseError("Command # 1 (SET) of pipeline caused error"), True]),
|
||||
_RaisingPipeline([add_error, True]),
|
||||
]
|
||||
failure_hook: Final = AsyncMock()
|
||||
added: Final = MagicMock()
|
||||
|
||||
with (
|
||||
patch.object(redis_cache, "init_async_client", return_value=client),
|
||||
patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook),
|
||||
pytest.raises(ResponseError) as raised,
|
||||
):
|
||||
await redis_cache.async_increment(
|
||||
key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert raised.value is add_error
|
||||
added.assert_not_called()
|
||||
failure_hook.assert_awaited_once()
|
||||
assert failure_hook.await_args.kwargs["error"] is add_error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"failure",
|
||||
[
|
||||
pytest.param(RedisConnectionError("Connection closed by server."), id="dropped-connection"),
|
||||
pytest.param(RedisTimeoutError("Timeout reading from socket"), id="stalled-node"),
|
||||
pytest.param(ClusterDownError("CLUSTERDOWN The cluster is down"), id="cluster-down"),
|
||||
],
|
||||
)
|
||||
async def test_async_increment_only_if_missing_sends_no_add_after_a_failure_the_client_already_retried(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, failure: Exception
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.return_value = _RaisingPipeline([failure, True])
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(type(failure)):
|
||||
await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True)
|
||||
|
||||
assert client.pipeline.call_count == 1
|
||||
|
||||
|
||||
class _Keyspace:
|
||||
"""String keys as Redis holds them, for pipelines that run each command in order and raise the first error after
|
||||
running the rest, as a non-transaction pipeline does. Commands named in ``refused`` are refused, a pipeline that
|
||||
has run ``writes_before_full`` write commands has the rest refused as Redis does past ``maxmemory``, and pipelines
|
||||
after the first ``reached`` never reach Redis, as when the connection drops."""
|
||||
|
||||
def __init__(
|
||||
self, refused: frozenset[str] = frozenset(), writes_before_full: int | None = None, reached: int | None = None
|
||||
) -> None:
|
||||
self.values: Final[dict[str, str]] = {} # mutable-ok: stands in for the Redis keyspace
|
||||
self.refused = refused
|
||||
self.writes_before_full = writes_before_full
|
||||
self.reached = reached
|
||||
self.pipelines = 0
|
||||
|
||||
def pipeline(self, **_: object) -> "_KeyspacePipeline":
|
||||
self.pipelines += 1
|
||||
return _KeyspacePipeline(self, lost=self.reached is not None and self.pipelines > self.reached)
|
||||
|
||||
|
||||
class _KeyspacePipeline(_RaisingPipeline):
|
||||
def __init__(self, keyspace: _Keyspace, lost: bool) -> None:
|
||||
super().__init__([])
|
||||
self.keyspace = keyspace
|
||||
self.lost = lost
|
||||
self.writes_left = keyspace.writes_before_full
|
||||
|
||||
def setnx(self, name: str, value: float) -> None:
|
||||
self.commands.append(("SETNX", name, value))
|
||||
|
||||
def _write(self, command: str) -> None:
|
||||
if self.writes_left is not None:
|
||||
if self.writes_left == 0:
|
||||
raise ResponseError("OOM command not allowed when used memory > 'maxmemory'.")
|
||||
self.writes_left -= 1
|
||||
|
||||
def _run(self, command: tuple[object, ...]) -> object:
|
||||
name: Final = str(command[1])
|
||||
held: Final = self.keyspace.values.get(name)
|
||||
if command[0] in self.keyspace.refused:
|
||||
raise NoPermissionError(f"this user has no permissions to run the '{str(command[0]).lower()}' command")
|
||||
if command[0] == "EXPIRE":
|
||||
return held is not None
|
||||
self._write(str(command[0]))
|
||||
if command[0] == "SET":
|
||||
if held is None:
|
||||
self.keyspace.values[name] = str(command[2])
|
||||
return None if held is None else held.encode()
|
||||
if command[0] == "SETNX":
|
||||
if held is None:
|
||||
self.keyspace.values[name] = repr(command[2])
|
||||
return held is None
|
||||
total: Final = float(held or 0) + float(cast(float, command[2]))
|
||||
self.keyspace.values[name] = repr(total)
|
||||
return total
|
||||
|
||||
async def execute(self) -> list[object]:
|
||||
if self.lost:
|
||||
raise RedisConnectionError("Connection closed by server.")
|
||||
replies: Final[list[object]] = []
|
||||
for number, command in enumerate(self.commands, start=1):
|
||||
try:
|
||||
replies.append(self._run(command))
|
||||
except RedisError as e:
|
||||
words: str = " ".join(str(part) for part in command)
|
||||
e.args = (f"Command # {number} ({words}) of pipeline caused error: {e.args}",)
|
||||
replies.append(e)
|
||||
first_error: Final = next((reply for reply in replies if isinstance(reply, Exception)), None)
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
return replies
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_increment_only_if_missing_leaves_a_missing_key_missing_when_the_set_is_refused_and_the_add_lost(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
keyspace: Final = _Keyspace(refused=frozenset(("SET", "SETNX")), reached=1)
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(RedisConnectionError):
|
||||
await redis_cache.async_increment(key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True)
|
||||
|
||||
assert keyspace.values == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_increment_only_if_missing_sets_a_missing_key_once_when_redis_runs_out_of_memory_after_the_set(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
keyspace: Final = _Keyspace(writes_before_full=1)
|
||||
added: Final = MagicMock()
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=keyspace):
|
||||
current: Final = await redis_cache.async_increment(
|
||||
key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added
|
||||
)
|
||||
|
||||
assert current == 6.0
|
||||
assert float(keyspace.values["ns:spend:key:k"]) == 6.0
|
||||
assert keyspace.pipelines == 1
|
||||
added.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_increment_only_if_missing_raises_a_refused_ttl_command_as_mains_add_does_without_adding_again(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
keyspace: Final = _Keyspace(refused=frozenset(("EXPIRE",)))
|
||||
added: Final = MagicMock()
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=keyspace), pytest.raises(NoPermissionError):
|
||||
await redis_cache.async_increment(
|
||||
key="spend:key:k", value=6.0, ttl=60, refresh_ttl=True, only_if_missing=True, on_added=added
|
||||
)
|
||||
|
||||
assert keyspace.values == {"ns:spend:key:k": "6"}
|
||||
assert keyspace.pipelines == 1
|
||||
added.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("held", ["not-a-number", "0x10", "6_0", " 6", "inf", "nan", "1e999"])
|
||||
async def test_async_increment_only_if_missing_leaves_a_value_it_cannot_read_like_incrbyfloat_to_mains_add(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, held: str
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
add: Final = _RaisingPipeline([7.5, True])
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.side_effect = [_RaisingPipeline([held.encode(), True]), add]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client):
|
||||
current: Final = await redis_cache.async_increment(
|
||||
key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True
|
||||
)
|
||||
|
||||
assert current == 7.5
|
||||
assert add.commands == [("INCRBYFLOAT", "ns:spend:key:k", 1.5), ("EXPIRE", "ns:spend:key:k", 60)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("value", [float("inf"), float("nan")])
|
||||
async def test_async_increment_only_if_missing_sends_a_value_incrbyfloat_refuses_as_mains_add(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object, value: float
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
add: Final = _RaisingPipeline([ResponseError("increment would produce NaN or Infinity"), True])
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.return_value = add
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client), pytest.raises(ResponseError):
|
||||
await redis_cache.async_increment(
|
||||
key="spend:key:k", value=value, ttl=60, refresh_ttl=True, only_if_missing=True
|
||||
)
|
||||
|
||||
assert client.pipeline.call_count == 1
|
||||
assert add.commands[0][0] == "INCRBYFLOAT"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_increment_only_if_missing_reports_a_redis_failure_like_an_increment(
|
||||
monkeypatch: pytest.MonkeyPatch, redis_no_ping: object
|
||||
):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache: Final = RedisCache(namespace="ns")
|
||||
client: Final = MagicMock()
|
||||
client.pipeline.side_effect = ConnectionError("connection reset")
|
||||
failure_hook: Final = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(redis_cache, "init_async_client", return_value=client),
|
||||
patch.object(redis_cache.service_logger_obj, "async_service_failure_hook", failure_hook),
|
||||
pytest.raises(ConnectionError),
|
||||
):
|
||||
await redis_cache.async_increment(key="spend:key:k", value=1.5, ttl=60, refresh_ttl=True, only_if_missing=True)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert failure_hook.await_args.kwargs["call_type"] == "async_increment"
|
||||
assert redis_cache._circuit_breaker._hard_failure_count == 1
|
||||
|
||||
|
||||
class _SetRecordingPipeline:
|
||||
def __init__(self) -> None:
|
||||
self.sets: list[tuple[str, str, timedelta | None]] = []
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue