mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(prometheus): keep the gauges right for pass-through rows, overlapping charges and slow DB spend writes
Pass-through rows carry the auth object, so a cold committed counter gets seeded there too. The gauges read a charge still waiting in the post-call pipeline as not counted, whatever other charges reached the counter. Committed counters outlive two DB spend writes. Cold seeds read the DB once per worker, and the global spend reset reads keys and teams in bounded pages.
This commit is contained in:
mateo-berri
parent
b4ab31e05f
commit
6046f1b015
13 changed files with 375 additions and 36 deletions
|
|
@ -461,6 +461,7 @@ class RequestRedisBatches:
|
|||
"_deadline",
|
||||
"_deadline_flush",
|
||||
"_post_call",
|
||||
"_post_call_sent",
|
||||
"post_call_deadline",
|
||||
"prefetched",
|
||||
)
|
||||
|
|
@ -468,6 +469,7 @@ class RequestRedisBatches:
|
|||
def __init__(self, post_call_deadline: float = POST_CALL_FLUSH_DEADLINE_SECONDS) -> None:
|
||||
self._batches: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
|
||||
self._post_call: Final[dict[object, RedisBatch]] = {} # mutable-ok: lazily filled per backend
|
||||
self._post_call_sent: bool = False
|
||||
self.post_call_deadline: Final = post_call_deadline
|
||||
self._deadline: asyncio.TimerHandle | None = None
|
||||
self._deadline_flush: asyncio.Task[None] | None = None
|
||||
|
|
@ -508,10 +510,15 @@ class RequestRedisBatches:
|
|||
if self._deadline is not None:
|
||||
self._deadline.cancel()
|
||||
self._deadline = None
|
||||
self._post_call_sent = self._post_call_sent or any(batch.pending for batch in self._post_call.values())
|
||||
await asyncio.gather(*(batch.flush() for batch in self._post_call.values() if batch.pending))
|
||||
if not any(batch.pending for batch in self._post_call.values()):
|
||||
_open_post_call.discard(self)
|
||||
|
||||
@property
|
||||
def post_call_sent(self) -> bool:
|
||||
return self._post_call_sent
|
||||
|
||||
@property
|
||||
def batches(self) -> tuple[RedisBatch, ...]:
|
||||
return tuple(self._batches.values())
|
||||
|
|
@ -542,6 +549,12 @@ def active_post_call_redis_batch(redis_cache: RedisCache) -> RedisBatch | None:
|
|||
return batches.post_call(redis_cache)
|
||||
|
||||
|
||||
def post_call_writes_unsent() -> bool:
|
||||
"""Inside a request scope that has sent no post-call pipeline yet, none of its post-call writes reached Redis."""
|
||||
batches: Final = _active_request_batches.get()
|
||||
return batches is not None and not batches.post_call_sent
|
||||
|
||||
|
||||
async def flush_post_call_redis_batches() -> None:
|
||||
"""Called where the success and failure callbacks of a request have all run."""
|
||||
batches: Final = _active_request_batches.get()
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from typing_extensions import ReadOnly, TypedDict, assert_never
|
|||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.caching.redis_batch import post_call_writes_unsent
|
||||
from litellm.constants import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
|
|
@ -405,12 +406,20 @@ async def _committed_spend_now(counter_keys: EntitySpendCounterKeys) -> Committe
|
|||
return CommittedBudgetSpend()
|
||||
|
||||
|
||||
def _own_charge_waits_in_post_call_pipeline() -> bool:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
return spend_counter_cache.redis_cache is not None and post_call_writes_unsent()
|
||||
|
||||
|
||||
def _committed_spend_before_request(
|
||||
at_auth: CommittedBudgetSpend, now: CommittedBudgetSpend, response_cost: float
|
||||
) -> CommittedBudgetSpend:
|
||||
def before(at_auth_spend: float | None, now_spend: float | None) -> float | None:
|
||||
if now_spend is None:
|
||||
return at_auth_spend
|
||||
if _own_charge_waits_in_post_call_pipeline():
|
||||
return now_spend
|
||||
if at_auth_spend is None:
|
||||
return now_spend - response_cost
|
||||
own_charge_counted: Final = now_spend - at_auth_spend >= response_cost * _OWN_CHARGE_MATCH_RATIO
|
||||
|
|
|
|||
|
|
@ -96,6 +96,10 @@ class SpendCounterReseed:
|
|||
_locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
|
||||
_registry_lock: ClassVar[asyncio.Lock | None] = None
|
||||
|
||||
@staticmethod
|
||||
async def counter_lock(counter_key: str) -> asyncio.Lock:
|
||||
return await SpendCounterReseed._get_lock(counter_key)
|
||||
|
||||
@staticmethod
|
||||
async def _get_lock(counter_key: str) -> asyncio.Lock:
|
||||
if SpendCounterReseed._registry_lock is None:
|
||||
|
|
|
|||
|
|
@ -2136,7 +2136,14 @@ class LiteLLMProxyRequestSetup:
|
|||
user_api_key_dict.matched_model_access_groups
|
||||
)
|
||||
# UserAPIKeyAuth object for MCP server access control
|
||||
data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict.model_copy(
|
||||
data[_metadata_variable_name]["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def auth_for_request_metadata(user_api_key_dict: UserAPIKeyAuth) -> UserAPIKeyAuth:
|
||||
return user_api_key_dict.model_copy(
|
||||
update={
|
||||
"metadata": strip_callback_config(user_api_key_dict.metadata),
|
||||
"team_metadata": strip_callback_config(user_api_key_dict.team_metadata),
|
||||
|
|
@ -2144,7 +2151,6 @@ class LiteLLMProxyRequestSetup:
|
|||
"organization_metadata": strip_callback_config(user_api_key_dict.organization_metadata),
|
||||
}
|
||||
)
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def add_management_endpoint_metadata_to_request_metadata(
|
||||
|
|
|
|||
|
|
@ -707,6 +707,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
|||
_metadata.pop(SESSION_ID_OMITTED_METADATA_KEY, None)
|
||||
_metadata["user_api_key"] = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
|
||||
_metadata["litellm_parent_otel_span"] = user_api_key_dict.parent_otel_span
|
||||
_metadata["user_api_key_auth"] = LiteLLMProxyRequestSetup.auth_for_request_metadata(user_api_key_dict)
|
||||
_metadata["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation
|
||||
_metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = user_api_key_dict.matched_model_access_groups
|
||||
# The per-model budget counters are keyed off these. get_sanitized_user_information_from_key
|
||||
|
|
|
|||
|
|
@ -927,6 +927,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
|
|||
from litellm.proxy.spend_tracking.committed_spend import (
|
||||
committed_spend_key,
|
||||
committed_spend_keys,
|
||||
is_committed_spend_key,
|
||||
seed_committed_spend,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_capture_rate import (
|
||||
|
|
@ -3099,6 +3100,7 @@ proxy_budget_rescheduler_min_time = PROXY_BUDGET_RESCHEDULER_MIN_TIME
|
|||
proxy_budget_rescheduler_max_time = PROXY_BUDGET_RESCHEDULER_MAX_TIME
|
||||
proxy_batch_polling_interval = PROXY_BATCH_POLLING_INTERVAL
|
||||
proxy_batch_write_at = PROXY_BATCH_WRITE_AT
|
||||
_BATCH_WRITE_JITTER_MAX_SECONDS: Final = 5
|
||||
proxy_config_reload_interval_seconds = PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
|
||||
litellm_master_key_hash = None
|
||||
disable_spend_logs = False
|
||||
|
|
@ -3874,9 +3876,25 @@ def _spend_counter_token(key_token: str) -> str:
|
|||
return hash_token(token=key_token) if isinstance(key_token, str) and key_token.startswith("sk-") else key_token
|
||||
|
||||
|
||||
def _committed_spend_ttl() -> int | None:
|
||||
"""A twin outlives two DB spend writes, or the next charge seeds it from a total still missing spend."""
|
||||
redis_cache: Final = spend_counter_cache.redis_cache
|
||||
redis_ttl: Final = redis_cache.default_ttl if redis_cache is not None else None
|
||||
if redis_ttl is None:
|
||||
return None
|
||||
return max(redis_ttl, 2 * (proxy_batch_write_at + _BATCH_WRITE_JITTER_MAX_SECONDS))
|
||||
|
||||
|
||||
def _spend_counter_ttl(counter_key: str, redis_ttl: int | None) -> int | None:
|
||||
return _committed_spend_ttl() if is_committed_spend_key(counter_key) else redis_ttl
|
||||
|
||||
|
||||
async def _committed_spend_seeded(counter_key: str) -> bool:
|
||||
return await _is_spend_counter_cache_warm(committed_spend_key(counter_key)) or await seed_committed_spend(
|
||||
spend_counter_cache=spend_counter_cache, prisma_client=prisma_client, counter_key=counter_key
|
||||
spend_counter_cache=spend_counter_cache,
|
||||
prisma_client=prisma_client,
|
||||
counter_key=counter_key,
|
||||
ttl=_committed_spend_ttl(),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -4301,7 +4319,9 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) ->
|
|||
ttl: Final = redis_cache.get_ttl()
|
||||
with service_target(SPEND_COUNTERS_TARGET):
|
||||
for item in pending:
|
||||
batch.increment(item.counter_key, item.increment, ttl).on_settled(_settle_spend_counter_increment(item))
|
||||
batch.increment(item.counter_key, item.increment, _spend_counter_ttl(item.counter_key, ttl)).on_settled(
|
||||
_settle_spend_counter_increment(item)
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
|
|
@ -4354,7 +4374,11 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) -
|
|||
)
|
||||
ttl: Final = redis_cache.get_ttl()
|
||||
increment_list: Final = [
|
||||
RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl)
|
||||
RedisPipelineIncrementOperation(
|
||||
key=item.counter_key,
|
||||
increment_value=item.increment,
|
||||
ttl=_spend_counter_ttl(item.counter_key, ttl),
|
||||
)
|
||||
for item in pending
|
||||
]
|
||||
results: Final = await redis_cache.async_increment_pipeline(increment_list=increment_list)
|
||||
|
|
@ -11367,7 +11391,7 @@ class ProxyStartupEvent:
|
|||
)
|
||||
|
||||
# Ensure minimum interval of 30 seconds for batch writing to prevent memory issues
|
||||
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, 5)
|
||||
batch_writing_interval: Final = proxy_batch_write_at + random.randint(0, _BATCH_WRITE_JITTER_MAX_SECONDS)
|
||||
|
||||
### PROXY WORKER HEARTBEAT ###
|
||||
worker_heartbeat: Final = ProxyWorkerHeartbeat(prisma_client=prisma_client)
|
||||
|
|
|
|||
|
|
@ -34,13 +34,14 @@ if TYPE_CHECKING:
|
|||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_COMMITTED_KEY_PREFIX: Final = "committed_"
|
||||
_BASE_COUNTER_PREFIXES: Final = ("spend:key:", "spend:team:", "spend:user:", "spend:org:", END_USER_COUNTER_PREFIX)
|
||||
_FETCHED_VALUES: Final = TypeAdapter(dict[str, float | None])
|
||||
_STORED_VALUE: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
|
||||
|
||||
|
||||
def committed_spend_key(counter_key: str) -> str:
|
||||
return f"committed_{counter_key}"
|
||||
return f"{_COMMITTED_KEY_PREFIX}{counter_key}"
|
||||
|
||||
|
||||
def _has_committed_twin(counter_key: str) -> bool:
|
||||
|
|
@ -51,6 +52,10 @@ def committed_spend_keys(counter_keys: Iterable[str]) -> tuple[str, ...]:
|
|||
return tuple(committed_spend_key(counter_key) for counter_key in counter_keys if _has_committed_twin(counter_key))
|
||||
|
||||
|
||||
def is_committed_spend_key(counter_key: str) -> bool:
|
||||
return counter_key.startswith(_COMMITTED_KEY_PREFIX)
|
||||
|
||||
|
||||
def committed_spend_tracked() -> bool:
|
||||
return litellm.logging_callback_manager.callback_is_active(callback_type=PrometheusLogger)
|
||||
|
||||
|
|
@ -142,13 +147,29 @@ async def drop_committed_spend(spend_counter_cache: DualCache, counter_keys: Seq
|
|||
"""After a direct write to the stored spend the twin holds the old total, so the next charge seeds it again."""
|
||||
try:
|
||||
await spend_counter_cache.async_delete_cache_keys(committed_spend_keys(counter_keys))
|
||||
except Exception as e: # noqa: BLE001 # the write already committed; a kept twin only lasts until its TTL
|
||||
except Exception as e: # noqa: BLE001 # the write already committed; a kept twin lasts until charges stop for its TTL
|
||||
verbose_proxy_logger.warning("committed spend drop failed: %s", e)
|
||||
|
||||
|
||||
@with_service_target(SPEND_COUNTERS_TARGET)
|
||||
async def seed_committed_spend(
|
||||
spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str
|
||||
spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str, ttl: int | None = None
|
||||
) -> bool:
|
||||
"""One DB read per worker for a cold twin: charges waiting on the lock find the twin the first one wrote."""
|
||||
twin_key: Final = committed_spend_key(counter_key)
|
||||
async with await SpendCounterReseed.counter_lock(twin_key):
|
||||
try:
|
||||
stored: Final = await _read_committed_values(spend_counter_cache=spend_counter_cache, twin_keys=(twin_key,))
|
||||
except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again
|
||||
verbose_proxy_logger.debug("committed spend seed read failed for %s: %s", twin_key, e)
|
||||
return False
|
||||
return stored.get(twin_key) is not None or await _seed_committed_spend_from_db(
|
||||
spend_counter_cache, prisma_client, counter_key, ttl
|
||||
)
|
||||
|
||||
|
||||
async def _seed_committed_spend_from_db(
|
||||
spend_counter_cache: DualCache, prisma_client: PrismaClient | None, counter_key: str, ttl: int | None
|
||||
) -> bool:
|
||||
db_spend: Final = await (
|
||||
SpendCounterReseed.end_user_from_db(prisma_client=prisma_client, counter_key=counter_key)
|
||||
|
|
@ -164,7 +185,7 @@ async def seed_committed_spend(
|
|||
spend_counter_cache.in_memory_cache.set_cache(key=twin_key, value=db_spend) # pyright: ignore[reportUnknownMemberType] # untyped cache API
|
||||
return True
|
||||
try:
|
||||
await redis_cache.async_set_max(key=twin_key, value=db_spend)
|
||||
await redis_cache.async_set_max(key=twin_key, value=db_spend, ttl=ttl)
|
||||
except Exception as e: # noqa: BLE001 # an unseeded twin only skips this charge; the next one seeds it again
|
||||
verbose_proxy_logger.debug("committed spend seed failed for %s: %s", twin_key, e)
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -3696,34 +3696,41 @@ async def view_spend_logs(
|
|||
)
|
||||
|
||||
|
||||
class _TokenRow(TypedDict):
|
||||
token: ReadOnly[str]
|
||||
class _EntityIdRow(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
|
||||
|
||||
class _TeamIdRow(TypedDict):
|
||||
team_id: ReadOnly[str]
|
||||
_COMMITTED_DROP_PAGE_SIZE: Final = 1000
|
||||
_TOKEN_PAGE_QUERY: Final = (
|
||||
'SELECT token AS id FROM "LiteLLM_VerificationToken" WHERE token > $1 ORDER BY token LIMIT $2'
|
||||
)
|
||||
_TEAM_PAGE_QUERY: Final = 'SELECT team_id AS id FROM "LiteLLM_TeamTable" WHERE team_id > $1 ORDER BY team_id LIMIT $2'
|
||||
|
||||
|
||||
async def _drop_committed_spend_page(
|
||||
prisma_client: PrismaClient, page_query: str, counter_prefix: str, after_id: str
|
||||
) -> str | None:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
rows: Final[Sequence[_EntityIdRow]] = await _query_raw(
|
||||
prisma_client, page_query, after_id, _COMMITTED_DROP_PAGE_SIZE
|
||||
)
|
||||
await drop_committed_spend(spend_counter_cache, tuple(f"{counter_prefix}{row['id']}" for row in rows))
|
||||
return rows[-1]["id"] if len(rows) == _COMMITTED_DROP_PAGE_SIZE else None
|
||||
|
||||
|
||||
async def _drop_committed_spend_pages(prisma_client: PrismaClient, page_query: str, counter_prefix: str) -> None:
|
||||
after_id: str | None = "" # rebind-ok: the keyset cursor moves one page at a time
|
||||
while after_id is not None:
|
||||
after_id = await _drop_committed_spend_page(prisma_client, page_query, counter_prefix, after_id)
|
||||
|
||||
|
||||
async def _drop_all_key_and_team_committed_spend(prisma_client: PrismaClient) -> None:
|
||||
from litellm.proxy.proxy_server import spend_counter_cache
|
||||
|
||||
try:
|
||||
token_rows: Final[Sequence[_TokenRow]] = await _query_raw(
|
||||
prisma_client, 'SELECT token FROM "LiteLLM_VerificationToken"'
|
||||
)
|
||||
team_rows: Final[Sequence[_TeamIdRow]] = await _query_raw(
|
||||
prisma_client, 'SELECT team_id FROM "LiteLLM_TeamTable"'
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin only lasts until its TTL
|
||||
await _drop_committed_spend_pages(prisma_client, _TOKEN_PAGE_QUERY, "spend:key:")
|
||||
await _drop_committed_spend_pages(prisma_client, _TEAM_PAGE_QUERY, "spend:team:")
|
||||
except Exception as e: # noqa: BLE001 # the reset already committed; a kept twin lasts until charges stop for its TTL
|
||||
verbose_proxy_logger.warning("committed spend drop after a global reset failed: %s", e)
|
||||
return
|
||||
await drop_committed_spend(
|
||||
spend_counter_cache,
|
||||
(
|
||||
*(f"spend:key:{row['token']}" for row in token_rows),
|
||||
*(f"spend:team:{row['team_id']}" for row in team_rows),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.caching.redis_batch import (
|
|||
active_post_call_redis_batch,
|
||||
drain_post_call_redis_batches,
|
||||
flush_post_call_redis_batches,
|
||||
post_call_writes_unsent,
|
||||
request_redis_batch_scope,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
|
@ -514,3 +515,16 @@ async def test_a_cancelled_post_call_flush_keeps_the_shared_spend_counter_and_co
|
|||
assert redis_cache.alone == [], "a cancel says nothing about the shared counter, so Redis keeps it"
|
||||
assert spend_cache.in_memory_cache.get_cache("spend:key:k1") == 3.5, "the local copy counts the cancelled spend"
|
||||
assert spend_cache.in_memory_cache.get_cache("spend:team:t1") is None, "an absent local copy is not seeded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_writes_count_as_unsent_until_a_pipeline_carried_them():
|
||||
assert post_call_writes_unsent() is False
|
||||
redis_cache = PostCallFakeRedisCache(FakeClient(_ok_replies))
|
||||
with request_redis_batch_scope():
|
||||
await flush_post_call_redis_batches()
|
||||
assert post_call_writes_unsent() is True
|
||||
active_post_call_redis_batch(redis_cache).increment("spend:key:k", 1.0, None)
|
||||
assert post_call_writes_unsent() is True
|
||||
await flush_post_call_redis_batches()
|
||||
assert post_call_writes_unsent() is False
|
||||
|
|
|
|||
|
|
@ -9,11 +9,18 @@ from datetime import datetime, timezone
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fakeredis import FakeAsyncRedis, FakeServer
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.caching.redis_batch import (
|
||||
active_post_call_redis_batch,
|
||||
flush_post_call_redis_batches,
|
||||
request_redis_batch_scope,
|
||||
)
|
||||
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker
|
||||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
|
|
@ -376,6 +383,103 @@ async def test_slow_committed_read_on_a_failed_request_is_bounded_and_keeps_the_
|
|||
prometheus_logger.litellm_remaining_org_budget_metric.labels().set.assert_called_once_with(pytest.approx(440.0))
|
||||
|
||||
|
||||
class _RedisCacheOnClient(RedisCache):
|
||||
def __init__(self, client: FakeAsyncRedis) -> None:
|
||||
self._injected_client = client
|
||||
self.namespace = None
|
||||
self._circuit_breaker = RedisCircuitBreaker(failure_threshold=5, recovery_timeout=30)
|
||||
self.service_logger_obj = ServiceLogging()
|
||||
self.default_ttl = None
|
||||
|
||||
def init_async_client(self) -> FakeAsyncRedis:
|
||||
return self._injected_client
|
||||
|
||||
|
||||
_TWIN_KEYS = tuple(
|
||||
committed_spend_key(counter_key)
|
||||
for counter_key in ("spend:key:hashed", "spend:team:t1", "spend:user:u1", "spend:org:o1")
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_spend_counter_cache():
|
||||
redis_cache = _RedisCacheOnClient(FakeAsyncRedis(server=FakeServer()))
|
||||
cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
|
||||
with patch( # test-quality-ok: prometheus reads the proxy_server spend counters at emit time, no injection seam
|
||||
"litellm.proxy.proxy_server.spend_counter_cache", cache
|
||||
):
|
||||
yield redis_cache
|
||||
|
||||
|
||||
async def _store_auth_time_spend_plus_a_two_dollar_charge(redis_cache: _RedisCacheOnClient) -> None:
|
||||
for twin_key, spend in zip(_TWIN_KEYS, (5.0, 32.0, 10.0, 62.0)):
|
||||
await redis_cache.init_async_client().set(twin_key, spend)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_add_the_own_charge_to_a_same_size_charge_of_another_request_while_the_own_waits_in_the_pipeline(
|
||||
prometheus_logger, getters, redis_spend_counter_cache
|
||||
):
|
||||
"""Another request's 2.0 reached the counters and this request's 2.0 is queued, so both count."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
await _store_auth_time_spend_plus_a_two_dollar_charge(redis_spend_counter_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
batch = active_post_call_redis_batch(redis_spend_counter_cache)
|
||||
for twin_key in _TWIN_KEYS:
|
||||
batch.increment(twin_key, 2.0, None)
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(3.0),
|
||||
"team": pytest.approx(266.0),
|
||||
"user": pytest.approx(38.0),
|
||||
"org": pytest.approx(436.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_do_not_add_the_own_charge_twice_once_the_pipeline_went_out(
|
||||
prometheus_logger, getters, redis_spend_counter_cache
|
||||
):
|
||||
"""The 2.0 over the auth time spend is this request's own charge, sent before the gauges read it."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
await _store_auth_time_spend_plus_a_two_dollar_charge(redis_spend_counter_cache)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
batch = active_post_call_redis_batch(redis_spend_counter_cache)
|
||||
batch.increment("unrelated-write", 1.0, None)
|
||||
await flush_post_call_redis_batches()
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_without_redis_treat_a_charge_equal_to_the_own_as_the_own(
|
||||
prometheus_logger, getters, spend_counter_cache
|
||||
):
|
||||
"""Without Redis the charge increments in memory at once, so a request scope says nothing about it."""
|
||||
_stub_gauges(prometheus_logger)
|
||||
_store_committed(spend_counter_cache, key=5.0, team=32.0, user=10.0, org=62.0)
|
||||
|
||||
with request_redis_batch_scope():
|
||||
await _emit(prometheus_logger, _request_metadata(_token_committed_at_auth()))
|
||||
|
||||
assert _remaining(prometheus_logger) == {
|
||||
"key": pytest.approx(5.0),
|
||||
"team": pytest.approx(268.0),
|
||||
"user": pytest.approx(40.0),
|
||||
"org": pytest.approx(438.0),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gauges_add_the_cost_to_the_auth_time_spend_before_the_charge_reaches_the_counters(
|
||||
prometheus_logger, getters, spend_counter_cache
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features
|
||||
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.types.proxy.carried_budget_state import CommittedBudgetSpend, EntitySpendCounterKeys
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
|
||||
VertexPassthroughLoggingHandler,
|
||||
|
|
@ -7829,6 +7830,31 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock:
|
|||
return increment_spend_counters
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passthrough_charge_moves_the_committed_spend_of_the_entities_auth_named_budgeted():
|
||||
"""A forged body value stays out, and the row's auth keeps no callback credentials for its loggers."""
|
||||
counter_keys: Final = EntitySpendCounterKeys.of(
|
||||
token="hashed-token", team_id="t1", user_id="u1", org_id=None, end_user_id=None
|
||||
)
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(
|
||||
api_key="hashed-token",
|
||||
user_id="u1",
|
||||
committed_budget_spend=CommittedBudgetSpend(counter_keys=counter_keys),
|
||||
metadata={"logging": [{"callback_name": "langfuse"}], "kept": "yes"},
|
||||
)
|
||||
|
||||
kwargs: Final = _passthrough_kwargs_for_reservation(
|
||||
user_api_key_dict,
|
||||
parsed_body={"litellm_metadata": {"user_api_key_auth": {"committed_budget_spend": None}}},
|
||||
)
|
||||
carried: Final = kwargs["litellm_params"]["metadata"]["user_api_key_auth"]
|
||||
increment_spend_counters: Final = await _track_cost_for_passthrough_kwargs(kwargs)
|
||||
|
||||
assert increment_spend_counters.await_args.kwargs["committed_spend_scope"] == counter_keys
|
||||
assert carried.metadata == {"kept": "yes"}
|
||||
assert user_api_key_dict.metadata["logging"] == [{"callback_name": "langfuse"}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_passthrough_success_reconciles_budget_reservation():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from fakeredis import FakeAsyncRedis, FakeServer
|
||||
|
|
@ -13,6 +15,7 @@ import litellm
|
|||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_batch import flush_post_call_redis_batches, request_redis_batch_scope
|
||||
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy.spend_tracking.committed_spend import (
|
||||
|
|
@ -23,7 +26,13 @@ from litellm.proxy.spend_tracking.committed_spend import (
|
|||
read_committed_budget_spend,
|
||||
seed_committed_spend,
|
||||
)
|
||||
from litellm.proxy.proxy_server import (
|
||||
_defer_spend_counter_increments,
|
||||
_prepare_committed_spend_increments,
|
||||
run_spend_counter_pipeline,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import (
|
||||
PendingSpendIncrement,
|
||||
active_spend_counter_batch,
|
||||
read_batched_spend_counter,
|
||||
spend_counter_batch_scope,
|
||||
|
|
@ -36,11 +45,11 @@ class _FakeRedisSettingMax(FakeAsyncRedis):
|
|||
"""fakeredis runs no Lua without lupa, so this stands in for the GET/compare/SET script of ``async_set_max``."""
|
||||
|
||||
async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> bytes:
|
||||
key, value, _ttl = keys_and_args
|
||||
key, value, ttl = keys_and_args
|
||||
current: Final = await self.get(key)
|
||||
if current is not None and float(current) >= float(value):
|
||||
return current
|
||||
await self.set(key, value)
|
||||
await self.set(key, value, ex=int(ttl) if int(ttl) > 0 else None)
|
||||
return value.encode()
|
||||
|
||||
|
||||
|
|
@ -59,8 +68,11 @@ class _RedisCacheOnClient(RedisCache):
|
|||
class _FindUniqueTable:
|
||||
def __init__(self, row: SimpleNamespace | None) -> None:
|
||||
self._row: Final = row
|
||||
self.reads = 0
|
||||
|
||||
async def find_unique(self, where: dict[str, str]) -> SimpleNamespace | None:
|
||||
self.reads += 1
|
||||
await asyncio.sleep(0)
|
||||
return self._row
|
||||
|
||||
|
||||
|
|
@ -142,6 +154,71 @@ async def test_seed_writes_the_db_spend_when_the_twin_is_absent(
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_charges_on_a_cold_twin_read_the_db_once(spend_counter_cache: DualCache):
|
||||
prisma_client: Final = _prisma_with_spend(12.0)
|
||||
counter_key: Final = f"spend:key:{uuid4()}"
|
||||
|
||||
seeded: Final = await asyncio.gather(
|
||||
*(seed_committed_spend(spend_counter_cache, prisma_client, counter_key) for _ in range(5))
|
||||
)
|
||||
|
||||
assert seeded == [True] * 5
|
||||
assert prisma_client.db.litellm_verificationtoken.reads == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_reports_no_twin_when_redis_cannot_be_read():
|
||||
redis_cache: Final = AsyncMock(spec=RedisCache)
|
||||
redis_cache.async_batch_get_cache.side_effect = ConnectionError("redis unreachable")
|
||||
spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
|
||||
|
||||
assert await seed_committed_spend(spend_counter_cache, _prisma_with_spend(12.0), "spend:team:t1") is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("batch_write_seconds", "expected_twin_ttl"),
|
||||
[
|
||||
(10, 60),
|
||||
(120, 250),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_twin_outlives_two_db_spend_writes_on_every_write_path(
|
||||
monkeypatch: pytest.MonkeyPatch, batch_write_seconds: int, expected_twin_ttl: int
|
||||
):
|
||||
"""A twin expiring before the DB write would be seeded from a total still missing its charges."""
|
||||
client: Final = _FakeRedisSettingMax(server=FakeServer())
|
||||
redis_cache: Final = _RedisCacheOnClient(client)
|
||||
redis_cache.default_ttl = 60
|
||||
spend_counter_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _prisma_with_spend(12.0))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_batch_write_at", batch_write_seconds)
|
||||
seeded_twin: Final = committed_spend_key("spend:key:seeded")
|
||||
piped_twin: Final = committed_spend_key("spend:key:piped")
|
||||
batched_twin: Final = committed_spend_key("spend:key:batched")
|
||||
|
||||
await _prepare_committed_spend_increments(("spend:key:seeded",), 1.0)
|
||||
await run_spend_counter_pipeline(
|
||||
(
|
||||
PendingSpendIncrement(counter_key=piped_twin, increment=1.0),
|
||||
PendingSpendIncrement(counter_key="spend:key:piped", increment=1.0),
|
||||
)
|
||||
)
|
||||
with request_redis_batch_scope():
|
||||
assert _defer_spend_counter_increments(
|
||||
(
|
||||
PendingSpendIncrement(counter_key=batched_twin, increment=1.0),
|
||||
PendingSpendIncrement(counter_key="spend:key:batched", increment=1.0),
|
||||
)
|
||||
)
|
||||
await flush_post_call_redis_batches()
|
||||
|
||||
assert [await client.ttl(twin) for twin in (seeded_twin, piped_twin, batched_twin)] == [expected_twin_ttl] * 3
|
||||
assert [await client.ttl(counter) for counter in ("spend:key:piped", "spend:key:batched")] == [60, 60]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_seed_never_lowers_a_twin_another_replica_already_advanced(spend_counter_cache: DualCache):
|
||||
await _store(spend_counter_cache, committed_spend_key("spend:team:t1"), 20.0)
|
||||
|
|
|
|||
|
|
@ -8284,8 +8284,9 @@ class _ResetDb:
|
|||
self.litellm_verificationtoken: Final = _ResetTableRows()
|
||||
self.litellm_teamtable: Final = _ResetTableRows()
|
||||
|
||||
async def query_raw(self, sql_query: str) -> list[dict[str, str]]:
|
||||
return [{"token": "hashed"}] if "VerificationToken" in sql_query else [{"team_id": "t1"}]
|
||||
async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]:
|
||||
rows: Final = ["hashed"] if "VerificationToken" in sql_query else ["t1"]
|
||||
return [{"id": row} for row in rows if row > after_id][:limit]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8313,6 +8314,38 @@ async def test_global_spend_reset_drops_the_committed_key_and_team_spend_the_gau
|
|||
assert [spend_counter_cache.in_memory_cache.get_cache(key=twin) for twin in twins] == [twin_after, twin_after]
|
||||
|
||||
|
||||
class _PagedResetDb(_ResetDb):
|
||||
def __init__(self, token_count: int) -> None:
|
||||
super().__init__()
|
||||
self._tokens: Final = sorted(f"token-{index}" for index in range(token_count))
|
||||
self.largest_page = 0
|
||||
|
||||
async def query_raw(self, sql_query: str, after_id: str, limit: int) -> list[dict[str, str]]:
|
||||
rows: Final = self._tokens if "VerificationToken" in sql_query else []
|
||||
page: Final = [{"id": row} for row in rows if row > after_id][:limit]
|
||||
self.largest_page = max(self.largest_page, len(page))
|
||||
return page
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("token_count", [999, 1000, 1001, 2500])
|
||||
async def test_global_spend_reset_drops_every_key_twin_reading_the_keys_in_bounded_pages(monkeypatch, token_count):
|
||||
db: Final = _PagedResetDb(token_count)
|
||||
spend_counter_cache: Final = DualCache()
|
||||
twins: Final = tuple(f"committed_spend:key:token-{index}" for index in range(token_count))
|
||||
for twin in twins:
|
||||
spend_counter_cache.in_memory_cache.set_cache(key=twin, value=7.0)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=db))
|
||||
monkeypatch.setattr(litellm, "callbacks", [MagicMock(spec=PrometheusLogger)])
|
||||
|
||||
response: Final = await global_spend_reset()
|
||||
|
||||
assert response["status"] == "success"
|
||||
assert [twin for twin in twins if spend_counter_cache.in_memory_cache.get_cache(key=twin) is not None] == []
|
||||
assert db.largest_page <= 1000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_global_spend_reset_still_succeeds_when_listing_the_committed_counters_fails(monkeypatch):
|
||||
db: Final = _ResetDb()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue