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

fix(proxy): backport auth cache freshness and spend updates

Browse source 
(cherry picked from commit b55024658c)
This commit is contained in:
Yuneng Jiang 2026-11-02 13:58:11 -08:00
parent 76e42d73bd
commit 5352ed94fb
No known key found for this signature in database
13 changed files with 665 additions and 742 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/caching/dual_cache.py
View file

@ -59,16 +59,6 @@ class PendingBatchRead:
previous_access_times: dict[str, float | None]
@dataclass(frozen=True, slots=True)
class DeclaredBatchRead:
"""A ``async_batch_get_cache`` split in two: the memory half done, the Redis half declared on a ``RedisBatch``
so it rides that batch's next round trip, resolved later with ``async_resolve_batch_get``."""
keys: tuple[str, ...]
pending: PendingBatchRead
result: BatchResult[Mapping[str, object]] | None
def _log_deferred_increment_failure(future: asyncio.Future[float]) -> None:
if future.cancelled():
return
@ -382,22 +372,6 @@ class DualCache(BaseCache):
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
return merged
async def declare_batch_get(self, keys: Sequence[str], batch: RedisBatch) -> DeclaredBatchRead:
pending: Final = await self._prepare_batch_get(
list(keys),
local_only=False,
throttle_redis=False,
)
return DeclaredBatchRead(
keys=tuple(keys),
pending=pending,
result=batch.mget(pending.redis_keys) if pending.redis_keys else None,
)
async def async_resolve_batch_get(self, declared: DeclaredBatchRead) -> list[object | None]:
redis_result: Final = None if declared.result is None else await declared.result
return await self._apply_batch_get(declared.pending, redis_result)
async def async_batch_get_cache(
self,
keys: list,

47
litellm/caching/redis_cache.py
View file

@ -134,7 +134,12 @@ _INCREMENT_WITH_FLOOR_LUA: Final = (
"return count"
)
_INCREMENT_IF_PRESENT_LUA: Final = (
"if redis.call('EXISTS', KEYS[1]) == 0 then return false end return redis.call('INCRBYFLOAT', KEYS[1], ARGV[1])"
)
_LUA_COUNT: Final = TypeAdapter(int)
_LUA_OPTIONAL_TOTAL: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
@ -675,7 +680,7 @@ class RedisCache(BaseCache):
# redis namespaces
self.namespace = namespace
# for high traffic, we store the redis results in memory and then batch write to redis
self.redis_batch_writing_buffer: list = []
self.redis_batch_writing_buffer: list[tuple[str, object]] = []
if redis_flush_size is None:
self.redis_flush_size: int = 100
else:
@ -1568,6 +1573,46 @@ class RedisCache(BaseCache):
count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl)
return _LUA_COUNT.validate_python(count)
@_redis_circuit_breaker_guard
async def async_increment_if_present(
self, key: str, value: float, parent_otel_span: Span | None = None
) -> float | None:
_redis_client: Final = self._async_commands()
namespaced_key: Final = self.check_and_fix_namespace(key=key)
start_time: Final = time.time()
try:
total: Final = _LUA_OPTIONAL_TOTAL.validate_python(
await _redis_client.eval(_INCREMENT_IF_PRESENT_LUA, 1, namespaced_key, value)
)
except Exception as e:
failed_at: Final = time.time()
asyncio.create_task(
self.service_logger_obj.async_service_failure_hook(
service=ServiceTypes.REDIS,
duration=failed_at - start_time,
error=e,
call_type="async_increment_if_present",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=failed_at,
parent_otel_span=parent_otel_span,
)
)
raise
end_time: Final = time.time()
asyncio.create_task(
self.service_logger_obj.async_service_success_hook(
service=ServiceTypes.REDIS,
duration=end_time - start_time,
call_type="async_increment_if_present",
caller=_get_call_stack_info(),
start_time=start_time,
end_time=end_time,
parent_otel_span=parent_otel_span,
)
)
return total
async def flush_cache_buffer(self):
print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}")
await self.async_set_cache_pipeline(self.redis_batch_writing_buffer)

2
litellm/integrations/otel/README.md
View file

@ -91,7 +91,7 @@ way (`pod_lock`, `budget_reset`, `spend_queue`, `health_check`, `scheduler_queue
rather than a bare `redis.set`; `tests/unit/test_internal_context.py` scans every
module under `litellm/` and `enterprise/` that calls a shared cache or declares a
read or write on the request batch (`reserve_redis_batch_reads`,
`declare_batch_get`, `batch.mget`, `batch.set`, `batch.script`) and fails when
`batch.mget`, `batch.set`, `batch.script`) and fails when
one has no declared family, with the process-local `InMemoryCache` callers listed
as the only exemptions. A batch op carries the family that was active when it was
declared, so the routing prefetch armed before deployment selection

1
litellm/integrations/otel/model/spans.py
View file

@ -214,6 +214,7 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = {
"increment_cache": "incr",
"async_increment": "incr",
"async_increment_pipeline": "incr",
"async_increment_if_present": "incr",
"delete_cache": "delete",
"async_delete_cache": "delete",
"async_rpush": "rpush",

14
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -503,7 +503,6 @@ class _ProxyDBLogger(CustomLogger):
increment_spend_counters,
proxy_logging_obj,
update_cache,
update_cache_read_keys,
)
accounting: Final = (
@ -618,13 +617,6 @@ class _ProxyDBLogger(CustomLogger):
request_tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
update_cache_read_keys=update_cache_read_keys(
user_id=user_id,
end_user_id=end_user_id,
team_id=team_id,
tags=tags,
response_cost=response_cost,
),
)
if not charged:
return
@ -954,7 +946,6 @@ async def _update_database_and_spend_counters(
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
update_cache_read_keys: Sequence[str] = (),
) -> bool:
"""The reservation is reconciled before the spend is persisted, from its own read. One spend counter batch then
spans the database write and the counter update, so the post-call counters are read with a single MGET after the
@ -996,7 +987,6 @@ async def _update_database_and_spend_counters(
request_tags=request_tags,
model_access_groups=model_access_groups,
project_id=project_id,
update_cache_read_keys=update_cache_read_keys,
)
@ -1017,10 +1007,7 @@ async def _update_database_and_spend_counters_in_batch(
request_tags: list[str] | None,
model_access_groups: Sequence[str] | None,
project_id: str | None,
update_cache_read_keys: Sequence[str],
) -> bool:
from litellm.proxy.proxy_server import arm_update_cache_read
try:
writer: Final = proxy_logging_obj.db_spend_update_writer
accounting: Final = (
@ -1060,7 +1047,6 @@ async def _update_database_and_spend_counters_in_batch(
await _release_budget_reservation(budget_reservation=budget_reservation)
return False
await arm_update_cache_read(update_cache_read_keys)
try:
await increment_spend_counters(
token=user_api_key,

352
litellm/proxy/proxy_server.py
View file

@ -144,7 +144,6 @@ from litellm.proxy._types import (
TransformRequestBody,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.callback_utils import (
is_sensitive_callback_key,
normalize_callback_names,
@ -302,10 +301,8 @@ from litellm import Router
from litellm._internal_context import service_target, with_service_target
from litellm._logging import redact_string, verbose_proxy_logger, verbose_router_logger
from litellm.caching.caching import DualCache, RedisCache
from litellm.caching.dual_cache import DeclaredBatchRead
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
active_request_redis_batches,
drain_post_call_redis_batches,
)
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure
@ -4153,75 +4150,61 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) -
return tuple(float(current_value) for current_value in results or ())
def update_cache_read_keys(
user_id: str | None,
end_user_id: str | None,
team_id: str | None,
tags: Sequence[object] | None,
response_cost: float | None,
) -> tuple[str, ...]:
if response_cost is None:
return ()
user_keys: tuple[str, ...] = (user_id, GLOBAL_PROXY_SPEND_CACHE_KEY) if user_id is not None else ()
end_user_keys: tuple[str, ...] = (end_user_cache_key(end_user_id),) if end_user_id is not None else ()
team_keys: tuple[str, ...] = (f"team_id:{team_id}",) if team_id is not None else ()
tag_keys: tuple[str, ...] = tuple(tag_cache_key(tag) for tag in tags or () if isinstance(tag, str) and tag)
return user_keys + end_user_keys + team_keys + tag_keys
_CachedSpendObjectT = TypeVar(
"_CachedSpendObjectT", LiteLLM_UserTable, LiteLLM_EndUserTable, LiteLLM_TeamTableCachedObj, LiteLLM_TagTable
)
_UPDATE_CACHE_PREFETCH_SLOT: Final = "update_cache_read"
async def arm_update_cache_read(keys: Sequence[str], cache: DualCache | None = None) -> None:
"""Declares the ``update_cache`` read on the request pipeline once the spend is persisted, so it rides the same
round trip as the post-call spend counter read instead of its own."""
request: Final = active_request_redis_batches()
target: Final = user_api_key_cache if cache is None else cache
if request is None or target.redis_cache is None or not keys:
def _add_spend_to_locally_cached_object(
cache: UserApiKeyCache, key: str, model_type: type[_CachedSpendObjectT], response_cost: float
) -> None:
current: Final = cache.get_cache(key=key, local_only=True, model_type=model_type)
if current is None:
return
with service_target(AUTH_OBJECTS_TARGET):
request.prefetched[_UPDATE_CACHE_PREFETCH_SLOT] = await target.declare_batch_get(
keys, request.batch(target.redis_cache)
)
cache.set_cache(
key=key,
value=current.model_copy(update={"spend": (current.spend or 0.0) + response_cost}),
local_only=True,
model_type=model_type,
ttl=get_management_object_ttl(cache),
)
async def _take_armed_update_cache_read(keys: Sequence[str], cache: DualCache) -> Mapping[str, object] | None:
request: Final = active_request_redis_batches()
if request is None:
async def _add_spend_to_shared_global_proxy_spend(
cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None
) -> float | None:
if cache.redis_cache is None:
return None
armed: Final = request.prefetched.pop(_UPDATE_CACHE_PREFETCH_SLOT, None)
if not isinstance(armed, DeclaredBatchRead) or armed.keys != tuple(keys):
return None
values: Final = await cache.async_resolve_batch_get(armed)
return MappingProxyType({key: value for key, value in zip(keys, values) if value is not None})
async def _read_update_cache_values(
keys: Sequence[str], parent_otel_span: Span | None, cache: DualCache | None = None
) -> Mapping[str, object]:
"""One batched read for every object ``update_cache`` refreshes; a failed read leaves them all untouched,
exactly as a failed per-object GET left that object untouched."""
if not keys:
return MappingProxyType({})
target: Final = user_api_key_cache if cache is None else cache
try:
armed: Final = await _take_armed_update_cache_read(keys, target)
if armed is not None:
return armed
values: Final = await target.async_batch_get_cache(
keys=list(keys), parent_otel_span=parent_otel_span, throttle_redis=False
return await cache.redis_cache.async_increment_if_present(
key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=response_cost, parent_otel_span=parent_otel_span
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to read cached spend objects. Budget enforcement may use stale spend values. "
"keys=%s - %s",
keys,
"Spend tracking - failed to add spend to the global proxy spend in Redis. "
"Budget enforcement may use stale spend values. response_cost=%s - %s",
response_cost,
str(e),
)
return MappingProxyType({})
if values is None:
return MappingProxyType({})
return MappingProxyType({key: value for key, value in zip(keys, values) if value is not None})
return None
async def _add_spend_to_global_proxy_spend(
cache: UserApiKeyCache, response_cost: float, parent_otel_span: Span | None
) -> None:
local_before: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True)
shared_total: Final = await _add_spend_to_shared_global_proxy_spend(cache, response_cost, parent_otel_span)
local_total: Final = cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True)
if not isinstance(local_total, (int, float)):
return
if shared_total is not None and local_total != local_before and shared_total < local_total:
return
cache.set_cache(
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
value=local_total + response_cost if shared_total is None else shared_total,
local_only=True,
ttl=get_management_object_ttl(cache),
)
async def update_cache(
@ -4232,6 +4215,7 @@ async def update_cache(
response_cost: float | None,
parent_otel_span: Span | None,
tags: list[str] | None = None,
cache: UserApiKeyCache | None = None,
):
"""
Use this to update the cache with new user spend.
@ -4239,14 +4223,7 @@ async def update_cache(
Put any alerting logic in here.
"""
values_to_update_in_cache: Final[list[tuple[str, object]]] = []
with service_target(AUTH_OBJECTS_TARGET):
cached_values: Final = await _read_update_cache_values(
keys=update_cache_read_keys(
user_id=user_id, end_user_id=end_user_id, team_id=team_id, tags=tags, response_cost=response_cost
),
parent_otel_span=parent_otel_span,
)
target: Final = user_api_key_cache if cache is None else cache
### UPDATE KEY SPEND ###
async def _update_key_cache(token: str, response_cost: float):
@ -4256,7 +4233,7 @@ async def update_cache(
else:
hashed_token = token
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
existing_spend_obj = await user_api_key_cache.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth)
existing_spend_obj = await target.async_get_cache(key=hashed_token, model_type=UserAPIKeyAuth)
verbose_proxy_logger.debug("_update_key_cache: existing_spend_obj=%s", existing_spend_obj)
if existing_spend_obj is None:
return
@ -4298,191 +4275,17 @@ async def update_cache(
)
# set cooldown on alert
### UPDATE USER SPEND ###
async def _update_user_cache():
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
if response_cost is None:
return
user_ids: Final = [user_id]
def _add_spend(key: str, model_type: type[_CachedSpendObjectT], cost: float) -> None:
try:
for _id in user_ids:
# Fetch the existing cost for the given user
if _id is None:
continue
cached_user = cached_values.get(_id)
if cached_user is None:
# do nothing if there is no cache value
return
existing_spend_obj = CacheCodec.deserialize(cached_user, LiteLLM_UserTable)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
"_update_user_db: existing spend: %s; response_cost: %s", existing_spend_obj, response_cost
)
existing_spend = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_UserTable),
)
)
## UPDATE GLOBAL PROXY ##
global_proxy_spend: Final = cached_values.get(GLOBAL_PROXY_SPEND_CACHE_KEY)
if not isinstance(global_proxy_spend, (int, float)):
# do nothing if not in cache
return
elif response_cost is not None:
increment: Final = global_proxy_spend + response_cost
values_to_update_in_cache.append((GLOBAL_PROXY_SPEND_CACHE_KEY, increment))
_add_spend_to_locally_cached_object(target, key, model_type, cost)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update user spend in cache. "
"Spend tracking - failed to update %s spend in cache. "
"Budget enforcement may use stale spend values. "
"user_id=%s, response_cost=%s - %s\n%s",
user_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE END-USER SPEND ###
async def _update_end_user_cache():
if end_user_id is None or response_cost is None:
return
_id: Final = end_user_cache_key(end_user_id)
try:
# Fetch the existing cost for the given user
cached_end_user: Final = cached_values.get(_id)
if cached_end_user is None:
# if user does not exist in LiteLLM_UserTable, create a new user
# do nothing if end-user not in api key cache
return
existing_spend_obj: Final = CacheCodec.deserialize(cached_end_user, LiteLLM_EndUserTable)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
"_update_end_user_db: existing spend: %s; response_cost: %s", existing_spend_obj, response_cost
)
existing_spend: Final = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend: Final = existing_spend + response_cost
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_EndUserTable),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update end user spend in cache. "
"Budget enforcement may use stale spend values. "
"end_user_id=%s, response_cost=%s - %s\n%s",
end_user_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE TEAM SPEND ###
async def _update_team_cache():
if team_id is None or response_cost is None:
return
_id: Final = f"team_id:{team_id}"
try:
cached_team: Final = cached_values.get(_id)
if cached_team is None:
# do nothing if team not in api key cache
return
existing_spend_obj: Final[LiteLLM_TeamTableCachedObj | None] = CacheCodec.deserialize(
cached_team, LiteLLM_TeamTableCachedObj
)
if existing_spend_obj is None:
return
verbose_proxy_logger.debug(
"_update_team_db: existing spend: %s; response_cost: %s", existing_spend_obj, response_cost
)
existing_spend: Final[float] = existing_spend_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend: Final = existing_spend + response_cost
existing_spend_obj.spend = new_spend
values_to_update_in_cache.append(
(
_id,
CacheCodec.serialize(existing_spend_obj, model_type=LiteLLM_TeamTableCachedObj),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update team spend in cache. "
"Budget enforcement may use stale spend values. "
"team_id=%s, response_cost=%s - %s\n%s",
team_id,
response_cost,
str(e),
traceback.format_exc(),
)
### UPDATE TAG SPEND ###
async def _update_tag_cache():
"""
Update the tag cache with the new spend.
"""
if tags is None or response_cost is None:
return
try:
for tag_name in tags:
if not tag_name or not isinstance(tag_name, str):
continue
cache_key = tag_cache_key(tag_name)
# Fetch the existing tag object from cache
cached_tag = cached_values.get(cache_key)
if cached_tag is None:
# do nothing if tag not in api key cache
continue
existing_tag_obj = CacheCodec.deserialize(cached_tag, LiteLLM_TagTable)
if existing_tag_obj is None:
continue
verbose_proxy_logger.debug(
"_update_tag_cache: existing spend for tag=%s: %s; response_cost: %s",
tag_name,
existing_tag_obj,
response_cost,
)
existing_spend = existing_tag_obj.spend or 0.0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
existing_tag_obj.spend = new_spend
values_to_update_in_cache.append(
(
cache_key,
CacheCodec.serialize(existing_tag_obj, model_type=LiteLLM_TagTable),
)
)
except Exception as e:
verbose_proxy_logger.warning(
"Spend tracking - failed to update tag spend in cache. "
"Budget enforcement may use stale spend values. "
"tags=%s, response_cost=%s - %s\n%s",
tags,
response_cost,
"key=%s, response_cost=%s - %s\n%s",
model_type.__name__,
key,
cost,
str(e),
traceback.format_exc(),
)
@ -4490,42 +4293,23 @@ async def update_cache(
with service_target(AUTH_OBJECTS_TARGET):
if token is not None and response_cost is not None:
await _update_key_cache(token=token, response_cost=response_cost)
if response_cost is None:
return
if user_id is not None:
await _update_user_cache()
if user_id is not None:
_add_spend(user_id, LiteLLM_UserTable, response_cost)
if end_user_id is not None:
_add_spend(end_user_cache_key(end_user_id), LiteLLM_EndUserTable, response_cost)
if team_id is not None:
_add_spend(f"team_id:{team_id}", LiteLLM_TeamTableCachedObj, response_cost)
for tag_name in tags or ():
if tag_name and isinstance(tag_name, str):
_add_spend(tag_cache_key(tag_name), LiteLLM_TagTable, response_cost)
if end_user_id is not None:
await _update_end_user_cache()
if team_id is not None:
await _update_team_cache()
if tags is not None:
await _update_tag_cache()
global_proxy_spend_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
local_object_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key)
shared_scalar_updates: Final = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key)
if local_object_updates:
with service_target(AUTH_OBJECTS_TARGET):
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=list(local_object_updates),
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
)
if shared_scalar_updates:
with service_target(SPEND_COUNTERS_TARGET):
asyncio.create_task(
user_api_key_cache.async_set_cache_pipeline(
cache_list=list(shared_scalar_updates),
ttl=get_management_object_ttl(user_api_key_cache),
litellm_parent_otel_span=parent_otel_span,
)
)
if litellm.max_budget <= 0:
return
with service_target(SPEND_COUNTERS_TARGET):
await _add_spend_to_global_proxy_spend(target, response_cost, parent_otel_span)
def run_ollama_serve():

58
tests/integration/spend/test_global_proxy_spend_counter.py Normal file
View file

@ -0,0 +1,58 @@
import os
import uuid
from typing import Final
import pytest
import litellm
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.proxy_server import update_cache
async def _add_request_cost(cache: UserApiKeyCache, response_cost: float) -> None:
await update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=response_cost,
parent_otel_span=None,
cache=cache,
)
@pytest.mark.asyncio
async def test_global_proxy_spend_adds_every_instance_cost_and_stays_cleared_after_a_budget_reset(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "max_budget", 100.0)
redis_cache: Final = RedisCache(
host=os.environ["REDIS_HOST"],
port=int(os.environ["REDIS_PORT"]),
namespace=f"global-proxy-spend-{uuid.uuid4().hex}",
)
gateway: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
peer: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
try:
await gateway.async_set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, ttl=60)
await _add_request_cost(peer, 0.5)
await _add_request_cost(gateway, 0.25)
before_reset: Final = (
gateway.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True),
await redis_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY),
)
ttl_before_reset: Final = await redis_cache.async_get_ttl(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
await gateway.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
await _add_request_cost(peer, 0.5)
await _add_request_cost(gateway, 0.25)
assert before_reset == (10.75, 10.75)
assert ttl_before_reset is not None and 0 < ttl_before_reset <= 60
assert await redis_cache.async_get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None
finally:
await redis_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)

143
tests/unit/caching/test_request_redis_batch_post_call.py
View file

@ -10,31 +10,27 @@ import datetime
import hashlib
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from unittest.mock import AsyncMock, MagicMock
import pytest
import litellm
from litellm._internal_context import current_service_target
from litellm.caching.caching import Cache
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_batch import (
active_post_call_redis_batch,
active_request_redis_batches,
drain_post_call_redis_batches,
flush_post_call_redis_batches,
request_redis_batch_scope,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PARALLEL_RELEASE_SCRIPT,
TOKEN_INCREMENT_SCRIPT,
ParallelSlotAcquisition,
RequestRateLimiterStash,
PROXY_MaxParallelRequestsHandler_v3,
RequestRateLimiterStash,
)
from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement
from litellm.proxy.utils import InternalUsageCache
@ -518,140 +514,3 @@ async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_co
assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it"
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend"
assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded"
@pytest.mark.asyncio
async def test_the_update_cache_read_armed_before_accounting_rides_the_pipeline_of_the_reconcile_read():
from litellm.proxy.proxy_server import _read_update_cache_values, arm_update_cache_read
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
cache = DualCache()
cache.attach_redis_cache(redis_cache)
keys = ["user-1", "team_id:t1"]
with request_redis_batch_scope() as request:
await arm_update_cache_read(keys, cache=cache)
assert client.pipelines == []
await request.batch(redis_cache).mget(["spend:key:k1"]) # the spend reconcile read of the same request
values = await _read_update_cache_values(keys, None, cache=cache)
assert len(client.pipelines) == 1
assert client.pipelines[0].commands == [("MGET", "user-1", "team_id:t1"), ("MGET", "spend:key:k1")]
assert values == {"user-1": {"spend": 1.0}, "team_id:t1": {"spend": 1.0}}
assert redis_cache.alone == []
assert active_request_redis_batches() is None
@pytest.mark.asyncio
async def test_the_armed_update_cache_read_is_declared_under_the_auth_objects_family():
"""The user, team and tag rows the accounting reads are auth objects, so the pipeline that carries
the armed read renders ``redis.pipeline auth_objects``, not a bare ``redis.pipeline``."""
from litellm.proxy.proxy_server import _read_update_cache_values, arm_update_cache_read
client = FakeClient(_ok_replies)
redis_cache = PostCallFakeRedisCache(client)
cache = DualCache()
cache.attach_redis_cache(redis_cache)
pipeline_targets: list[str | None] = [] # mutable-ok: filled by the recording hook
async def record(**kwargs: object) -> None:
pipeline_targets.append(current_service_target())
redis_cache.service_logger_obj.async_service_success_hook = record # pyright: ignore[reportAttributeAccessIssue] # fake, records the hook call
with request_redis_batch_scope():
await arm_update_cache_read(["user-1", "team_id:t1"], cache=cache)
await _read_update_cache_values(["user-1", "team_id:t1"], None, cache=cache)
await asyncio.gather(*(t for t in asyncio.all_tasks() if t is not asyncio.current_task()))
assert pipeline_targets == [AUTH_OBJECTS_TARGET]
@pytest.mark.asyncio
async def test_an_update_cache_read_armed_for_other_keys_is_ignored_and_the_read_happens_as_before():
from litellm.proxy.proxy_server import _read_update_cache_values, arm_update_cache_read
redis_cache = PostCallFakeRedisCache(FakeClient(_ok_replies))
redis_cache.store["team_id:t1"] = {"spend": 2.0}
cache = DualCache()
cache.attach_redis_cache(redis_cache)
with request_redis_batch_scope():
await arm_update_cache_read(["user-1"], cache=cache)
values = await _read_update_cache_values(["team_id:t1"], None, cache=cache)
assert values == {"team_id:t1": {"spend": 2.0}}
assert ("MGET", ("team_id:t1",)) in redis_cache.alone
@pytest.mark.asyncio
async def test_the_update_cache_read_sees_a_cached_spend_written_while_the_spend_was_persisted(monkeypatch):
from litellm.proxy import proxy_server
from litellm.proxy.hooks.proxy_track_cost_callback import _update_database_and_spend_counters
cached_user_spend = {"user-1": 1.0}
def replies(command: tuple[object, ...]) -> object:
if command[0] == "MGET":
return [
json.dumps({"spend": cached_user_spend[key]}) if key in cached_user_spend else b"0.5"
for key in command[1:]
]
return _ok_replies(command)
client = FakeClient(replies)
redis_cache = PostCallFakeRedisCache(client)
spend_cache = DualCache()
spend_cache.attach_redis_cache(redis_cache)
user_cache = DualCache()
user_cache.attach_redis_cache(redis_cache)
monkeypatch.setattr(proxy_server, "spend_counter_cache", spend_cache)
monkeypatch.setattr(proxy_server, "user_api_key_cache", user_cache)
async def _read_on_the_request_pipeline_then_a_concurrent_callback_writes_the_user(**kwargs: object) -> bool:
request = active_request_redis_batches()
assert request is not None
await request.batch(redis_cache).mget(["key-object"])
cached_user_spend["user-1"] = 5.0
return True
proxy_logging_obj = MagicMock()
proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(
side_effect=_read_on_the_request_pipeline_then_a_concurrent_callback_writes_the_user
)
reservation = {
"reserved_cost": 0.5,
"entries": [
{
"counter_key": "spend:key:k1",
"entity_type": "Key",
"entity_id": "k1",
"reserved_cost": 0.5,
"applied_adjustment": 0.0,
}
],
"finalized": False,
}
with request_redis_batch_scope():
charged = await _update_database_and_spend_counters(
proxy_logging_obj=proxy_logging_obj,
increment_spend_counters=proxy_server.increment_spend_counters,
user_api_key="k1",
user_id="user-1",
end_user_id=None,
team_id=None,
org_id=None,
kwargs={},
completion_response=None,
start_time=datetime.datetime.now(),
end_time=datetime.datetime.now(),
response_cost=0.2,
budget_reservation=reservation,
update_cache_read_keys=("user-1",),
)
values = await proxy_server._read_update_cache_values(("user-1",), None)
assert charged is True
assert values == {"user-1": {"spend": 5.0}}, client.pipelines

1
tests/unit/integrations/otel/test_otel_v2_logger.py
View file

@ -1885,6 +1885,7 @@ def test_service_span_is_named_by_purpose_when_the_producer_declares_a_target():
("async_batch_get_cache", "redis.mget auth_objects", "redis.mget"),
("async_set_cache_pipeline_with_ttls", "redis.set auth_objects", "redis.set"),
("async_increment_pipeline", "redis.incr auth_objects", "redis.incr"),
("async_increment_if_present", "redis.incr auth_objects", "redis.incr"),
("async_delete_cache", "redis.delete auth_objects", "redis.delete"),
("async_scan_iter", "redis.scan auth_objects", "redis.scan"),
("request_redis_batch", "redis.pipeline auth_objects", "redis.pipeline"),

469
tests/unit/proxy/proxy_server/test_spend_counters.py
View file

@ -22,13 +22,34 @@ Pins covered:
from __future__ import annotations
import asyncio
from collections.abc import Awaitable, Callable
from datetime import datetime
from typing import Final, cast
from unittest.mock import AsyncMock, MagicMock
from typing import Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, call
import pytest
from fakeredis import FakeAsyncRedis, FakeServer
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import StatusCode
import litellm
import litellm.proxy.proxy_server as ps
from litellm._service_logger import ServiceLogging
from litellm._internal_context import current_service_target
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.models.end_user import LiteLLM_EndUserTable
from litellm.models.tag import LiteLLM_TagTable
from litellm.models.team import LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj
from litellm.models.user import LiteLLM_UserTable
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_local
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, end_user_cache_key, tag_cache_key
from litellm.proxy.management_endpoints.team_endpoints import refresh_cached_team
from .conftest import normalize
@ -1554,32 +1575,6 @@ async def test_invalidate_spend_counter_swallows_redis_failure_no_raise(monkeypa
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_update_cache_no_cached_entities_schedules_pipeline_flush(monkeypatch):
fake_user_cache = _make_user_api_key_cache(get_value=None)
monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache)
await ps.update_cache(
token=None,
user_id="u1",
end_user_id="eu1",
team_id="t1",
response_cost=1.0,
parent_otel_span=None,
tags=["x"],
)
assert fake_user_cache.async_get_cache.await_count == 0
fake_user_cache.async_batch_get_cache.assert_awaited_once()
assert fake_user_cache.async_batch_get_cache.await_args.kwargs["keys"] == [
"u1",
f"{ps.litellm_proxy_admin_name}:spend",
"end_user_id:eu1",
"team_id:t1",
"tag:x",
]
@pytest.mark.asyncio
async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkeypatch):
"""An inner _update_user_cache raising must not propagate — update_cache
@ -1600,3 +1595,421 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey
)
assert result is None
async def _drain_spend_cache_writes() -> None:
for _ in range(3):
await asyncio.sleep(0)
async def _update_team_spend(team_id: str, response_cost: float) -> None:
await ps.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=team_id,
response_cost=response_cost,
parent_otel_span=None,
)
@pytest.mark.asyncio
async def test_update_cache_adds_spend_to_the_cached_team_only(monkeypatch: pytest.MonkeyPatch) -> None:
cache: Final = UserApiKeyCache()
cache.set_cache(key="team_id:warm", value=LiteLLM_TeamTableCachedObj(team_id="warm", spend=5.0, blocked=True))
monkeypatch.setattr(ps, "user_api_key_cache", cache)
await _update_team_spend(team_id="warm", response_cost=2.5)
await _update_team_spend(team_id="cold", response_cost=2.5)
await _drain_spend_cache_writes()
warm: Final = cache.get_cache(key="team_id:warm", model_type=LiteLLM_TeamTableCachedObj)
assert warm is not None
assert (warm.spend, warm.blocked) == (7.5, True)
assert cache.get_cache(key="team_id:cold") is None
async def _block_on_this_worker(cache: UserApiKeyCache, team_id: str) -> None:
await refresh_cached_team(
team_row=LiteLLM_TeamTable(team_id=team_id, blocked=True),
user_api_key_cache=cache,
proxy_logging_obj=ps.proxy_logging_obj,
)
async def _invalidated_by_a_peer(cache: UserApiKeyCache, team_id: str) -> None:
evict_local(cache_keys=(f"team_id:{team_id}",), user_api_key_cache=cache)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("management_write", "expected_blocked"),
[(_block_on_this_worker, True), (_invalidated_by_a_peer, None)],
)
async def test_update_cache_never_restores_a_team_a_management_write_replaced(
management_write: Callable[[UserApiKeyCache, str], Awaitable[None]],
expected_blocked: bool | None,
monkeypatch: pytest.MonkeyPatch,
) -> None:
team_id: Final = "kill-switch"
cache: Final = UserApiKeyCache()
cache.set_cache(key=f"team_id:{team_id}", value=LiteLLM_TeamTableCachedObj(team_id=team_id, blocked=False))
monkeypatch.setattr(ps, "user_api_key_cache", cache)
await _update_team_spend(team_id=team_id, response_cost=1.0)
await management_write(cache, team_id)
await _drain_spend_cache_writes()
seen_by_auth: Final = cache.get_cache(key=f"team_id:{team_id}", model_type=LiteLLM_TeamTableCachedObj)
assert (seen_by_auth.blocked if seen_by_auth is not None else None) == expected_blocked
_CachedSpendObject = LiteLLM_UserTable | LiteLLM_EndUserTable | LiteLLM_TeamTableCachedObj | LiteLLM_TagTable
def _cached_spend_object(
kind: Literal["user", "end_user", "team", "tag"], models: list[str]
) -> tuple[str, _CachedSpendObject]:
match kind:
case "user":
return "race-user", LiteLLM_UserTable(user_id="race-user", models=models)
case "end_user":
return end_user_cache_key("race-end-user"), LiteLLM_EndUserTable(
user_id="race-end-user", blocked=False, models=models
)
case "team":
return "team_id:race-team", LiteLLM_TeamTableCachedObj(team_id="race-team", models=models)
case "tag":
return tag_cache_key("race-tag"), LiteLLM_TagTable(tag_name="race-tag", models=models)
class _RunsOnKeyRead(InMemoryCache):
def __init__(self, on_read: Callable[[], Awaitable[None]]) -> None:
super().__init__()
self._on_read: Final = on_read
async def async_get_cache(self, key: str, **kwargs: object) -> object:
await self._on_read()
return await super().async_get_cache(key, **kwargs)
def _cache_that_runs_on_key_read(on_read: Callable[[], Awaitable[None]]) -> UserApiKeyCache:
return UserApiKeyCache(in_memory_cache=InMemoryCache(), key_object_in_memory_cache=_RunsOnKeyRead(on_read))
@pytest.fixture
def proxy_max_budget(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "max_budget", 100.0)
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ["user", "end_user", "team", "tag"])
@pytest.mark.parametrize(
("management_write", "expected"),
[("none", (["gpt-a"], 0.25)), ("evict", None), ("refresh", (["gpt-b"], 0.25))],
)
async def test_update_cache_adds_spend_to_the_entry_a_management_write_left_mid_flight(
kind: Literal["user", "end_user", "team", "tag"],
management_write: Literal["none", "evict", "refresh"],
expected: tuple[list[str], float] | None,
) -> None:
key, seeded = _cached_spend_object(kind, models=["gpt-a"])
_, refreshed = _cached_spend_object(kind, models=["gpt-b"])
async def _write() -> None:
if management_write == "none":
return
evict_local(cache_keys=(key,), user_api_key_cache=cache)
if management_write == "refresh":
cache.set_cache(key=key, value=refreshed, local_only=True)
cache: Final = _cache_that_runs_on_key_read(on_read=_write)
cache.set_cache(key=key, value=seeded, local_only=True)
await ps.update_cache(
token="sk-race",
user_id="race-user",
end_user_id="race-end-user",
team_id="race-team",
response_cost=0.25,
parent_otel_span=None,
tags=["race-tag"],
cache=cache,
)
await _drain_spend_cache_writes()
stored: Final = cache.get_cache(key=key, local_only=True, model_type=type(seeded))
assert ((stored.models, stored.spend) if stored is not None else None) == expected
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
@pytest.mark.parametrize("user", ["cached_on_this_worker", "not_cached_on_this_worker", "none"])
async def test_update_cache_adds_spend_to_the_global_proxy_spend_for_every_request(
user: Literal["cached_on_this_worker", "not_cached_on_this_worker", "none"],
) -> None:
cache: Final = UserApiKeyCache()
cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0)
if user == "cached_on_this_worker":
cache.set_cache(key="spender", value=LiteLLM_UserTable(user_id="spender", spend=1.0))
await ps.update_cache(
token=None,
user_id=None if user == "none" else "spender",
end_user_id=None,
team_id=None,
response_cost=0.25,
parent_otel_span=None,
cache=cache,
)
await _drain_spend_cache_writes()
assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) == 10.25
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_never_brings_back_a_global_proxy_spend_a_budget_reset_cleared() -> None:
async def _budget_reset() -> None:
await cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
cache: Final = _cache_that_runs_on_key_read(on_read=_budget_reset)
cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0)
await ps.update_cache(
token="sk-spender",
user_id="spender",
end_user_id=None,
team_id=None,
response_cost=0.25,
parent_otel_span=None,
cache=cache,
)
await _drain_spend_cache_writes()
assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY) is None
class _FakeRedisRunningIncrementIfPresent(FakeAsyncRedis):
def __init__(self, server: FakeServer, after_increment: Callable[[], object] = lambda: None) -> None:
super().__init__(server=server)
self._after_increment: Final = after_increment
async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None:
key, increment = keys_and_args
if not await self.exists(key):
return None
total: Final = await self.incrbyfloat(key, increment)
self._after_increment()
return str(total).encode()
class _RedisCacheOnClient(RedisCache):
def __init__(self, client: FakeAsyncRedis) -> None:
self._injected_client: Final = client
self.namespace = None
self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30)
self.service_logger_obj = ServiceLogging()
self.default_ttl = None
def init_async_client(self) -> FakeAsyncRedis:
return self._injected_client
class _FakeRedisHoldingOneReply(FakeAsyncRedis):
def __init__(self, server: FakeServer, held_increment: float, applied: asyncio.Event, released: asyncio.Event) -> None:
super().__init__(server=server)
self._held_increment: Final = held_increment
self._applied: Final = applied
self._released: Final = released
async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None:
key, increment = keys_and_args
total: Final = await self.incrbyfloat(key, increment)
if float(increment) == self._held_increment:
self._applied.set()
await self._released.wait()
return str(total).encode()
_SharedGlobalSpendOutcome = Literal["peers_added_spend", "reset_elsewhere", "reset_while_redis_answers", "redis_down"]
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
@pytest.mark.parametrize(
("outcome", "shared_total", "expected"),
[
("peers_added_spend", 11.75, 12.0),
("reset_elsewhere", None, 10.25),
("reset_while_redis_answers", 9.75, None),
("redis_down", 10.0, 10.25),
],
)
async def test_update_cache_moves_this_workers_global_proxy_spend_to_the_shared_total(
outcome: _SharedGlobalSpendOutcome, shared_total: float | None, expected: float | None
) -> None:
def _reset_this_worker() -> None:
if outcome == "reset_while_redis_answers":
cache.in_memory_cache.delete_cache(GLOBAL_PROXY_SPEND_CACHE_KEY)
server: Final = FakeServer()
redis: Final = _FakeRedisRunningIncrementIfPresent(server, after_increment=_reset_this_worker)
if shared_total is not None:
await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, shared_total)
server.connected = outcome != "redis_down"
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, local_only=True)
await ps.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=0.25,
parent_otel_span=None,
cache=cache,
)
assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) == expected
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_keeps_the_newest_global_proxy_spend_when_redis_replies_out_of_order() -> None:
held_increment_applied: Final = asyncio.Event()
held_reply_released: Final = asyncio.Event()
server: Final = FakeServer()
redis: Final = _FakeRedisHoldingOneReply(
server, held_increment=0.25, applied=held_increment_applied, released=held_reply_released
)
await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0)
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0, local_only=True)
async def _request(response_cost: float) -> None:
await ps.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=response_cost,
parent_otel_span=None,
cache=cache,
)
earlier: Final = asyncio.create_task(_request(0.25))
await held_increment_applied.wait()
await _request(0.5)
held_reply_released.set()
await earlier
assert cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, local_only=True) == 10.75
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_adds_the_global_spend_scalar_in_redis_as_spend_counters() -> None:
target_during_increment: Final = MagicMock()
redis: Final = _FakeRedisRunningIncrementIfPresent(
FakeServer(), after_increment=lambda: target_during_increment(current_service_target())
)
await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0)
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
await ps.update_cache(
token=None,
user_id="u1",
end_user_id=None,
team_id=None,
response_cost=2.0,
parent_otel_span=None,
cache=cache,
)
assert target_during_increment.call_args_list == [call("spend_counters")]
assert float(await redis.get(GLOBAL_PROXY_SPEND_CACHE_KEY)) == 12.0
assert current_service_target() is None
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
@pytest.mark.parametrize(("redis_up", "expected_status"), [(True, StatusCode.OK), (False, StatusCode.ERROR)])
async def test_update_cache_traces_the_global_spend_redis_call_under_the_request_span(
monkeypatch: pytest.MonkeyPatch, redis_up: bool, expected_status: StatusCode
) -> None:
exporter: Final = InMemorySpanExporter()
provider: Final = TracerProvider()
provider.add_span_processor(SimpleSpanProcessor(exporter))
monkeypatch.setattr(
litellm,
"service_callback",
[OpenTelemetry(config=OpenTelemetryConfig(exporter=exporter), tracer_provider=provider)],
)
server: Final = FakeServer()
redis: Final = _FakeRedisRunningIncrementIfPresent(server)
await redis.set(GLOBAL_PROXY_SPEND_CACHE_KEY, 10.0)
server.connected = redis_up
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis))
request_span: Final = provider.get_tracer("test").start_span("request")
await ps.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=2.0,
parent_otel_span=request_span,
cache=cache,
)
await _drain_spend_cache_writes()
request_span.end()
redis_spans: Final = tuple(
(span.attributes.get("call_type"), span.status.status_code, span.parent.span_id)
for span in exporter.get_finished_spans()
if span.name == "redis" and span.attributes is not None and span.parent is not None
)
assert redis_spans == (
("async_increment_if_present", expected_status, request_span.get_span_context().span_id),
)
class _UnwritableKeyUserApiKeyCache(UserApiKeyCache):
def __init__(self, unwritable_key: str) -> None:
super().__init__(in_memory_cache=InMemoryCache())
self._unwritable_key: Final = unwritable_key
def set_cache(self, key: str | None, value: object, local_only: bool = False, **kwargs: object):
if key == self._unwritable_key:
raise RuntimeError("cache write failed")
return super().set_cache(key, value, local_only=local_only, **kwargs)
@pytest.mark.asyncio
@pytest.mark.usefixtures("proxy_max_budget")
async def test_update_cache_still_adds_spend_to_the_others_when_one_cached_object_fails() -> None:
cache: Final = _UnwritableKeyUserApiKeyCache(unwritable_key="broken-user")
cache.in_memory_cache.set_cache(
key="broken-user",
value=CacheCodec.serialize(LiteLLM_UserTable(user_id="broken-user", spend=1.0), model_type=LiteLLM_UserTable),
)
cache.set_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY, value=10.0)
cache.set_cache(key="team_id:still-billed", value=LiteLLM_TeamTableCachedObj(team_id="still-billed", spend=1.0))
await ps.update_cache(
token=None,
user_id="broken-user",
end_user_id=None,
team_id="still-billed",
response_cost=0.25,
parent_otel_span=None,
cache=cache,
)
await _drain_spend_cache_writes()
team: Final = cache.get_cache(key="team_id:still-billed", model_type=LiteLLM_TeamTableCachedObj)
assert (team.spend if team is not None else None, cache.get_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)) == (
1.25,
10.25,
)

19
tests/unit/proxy/spend_tracking/test_spend_counter_batch.py
View file

@ -737,22 +737,3 @@ async def test_a_resized_reservation_is_carried_at_its_resized_cost_to_the_count
assert reservation["reserved_cost"] == pytest.approx(0.1)
assert redis.store["spend:key:hashed"] == pytest.approx(9.9)
assert redis.store["spend:team:team"] == pytest.approx(2.1)
@pytest.mark.asyncio
async def test_update_cache_reads_an_object_redis_gained_right_after_a_batch_read_missed_it(monkeypatch):
"""DualCache throttles repeated batch reads of a key that just missed; the per-object GET update_cache used to
issue never did, so its batched read must not either."""
from litellm.caching.dual_cache import DualCache
redis = CountingRedis()
cache = DualCache(redis_cache=redis)
monkeypatch.setattr(ps, "user_api_key_cache", cache)
assert await cache.async_batch_get_cache(keys=["team_id:team"]) == [None]
redis.store["team_id:team"] = {"spend": 1.0}
assert await cache.async_batch_get_cache(keys=["team_id:team"]) == [None]
assert await ps._read_update_cache_values(keys=["team_id:team"], parent_otel_span=None) == {
"team_id:team": {"spend": 1.0}
}
assert redis.commands.count("MGET team_id:team") == 2

271
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
View file

@ -6189,53 +6189,24 @@ async def test_tag_cache_update_called():
"""
Test that update_cache updates tag cache when tags are provided.
"""
from litellm.caching.caching import DualCache
from litellm.proxy.proxy_server import user_api_key_cache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache = DualCache()
cache = UserApiKeyCache(in_memory_cache=InMemoryCache())
cache.in_memory_cache.set_cache("tag:test-tag", {"tag_name": "test-tag", "spend": 10.0})
setattr(
litellm.proxy.proxy_server,
"user_api_key_cache",
cache,
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["test-tag"],
cache=cache,
)
mock_tag_obj = {
"tag_name": "test-tag",
"spend": 10.0,
}
set_cache_called: Final = asyncio.Event()
async def _record_set_cache(*args: object, **kwargs: object) -> None:
set_cache_called.set()
with patch.object(cache, "async_batch_get_cache", new=AsyncMock(return_value=[mock_tag_obj])) as mock_get_cache:
with patch.object(
cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_record_set_cache)
) as mock_set_cache:
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["test-tag"],
)
assert await asyncio.wait_for(set_cache_called.wait(), timeout=5)
mock_get_cache.assert_awaited_once_with(keys=["tag:test-tag"], parent_otel_span=None, throttle_redis=False)
mock_set_cache.assert_awaited_once()
call_args = mock_set_cache.call_args
cache_list = call_args.kwargs["cache_list"]
assert len(cache_list) == 1
cache_key, cache_value = cache_list[0]
assert cache_key == "tag:test-tag"
assert cache_value["spend"] == 15.0
assert cache.in_memory_cache.get_cache("tag:test-tag")["spend"] == 15.0
@pytest.mark.asyncio
@ -6243,59 +6214,26 @@ async def test_tag_cache_update_multiple_tags():
"""
Test that multiple tags are updated in cache.
"""
from litellm.caching.caching import DualCache
from litellm.proxy.proxy_server import user_api_key_cache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
cache = DualCache()
cache = UserApiKeyCache(in_memory_cache=InMemoryCache())
cache.in_memory_cache.set_cache("tag:tag1", {"tag_name": "tag1", "spend": 10.0})
cache.in_memory_cache.set_cache("tag:tag2", {"tag_name": "tag2", "spend": 20.0})
setattr(
litellm.proxy.proxy_server,
"user_api_key_cache",
cache,
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["tag1", "tag2"],
cache=cache,
)
mock_tag1_obj = {"tag_name": "tag1", "spend": 10.0}
mock_tag2_obj = {"tag_name": "tag2", "spend": 20.0}
async def mock_get_cache_side_effect(keys, **kwargs):
return [{"tag:tag1": mock_tag1_obj, "tag:tag2": mock_tag2_obj}.get(key) for key in keys]
set_cache_called: Final = asyncio.Event()
async def _record_set_cache(*args: object, **kwargs: object) -> None:
set_cache_called.set()
with patch.object(
cache, "async_batch_get_cache", new=AsyncMock(side_effect=mock_get_cache_side_effect)
) as mock_get_cache:
with patch.object(
cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_record_set_cache)
) as mock_set_cache:
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["tag1", "tag2"],
)
assert await asyncio.wait_for(set_cache_called.wait(), timeout=5)
mock_get_cache.assert_awaited_once_with(keys=["tag:tag1", "tag:tag2"], parent_otel_span=None, throttle_redis=False)
mock_set_cache.assert_awaited_once()
call_args = mock_set_cache.call_args
cache_list = call_args.kwargs["cache_list"]
assert len(cache_list) == 2
tag_updates = {cache_key: cache_value for cache_key, cache_value in cache_list}
assert "tag:tag1" in tag_updates
assert "tag:tag2" in tag_updates
assert tag_updates["tag:tag1"]["spend"] == 15.0
assert tag_updates["tag:tag2"]["spend"] == 25.0
assert cache.in_memory_cache.get_cache("tag:tag1")["spend"] == 15.0
assert cache.in_memory_cache.get_cache("tag:tag2")["spend"] == 25.0
@pytest.mark.asyncio
@ -6306,41 +6244,25 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl():
a hardcoded 60s, otherwise every priced request resets an active key's cache
entry back to 60s and the configured TTL is never observed.
"""
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
original_cache = litellm.proxy.proxy_server.user_api_key_cache
cache = DualCache(default_in_memory_ttl=300)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
set_cache_called: Final = asyncio.Event()
cache = UserApiKeyCache(in_memory_cache=InMemoryCache(clock=lambda: 1000.0), default_in_memory_ttl=300)
cache.in_memory_cache.cache_dict["tag:active-tag"] = {"tag_name": "active-tag", "spend": 1.0}
async def _record_set_cache(*args: object, **kwargs: object) -> None:
set_cache_called.set()
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["active-tag"],
cache=cache,
)
try:
with patch.object(
cache,
"async_batch_get_cache",
new=AsyncMock(return_value=[{"tag_name": "active-tag", "spend": 1.0}]),
):
with patch.object(
cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_record_set_cache)
) as mock_set_cache:
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id=None,
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
tags=["active-tag"],
)
assert await asyncio.wait_for(set_cache_called.wait(), timeout=5)
mock_set_cache.assert_awaited_once()
assert mock_set_cache.call_args.kwargs["ttl"] == 300
finally:
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
assert cache.in_memory_cache.ttl_dict["tag:active-tag"] == 1300.0
assert cache.in_memory_cache.get_cache("tag:active-tag")["spend"] == 6.0
@pytest.mark.asyncio
@ -6393,63 +6315,6 @@ async def test_spend_tracking_never_writes_the_auth_object_back():
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
@pytest.mark.asyncio
async def test_update_cache_global_proxy_spend_scalar_stays_shared():
"""
The proxy-wide spend estimate must keep flowing to Redis when the spend
writeback goes per-pod: the global max_budget check reads the
``{litellm_proxy_admin_name}:spend`` cache entry between authoritative DB
reloads, so keeping it pod-local would let traffic spread across replicas
exceed the proxy budget by roughly a factor of the replica count within a
cache TTL. Sharing this scalar is safe because it carries no limits or
permissions, so it cannot resurrect an invalidated auth blob.
"""
from litellm.caching.caching import DualCache
admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name
global_key = "{}:spend".format(admin_name)
def fake_get(key):
if key == "user-lit":
return {"user_id": "user-lit", "spend": 1.0}
if key == global_key:
return 10.0
return None
async def fake_batch_get(keys, **kwargs):
return [fake_get(key) for key in keys]
original_cache = litellm.proxy.proxy_server.user_api_key_cache
cache = DualCache(default_in_memory_ttl=300)
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
try:
with patch.object(cache, "async_batch_get_cache", new=AsyncMock(side_effect=fake_batch_get)):
with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache:
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id="user-lit",
end_user_id=None,
team_id=None,
response_cost=5.0,
parent_otel_span=None,
)
pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
if pending:
await asyncio.wait(pending, timeout=5)
calls = mock_set_cache.await_args_list
local_keys = [k for c in calls if c.kwargs.get("local_only") is True for k, _ in c.kwargs["cache_list"]]
shared_keys = [
k for c in calls if c.kwargs.get("local_only") is not True for k, _ in c.kwargs["cache_list"]
]
assert "user-lit" in local_keys
assert global_key not in local_keys
assert shared_keys == [global_key]
finally:
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
@pytest.mark.asyncio
async def test_init_sso_settings_in_db():
"""
@ -15661,47 +15526,3 @@ async def test_spend_capture_rate_check_job_clears_the_gauge_once_the_setting_is
call(api_provider="openai", capture_rate=0.97),
call(api_provider="openai", capture_rate=None),
]
@pytest.mark.asyncio
async def test_update_cache_reads_and_writes_declare_the_auth_objects_key_family():
"""The post-call spend write-back reads and rewrites the cached auth objects, so its
Redis spans must read ``redis.mget auth_objects`` / ``redis.set auth_objects`` (the key
family the auth phase declares) and the global spend scalar ``redis.set spend_counters``,
never a bare ``redis.mget`` with no owner."""
from litellm.caching.caching import DualCache
original_cache = litellm.proxy.proxy_server.user_api_key_cache
cache = DualCache()
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
seen: list[tuple[str, str | None]] = []
async def _mget(keys, **_kwargs):
seen.append(("mget", current_service_target()))
return [{"user_id": "u1", "spend": 1.0} for _ in keys]
async def _set_pipeline(**_kwargs):
seen.append(("set", current_service_target()))
try:
with (
patch.object(cache, "async_batch_get_cache", new=AsyncMock(side_effect=_mget)),
patch.object(cache, "async_set_cache_pipeline", new=AsyncMock(side_effect=_set_pipeline)),
):
await litellm.proxy.proxy_server.update_cache(
token=None,
user_id="u1",
end_user_id=None,
team_id=None,
response_cost=2.0,
parent_otel_span=None,
)
pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
if pending:
await asyncio.wait(pending, timeout=5)
finally:
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
assert seen, "update_cache must touch the cache for a priced user request"
assert {target for _, target in seen} == {"auth_objects"}
assert current_service_target() is None

4
tests/unit/test_internal_context.py
View file

@ -87,10 +87,10 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset(
)
_CACHE_CALL: Final = re.compile(
r"\.(?:async_)?(?:get_cache|set_cache|batch_get_cache|batch_get_cache_shared|increment_cache|increment"
r"\.(?:async_)?(?:get_cache|set_cache|batch_get_cache|batch_get_cache_shared|increment_cache|increment_if_present|increment"
r"|set_cache_pipeline|set_cache_pipeline_with_ttls|set_cache_sadd|delete_cache|batch_set_cache|increment_pipeline"
r"|rpush|lpop|scan_iter|get_ttl|mget)\("
r"|\b(?:reserve_redis_batch_reads|declare_batch_get|_prepare_batch_get)\("
r"|\b(?:reserve_redis_batch_reads|_prepare_batch_get)\("
r"|\bbatch\.(?:set|delete|script|increment)\("
)
_DECLARES_TARGET: Final = re.compile(r"\b(?:with_service_target|service_target|response_cache_phase)\(")