mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
fix(proxy): track global proxy spend with the shared spend-counter machinery
The previous iterations hand-rolled a separate cache path for the global proxy spend scalar, and each variant reintroduced a race the shared spend counter system was built to avoid. This routes global proxy spend through that machinery instead: - Write: _init_and_increment_spend_counter seeds a cold counter from the authoritative litellm-proxy-budget row (cross-pod SET NX) and then increments atomically, keyed spend:user:litellm-proxy-budget. - Read: get_current_spend reads Redis-first, reseeds from the DB on a clean miss, and only floors to the DB row when the freshly loaded fallback exceeds the counter, repaired with the monotonic async_set_max. - Reset: the existing user-reset path already zeroes the counter for the proxy-budget row; the job now also invalidates the cached row-spend hint so no worker acts on a pre-reset fallback. Removes the bespoke coordinator, per-worker DB-floor marker, and delete-on- reset invalidation.
This commit is contained in:
lostmartian
parent
d66e7848cf
commit
3c3d7fd634
8 changed files with 181 additions and 535 deletions
|
|
@ -1449,7 +1449,7 @@ LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE: Final = int(
|
|||
)
|
||||
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
|
||||
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
|
||||
GLOBAL_PROXY_SPEND_COUNTER_KEY: Final = f"spend:user:{LITELLM_PROXY_BUDGET_NAME}"
|
||||
|
||||
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
|
||||
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ import litellm
|
|||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.constants import (
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
GLOBAL_PROXY_SPEND_COUNTER_KEY,
|
||||
LITELLM_PROXY_BUDGET_NAME,
|
||||
LITELLM_PROXY_MASTER_KEY_ALIAS,
|
||||
)
|
||||
|
|
@ -76,7 +76,6 @@ from litellm.proxy.auth.resolvers import CredentialRef, Principal
|
|||
from litellm.proxy.auth.resolvers.store import IdentityStore
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.trusted_proxy_utils import get_trusted_proxy_cidrs
|
||||
from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_read_request_body,
|
||||
_safe_get_request_headers,
|
||||
|
|
@ -491,103 +490,57 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use
|
|||
return valid_token
|
||||
|
||||
|
||||
# Reusable coordinator for global spend to prevent cache stampede
|
||||
_global_spend_coordinator: Final = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEND]")
|
||||
|
||||
# Cache the authoritative DB spend read for a few seconds per worker so the
|
||||
# global-proxy budget floor check does not query the DB on every request.
|
||||
_GLOBAL_PROXY_SPEND_DB_FLOOR_TTL_SECONDS: Final = 5
|
||||
_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX: Final = "global_proxy_spend_db_floor:"
|
||||
async def _read_global_proxy_spend_hint_from_cache(user_api_key_cache: UserApiKeyCache) -> float | None:
|
||||
redis_cache: Final = user_api_key_cache.redis_cache
|
||||
if redis_cache is not None:
|
||||
try:
|
||||
cached: Final = await redis_cache.async_get_cache(key=LITELLM_PROXY_BUDGET_NAME)
|
||||
except Exception as e: # noqa: BLE001 # hint read failure must not break auth; in-memory fallback below
|
||||
verbose_proxy_logger.debug("Global proxy spend hint: Redis read failed, falling back to in-memory: %s", e)
|
||||
else:
|
||||
return float(cached) if cached is not None else None
|
||||
cached_in_memory: Final = user_api_key_cache.in_memory_cache.get_cache(key=LITELLM_PROXY_BUDGET_NAME)
|
||||
return float(cached_in_memory) if cached_in_memory is not None else None
|
||||
|
||||
|
||||
def invalidate_global_proxy_spend_db_floor_marker(user_api_key_cache: UserApiKeyCache, cache_key: str) -> None:
|
||||
"""Drop the worker-local DB-floor marker for ``cache_key``.
|
||||
|
||||
ResetBudgetJob calls this after zeroing the aggregate DB row so a stale
|
||||
pre-reset marker is not used to write previous-window spend back into
|
||||
Redis during the next auth-time floor check.
|
||||
"""
|
||||
marker_key: Final = f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}{cache_key}"
|
||||
user_api_key_cache.in_memory_cache.delete_cache(key=marker_key)
|
||||
|
||||
|
||||
async def _fetch_global_spend_with_event_coordination(
|
||||
cache_key: str,
|
||||
async def _load_global_proxy_spend_hint(
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
prisma_client: PrismaClient,
|
||||
) -> float | None:
|
||||
"""Spend of the proxy-budget aggregate user row, cached in the management cache.
|
||||
|
||||
This is only the fallback/staleness signal for get_current_spend; the
|
||||
enforced value is the shared spend counter. Read Redis-first so the reset
|
||||
job's invalidation is visible to every worker immediately.
|
||||
"""
|
||||
Fetch global spend with event-driven coordination to prevent cache stampede.
|
||||
Uses EventDrivenCacheCoordinator: first request queries DB and signals others when done.
|
||||
|
||||
Reads the proxy budget aggregate user row, which accrues proxy-wide spend
|
||||
per request and is zeroed by ResetBudgetJob every ``litellm.budget_duration``.
|
||||
|
||||
With Redis configured, the shared scalar is read straight from Redis: the
|
||||
global spend is incremented by every worker, so a per-worker in-memory copy
|
||||
can go stale across resets and falsely reject requests at auth time. Reading
|
||||
the Redis value directly keeps the budget check on the shared total. On a
|
||||
Redis miss the coordinator loads the authoritative DB spend once and
|
||||
repopulates Redis. Without Redis (single process), the normal DualCache
|
||||
coordinator path is used.
|
||||
"""
|
||||
|
||||
async def _load_global_spend() -> float | None:
|
||||
proxy_budget_row: Final = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": LITELLM_PROXY_BUDGET_NAME}
|
||||
)
|
||||
return float(proxy_budget_row.spend) if proxy_budget_row is not None else None
|
||||
|
||||
async def _authoritative_db_spend() -> float | None:
|
||||
# Cached per-worker for a few seconds to avoid a DB query per request.
|
||||
marker_key: Final = f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}{cache_key}"
|
||||
cached = user_api_key_cache.in_memory_cache.get_cache(key=marker_key)
|
||||
if cached is not None:
|
||||
return float(cached)
|
||||
db_spend = await _load_global_spend()
|
||||
if db_spend is not None:
|
||||
user_api_key_cache.in_memory_cache.set_cache(
|
||||
key=marker_key,
|
||||
value=db_spend,
|
||||
ttl=_GLOBAL_PROXY_SPEND_DB_FLOOR_TTL_SECONDS,
|
||||
)
|
||||
return db_spend
|
||||
|
||||
spend_cache: Final = (
|
||||
user_api_key_cache.redis_cache if user_api_key_cache.redis_cache is not None else user_api_key_cache
|
||||
cached: Final = await _read_global_proxy_spend_hint_from_cache(user_api_key_cache)
|
||||
if cached is not None:
|
||||
return cached
|
||||
proxy_budget_row: Final = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": LITELLM_PROXY_BUDGET_NAME}
|
||||
)
|
||||
value = await _global_spend_coordinator.get_or_load(
|
||||
cache_key=cache_key,
|
||||
cache=spend_cache, # pyright: ignore[reportArgumentType]
|
||||
load_fn=_load_global_spend,
|
||||
)
|
||||
if value is None:
|
||||
if proxy_budget_row is None:
|
||||
return None
|
||||
spend: Final = float(proxy_budget_row.spend or 0.0)
|
||||
await user_api_key_cache.async_set_cache(key=LITELLM_PROXY_BUDGET_NAME, value=spend)
|
||||
return spend
|
||||
|
||||
try:
|
||||
numeric_value = float(value)
|
||||
except (TypeError, ValueError):
|
||||
numeric_value = 0.0
|
||||
|
||||
# Floor to the authoritative DB spend when Redis is in use. A Redis counter
|
||||
# that was clobbered by a reset race, restarted from an old RDB snapshot,
|
||||
# or recreated from zero after expiry can under-report current-period spend;
|
||||
# capping the value to the DB row prevents requests from leaking past the
|
||||
# global cap. Without Redis the coordinator already loaded the DB row.
|
||||
if user_api_key_cache.redis_cache is not None:
|
||||
db_spend = await _authoritative_db_spend()
|
||||
if db_spend is not None and db_spend > numeric_value:
|
||||
# Use a monotonic write so a concurrent increment that already
|
||||
# pushed Redis above the DB value is not overwritten by a repair
|
||||
# carrying a slightly-stale DB total. Use the resulting Redis value
|
||||
# so the enforced total reflects any concurrent increment.
|
||||
repaired_value: Final[float | None] = await user_api_key_cache.redis_cache.async_set_max(
|
||||
key=cache_key, value=db_spend
|
||||
)
|
||||
numeric_value = repaired_value if repaired_value is not None else db_spend
|
||||
user_api_key_cache.in_memory_cache.set_cache(key=cache_key, value=numeric_value)
|
||||
async def _get_global_proxy_spend_value(
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
prisma_client: PrismaClient,
|
||||
) -> float:
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
return numeric_value
|
||||
hint: Final = await _load_global_proxy_spend_hint(
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
return await get_current_spend(
|
||||
counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY,
|
||||
fallback_spend=hint if hint is not None else 0.0,
|
||||
max_budget=litellm.max_budget,
|
||||
)
|
||||
|
||||
|
||||
async def get_global_proxy_spend(
|
||||
|
|
@ -599,10 +552,7 @@ async def get_global_proxy_spend(
|
|||
) -> float | None:
|
||||
global_proxy_spend = None
|
||||
if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget
|
||||
# Use event-driven coordination to prevent cache stampede
|
||||
cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key=cache_key,
|
||||
global_proxy_spend = await _get_global_proxy_spend_value(
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
|
@ -2112,10 +2062,8 @@ async def _user_api_key_auth_builder(
|
|||
|
||||
global_proxy_spend = None
|
||||
if litellm.max_budget > 0 and prisma_client is not None: # user set proxy max budget
|
||||
cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
with tracer.trace("litellm.proxy.auth.get_global_proxy_spend"):
|
||||
global_proxy_spend = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key=cache_key,
|
||||
global_proxy_spend = await _get_global_proxy_spend_value(
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.constants import (
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
LITELLM_PROXY_BUDGET_NAME,
|
||||
RESET_BUDGET_JOB_BATCH_SIZE,
|
||||
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN,
|
||||
|
|
@ -262,28 +261,14 @@ class ResetBudgetJob:
|
|||
|
||||
@staticmethod
|
||||
async def _invalidate_global_proxy_spend_cache() -> None:
|
||||
"""Drop the cached global-proxy spend accumulator after the proxy
|
||||
budget aggregate row is reset so the next auth-time read reloads the
|
||||
zeroed DB row. Reads and writes are atomic now (increment-if-present on
|
||||
Redis), so deleting cannot resurrect a stale pre-reset value, and it
|
||||
cannot clobber a concurrent increment the way a plain SET to zero could.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
invalidate_global_proxy_spend_db_floor_marker,
|
||||
)
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
"""Drop the cached proxy-budget row spend after the row is reset.
|
||||
|
||||
await user_api_key_cache.async_delete_cache(key=GLOBAL_PROXY_SPEND_CACHE_KEY)
|
||||
invalidate_global_proxy_spend_db_floor_marker(
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
cache_key=GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # cache invalidation failure must not break budget reset
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to invalidate global proxy spend cache after budget reset: %s",
|
||||
e,
|
||||
)
|
||||
Auth uses it only as the fallback signal for the shared spend counter's
|
||||
DB floor, so every worker must reload the zeroed row instead of acting
|
||||
on a pre-reset value. The counter itself is zeroed by the standard
|
||||
``_invalidate_spend_counter`` path for the proxy-budget user row.
|
||||
"""
|
||||
await ResetBudgetJob._invalidate_user_api_key_cache_entry(LITELLM_PROXY_BUDGET_NAME)
|
||||
|
||||
@staticmethod
|
||||
async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
|
||||
|
|
|
|||
|
|
@ -235,7 +235,7 @@ from litellm.constants import (
|
|||
DAYS_IN_A_MONTH,
|
||||
DEFAULT_HEALTH_CHECK_INTERVAL,
|
||||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
GLOBAL_PROXY_SPEND_COUNTER_KEY,
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
LITELLM_PROXY_BUDGET_NAME,
|
||||
MONTHLY_SPEND_REPORT_JOB_ID,
|
||||
|
|
@ -292,7 +292,6 @@ from litellm.proxy.auth.model_checks import (
|
|||
get_team_models,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
user_api_key_auth,
|
||||
user_api_key_auth_websocket,
|
||||
)
|
||||
|
|
@ -3123,27 +3122,16 @@ async def update_cache(
|
|||
|
||||
### UPDATE GLOBAL PROXY SPEND ###
|
||||
async def _update_global_proxy_spend() -> None:
|
||||
if response_cost is None:
|
||||
if response_cost is None or litellm.max_budget <= 0:
|
||||
return
|
||||
try:
|
||||
# Atomically accumulate the shared global-proxy scalar in Redis.
|
||||
# Authorization reads this Redis scalar directly, so there is no
|
||||
# per-worker in-memory copy that can go stale across resets.
|
||||
if user_api_key_cache.redis_cache is not None:
|
||||
await user_api_key_cache.redis_cache.async_increment(
|
||||
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
value=response_cost,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
refresh_ttl=True,
|
||||
)
|
||||
else:
|
||||
# No Redis: single-process, so an in-memory increment is coherent.
|
||||
await user_api_key_cache.async_increment_cache(
|
||||
key=GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
value=response_cost,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
refresh_ttl=True,
|
||||
)
|
||||
# Same spend-counter path as every other budget entity: seed a cold
|
||||
# counter from the authoritative DB row (SET NX), then increment.
|
||||
await _init_and_increment_spend_counter(
|
||||
counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY,
|
||||
source_cache_key=LITELLM_PROXY_BUDGET_NAME,
|
||||
increment=response_cost,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # spend update must not fail the request
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - failed to update global proxy spend in cache. "
|
||||
|
|
@ -8613,11 +8601,9 @@ class ProxyStartupEvent:
|
|||
) -> None:
|
||||
"""Warm global spend cache once at startup to reduce impact of first wave of requests."""
|
||||
try:
|
||||
cache_key: Final = GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
await _fetch_global_spend_with_event_coordination(
|
||||
cache_key=cache_key,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
prisma_client=prisma_client,
|
||||
await get_current_spend(
|
||||
counter_key=GLOBAL_PROXY_SPEND_COUNTER_KEY,
|
||||
fallback_spend=0.0,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Global spend cache warm-up at startup skipped or failed: %s", e)
|
||||
|
|
|
|||
|
|
@ -6068,16 +6068,14 @@ class TestCheckKeyModelBudgetWithFallback:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
|
||||
async def test_global_proxy_spend_hint_reads_resettable_proxy_budget_row():
|
||||
"""Regression for the global proxy budget ignoring budget_duration
|
||||
(LIT-4309 / gh#31292): the enforced global spend must be loaded from the
|
||||
(LIT-4309 / gh#31292): the fallback spend must be loaded from the
|
||||
"litellm-proxy-budget" user row, which the spend writer increments per
|
||||
request and ResetBudgetJob zeroes every budget_duration. It must NOT be
|
||||
loaded from the MonthlyGlobalSpend view, whose window is hardcoded to a
|
||||
trailing 30 days and never resets on the configured duration."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import _load_global_proxy_spend_hint
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
|
|
@ -6088,8 +6086,7 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
|
|||
side_effect=AssertionError("global spend must not be loaded from the fixed-30d MonthlyGlobalSpend view")
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
result = await _load_global_proxy_spend_hint(
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
|
@ -6101,19 +6098,16 @@ async def test_global_proxy_spend_reads_resettable_proxy_budget_row():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_none_when_proxy_budget_row_missing():
|
||||
"""Before the startup upsert creates the aggregate row, enforcement must
|
||||
see None (no cap applied) rather than raising."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
async def test_global_proxy_spend_hint_none_when_proxy_budget_row_missing():
|
||||
"""Before the startup upsert creates the aggregate row, the fallback hint
|
||||
is None so the counter falls back to 0 rather than raising."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _load_global_proxy_spend_hint
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
result = await _load_global_proxy_spend_hint(
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
|
@ -6122,209 +6116,61 @@ async def test_global_proxy_spend_none_when_proxy_budget_row_missing():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_reads_shared_redis_scalar_when_configured():
|
||||
"""With Redis, the global spend read goes straight to the shared Redis
|
||||
scalar instead of a per-worker in-memory copy, which can go stale across
|
||||
resets and falsely reject requests at auth time."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
async def test_global_proxy_spend_hint_reads_redis_first():
|
||||
"""The hint is read from the shared Redis tier first so a reset
|
||||
invalidation is visible to every worker immediately; a per-worker
|
||||
in-memory copy is only a Redis-outage fallback."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _load_global_proxy_spend_hint
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value=8.5)
|
||||
redis_cache.async_set_cache = AsyncMock()
|
||||
cache.redis_cache = redis_cache
|
||||
cache.in_memory_cache.set_cache(key="litellm-proxy-budget", value=999.0)
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 5.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
side_effect=AssertionError("DB must not be queried when Redis has the hint")
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
result = await _load_global_proxy_spend_hint(
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 8.5
|
||||
redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend")
|
||||
redis_cache.async_set_cache.assert_not_called()
|
||||
redis_cache.async_get_cache.assert_awaited_once_with(key="litellm-proxy-budget")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_floors_stale_low_redis_to_db():
|
||||
"""A Redis counter that is stale-low (e.g., recreated from zero after a
|
||||
reset or Redis restart) must be floored to the authoritative DB spend so
|
||||
requests cannot leak past the global cap."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
async def test_global_proxy_spend_value_uses_shared_spend_counter(monkeypatch):
|
||||
"""The enforced global proxy spend comes from the shared spend counter
|
||||
(spend:user:litellm-proxy-budget) via get_current_spend, with the DB row
|
||||
spend as the fallback signal and the configured cap as max_budget."""
|
||||
import litellm
|
||||
from litellm.proxy.auth.user_api_key_auth import _get_global_proxy_spend_value
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="2.0")
|
||||
redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 9.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
)
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=proxy_budget_row)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
get_current_spend = AsyncMock(return_value=9.0)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.get_current_spend", get_current_spend)
|
||||
monkeypatch.setattr(litellm, "max_budget", 100.0)
|
||||
|
||||
result = await _get_global_proxy_spend_value(
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 9.0
|
||||
redis_cache.async_get_cache.assert_awaited_once_with(key="default_user_id:spend")
|
||||
redis_cache.async_set_max.assert_awaited_once_with(
|
||||
key="default_user_id:spend", value=9.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_floor_does_not_overwrite_concurrent_increment():
|
||||
"""When Redis already holds a higher value than the DB floor (because a
|
||||
concurrent request incremented it after our read), the repair must use a
|
||||
monotonic write so the concurrent increment is not discarded."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="2.0")
|
||||
redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 1.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
# DB (1.0) is not greater than Redis (2.0), so no repair is attempted.
|
||||
assert result == 2.0
|
||||
redis_cache.async_set_max.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_marker_cached_db_value_used_without_db_query():
|
||||
"""The DB-floor marker should be reused for subsequent reads so we do not
|
||||
query Postgres on every request."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="2.0")
|
||||
redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
# Seed the marker cache with a DB value greater than Redis.
|
||||
marker_key = f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}default_user_id:spend"
|
||||
cache.in_memory_cache.set_cache(key=marker_key, value=9.0)
|
||||
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
side_effect=AssertionError("DB must not be queried when marker is present")
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 9.0
|
||||
prisma_client.db.litellm_usertable.find_unique.assert_not_called()
|
||||
redis_cache.async_set_max.assert_awaited_once_with(
|
||||
key="default_user_id:spend", value=9.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_floor_falls_back_when_set_max_returns_none():
|
||||
"""If the monotonic Redis repair returns None (e.g. connection hiccup),
|
||||
the enforced total must still fall back to the authoritative DB floor."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="2.0")
|
||||
redis_cache.async_set_max = AsyncMock(return_value=None)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 9.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 9.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_proxy_spend_treats_unparseable_redis_as_zero():
|
||||
"""A corrupted/non-numeric Redis value is treated as 0.0 and then floored
|
||||
to the DB total so enforcement remains safe."""
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_fetch_global_spend_with_event_coordination,
|
||||
)
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="not-a-number")
|
||||
redis_cache.async_set_max = AsyncMock(side_effect=lambda key, value: value)
|
||||
cache.redis_cache = redis_cache
|
||||
|
||||
proxy_budget_row = MagicMock()
|
||||
proxy_budget_row.spend = 7.0
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_usertable.find_unique = AsyncMock(
|
||||
return_value=proxy_budget_row
|
||||
)
|
||||
|
||||
result = await _fetch_global_spend_with_event_coordination(
|
||||
cache_key="default_user_id:spend",
|
||||
user_api_key_cache=cache,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
assert result == 7.0
|
||||
redis_cache.async_set_max.assert_awaited_once_with(
|
||||
key="default_user_id:spend", value=7.0
|
||||
get_current_spend.assert_awaited_once_with(
|
||||
counter_key="spend:user:litellm-proxy-budget",
|
||||
fallback_spend=9.0,
|
||||
max_budget=100.0,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -863,48 +863,6 @@ def test_reset_budget_windows_uses_is_not_null_filter(monkeypatch):
|
|||
assert "budget_limits IS NOT NULL" in team_query
|
||||
|
||||
|
||||
def test_invalidate_global_proxy_spend_cache_deletes_key_and_db_floor_marker(monkeypatch):
|
||||
"""ResetBudgetJob must drop both the shared Redis spend key and the worker-
|
||||
local DB-floor marker. Otherwise a stale pre-reset marker can be used to
|
||||
write previous-window spend back into Redis at the next auth check."""
|
||||
from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX,
|
||||
)
|
||||
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_delete_cache = AsyncMock()
|
||||
user_api_key_cache.in_memory_cache.delete_cache = MagicMock()
|
||||
|
||||
fake_module = types.ModuleType("litellm.proxy.proxy_server")
|
||||
fake_module.user_api_key_cache = user_api_key_cache
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
|
||||
|
||||
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=MagicMock())
|
||||
asyncio.run(job._invalidate_global_proxy_spend_cache())
|
||||
|
||||
user_api_key_cache.async_delete_cache.assert_awaited_once_with(
|
||||
key=GLOBAL_PROXY_SPEND_CACHE_KEY
|
||||
)
|
||||
user_api_key_cache.in_memory_cache.delete_cache.assert_called_once_with(
|
||||
key=f"{_GLOBAL_PROXY_SPEND_DB_FLOOR_MARKER_PREFIX}{GLOBAL_PROXY_SPEND_CACHE_KEY}"
|
||||
)
|
||||
|
||||
|
||||
def test_invalidate_global_proxy_spend_cache_swallows_cache_failure(monkeypatch):
|
||||
"""A failure to invalidate the cache must not break the budget reset job."""
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
|
||||
fake_module = types.ModuleType("litellm.proxy.proxy_server")
|
||||
fake_module.user_api_key_cache = user_api_key_cache
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
|
||||
|
||||
job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=MagicMock())
|
||||
# Should not raise.
|
||||
asyncio.run(job._invalidate_global_proxy_spend_cache())
|
||||
|
||||
|
||||
def test_reset_budget_windows_resets_expired_key_window(monkeypatch):
|
||||
"""A key whose window's `reset_at` has passed gets an update with a new
|
||||
`reset_at` in the future, and the in-memory spend counter is cleared."""
|
||||
|
|
@ -1134,8 +1092,9 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock
|
|||
def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache(
|
||||
reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
"""Resetting the proxy-wide budget aggregate row invalidates the cached
|
||||
global-spend accumulator so the next read reloads the zeroed DB row."""
|
||||
"""Resetting the proxy-wide budget aggregate row zeroes its spend counter
|
||||
and invalidates the cached row-spend hint so the next auth reloads the
|
||||
zeroed DB row instead of acting on a pre-reset value."""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
|
|
@ -1155,16 +1114,19 @@ def test_reset_budget_for_proxy_budget_row_resets_global_spend_cache(
|
|||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
|
||||
|
||||
counter_cache.in_memory_cache.set_cache.assert_any_call(
|
||||
key="spend:user:litellm-proxy-budget", value=0.0, ttl=60
|
||||
)
|
||||
counter_cache.user_api_key_cache.async_delete_cache.assert_any_call(
|
||||
key="default_user_id:spend"
|
||||
key="litellm-proxy-budget"
|
||||
)
|
||||
|
||||
|
||||
def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
|
||||
reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
"""The global-spend accumulator must only be invalidated when the proxy
|
||||
budget aggregate row itself resets, not on every user reset."""
|
||||
"""The proxy-budget hint must only be invalidated when the proxy budget
|
||||
aggregate row itself resets, not on every user reset."""
|
||||
counter_cache = _make_counter_invalidation_job(monkeypatch)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
|
|
@ -1185,7 +1147,7 @@ def test_reset_budget_for_ordinary_user_does_not_touch_global_spend_cache(
|
|||
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
|
||||
|
||||
assert not any(
|
||||
call.kwargs.get("key") == "default_user_id:spend"
|
||||
call.kwargs.get("key") == "litellm-proxy-budget"
|
||||
for call in counter_cache.user_api_key_cache.async_delete_cache.call_args_list
|
||||
)
|
||||
assert not any(
|
||||
|
|
|
|||
|
|
@ -1273,133 +1273,79 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_global_proxy_spend_starts_from_zero_after_reset(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Global proxy spend starts from zero after the cache key is reset."""
|
||||
async def test_update_cache_global_proxy_spend_uses_spend_counter(monkeypatch):
|
||||
"""Global proxy spend goes through the shared spend-counter path so cold
|
||||
counters seed from the authoritative DB row and increments stay atomic."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
cache = DualCache(default_in_memory_ttl=300)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", cache)
|
||||
init_and_increment = AsyncMock()
|
||||
monkeypatch.setattr(ps, "_init_and_increment_spend_counter", init_and_increment)
|
||||
monkeypatch.setattr(ps.litellm, "max_budget", 100.0)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", DualCache(default_in_memory_ttl=300))
|
||||
|
||||
await cache.async_set_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, value=50.0)
|
||||
await cache.async_delete_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
|
||||
await cache.async_set_cache(
|
||||
key="user-lit",
|
||||
value={"user_id": "user-lit", "spend": 0.0},
|
||||
)
|
||||
await ps.update_cache(
|
||||
token=None,
|
||||
user_id="user-lit",
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=1.05,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
final_value = await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
|
||||
assert final_value == 1.05
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_global_proxy_spend_concurrent_increments_after_reset(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Concurrent spend updates after a reset accumulate from zero."""
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
cache = DualCache(default_in_memory_ttl=300)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", cache)
|
||||
|
||||
await cache.async_set_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY, value=100.0)
|
||||
await cache.async_delete_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
|
||||
await cache.async_set_cache(
|
||||
key="user-concurrent",
|
||||
value={"user_id": "user-concurrent", "spend": 0.0},
|
||||
)
|
||||
|
||||
async def _update(cost: float) -> None:
|
||||
await ps.update_cache(
|
||||
token=None,
|
||||
user_id="user-concurrent",
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=cost,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
await asyncio.gather(*[_update(0.01) for _ in range(50)])
|
||||
|
||||
final_value = await cache.async_get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY)
|
||||
assert final_value == pytest.approx(0.5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_global_proxy_spend_does_not_mirror_to_in_memory(
|
||||
monkeypatch,
|
||||
):
|
||||
"""Authorization reads the shared Redis scalar directly, so the increment
|
||||
path must not maintain a per-worker in-memory copy that could go stale
|
||||
across resets and confuse budget checks."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_increment = AsyncMock(return_value=3.5)
|
||||
cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", cache)
|
||||
|
||||
await cache.async_set_cache(
|
||||
key="user-no-mirror",
|
||||
value={"user_id": "user-no-mirror", "spend": 0.0},
|
||||
)
|
||||
await ps.update_cache(
|
||||
token=None,
|
||||
user_id="user-no-mirror",
|
||||
user_id=None,
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=1.5,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
redis_cache.async_increment.assert_awaited_once_with(
|
||||
key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
value=1.5,
|
||||
ttl=ps.get_management_object_ttl(cache),
|
||||
refresh_ttl=True,
|
||||
init_and_increment.assert_awaited_once_with(
|
||||
counter_key="spend:user:litellm-proxy-budget",
|
||||
source_cache_key="litellm-proxy-budget",
|
||||
increment=1.5,
|
||||
)
|
||||
assert cache.in_memory_cache.get_cache(key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_global_proxy_spend_swallows_redis_failure(monkeypatch):
|
||||
"""A Redis failure on the global spend increment must not fail the request."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
async def test_update_cache_global_proxy_spend_skipped_without_max_budget(monkeypatch):
|
||||
"""No global cap configured means the proxy-budget row is never written by
|
||||
the DB flusher either, so the counter must not accumulate spend it could
|
||||
not be reseeded from."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=300)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", cache)
|
||||
init_and_increment = AsyncMock()
|
||||
monkeypatch.setattr(ps, "_init_and_increment_spend_counter", init_and_increment)
|
||||
monkeypatch.setattr(ps.litellm, "max_budget", 0)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", DualCache(default_in_memory_ttl=300))
|
||||
|
||||
await cache.async_set_cache(
|
||||
key="user-redis-down",
|
||||
value={"user_id": "user-redis-down", "spend": 0.0},
|
||||
)
|
||||
await ps.update_cache(
|
||||
token=None,
|
||||
user_id="user-redis-down",
|
||||
user_id=None,
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=1.5,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
init_and_increment.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_global_proxy_spend_swallows_failure(monkeypatch):
|
||||
"""A failure on the global spend increment must not fail the request."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
init_and_increment = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
monkeypatch.setattr(ps, "_init_and_increment_spend_counter", init_and_increment)
|
||||
monkeypatch.setattr(ps.litellm, "max_budget", 100.0)
|
||||
monkeypatch.setattr(ps, "user_api_key_cache", DualCache(default_in_memory_ttl=300))
|
||||
|
||||
await ps.update_cache(
|
||||
token=None,
|
||||
user_id=None,
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=1.0,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
redis_cache.async_increment.assert_awaited_once_with(
|
||||
key=ps.GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
value=1.0,
|
||||
ttl=ps.get_management_object_ttl(cache),
|
||||
refresh_ttl=True,
|
||||
)
|
||||
init_and_increment.assert_awaited_once()
|
||||
|
|
|
|||
|
|
@ -5273,65 +5273,38 @@ async def test_spend_tracking_never_writes_the_auth_object_back():
|
|||
@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. The scalar is updated via an atomic cache increment, which
|
||||
remains shared across pods and avoids resurrecting a stale pre-reset value.
|
||||
The proxy-wide spend estimate must live in the shared spend counter
|
||||
(``spend:user:litellm-proxy-budget``), which seeds from the authoritative
|
||||
DB row when cold and increments atomically. A per-pod cache key would let
|
||||
traffic spread across replicas exceed the proxy budget by roughly a factor
|
||||
of the replica count within a cache TTL.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name
|
||||
global_key = f"{admin_name}:spend"
|
||||
|
||||
async def fake_get(key, **kwargs):
|
||||
if key == "user-lit":
|
||||
return {"user_id": "user-lit", "spend": 1.0}
|
||||
return None
|
||||
|
||||
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)
|
||||
original_max_budget = litellm.max_budget
|
||||
init_and_increment = AsyncMock()
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", DualCache(default_in_memory_ttl=300))
|
||||
setattr(litellm, "max_budget", 100.0)
|
||||
setattr(litellm.proxy.proxy_server, "_init_and_increment_spend_counter", init_and_increment)
|
||||
try:
|
||||
with patch.object(cache, "async_get_cache", new=AsyncMock(side_effect=fake_get)):
|
||||
with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache:
|
||||
with patch.object(cache, "async_increment_cache", new=AsyncMock()) as mock_increment:
|
||||
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,
|
||||
)
|
||||
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,
|
||||
)
|
||||
|
||||
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 == []
|
||||
|
||||
increment_calls = [
|
||||
c
|
||||
for c in mock_increment.await_args_list
|
||||
if (c.args[0] if c.args else c.kwargs.get("key")) == global_key
|
||||
]
|
||||
assert len(increment_calls) == 1
|
||||
call = increment_calls[0]
|
||||
value = call.args[1] if len(call.args) >= 2 else call.kwargs.get("value")
|
||||
assert value == 5.0
|
||||
assert call.kwargs.get("refresh_ttl") is True
|
||||
init_and_increment.assert_awaited_once_with(
|
||||
counter_key="spend:user:litellm-proxy-budget",
|
||||
source_cache_key="litellm-proxy-budget",
|
||||
increment=5.0,
|
||||
)
|
||||
finally:
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
|
||||
setattr(litellm, "max_budget", original_max_budget)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue