mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
b55024658c
* fix(proxy): stop update_cache from resurrecting evicted auth cache entries Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): update update_cache pins for synchronous write-back Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): wrap long assert in update_cache tag test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover refill race and pin TTL with injected clock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): count global proxy spend for every request with a user, cached on this worker or not * fix(proxy): add global proxy spend with an atomic Redis increment that never re-creates a reset counter The global proxy spend bump read the total, awaited, and wrote old + cost back from a background task, so a budget reset that deleted the key in between got the pre-reset total written back, and workers overwrote each other's totals. RedisCache.async_increment_if_present runs INCRBYFLOAT only when the key exists, keeping its TTL, and this worker's copy is then set from the shared total with no await between its read and write. The bump now runs whenever a proxy max_budget is set, matching the DB writer, and the armed post-call read and the batch-read declaration it used are removed. * fix(proxy): report the global spend Redis increment to service metrics and request traces The global proxy spend increment now emits Redis service success and failure events under the spend_counters target and passes the request span through, so it shows as redis.incr spend_counters like the write it replaced. The unit tests run the real RedisCache increment against fakeredis instead of mocks, and a new test checks the span lands under the request span on success and on a Redis failure. * test(proxy): keep the spend counter Redis fake off the network * test(proxy): count increment_if_present as a shared-cache call in the key-family guard * fix(proxy): keep the newer global spend total when Redis replies arrive out of order --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
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)
|