diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index aa458c9c926..87245023633 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -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() diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 28f505370e2..9e0895b4b51 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -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 diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index f30b351eb5a..d3f26f17371 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -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: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index edaacd53d22..91ffc3d2ef6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index cfac28db1a9..bbfb8319789 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dc1133cf596..850e1498608 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/spend_tracking/committed_spend.py b/litellm/proxy/spend_tracking/committed_spend.py index 0f542040d8e..2b69d9cf4a5 100644 --- a/litellm/proxy/spend_tracking/committed_spend.py +++ b/litellm/proxy/spend_tracking/committed_spend.py @@ -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 diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 6c3272c74bb..8a89df55cec 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -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( diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 4b73bb6dd34..bbc32b3fe9d 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -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 diff --git a/tests/unit/integrations/test_prometheus_carried_budget_state.py b/tests/unit/integrations/test_prometheus_carried_budget_state.py index 1a746de588e..351d5be1c8f 100644 --- a/tests/unit/integrations/test_prometheus_carried_budget_state.py +++ b/tests/unit/integrations/test_prometheus_carried_budget_state.py @@ -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 diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index a9098ac3d59..738a7c709c3 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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(): """ diff --git a/tests/unit/proxy/spend_tracking/test_committed_spend.py b/tests/unit/proxy/spend_tracking/test_committed_spend.py index a6beb006d0c..1d7c8fdf304 100644 --- a/tests/unit/proxy/spend_tracking/test_committed_spend.py +++ b/tests/unit/proxy/spend_tracking/test_committed_spend.py @@ -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) diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index c82dcb57bfd..057f512ab0d 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -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()