diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index bef04a5c23c..e6214a3650c 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -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, diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 461ac5d947f..eb17facfb12 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -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) diff --git a/litellm/integrations/otel/README.md b/litellm/integrations/otel/README.md index 7414a451038..06c2da117b5 100644 --- a/litellm/integrations/otel/README.md +++ b/litellm/integrations/otel/README.md @@ -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 diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 76fc969cfc0..d8be607150c 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -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", diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index d585880e333..ea6a26174ac 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -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, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3fb7043aa2c..6dcde1aacae 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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(): diff --git a/tests/integration/spend/test_global_proxy_spend_counter.py b/tests/integration/spend/test_global_proxy_spend_counter.py new file mode 100644 index 00000000000..0185e48ffbd --- /dev/null +++ b/tests/integration/spend/test_global_proxy_spend_counter.py @@ -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) diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 5de785e44db..4b73bb6dd34 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -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 diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 6a9a963948e..a18e19e16a7 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -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"), diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 01400384842..d273ef8f42f 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -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, + ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index 1fddfaaa766..8ba79a45ff3 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -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 diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index e824c5017c6..114f5123dd6 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -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 diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 8d3e341096c..60a75c688c4 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -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)\(")