mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-18 04:28:19 +00:00
Reseed enforcement read path from DB on counter miss (#26459)
Co-authored-by: Michael Riad Zaky <michaelr@Michaels-MacBook-Air.local>
This commit is contained in:
Michael-RZ-Berri
co-authored by
Michael Riad Zaky
Michael Riad Zaky
parent
b021d5c109
commit
3ef16098f2
@@ -1429,6 +1429,9 @@ SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
|
||||
SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100))
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL = float(os.getenv("SPEND_LOG_QUEUE_POLL_INTERVAL", 2.0))
|
||||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = int(
|
||||
os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000)
|
||||
)
|
||||
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS = int(
|
||||
os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)
|
||||
) # 1 minute
|
||||
|
||||
@@ -0,0 +1,159 @@
|
||||
"""
|
||||
Coalesced reseed of spend counters from the authoritative DB.
|
||||
|
||||
When a Redis spend counter expires (or is missing on a fresh pod), enforcement
|
||||
must read the current spend from somewhere. The in-process management cache
|
||||
(`user_api_key_cache.team_membership.spend`, etc.) is per-pod and lags DB
|
||||
writes from other pods, so trusting it allows budget bypass in multi-pod
|
||||
deployments. This module reseeds from the authoritative DB instead.
|
||||
|
||||
A per-counter singleflight lock collapses concurrent reseeds on the same pod
|
||||
to one DB query per cold-cache window. The lock dict is bounded LRU to cap
|
||||
memory in long-lived deployments.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING, ClassVar, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class SpendCounterReseed:
|
||||
"""
|
||||
Reseeds spend counters from the authoritative DB and warms the cache,
|
||||
coalesced via per-counter singleflight locks.
|
||||
|
||||
Counter key prefixes map to DB tables:
|
||||
spend:key:{token} -> LiteLLM_VerificationToken.spend
|
||||
spend:team:{team_id} -> LiteLLM_TeamTable.spend
|
||||
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
|
||||
spend:user:{user_id} -> LiteLLM_UserTable.spend
|
||||
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
|
||||
"""
|
||||
|
||||
_locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
|
||||
_registry_lock: ClassVar[Optional[asyncio.Lock]] = None
|
||||
|
||||
@staticmethod
|
||||
async def _get_lock(counter_key: str) -> asyncio.Lock:
|
||||
if SpendCounterReseed._registry_lock is None:
|
||||
SpendCounterReseed._registry_lock = asyncio.Lock()
|
||||
async with SpendCounterReseed._registry_lock:
|
||||
lock = SpendCounterReseed._locks.get(counter_key)
|
||||
if lock is not None:
|
||||
SpendCounterReseed._locks.move_to_end(counter_key)
|
||||
return lock
|
||||
lock = asyncio.Lock()
|
||||
SpendCounterReseed._locks[counter_key] = lock
|
||||
if len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE:
|
||||
SpendCounterReseed._locks.popitem(last=False)
|
||||
return lock
|
||||
|
||||
@staticmethod
|
||||
async def from_db(
|
||||
prisma_client: Optional["PrismaClient"], counter_key: str
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
Read the authoritative spend for a counter from the DB.
|
||||
|
||||
Returns the spend value (including 0.0) when the DB is reachable
|
||||
and the row exists. Returns None when prisma is unavailable, the
|
||||
row is missing, the key format is unrecognized, or the query
|
||||
raises. Callers use None to fall back to a caller-supplied source.
|
||||
"""
|
||||
if prisma_client is None:
|
||||
return None
|
||||
# Per-window counters share prefixes with primary counters but
|
||||
# don't correspond to a DB row.
|
||||
if ":window:" in counter_key:
|
||||
return None
|
||||
try:
|
||||
if counter_key.startswith("spend:key:"):
|
||||
token = counter_key[len("spend:key:") :]
|
||||
row = await prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
where={"token": token}
|
||||
)
|
||||
elif counter_key.startswith("spend:team_member:"):
|
||||
suffix = counter_key[len("spend:team_member:") :]
|
||||
if ":" not in suffix:
|
||||
return None
|
||||
user_id, team_id = suffix.rsplit(":", 1)
|
||||
row = await prisma_client.db.litellm_teammembership.find_unique(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}
|
||||
)
|
||||
elif counter_key.startswith("spend:team:"):
|
||||
team_id = counter_key[len("spend:team:") :]
|
||||
row = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
elif counter_key.startswith("spend:user:"):
|
||||
user_id = counter_key[len("spend:user:") :]
|
||||
row = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_id}
|
||||
)
|
||||
elif counter_key.startswith("spend:org:"):
|
||||
org_id = counter_key[len("spend:org:") :]
|
||||
row = await prisma_client.db.litellm_organizationtable.find_unique(
|
||||
where={"organization_id": org_id}
|
||||
)
|
||||
else:
|
||||
return None
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"SpendCounterReseed.from_db: failed for %s", counter_key
|
||||
)
|
||||
return None
|
||||
if row is None:
|
||||
return None
|
||||
return float(getattr(row, "spend", 0.0) or 0.0)
|
||||
|
||||
@staticmethod
|
||||
async def coalesced(
|
||||
prisma_client: Optional["PrismaClient"],
|
||||
spend_counter_cache: "DualCache",
|
||||
counter_key: str,
|
||||
) -> Optional[float]:
|
||||
"""
|
||||
Reseed a cold spend counter from the DB and warm the cache,
|
||||
coalesced via a per-counter lock so concurrent callers (read path
|
||||
+ write path) collapse to one DB query per cold-cache window.
|
||||
|
||||
Returns the spend value (including 0.0 from a fresh budget reset)
|
||||
when the DB read succeeds, or None when the DB is unavailable.
|
||||
"""
|
||||
lock = await SpendCounterReseed._get_lock(counter_key)
|
||||
async with lock:
|
||||
# Re-check after acquiring the lock - another waiter may have warmed it.
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
try:
|
||||
val = await spend_counter_cache.redis_cache.async_get_cache(
|
||||
key=counter_key
|
||||
)
|
||||
if val is not None:
|
||||
return float(val)
|
||||
except Exception:
|
||||
pass
|
||||
val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
|
||||
if val is not None:
|
||||
return float(val)
|
||||
|
||||
db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key)
|
||||
if db_spend is None:
|
||||
return None
|
||||
# Warm even when 0 so subsequent reads hit cache, not DB.
|
||||
try:
|
||||
await spend_counter_cache.async_increment_cache(
|
||||
key=counter_key, value=db_spend
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"SpendCounterReseed.coalesced: failed to warm counter %s",
|
||||
counter_key,
|
||||
)
|
||||
return db_spend
|
||||
@@ -323,6 +323,7 @@ from litellm.proxy.container_endpoints.endpoints import router as container_rout
|
||||
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router
|
||||
from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config
|
||||
@@ -1775,7 +1776,8 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
|
||||
Fallback chain:
|
||||
1. Redis counter (cross-pod, authoritative)
|
||||
2. In-memory counter (single-instance or Redis failure)
|
||||
3. Cached object's .spend from DB (cold start, no counter yet)
|
||||
3. Reseed from authoritative DB spend (counter expired, cross-pod stale)
|
||||
4. Caller-supplied fallback (DB unavailable, cold start)
|
||||
"""
|
||||
# 1. Try Redis first (cross-pod authoritative)
|
||||
if spend_counter_cache.redis_cache is not None:
|
||||
@@ -1795,7 +1797,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float:
|
||||
if val is not None:
|
||||
return float(val)
|
||||
|
||||
# 3. Final fallback: cached object's spend from DB
|
||||
# 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass.
|
||||
db_spend = await SpendCounterReseed.coalesced(
|
||||
prisma_client=prisma_client,
|
||||
spend_counter_cache=spend_counter_cache,
|
||||
counter_key=counter_key,
|
||||
)
|
||||
if db_spend is not None:
|
||||
return db_spend
|
||||
|
||||
# 4. Caller-supplied fallback (DB unavailable).
|
||||
return fallback_spend
|
||||
|
||||
|
||||
@@ -1906,69 +1917,6 @@ async def increment_spend_counters(
|
||||
)
|
||||
|
||||
|
||||
async def _reseed_spend_from_db(counter_key: str) -> float:
|
||||
"""
|
||||
Read the authoritative spend for a missing counter from the DB. The
|
||||
counter_key prefix encodes the table to query:
|
||||
|
||||
spend:key:{token} -> LiteLLM_VerificationToken.spend
|
||||
spend:team:{team_id} -> LiteLLM_TeamTable.spend
|
||||
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
|
||||
spend:user:{user_id} -> LiteLLM_UserTable.spend
|
||||
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
|
||||
|
||||
Returns 0.0 if prisma is unavailable, the row is missing, or the
|
||||
key format is unrecognized. On failure, logs and returns 0.0 rather
|
||||
than raising so the caller can still record the current increment.
|
||||
"""
|
||||
if prisma_client is None:
|
||||
return 0.0
|
||||
# Per-window counters (spend:*:window:{duration}) share prefixes with
|
||||
# primary counters but don't correspond to a DB row; their ambiguity
|
||||
# would otherwise be silently parsed as a regular counter and miss.
|
||||
if ":window:" in counter_key:
|
||||
return 0.0
|
||||
try:
|
||||
if counter_key.startswith("spend:key:"):
|
||||
token = counter_key[len("spend:key:") :]
|
||||
row = await prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
where={"token": token}
|
||||
)
|
||||
elif counter_key.startswith("spend:team_member:"):
|
||||
suffix = counter_key[len("spend:team_member:") :]
|
||||
if ":" not in suffix:
|
||||
return 0.0
|
||||
user_id, team_id = suffix.rsplit(":", 1)
|
||||
row = await prisma_client.db.litellm_teammembership.find_unique(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}
|
||||
)
|
||||
elif counter_key.startswith("spend:team:"):
|
||||
team_id = counter_key[len("spend:team:") :]
|
||||
row = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
elif counter_key.startswith("spend:user:"):
|
||||
user_id = counter_key[len("spend:user:") :]
|
||||
row = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_id}
|
||||
)
|
||||
elif counter_key.startswith("spend:org:"):
|
||||
org_id = counter_key[len("spend:org:") :]
|
||||
row = await prisma_client.db.litellm_organizationtable.find_unique(
|
||||
where={"organization_id": org_id}
|
||||
)
|
||||
else:
|
||||
return 0.0
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception(
|
||||
"Failed to reseed spend counter %s from DB", counter_key
|
||||
)
|
||||
return 0.0
|
||||
if row is None:
|
||||
return 0.0
|
||||
return float(getattr(row, "spend", 0.0) or 0.0)
|
||||
|
||||
|
||||
async def _init_and_increment_spend_counter(
|
||||
counter_key: str,
|
||||
source_cache_key: str,
|
||||
@@ -1980,9 +1928,9 @@ async def _init_and_increment_spend_counter(
|
||||
|
||||
On first access per pod:
|
||||
1. Check spend_counter_cache (in-memory -> Redis via DualCache)
|
||||
2. If not found, reseed from the DB (`_reseed_spend_from_db`). Falls
|
||||
back to the cached object's `.spend` via user_api_key_cache only
|
||||
if prisma is unavailable, since that value can lag the flusher.
|
||||
2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`.
|
||||
Falls back to the cached object's `.spend` via user_api_key_cache
|
||||
only if prisma is unavailable, since that value can lag the flusher.
|
||||
3. Seed counter via async_increment_cache (not async_set_cache) to avoid a
|
||||
check-then-set race: if two pods cold-start simultaneously, both may see
|
||||
the counter as absent and seed it. Using increment means the worst case
|
||||
@@ -1992,20 +1940,25 @@ async def _init_and_increment_spend_counter(
|
||||
"""
|
||||
current = await spend_counter_cache.async_get_cache(key=counter_key)
|
||||
if current is None:
|
||||
base_spend = await _reseed_spend_from_db(counter_key)
|
||||
if prisma_client is None:
|
||||
# Best-effort fallback when prisma is unavailable (tests or
|
||||
# early-startup paths). May be stale but avoids resetting to 0.
|
||||
# Shares the per-counter lock with get_current_spend.
|
||||
db_spend = await SpendCounterReseed.coalesced(
|
||||
prisma_client=prisma_client,
|
||||
spend_counter_cache=spend_counter_cache,
|
||||
counter_key=counter_key,
|
||||
)
|
||||
if db_spend is None:
|
||||
# DB unavailable - fall back to in-process cache (may be stale).
|
||||
source = await user_api_key_cache.async_get_cache(key=source_cache_key)
|
||||
base_spend: float = 0.0
|
||||
if source is not None:
|
||||
if isinstance(source, dict):
|
||||
base_spend = source.get("spend", 0.0) or 0.0
|
||||
else:
|
||||
base_spend = getattr(source, "spend", 0.0) or 0.0
|
||||
if base_spend > 0:
|
||||
await spend_counter_cache.async_increment_cache(
|
||||
key=counter_key, value=base_spend
|
||||
)
|
||||
if base_spend > 0:
|
||||
await spend_counter_cache.async_increment_cache(
|
||||
key=counter_key, value=base_spend
|
||||
)
|
||||
|
||||
await spend_counter_cache.async_increment_cache(key=counter_key, value=increment)
|
||||
|
||||
|
||||
@@ -5034,8 +5034,7 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss(
|
||||
async def test_reseed_spend_from_db_user_and_org_prefixes():
|
||||
"""User and org counters must reseed from their own DB tables, not
|
||||
fall through to 0.0 like the other counters do today."""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.proxy_server import _reseed_spend_from_db
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
|
||||
user_row = MagicMock()
|
||||
user_row.spend = 17.0
|
||||
@@ -5048,20 +5047,15 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
|
||||
return_value=org_row
|
||||
)
|
||||
|
||||
orig_prisma = ps.prisma_client
|
||||
ps.prisma_client = fake_prisma
|
||||
try:
|
||||
assert await _reseed_spend_from_db("spend:user:alice") == 17.0
|
||||
fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with(
|
||||
where={"user_id": "alice"}
|
||||
)
|
||||
assert await SpendCounterReseed.from_db(fake_prisma, "spend:user:alice") == 17.0
|
||||
fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with(
|
||||
where={"user_id": "alice"}
|
||||
)
|
||||
|
||||
assert await _reseed_spend_from_db("spend:org:acme") == 305.0
|
||||
fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(
|
||||
where={"organization_id": "acme"}
|
||||
)
|
||||
finally:
|
||||
ps.prisma_client = orig_prisma
|
||||
assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0
|
||||
fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(
|
||||
where={"organization_id": "acme"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -5069,19 +5063,391 @@ async def test_reseed_spend_from_db_skips_window_variant_keys():
|
||||
"""Window counters (spend:*:window:{duration}) share prefixes with
|
||||
primary counters but don't correspond to a DB row. The guard must
|
||||
short-circuit without querying the DB."""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.proxy_server import _reseed_spend_from_db
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_verificationtoken.find_unique = AsyncMock()
|
||||
fake_prisma.db.litellm_teamtable.find_unique = AsyncMock()
|
||||
|
||||
orig_prisma = ps.prisma_client
|
||||
assert (
|
||||
await SpendCounterReseed.from_db(fake_prisma, "spend:key:sk-abc:window:1h")
|
||||
is None
|
||||
)
|
||||
assert (
|
||||
await SpendCounterReseed.from_db(fake_prisma, "spend:team:team-1:window:1d")
|
||||
is None
|
||||
)
|
||||
fake_prisma.db.litellm_verificationtoken.find_unique.assert_not_awaited()
|
||||
fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_reseeds_from_db_when_counter_missing():
|
||||
"""
|
||||
When both the Redis and in-memory counters are missing, the enforcement
|
||||
read path must reseed from the authoritative DB, not fall back to the
|
||||
caller-supplied stale value. Otherwise, every Redis TTL expiry lets a
|
||||
request through against a stale in-process `team_membership.spend`.
|
||||
"""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
counter_cache = DualCache()
|
||||
recorded_warms: list = []
|
||||
|
||||
async def record_increment(key, value, ttl=None, **kwargs):
|
||||
recorded_warms.append({"key": key, "value": value})
|
||||
return value
|
||||
|
||||
fake_redis = AsyncMock()
|
||||
fake_redis.async_increment = AsyncMock(side_effect=record_increment)
|
||||
fake_redis.async_get_cache = AsyncMock(return_value=None)
|
||||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
# DB has authoritative spend=362.0; caller hands us stale fallback=30.0
|
||||
# (the in-process team_membership.spend that hasn't caught up to DB).
|
||||
db_row = MagicMock()
|
||||
db_row.spend = 362.0
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
|
||||
ps.spend_counter_cache = counter_cache
|
||||
ps.prisma_client = fake_prisma
|
||||
try:
|
||||
assert await _reseed_spend_from_db("spend:key:sk-abc:window:1h") == 0.0
|
||||
assert await _reseed_spend_from_db("spend:team:team-1:window:1d") == 0.0
|
||||
fake_prisma.db.litellm_verificationtoken.find_unique.assert_not_awaited()
|
||||
fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited()
|
||||
spend = await get_current_spend(
|
||||
counter_key="spend:team_member:user-1:team-1",
|
||||
fallback_spend=30.0,
|
||||
)
|
||||
assert spend == 362.0, (
|
||||
f"expected DB reseed to return 362.0, got {spend} "
|
||||
f"(fallback would have returned 30.0 and caused bypass)"
|
||||
)
|
||||
# Counter warmed so subsequent reads are fast
|
||||
assert ("spend:team_member:user-1:team-1", 362.0) in [
|
||||
(w["key"], w["value"]) for w in recorded_warms
|
||||
]
|
||||
finally:
|
||||
ps.spend_counter_cache = orig_counter
|
||||
ps.prisma_client = orig_prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_uses_fallback_when_db_unavailable():
|
||||
"""
|
||||
If prisma is unavailable and both counters are missing, the read path
|
||||
must degrade to the caller-supplied fallback rather than raising.
|
||||
"""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
counter_cache = DualCache()
|
||||
fake_redis = AsyncMock()
|
||||
fake_redis.async_get_cache = AsyncMock(return_value=None)
|
||||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
|
||||
ps.spend_counter_cache = counter_cache
|
||||
ps.prisma_client = None # simulate prisma unavailable
|
||||
try:
|
||||
spend = await get_current_spend(
|
||||
counter_key="spend:team_member:user-1:team-1",
|
||||
fallback_spend=15.5,
|
||||
)
|
||||
assert spend == 15.5
|
||||
finally:
|
||||
ps.spend_counter_cache = orig_counter
|
||||
ps.prisma_client = orig_prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_coalesces_concurrent_reseeds():
|
||||
"""
|
||||
When several concurrent calls hit a cold counter on the same pod,
|
||||
only one DB query should fire. The rest should wait for the lock,
|
||||
re-check the warmed counter, and return without hitting the DB.
|
||||
"""
|
||||
import asyncio as _asyncio
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
counter_cache = DualCache()
|
||||
counter_key = "spend:team_member:user-1:team-coalesce"
|
||||
|
||||
# Track DB query calls and inject a small delay so the concurrent
|
||||
# callers actually overlap in the lock-acquire window.
|
||||
db_call_count = 0
|
||||
|
||||
async def slow_find_unique(**kwargs):
|
||||
nonlocal db_call_count
|
||||
db_call_count += 1
|
||||
await _asyncio.sleep(0.05)
|
||||
row = MagicMock()
|
||||
row.spend = 100.0
|
||||
return row
|
||||
|
||||
fake_redis = AsyncMock()
|
||||
redis_store: dict = {}
|
||||
|
||||
async def redis_get(key, **_):
|
||||
return redis_store.get(key)
|
||||
|
||||
async def redis_increment(key, value, **_):
|
||||
redis_store[key] = (redis_store.get(key) or 0.0) + value
|
||||
return redis_store[key]
|
||||
|
||||
fake_redis.async_get_cache = AsyncMock(side_effect=redis_get)
|
||||
fake_redis.async_increment = AsyncMock(side_effect=redis_increment)
|
||||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(
|
||||
side_effect=slow_find_unique
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
|
||||
ps.spend_counter_cache = counter_cache
|
||||
ps.prisma_client = fake_prisma
|
||||
try:
|
||||
results = await _asyncio.gather(
|
||||
*[
|
||||
get_current_spend(counter_key=counter_key, fallback_spend=0.0)
|
||||
for _ in range(5)
|
||||
]
|
||||
)
|
||||
assert results == [100.0] * 5, f"all callers should see DB value, got {results}"
|
||||
assert (
|
||||
db_call_count == 1
|
||||
), f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}"
|
||||
finally:
|
||||
ps.spend_counter_cache = orig_counter
|
||||
ps.prisma_client = orig_prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_current_spend_uses_db_zero_over_stale_fallback():
|
||||
"""
|
||||
When DB returns spend=0 (e.g. just after a budget period reset), the
|
||||
authoritative DB value must win over a stale non-zero fallback. The
|
||||
fallback in production is the in-process team_membership.spend, which
|
||||
can still hold the pre-reset value across pods.
|
||||
"""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
counter_cache = DualCache()
|
||||
fake_redis = AsyncMock()
|
||||
fake_redis.async_get_cache = AsyncMock(return_value=None)
|
||||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
db_row = MagicMock()
|
||||
db_row.spend = 0.0
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
|
||||
ps.spend_counter_cache = counter_cache
|
||||
ps.prisma_client = fake_prisma
|
||||
try:
|
||||
spend = await get_current_spend(
|
||||
counter_key="spend:team_member:user-1:team-after-reset",
|
||||
fallback_spend=42.0,
|
||||
)
|
||||
assert (
|
||||
spend == 0.0
|
||||
), f"DB authoritative 0 must override stale fallback 42, got {spend}"
|
||||
finally:
|
||||
ps.spend_counter_cache = orig_counter
|
||||
ps.prisma_client = orig_prisma
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_read_and_write_paths_share_one_db_query():
|
||||
"""
|
||||
The read path (`get_current_spend`) and the write path
|
||||
(`_init_and_increment_spend_counter`) both reseed cold counters from
|
||||
the DB. They must share the per-counter lock so a concurrent pre-call
|
||||
enforcement read and post-call increment for the same counter collapse
|
||||
to one DB query, not two.
|
||||
"""
|
||||
import asyncio as _asyncio
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.proxy_server import (
|
||||
_init_and_increment_spend_counter,
|
||||
get_current_spend,
|
||||
)
|
||||
|
||||
counter_cache = DualCache()
|
||||
counter_key = "spend:team_member:user-1:team-cross-path"
|
||||
|
||||
db_call_count = 0
|
||||
|
||||
async def slow_find_unique(**kwargs):
|
||||
nonlocal db_call_count
|
||||
db_call_count += 1
|
||||
await _asyncio.sleep(0.05)
|
||||
row = MagicMock()
|
||||
row.spend = 50.0
|
||||
return row
|
||||
|
||||
redis_store: dict = {}
|
||||
|
||||
async def redis_get(key, **_):
|
||||
return redis_store.get(key)
|
||||
|
||||
async def redis_increment(key, value, **_):
|
||||
redis_store[key] = (redis_store.get(key) or 0.0) + value
|
||||
return redis_store[key]
|
||||
|
||||
fake_redis = AsyncMock()
|
||||
fake_redis.async_get_cache = AsyncMock(side_effect=redis_get)
|
||||
fake_redis.async_increment = AsyncMock(side_effect=redis_increment)
|
||||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(
|
||||
side_effect=slow_find_unique
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
orig_counter, orig_prisma, orig_user = (
|
||||
ps.spend_counter_cache,
|
||||
ps.prisma_client,
|
||||
ps.user_api_key_cache,
|
||||
)
|
||||
ps.spend_counter_cache = counter_cache
|
||||
ps.prisma_client = fake_prisma
|
||||
ps.user_api_key_cache = DualCache()
|
||||
try:
|
||||
results = await _asyncio.gather(
|
||||
get_current_spend(counter_key=counter_key, fallback_spend=0.0),
|
||||
_init_and_increment_spend_counter(
|
||||
counter_key=counter_key,
|
||||
source_cache_key="ignored",
|
||||
increment=1.5,
|
||||
),
|
||||
get_current_spend(counter_key=counter_key, fallback_spend=0.0),
|
||||
)
|
||||
assert (
|
||||
db_call_count == 1
|
||||
), f"expected 1 DB query for concurrent read+write+read, got {db_call_count}"
|
||||
# Read-path callers see the warmed counter; the write path's
|
||||
# increment may or may not have landed by then, so accept either
|
||||
# the seeded value or seeded+increment.
|
||||
assert results[0] in (50.0, 51.5), f"got {results[0]}"
|
||||
assert results[2] in (50.0, 51.5), f"got {results[2]}"
|
||||
finally:
|
||||
ps.spend_counter_cache = orig_counter
|
||||
ps.prisma_client = orig_prisma
|
||||
ps.user_api_key_cache = orig_user
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reseed_locks_dict_is_bounded():
|
||||
"""
|
||||
`SpendCounterReseed._locks` is an LRU bounded at
|
||||
`SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE` to prevent unbounded growth in
|
||||
long-lived deployments with high counter-key churn. Inserting more
|
||||
than the cap evicts the oldest entries.
|
||||
"""
|
||||
import litellm.constants as constants
|
||||
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
|
||||
|
||||
orig_locks = SpendCounterReseed._locks.copy()
|
||||
SpendCounterReseed._locks.clear()
|
||||
orig_max = constants.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
|
||||
constants.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = 5
|
||||
# The class reads the constant via module-level import, so patch the
|
||||
# module-level name on the spend_counter_reseed module too.
|
||||
import litellm.proxy.db.spend_counter_reseed as scr
|
||||
|
||||
orig_module_max = scr.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
|
||||
scr.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = 5
|
||||
try:
|
||||
for i in range(7):
|
||||
await SpendCounterReseed._get_lock(f"spend:key:test-key-{i}")
|
||||
assert (
|
||||
len(SpendCounterReseed._locks) == 5
|
||||
), f"got {len(SpendCounterReseed._locks)}"
|
||||
# Oldest two evicted
|
||||
assert "spend:key:test-key-0" not in SpendCounterReseed._locks
|
||||
assert "spend:key:test-key-1" not in SpendCounterReseed._locks
|
||||
# Most recent retained
|
||||
assert "spend:key:test-key-6" in SpendCounterReseed._locks
|
||||
finally:
|
||||
constants.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = orig_max
|
||||
scr.SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE = orig_module_max
|
||||
SpendCounterReseed._locks.clear()
|
||||
SpendCounterReseed._locks.update(orig_locks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reseed_warms_cache_even_on_zero_db_spend():
|
||||
"""
|
||||
When DB returns 0.0 (fresh entity / just after reset), the cache must
|
||||
still be warmed so subsequent reads hit the cache instead of issuing
|
||||
another DB query. Skipping the warm causes O(requests) DB load on
|
||||
zero-spend entities.
|
||||
"""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
counter_cache = DualCache()
|
||||
counter_key = "spend:team_member:user-1:team-zero-warm"
|
||||
redis_store: dict = {}
|
||||
|
||||
async def redis_get(key, **_):
|
||||
return redis_store.get(key)
|
||||
|
||||
async def redis_increment(key, value, **_):
|
||||
redis_store[key] = (redis_store.get(key) or 0.0) + value
|
||||
return redis_store[key]
|
||||
|
||||
fake_redis = AsyncMock()
|
||||
fake_redis.async_get_cache = AsyncMock(side_effect=redis_get)
|
||||
fake_redis.async_increment = AsyncMock(side_effect=redis_increment)
|
||||
counter_cache.redis_cache = fake_redis
|
||||
|
||||
db_call_count = 0
|
||||
|
||||
async def find_unique(**kwargs):
|
||||
nonlocal db_call_count
|
||||
db_call_count += 1
|
||||
row = MagicMock()
|
||||
row.spend = 0.0
|
||||
return row
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(
|
||||
side_effect=find_unique
|
||||
)
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client
|
||||
ps.spend_counter_cache = counter_cache
|
||||
ps.prisma_client = fake_prisma
|
||||
try:
|
||||
# First call: cold cache, hits DB, returns 0.
|
||||
spend1 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
|
||||
# Second call: cache should be warmed at 0, no second DB query.
|
||||
spend2 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0)
|
||||
assert spend1 == 0.0 and spend2 == 0.0
|
||||
assert (
|
||||
db_call_count == 1
|
||||
), f"second read should hit warmed cache, got {db_call_count} DB queries"
|
||||
assert redis_store.get(counter_key) == 0.0, "cache must be warmed at 0"
|
||||
finally:
|
||||
ps.spend_counter_cache = orig_counter
|
||||
ps.prisma_client = orig_prisma
|
||||
|
||||
Reference in New Issue
Block a user