mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(proxy): invalidate stale key/user/team caches and retry Redis spend counter resets on budget reset
Fixes three of the phantom Redis spend-counter inflation mechanisms reported in #30460 for a multi-pod deployment with intermittent Redis timeouts. Path 2: reset_budget_for_litellm_keys/_users/_teams (the per-row budget_duration reset path) invalidated the Redis spend counter but never dropped the corresponding user_api_key_cache entry, unlike the budget-table cascade path. A stale cached object could get read back by _get_source_cache_base_spend and re-seed the counter with the pre-reset spend. Also route that invalidation through the existing cross-pod broadcast (evict_and_broadcast/LIT-3803) instead of a local-only delete, since only one pod runs the reset job per tick and every other pod's cache needs to drop its copy too. Path 3: a failed Redis SET-to-zero during a budget reset only logged a warning and left the inflated pre-reset counter in place until its TTL. Now retries the SET a bounded number of times with a short backoff, and falls back to deleting the key (which reads as cold and reseeds from the DB) if every attempt fails, logging at ERROR instead of WARNING. Path 1 (narrow): release_budget_reservation_on_cancel silently swallowed a reconcile failure with no invalidation fallback, unlike its sibling release_or_invalidate_budget_reservation. Added the same invalidate-and-mark-finalized fallback so a Redis timeout during the cancel-path reconcile cannot leave a pre-charge stuck in the counter. The broader pre-charge/reconcile path already retries via a delete-and-reseed pattern (increment_spend_counters_pipeline / _reconcile_budget_reservation_for_counter_update), so no further change was made there. Relates to #30460
This commit is contained in:
Abhyuday
parent
b3882d8e43
commit
fba32eea20
5 changed files with 377 additions and 18 deletions
|
|
@ -1696,6 +1696,12 @@ RESET_BUDGET_JOB_NAME: Final = "reset_budget_job"
|
|||
# leader keeps the lease across its own run, and a crashed one strands the sweep for
|
||||
# at most a single tick.
|
||||
RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS: Final = max(
|
||||
1, int(os.getenv("RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS", "3"))
|
||||
)
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS: Final = float(
|
||||
os.getenv("RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS", "0.25")
|
||||
)
|
||||
PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
|
||||
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from typing_extensions import assert_never
|
|||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.constants import (
|
||||
GLOBAL_PROXY_SPEND_CACHE_KEY,
|
||||
LITELLM_PROXY_BUDGET_NAME,
|
||||
|
|
@ -21,6 +22,8 @@ from litellm.constants import (
|
|||
RESET_BUDGET_JOB_LOCK_TTL_SECONDS,
|
||||
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN,
|
||||
RESET_BUDGET_JOB_NAME,
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS,
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS,
|
||||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -32,6 +35,9 @@ from litellm.proxy._types import (
|
|||
LiteLLM_UserTable,
|
||||
LiteLLM_VerificationToken,
|
||||
)
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
evict_and_broadcast,
|
||||
)
|
||||
from litellm.proxy.common_utils.timezone_utils import (
|
||||
BudgetResetSettings,
|
||||
compute_budget_reset_at,
|
||||
|
|
@ -538,24 +544,61 @@ class ResetBudgetJob:
|
|||
Call AFTER the DB write commits. Clearing Redis before the DB
|
||||
commit opens a window where get_current_spend reads 0 from Redis
|
||||
while the DB still holds the pre-reset value, allowing bypass.
|
||||
|
||||
A SET that keeps failing after retrying falls back to deleting the key rather
|
||||
than leaving the pre-reset (possibly far higher) value authoritative in Redis
|
||||
until its TTL expires: a missing counter reads as cold and reseeds from the
|
||||
DB on the next request (_ensure_spend_counter_initialized), which is always
|
||||
closer to the truth than the stale value a failed SET would otherwise leave behind.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_spend, ttl=60)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_spend, ttl=60)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to reset spend counter %s in Redis: %s. "
|
||||
"Budget may be over-enforced until counter expires.",
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
redis_cache = spend_counter_cache.redis_cache
|
||||
if redis_cache is None:
|
||||
return
|
||||
if await ResetBudgetJob._reset_redis_spend_counter(
|
||||
redis_cache=redis_cache, counter_key=counter_key, new_spend=new_spend
|
||||
):
|
||||
return
|
||||
verbose_proxy_logger.error(
|
||||
"Failed to reset spend counter %s in Redis after %d attempts; deleting it instead so the "
|
||||
"next read reseeds from the DB rather than keeping the inflated pre-reset value authoritative.",
|
||||
counter_key,
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS,
|
||||
)
|
||||
try:
|
||||
await redis_cache.async_delete_cache(key=counter_key)
|
||||
except Exception as delete_err:
|
||||
verbose_proxy_logger.error(
|
||||
"Failed to delete spend counter %s in Redis after its reset also failed; the "
|
||||
"pre-reset value stays authoritative until its TTL expires: %s",
|
||||
counter_key,
|
||||
delete_err,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning("Failed to reset spend counter %s: %s", counter_key, e)
|
||||
|
||||
@staticmethod
|
||||
async def _reset_redis_spend_counter(redis_cache: RedisCache, counter_key: str, new_spend: float) -> bool:
|
||||
for attempt in range(RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS):
|
||||
try:
|
||||
await redis_cache.async_set_cache(key=counter_key, value=new_spend, ttl=60)
|
||||
return True
|
||||
except Exception as redis_err:
|
||||
is_last_attempt = attempt == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS - 1
|
||||
verbose_proxy_logger.warning(
|
||||
"Attempt %d/%d to reset spend counter %s in Redis failed: %s",
|
||||
attempt + 1,
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS,
|
||||
counter_key,
|
||||
redis_err,
|
||||
)
|
||||
if not is_last_attempt:
|
||||
await asyncio.sleep(RESET_BUDGET_SPEND_COUNTER_RESET_RETRY_DELAY_SECONDS)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def _invalidate_global_proxy_spend_cache() -> None:
|
||||
"""Drop the cached global-proxy spend accumulator after the proxy
|
||||
|
|
@ -566,20 +609,25 @@ class ResetBudgetJob:
|
|||
|
||||
@staticmethod
|
||||
async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
|
||||
"""Drop a stale management-cache entry so the next read fetches from DB.
|
||||
"""Drop a stale management-cache entry, on every pod, so the next read
|
||||
fetches from DB.
|
||||
|
||||
Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``;
|
||||
for those, when the spend counter expires the budget check falls back
|
||||
to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded
|
||||
from the DB, but auth still may consult ``user_api_key_cache`` objects
|
||||
whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache
|
||||
entry forces the next auth-time fetch to reload the zeroed row from
|
||||
Postgres.
|
||||
whose ``.spend`` field can lag a cross-pod DB reset.
|
||||
|
||||
Only one pod runs this job per tick (see ``_acquire_lease``), so a local-only
|
||||
delete would leave every other pod's copy stale until its TTL: ``evict_and_broadcast``
|
||||
is the same LIT-3803 cross-pod eviction every other cache-mutating endpoint already
|
||||
uses (e.g. auth_checks.delete_cache_team_object), broadcasting the delete over Redis
|
||||
pub/sub so every pod's copy is dropped, not just the one that ran the reset.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
await user_api_key_cache.async_delete_cache(key=cache_key)
|
||||
await evict_and_broadcast(cache_keys=(cache_key,), user_api_key_cache=user_api_key_cache)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to invalidate user_api_key_cache entry %s: %s",
|
||||
|
|
@ -988,6 +1036,7 @@ class ResetBudgetJob:
|
|||
token = getattr(k, "token", None)
|
||||
if token:
|
||||
await self._invalidate_spend_counter(f"spend:key:{token}", new_spend=k.spend or 0.0)
|
||||
await self._invalidate_user_api_key_cache_entry(token)
|
||||
|
||||
end_time = time.time()
|
||||
outcome: Final = _ChunkOutcome(
|
||||
|
|
@ -1093,6 +1142,7 @@ class ResetBudgetJob:
|
|||
user_id = getattr(u, "user_id", None)
|
||||
if user_id:
|
||||
await self._invalidate_spend_counter(f"spend:user:{user_id}", new_spend=u.spend or 0.0)
|
||||
await self._invalidate_user_api_key_cache_entry(user_id)
|
||||
if user_id == LITELLM_PROXY_BUDGET_NAME:
|
||||
await self._invalidate_global_proxy_spend_cache()
|
||||
|
||||
|
|
@ -1202,6 +1252,7 @@ class ResetBudgetJob:
|
|||
team_id = getattr(t, "team_id", None)
|
||||
if team_id:
|
||||
await self._invalidate_spend_counter(f"spend:team:{team_id}", new_spend=t.spend or 0.0)
|
||||
await self._invalidate_user_api_key_cache_entry(f"team_id:{team_id}")
|
||||
|
||||
end_time = time.time()
|
||||
outcome: Final = _ChunkOutcome(
|
||||
|
|
|
|||
|
|
@ -371,6 +371,12 @@ async def release_budget_reservation_on_cancel(
|
|||
surrounding task is being cancelled. The `finalized` guard makes this a no-op
|
||||
when success/failure handling already reconciled, so calling it on every
|
||||
cancellation path is safe.
|
||||
|
||||
A reconcile failure here (e.g. a Redis timeout) falls back to
|
||||
invalidate_budget_reservation_counters, mirroring release_or_invalidate_budget_reservation's
|
||||
handling of the same failure on the non-cancel release path: dropping the reserved counters
|
||||
forces the next read to reseed from the DB instead of leaving the pre-charge stuck in Redis
|
||||
with nothing left to correct it, since this is the terminal handler for a cancelled request.
|
||||
"""
|
||||
if not budget_reservation or budget_reservation.get("finalized") is True:
|
||||
return
|
||||
|
|
@ -379,8 +385,20 @@ async def release_budget_reservation_on_cancel(
|
|||
await asyncio.shield(
|
||||
reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost)
|
||||
)
|
||||
except (asyncio.CancelledError, Exception):
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to reconcile budget reservation on cancel; invalidating reserved counters"
|
||||
)
|
||||
try:
|
||||
await invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to invalidate budget reservation counters after cancel-path reconcile failed"
|
||||
)
|
||||
finally:
|
||||
budget_reservation["finalized"] = True
|
||||
|
||||
|
||||
async def invalidate_budget_reservation_counters(
|
||||
|
|
|
|||
|
|
@ -10,17 +10,22 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import httpx
|
||||
import prisma
|
||||
import pytest
|
||||
from redis.asyncio import Redis
|
||||
|
||||
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module
|
||||
from litellm.constants import (
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME,
|
||||
RESET_BUDGET_JOB_LOCK_TTL_SECONDS,
|
||||
RESET_BUDGET_JOB_NAME,
|
||||
RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS,
|
||||
)
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
AuthCacheInvalidationSubscriber,
|
||||
)
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
|
||||
# Mock classes for testing
|
||||
|
|
@ -3243,3 +3248,248 @@ def test_window_reset_zeroes_counter_when_rollover_disabled(monkeypatch):
|
|||
|
||||
spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-off:window:1d", value=0.0)
|
||||
spend_counter_cache.async_get_cache.assert_not_awaited()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path 2 (#30460): resetting a key/user/team's own budget must also drop its
|
||||
# user_api_key_cache entry, not just the Redis spend counter, and that
|
||||
# invalidation must reach every pod, not only the one that ran the sweep.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _direct_reset_fake_module(monkeypatch, user_api_key_cache, redis_usage_cache=None):
|
||||
spend_counter_cache = MagicMock()
|
||||
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
|
||||
spend_counter_cache.redis_cache = None
|
||||
|
||||
fake_module = types.ModuleType("litellm.proxy.proxy_server")
|
||||
fake_module.spend_counter_cache = spend_counter_cache
|
||||
fake_module.user_api_key_cache = user_api_key_cache
|
||||
fake_module.redis_usage_cache = redis_usage_cache
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
|
||||
return spend_counter_cache
|
||||
|
||||
|
||||
def test_reset_budget_for_keys_invalidates_user_api_key_cache(reset_budget_job, mock_prisma_client, monkeypatch):
|
||||
"""
|
||||
A stale user_api_key_cache entry must not survive a key's own budget
|
||||
reset. Before the fix, _get_source_cache_base_spend could read this
|
||||
cached object's pre-reset .spend straight back and re-seed the Redis
|
||||
counter from it on the next request whose DB read fails, reinflating a
|
||||
just-reset budget with no corresponding spend log (#30460 Path 2).
|
||||
"""
|
||||
user_api_key_cache = UserApiKeyCache()
|
||||
user_api_key_cache.set_cache(
|
||||
"sk-abc", UserAPIKeyAuth(token="sk-abc", spend=100.0, max_budget=50.0), model_type=UserAPIKeyAuth
|
||||
)
|
||||
assert user_api_key_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is not None
|
||||
_direct_reset_fake_module(monkeypatch, user_api_key_cache)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
mock_prisma_client.data["key"] = [
|
||||
type(
|
||||
"Key",
|
||||
(),
|
||||
{"spend": 100.0, "budget_duration": "30d", "budget_reset_at": now, "id": "key-1", "token": "sk-abc"},
|
||||
)
|
||||
]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
|
||||
|
||||
assert user_api_key_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None
|
||||
|
||||
|
||||
def test_reset_budget_for_users_invalidates_user_api_key_cache(reset_budget_job, mock_prisma_client, monkeypatch):
|
||||
"""Same Path 2 gap on the user-scoped reset chunk."""
|
||||
from litellm.proxy._types import LiteLLM_UserTable
|
||||
|
||||
user_api_key_cache = UserApiKeyCache()
|
||||
user_api_key_cache.set_cache(
|
||||
"alice", LiteLLM_UserTable(user_id="alice", spend=50.0, max_budget=20.0), model_type=LiteLLM_UserTable
|
||||
)
|
||||
assert user_api_key_cache.get_cache("alice", model_type=LiteLLM_UserTable) is not None
|
||||
_direct_reset_fake_module(monkeypatch, user_api_key_cache)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
mock_prisma_client.data["user"] = [
|
||||
type(
|
||||
"User",
|
||||
(),
|
||||
{"spend": 50.0, "budget_duration": "7d", "budget_reset_at": now, "id": "user-1", "user_id": "alice"},
|
||||
)
|
||||
]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_users())
|
||||
|
||||
assert user_api_key_cache.get_cache("alice", model_type=LiteLLM_UserTable) is None
|
||||
|
||||
|
||||
def test_reset_budget_for_teams_invalidates_user_api_key_cache(reset_budget_job, mock_prisma_client, monkeypatch):
|
||||
"""Same Path 2 gap on the team-scoped reset chunk; teams cache under
|
||||
``team_id:{team_id}`` (see auth_checks.get_team_object), not the bare id."""
|
||||
from litellm.proxy._types import LiteLLM_TeamTable
|
||||
|
||||
user_api_key_cache = UserApiKeyCache()
|
||||
user_api_key_cache.set_cache(
|
||||
"team_id:team-x",
|
||||
LiteLLM_TeamTable(team_id="team-x", spend=200.0, max_budget=100.0),
|
||||
model_type=LiteLLM_TeamTable,
|
||||
)
|
||||
assert user_api_key_cache.get_cache("team_id:team-x", model_type=LiteLLM_TeamTable) is not None
|
||||
_direct_reset_fake_module(monkeypatch, user_api_key_cache)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
mock_prisma_client.data["team"] = [
|
||||
type(
|
||||
"Team",
|
||||
(),
|
||||
{"spend": 200.0, "budget_duration": "1mo", "budget_reset_at": now, "id": "team-1", "team_id": "team-x"},
|
||||
)
|
||||
]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_teams())
|
||||
|
||||
assert user_api_key_cache.get_cache("team_id:team-x", model_type=LiteLLM_TeamTable) is None
|
||||
|
||||
|
||||
class _LoopbackPubSubRedisClient(Redis):
|
||||
"""One fake client good enough to both publish() and pubsub() against the
|
||||
same in-process queue, so a test can prove a message one pod publishes is
|
||||
actually delivered to another pod's subscriber."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._queue: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
async def publish(self, channel: str, message: str) -> int:
|
||||
self._queue.put_nowait({"type": "message", "data": message.encode()})
|
||||
return 1
|
||||
|
||||
def pubsub(self):
|
||||
return self
|
||||
|
||||
async def subscribe(self, *channels: str) -> None:
|
||||
pass
|
||||
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float):
|
||||
try:
|
||||
return await asyncio.wait_for(self._queue.get(), timeout)
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def aclose(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _FakeCoordinationRedisCache:
|
||||
def __init__(self, client: object) -> None:
|
||||
self._client = client
|
||||
self.namespace = None
|
||||
|
||||
def init_async_client(self) -> object:
|
||||
return self._client
|
||||
|
||||
|
||||
async def test_reset_budget_for_keys_broadcasts_cache_invalidation_to_other_pods(
|
||||
reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
"""
|
||||
Only one pod runs the reset job per tick (_acquire_lease elects one
|
||||
sweeper), so a local-only cache delete would leave every other pod's
|
||||
user_api_key_cache serving the pre-reset spend until its TTL. This proves
|
||||
the invalidation is actually broadcast (LIT-3803) to a second pod that
|
||||
never ran the sweep, closing the multi-pod half of #30460 Path 2.
|
||||
"""
|
||||
shared_redis_cache = _FakeCoordinationRedisCache(client=_LoopbackPubSubRedisClient())
|
||||
|
||||
leader_cache = UserApiKeyCache()
|
||||
follower_cache = UserApiKeyCache()
|
||||
for cache in (leader_cache, follower_cache):
|
||||
cache.set_cache(
|
||||
"sk-abc", UserAPIKeyAuth(token="sk-abc", spend=100.0, max_budget=50.0), model_type=UserAPIKeyAuth
|
||||
)
|
||||
|
||||
_direct_reset_fake_module(monkeypatch, leader_cache, redis_usage_cache=shared_redis_cache)
|
||||
|
||||
subscriber = AuthCacheInvalidationSubscriber(redis_cache=shared_redis_cache, user_api_key_cache=follower_cache)
|
||||
subscriber.start()
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
mock_prisma_client.data["key"] = [
|
||||
type(
|
||||
"Key",
|
||||
(),
|
||||
{"spend": 100.0, "budget_duration": "30d", "budget_reset_at": now, "id": "key-1", "token": "sk-abc"},
|
||||
)
|
||||
]
|
||||
|
||||
try:
|
||||
await reset_budget_job.reset_budget_for_litellm_keys()
|
||||
|
||||
for _ in range(200):
|
||||
if follower_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
finally:
|
||||
await subscriber.stop()
|
||||
|
||||
assert leader_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None
|
||||
assert follower_cache.get_cache("sk-abc", model_type=UserAPIKeyAuth) is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Path 3 (#30460): a failed Redis SET-to-zero on budget reset must not
|
||||
# silently leave the inflated pre-reset counter authoritative in Redis.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_invalidate_spend_counter_retries_then_deletes_on_persistent_redis_failure(monkeypatch):
|
||||
"""
|
||||
Every reset-to-zero attempt fails: the fix must retry a bounded number of
|
||||
times, then fall back to deleting the counter (a missing counter reads as
|
||||
cold and reseeds from the DB, see _ensure_spend_counter_initialized)
|
||||
instead of leaving the old, inflated value authoritative until its TTL.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
|
||||
spend_counter_cache = MagicMock()
|
||||
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
|
||||
spend_counter_cache.redis_cache = MagicMock()
|
||||
spend_counter_cache.redis_cache.async_set_cache = AsyncMock(side_effect=RuntimeError("elasticache timeout"))
|
||||
spend_counter_cache.redis_cache.async_delete_cache = AsyncMock()
|
||||
|
||||
fake_module = types.ModuleType("litellm.proxy.proxy_server")
|
||||
fake_module.spend_counter_cache = spend_counter_cache
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
|
||||
|
||||
with patch("litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.error") as mock_error, patch(
|
||||
"asyncio.sleep", new=AsyncMock()
|
||||
):
|
||||
asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-inflated", new_spend=0.0))
|
||||
|
||||
assert spend_counter_cache.redis_cache.async_set_cache.await_count == RESET_BUDGET_SPEND_COUNTER_RESET_MAX_ATTEMPTS
|
||||
spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:key:sk-inflated")
|
||||
assert mock_error.called
|
||||
|
||||
|
||||
def test_invalidate_spend_counter_recovers_after_a_transient_redis_failure(monkeypatch):
|
||||
"""A SET that fails once and then succeeds must not fall back to delete:
|
||||
the counter ends up reset to the real value, not merely absent."""
|
||||
from unittest.mock import patch
|
||||
|
||||
spend_counter_cache = MagicMock()
|
||||
spend_counter_cache.in_memory_cache.set_cache = MagicMock()
|
||||
spend_counter_cache.redis_cache = MagicMock()
|
||||
spend_counter_cache.redis_cache.async_set_cache = AsyncMock(
|
||||
side_effect=[RuntimeError("elasticache timeout"), None]
|
||||
)
|
||||
spend_counter_cache.redis_cache.async_delete_cache = AsyncMock()
|
||||
|
||||
fake_module = types.ModuleType("litellm.proxy.proxy_server")
|
||||
fake_module.spend_counter_cache = spend_counter_cache
|
||||
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module)
|
||||
|
||||
with patch("asyncio.sleep", new=AsyncMock()):
|
||||
asyncio.run(ResetBudgetJob._invalidate_spend_counter("spend:key:sk-recovered", new_spend=0.0))
|
||||
|
||||
assert spend_counter_cache.redis_cache.async_set_cache.await_count == 2
|
||||
spend_counter_cache.redis_cache.async_delete_cache.assert_not_awaited()
|
||||
|
|
|
|||
|
|
@ -2771,6 +2771,40 @@ async def test_release_budget_reservation_on_cancel_swallows_release_errors():
|
|||
await release_budget_reservation_on_cancel(reservation)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_budget_reservation_on_cancel_invalidates_counter_when_reconcile_fails(
|
||||
spend_counter_state,
|
||||
):
|
||||
"""
|
||||
Regression for #30460 Path 1: if reconcile itself fails on the cancel path
|
||||
(e.g. a Redis timeout), the pre-charge must not be left stuck in the
|
||||
counter with nothing left to correct it. This mirrors what
|
||||
release_or_invalidate_budget_reservation already does on the non-cancel
|
||||
release path: fall back to invalidate_budget_reservation_counters so the
|
||||
next read reseeds from the DB instead of enforcing the stale reservation
|
||||
forever.
|
||||
"""
|
||||
counter_cache, _key_cache = spend_counter_state
|
||||
counter_key = "spend:key:key-cancel-redis-down"
|
||||
counter_cache.in_memory_cache.set_cache(key=counter_key, value=3.0)
|
||||
|
||||
reservation = {
|
||||
"reserved_cost": 3.0,
|
||||
"input_cost": 0.5,
|
||||
"finalized": False,
|
||||
"entries": [{"counter_key": counter_key, "reserved_cost": 3.0, "applied_adjustment": 0.0}],
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation",
|
||||
new=AsyncMock(side_effect=RuntimeError("redis down")),
|
||||
):
|
||||
await release_budget_reservation_on_cancel(reservation)
|
||||
|
||||
assert counter_cache.in_memory_cache.get_cache(key=counter_key) is None
|
||||
assert reservation["finalized"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_cancel_in_slow_path_before_yield_refunds(spend_counter_state):
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue