From f6be848368ebeb9b3f70b36d193380879390d8c9 Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Sun, 11 Oct 2026 21:06:04 +0000 Subject: [PATCH 01/11] fix(proxy): give spend counters a Redis TTL that outlives batched DB spend writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../proxy/common_utils/reset_budget_job.py | 6 ++- litellm/proxy/db/spend_counter_reseed.py | 7 ++++ .../key_management_endpoints.py | 16 +++++-- litellm/proxy/proxy_server.py | 22 ++++++++-- .../spend_tracking/budget_reservation.py | 4 +- .../test_key_management_endpoints.py | 6 +-- .../proxy/proxy_server/test_spend_counters.py | 42 +++++++++++++++++-- .../test_spend_counter_batch.py | 8 +++- tests/unit/proxy/test_budget_reservation.py | 21 +++++++--- 10 files changed, 108 insertions(+), 25 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 38d16d61d7b..17633df07d8 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1894,6 +1894,7 @@ _background_interaction_cost_polling_env: Final = os.getenv( ).lower() BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true" PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) +SPEND_COUNTER_REDIS_TTL_SECONDS: Final = int(os.getenv("SPEND_COUNTER_REDIS_TTL_SECONDS", "3600")) PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 61330d12ac3..99ce358959c 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1415,8 +1415,12 @@ class ResetBudgetJob: new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_cache) spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_value) if spend_counter_cache.redis_cache is not None: + from litellm.proxy.proxy_server import spend_counter_redis_ttl + try: - await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_value) + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=new_value, ttl=spend_counter_redis_ttl() + ) except Exception as redis_err: verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err) budget_duration: Final = window["budget_duration"] diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 721999b1868..ab7ec7d3b1d 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -267,10 +267,13 @@ class SpendCounterReseed: current_value: float = float(db_spend) try: if spend_counter_cache.redis_cache is not None: + from litellm.proxy.proxy_server import spend_counter_redis_ttl + seeded: Final = await spend_counter_cache.redis_cache.async_set_cache( key=counter_key, value=db_spend, nx=True, + ttl=spend_counter_redis_ttl(), ) if seeded: current_value = float(db_spend) @@ -462,10 +465,13 @@ class SpendCounterReseed: return None try: if spend_counter_cache.redis_cache is not None: + from litellm.proxy.proxy_server import spend_counter_redis_ttl + seeded: Final = await spend_counter_cache.redis_cache.async_set_cache( key=counter_key, value=window_spend, nx=True, + ttl=spend_counter_redis_ttl(), ) if seeded: current_value = window_spend @@ -475,6 +481,7 @@ class SpendCounterReseed: current_value = await spend_counter_cache.redis_cache.async_increment( key=counter_key, value=window_spend, + ttl=spend_counter_redis_ttl(), ) else: current_value = float(current_cached_value) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 782a7e1db3c..6cb761174c9 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -3799,14 +3799,16 @@ async def update_key_fn( ) if data.spend is not None: - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import spend_counter_cache, spend_counter_redis_ttl counter_key: Final = f"spend:key:{hash_token_if_needed(key)}" spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=data.spend, ttl=60) if spend_counter_cache.redis_cache is not None: try: with service_target(SPEND_COUNTERS_TARGET): - await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=data.spend, ttl=60) + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=data.spend, ttl=spend_counter_redis_ttl() + ) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to update spend counter %s in Redis after key spend update: %s. " @@ -6350,12 +6352,18 @@ async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: f Without the broadcast, a worker that already cached the pre-reset key object or floor marker keeps enforcing against it until its own TTL expires. """ - from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache + from litellm.proxy.proxy_server import ( + SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + spend_counter_cache, + spend_counter_redis_ttl, + ) spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=value, ttl=60) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=value, ttl=60) + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=value, ttl=spend_counter_redis_ttl() + ) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to update spend counter %s in Redis: %s. " diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 089647c6297..c058de8b782 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -306,6 +306,7 @@ from litellm.constants import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG, SPEND_CAPTURE_RATE_CHECK_JOB_ID, + SPEND_COUNTER_REDIS_TTL_SECONDS, USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, ) @@ -2914,6 +2915,16 @@ async def _drain_spend_event_producer_on_shutdown() -> None: # Bounds authoritative DB re-reads when enforcing a budget against a # stale-low spend counter: at most one DB read per counter per window. SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5 +_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE: Final = 10 + + +def spend_counter_redis_ttl() -> int: + """Redis TTL for spend counters. A counter that expires before its spend is flushed to the DB reseeds from a + stale-low row and re-admits a key that is over budget, so it must outlive many batch-write intervals.""" + batch_write_at: Final = ( + proxy_batch_write_at if isinstance(proxy_batch_write_at, (int, float)) else PROXY_BATCH_WRITE_AT + ) + return max(SPEND_COUNTER_REDIS_TTL_SECONDS, int(_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE * batch_write_at)) def _fail_closed_budget_enforcement() -> bool: @@ -3046,7 +3057,9 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None if spend_counter_cache.redis_cache is not None: forget_spend_counter(counter_key) try: - repaired: Final = await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend) + repaired: Final = await spend_counter_cache.redis_cache.async_set_max( + key=counter_key, value=db_spend, ttl=spend_counter_redis_ttl() + ) except Exception: verbose_proxy_logger.debug( "Unable to repair stale spend counter %s in Redis", @@ -3879,7 +3892,7 @@ async def refresh_spend_counter_ttl(counter_key: str) -> bool: if spend_counter_cache.redis_cache is None: return False try: - return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key, ttl=spend_counter_redis_ttl()) except Exception as e: verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e) return False @@ -3892,6 +3905,7 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f current_value: Final = await spend_counter_cache.redis_cache.async_increment( key=counter_key, value=increment, + ttl=spend_counter_redis_ttl(), refresh_ttl=True, ) except Exception: @@ -3951,7 +3965,7 @@ def _defer_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> batch: Final = active_post_call_redis_batch(redis_cache) if batch is None: return False - ttl: Final = redis_cache.get_ttl() + ttl: Final = spend_counter_redis_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)) @@ -4005,7 +4019,7 @@ async def run_spend_counter_pipeline(pending: Sequence[PendingSpendIncrement]) - for item in pending ] ) - ttl: Final = redis_cache.get_ttl() + ttl: Final = spend_counter_redis_ttl() increment_list: Final = [ RedisPipelineIncrementOperation(key=item.counter_key, increment_value=item.increment, ttl=ttl) for item in pending diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4c9df95cd8d..4e45adde036 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -119,7 +119,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c """A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL while the request is in flight so a request longer than the TTL does not drop its reservation and admit concurrent requests against the DB floor on any worker.""" - from litellm.proxy.proxy_server import spend_counter_cache + from litellm.proxy.proxy_server import spend_counter_cache, spend_counter_redis_ttl if spend_counter_cache.redis_cache is None or not counter_keys: return @@ -127,7 +127,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c _renew_reservation_lease( budget_reservation=budget_reservation, counter_keys=counter_keys, - interval=spend_counter_cache.redis_cache.default_ttl / 2, + interval=spend_counter_redis_ttl() / 2, request_task=asyncio.current_task(), ) ) diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 71334e2321b..4c2c8c7b509 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -95,7 +95,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_key_list_check, validate_key_team_change, ) -from litellm.proxy.proxy_server import app +from litellm.proxy.proxy_server import app, spend_counter_redis_ttl from litellm.types.proxy.management_endpoints.key_management_endpoints import ( BulkUpdateKeyRequest, BulkUpdateKeyResponse, @@ -8449,7 +8449,7 @@ async def test_reset_key_spend_resets_budget_windows(monkeypatch): key=window_counter_key, value=0.0, ttl=60 ) mock_spend_counter_cache.redis_cache.async_set_cache.assert_any_call( - key=window_counter_key, value=0.0, ttl=60 + key=window_counter_key, value=0.0, ttl=spend_counter_redis_ttl() ) mock_spend_counter_cache.in_memory_cache.set_cache.assert_any_call( key=f"spend_db_floor:{window_counter_key}", value=0.0, ttl=5 @@ -8658,7 +8658,7 @@ async def test_update_key_spend_updates_counter(monkeypatch): key=f"spend:key:{hashed_key}", value=0.0, ttl=60 ) mock_spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( - key=f"spend:key:{hashed_key}", value=0.0, ttl=60 + key=f"spend:key:{hashed_key}", value=0.0, ttl=spend_counter_redis_ttl() ) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index ad86c3c5267..7ede9185043 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -134,7 +134,9 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # the stale counter is repaired up to the authoritative DB value via a # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:key:abc", value=12.0) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key="spend:key:abc", value=12.0, ttl=ps.spend_counter_redis_ttl() + ) @pytest.mark.asyncio @@ -282,7 +284,9 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk ) assert result == 12.0 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key="spend:end_user:customer-42", value=12.0) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key="spend:end_user:customer-42", value=12.0, ttl=ps.spend_counter_redis_ttl() + ) @pytest.mark.asyncio @@ -326,7 +330,9 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key=counter_key, value=15.0, ttl=ps.spend_counter_redis_ttl() + ) def _make_window_spend_prisma(row=None, spend_logs_total=0.0): @@ -367,7 +373,9 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() - fake_cache.redis_cache.async_set_max.assert_awaited_once_with(key=counter_key, value=15.0) + fake_cache.redis_cache.async_set_max.assert_awaited_once_with( + key=counter_key, value=15.0, ttl=ps.spend_counter_redis_ttl() + ) @pytest.mark.asyncio @@ -830,6 +838,32 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 +@pytest.mark.asyncio +@pytest.mark.parametrize("batch_write_at", [60, 900]) +async def test_spend_counter_redis_ttl_outlives_db_batch_write_interval(monkeypatch, batch_write_at): + """A counter that expires before its buffered spend is written to the DB reseeds from a stale-low + row and re-admits an over-budget key, so every counter write must outlive the batch interval""" + counter_cache = ps.DualCache() + fake_redis = AsyncMock() + fake_redis.async_increment = AsyncMock(return_value=1.0) + fake_redis.async_increment_pipeline = AsyncMock(return_value=[1.0]) + fake_redis.async_refresh_ttl = AsyncMock(return_value=True) + fake_redis.get_ttl = MagicMock(return_value=60) + fake_redis.default_ttl = 60 + counter_cache.redis_cache = fake_redis + monkeypatch.setattr(ps, "spend_counter_cache", counter_cache) + monkeypatch.setattr(ps, "proxy_batch_write_at", batch_write_at) + max_flush_delay: Final = batch_write_at + 5 + + await ps.increment_spend_counter_cache(counter_key="spend:key:direct", increment=0.5) + await ps.run_spend_counter_pipeline((ps.PendingSpendIncrement(counter_key="spend:key:batched", increment=0.5),)) + await ps.refresh_spend_counter_ttl(counter_key="spend:key:reserved") + + assert fake_redis.async_increment.await_args.kwargs["ttl"] > max_flush_delay + assert fake_redis.async_increment_pipeline.await_args.kwargs["increment_list"][0]["ttl"] > max_flush_delay + assert fake_redis.async_refresh_ttl.await_args.kwargs["ttl"] > max_flush_delay + + @pytest.mark.asyncio async def test_increment_spend_counters_pipeline_failure_invalidates_all_counters( monkeypatch, diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index 1fddfaaa766..bdb430e1945 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -351,7 +351,9 @@ async def test_reseed_treats_a_batched_clean_miss_as_authoritative_and_seeds_fro assert value == 2.25 assert [c for c in redis.commands if c.startswith("GET")] == [] - redis.async_set_cache.assert_awaited_once_with(key="spend:key:hashed", value=2.25, nx=True) + redis.async_set_cache.assert_awaited_once_with( + key="spend:key:hashed", value=2.25, nx=True, ttl=ps.spend_counter_redis_ttl() + ) @pytest.mark.asyncio @@ -404,7 +406,9 @@ async def test_post_call_cold_counters_seed_from_the_mget_miss_without_a_second_ await ps.increment_spend_counters(token="hashed", team_id="team", user_id=None, response_cost=0.5) assert [c.split()[0] for c in redis.commands] == ["MGET", "PIPELINE"], redis.commands - redis.async_set_cache.assert_awaited_once_with(key="spend:team:team", value=4.0, nx=True) + redis.async_set_cache.assert_awaited_once_with( + key="spend:team:team", value=4.0, nx=True, ttl=ps.spend_counter_redis_ttl() + ) assert redis.store["spend:key:hashed"] == 1.5 diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 4635e145914..eb16863662e 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -2347,6 +2347,9 @@ class _ExpiringRedisCache: self.refresh_count = 0 self.fail_first_refresh = fail_first_refresh + def _expiry(self, ttl: object) -> float: + return time.monotonic() + (float(ttl) if isinstance(ttl, (int, float)) else self.default_ttl) + def _evict_expired(self, key: str) -> None: if self.expires_at.get(key, float("inf")) <= time.monotonic(): self.store.pop(key, None) @@ -2359,18 +2362,18 @@ class _ExpiringRedisCache: async def async_increment(self, key: str, value: float, **kwargs: object) -> float: self._evict_expired(key) self.store[key] = self.store.get(key, 0.0) + float(value) - self.expires_at[key] = time.monotonic() + self.default_ttl + self.expires_at[key] = self._expiry(kwargs.get("ttl")) return self.store[key] async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: self._evict_expired(key) self.store[key] = max(self.store.get(key, float("-inf")), float(value)) - self.expires_at[key] = time.monotonic() + self.default_ttl + self.expires_at[key] = self._expiry(kwargs.get("ttl")) return self.store[key] async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: self.store[key] = float(value) - self.expires_at[key] = time.monotonic() + self.default_ttl + self.expires_at[key] = self._expiry(kwargs.get("ttl")) return True async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: @@ -2385,13 +2388,15 @@ class _ExpiringRedisCache: if key not in self.store: return False self.refresh_count += 1 - self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) + self.expires_at[key] = self._expiry(ttl) return True async def async_increment_pipeline( self, increment_list: Sequence[RedisPipelineIncrementOperation], **kwargs: object ) -> list[float]: - return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list] + return [ + await self.async_increment(op["key"], op["increment_value"], ttl=op.get("ttl")) for op in increment_list + ] def get_ttl(self, **kwargs: object) -> int | None: return int(self.default_ttl) @@ -2400,11 +2405,13 @@ class _ExpiringRedisCache: @pytest.mark.asyncio async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( spend_counter_state, + monkeypatch, ): """A request that runs longer than the counter TTL must keep its reservation in Redis (so a concurrent request on any worker still sees it), and renewal must stop once the reservation is reconciled so an idle counter still expires on its own.""" counter_cache, key_cache = spend_counter_state + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 0.2) redis_cache = _ExpiringRedisCache(default_ttl=0.2) counter_cache.redis_cache = redis_cache proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) @@ -2432,10 +2439,12 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( @pytest.mark.asyncio async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( spend_counter_state, + monkeypatch, ): """One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the rest of the request.""" counter_cache, key_cache = spend_counter_state + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 0.2) redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True) counter_cache.redis_cache = redis_cache proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) @@ -2452,11 +2461,13 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( @pytest.mark.asyncio async def test_reservation_lease_stops_when_request_task_ends_without_reconciling( spend_counter_state, + monkeypatch, ): """A request whose task ends without reconciling (client disconnect path that skips the cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of pinning the reservation until the request timeout.""" counter_cache, key_cache = spend_counter_state + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 0.2) redis_cache = _ExpiringRedisCache(default_ttl=0.2) counter_cache.redis_cache = redis_cache proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) From 495a727d93538f71055dd893f08d6d1bb431b194 Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Sun, 11 Oct 2026 21:10:12 +0000 Subject: [PATCH 02/11] fix(proxy): keep the spend counter TTL a constant instead of a new env var Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index 17633df07d8..4e67f31eae4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1894,7 +1894,7 @@ _background_interaction_cost_polling_env: Final = os.getenv( ).lower() BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true" PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) -SPEND_COUNTER_REDIS_TTL_SECONDS: Final = int(os.getenv("SPEND_COUNTER_REDIS_TTL_SECONDS", "3600")) +SPEND_COUNTER_REDIS_TTL_SECONDS: Final = 3600 PROXY_BATCH_WRITE_AT: Final = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) From 1985f7890851f2b0c90172b681b2321318d09e3f Mon Sep 17 00:00:00 2001 From: Mubashir Osmani Date: Sun, 11 Oct 2026 21:27:47 +0000 Subject: [PATCH 03/11] fix(proxy): keep lease renewal polling short and never shorten a configured Redis TTL Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 7 +++++- .../spend_tracking/budget_reservation.py | 4 ++-- .../proxy/proxy_server/test_spend_counters.py | 8 +++++++ tests/unit/proxy/test_budget_reservation.py | 23 +++++++++++++++++++ ...test_proxy_server_endpoints_and_startup.py | 5 +++- 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c058de8b782..6c09cbeccf4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2924,7 +2924,12 @@ def spend_counter_redis_ttl() -> int: batch_write_at: Final = ( proxy_batch_write_at if isinstance(proxy_batch_write_at, (int, float)) else PROXY_BATCH_WRITE_AT ) - return max(SPEND_COUNTER_REDIS_TTL_SECONDS, int(_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE * batch_write_at)) + configured_redis_ttl: Final = int(litellm.default_redis_ttl) if litellm.default_redis_ttl is not None else 0 + return max( + SPEND_COUNTER_REDIS_TTL_SECONDS, + int(_SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE * batch_write_at), + configured_redis_ttl, + ) def _fail_closed_budget_enforcement() -> bool: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4e45adde036..4c9df95cd8d 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -119,7 +119,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c """A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL while the request is in flight so a request longer than the TTL does not drop its reservation and admit concurrent requests against the DB floor on any worker.""" - from litellm.proxy.proxy_server import spend_counter_cache, spend_counter_redis_ttl + from litellm.proxy.proxy_server import spend_counter_cache if spend_counter_cache.redis_cache is None or not counter_keys: return @@ -127,7 +127,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c _renew_reservation_lease( budget_reservation=budget_reservation, counter_keys=counter_keys, - interval=spend_counter_redis_ttl() / 2, + interval=spend_counter_cache.redis_cache.default_ttl / 2, request_task=asyncio.current_task(), ) ) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 7ede9185043..64c730e4f81 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -838,6 +838,14 @@ async def test_increment_spend_counters_pipelines_all_scopes_in_one_redis_call( assert counter_cache.in_memory_cache.get_cache(key=key) == 1.5 +def test_spend_counter_redis_ttl_never_shortens_configured_redis_ttl(monkeypatch): + """Deployments that configured a longer Redis default must not see spend counters expire sooner""" + monkeypatch.setattr(ps.litellm, "default_redis_ttl", 86400) + monkeypatch.setattr(ps, "proxy_batch_write_at", 60) + + assert ps.spend_counter_redis_ttl() == 86400 + + @pytest.mark.asyncio @pytest.mark.parametrize("batch_write_at", [60, 900]) async def test_spend_counter_redis_ttl_outlives_db_batch_write_interval(monkeypatch, batch_write_at): diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index eb16863662e..9fe30f4c696 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -40,6 +40,7 @@ from litellm.proxy.common_utils.user_api_key_cache import ( ) from litellm.proxy.spend_tracking.budget_reservation import ( _get_model_access_group_budget_counters, + _lease_renewals, estimate_request_max_cost, get_budget_window_start, invalidate_budget_reservation_counters, @@ -2436,6 +2437,28 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( assert await redis_cache.async_get_cache(key=counter_key) is None +@pytest.mark.asyncio +async def test_reservation_lease_renewal_stops_promptly_with_long_spend_counter_ttl( + spend_counter_state, + monkeypatch, +): + """A long spend-counter TTL must not stretch how often renewal checks for completion, or finished + requests keep their renewal tasks alive for half that TTL""" + counter_cache, key_cache = spend_counter_state + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_redis_ttl", lambda: 3600) + counter_cache.redis_cache = _ExpiringRedisCache(default_ttl=0.2) + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-long-ttl", spend=0.0, max_budget=1.0) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + assert len(_lease_renewals) == 1 + + await release_budget_reservation(reservation) + await asyncio.sleep(0.3) + assert len(_lease_renewals) == 0 + + @pytest.mark.asyncio async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( spend_counter_state, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 920f9d90bb8..e7df36f1e3b 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -8917,7 +8917,9 @@ async def test_prepare_spend_counter_increment_reseeds_from_db_on_counter_miss() fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. # Only the per-request delta (1.5) goes through INCRBYFLOAT. - fake_redis.async_set_cache.assert_awaited_once_with(key="spend:team:team-9", value=42.0, nx=True) + fake_redis.async_set_cache.assert_awaited_once_with( + key="spend:team:team-9", value=42.0, nx=True, ttl=proxy_server_module.spend_counter_redis_ttl() + ) writes = [(c["key"], c["value"]) for c in recorded_increments] assert writes == [("spend:team:team-9", 1.5)] finally: @@ -9346,6 +9348,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() key=counter_key, value=2.25, nx=True, + ttl=proxy_server_module.spend_counter_redis_ttl(), ) assert redis_store[counter_key] == pytest.approx(3.25) assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(3.25) From 21aa23caaea4c93cec6a44125112ff6ccadc08c8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:04:20 -0700 Subject: [PATCH 04/11] fix(proxy): give the team member spend reset the spend counter Redis TTL --- litellm/proxy/auth/auth_checks.py | 10 ++++++++-- .../auth/test_auth_checks_object_access_and_lookup.py | 5 ++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 56773789059..ab2729c26b5 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3198,7 +3198,11 @@ async def invalidate_team_member_spend_state( await asyncio.wait((inflight,)) if new_spend is not None: - from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache + from litellm.proxy.proxy_server import ( + SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + spend_counter_cache, + spend_counter_redis_ttl, + ) spend_counter_key: Final = f"spend:team_member:{user_id}:{team_id}" spend_db_floor_key: Final = f"spend_db_floor:{spend_counter_key}" @@ -3206,7 +3210,9 @@ async def invalidate_team_member_spend_state( spend_counter_cache.in_memory_cache.set_cache(key=spend_counter_key, value=new_spend, ttl=60) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache(key=spend_counter_key, value=new_spend, ttl=60) + await spend_counter_cache.redis_cache.async_set_cache( + key=spend_counter_key, value=new_spend, ttl=spend_counter_redis_ttl() + ) except Exception as e: # noqa: BLE001 # fall back to deleting the stale entry before giving up verbose_proxy_logger.warning( "Failed to set spend counter %s in Redis after reset: %s; deleting it instead so the next " diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 25022621edc..de8ff5c0545 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10011,6 +10011,7 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele caller passing new_spend must match that precedent, not merely delete the counter.""" from litellm.caching.dual_cache import DualCache from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.proxy_server import spend_counter_redis_ttl real_cache = UserApiKeyCache() real_spend_counter_cache = DualCache() @@ -10029,7 +10030,9 @@ async def test_invalidate_team_member_spend_state_sets_new_spend_instead_of_dele ) assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == 2.5 - fake_redis_cache.async_set_cache.assert_awaited_once_with(key="spend:team_member:user-1:team-1", value=2.5, ttl=60) + fake_redis_cache.async_set_cache.assert_awaited_once_with( + key="spend:team_member:user-1:team-1", value=2.5, ttl=spend_counter_redis_ttl() + ) @pytest.mark.asyncio From 2cc0c79d074761b7f7ada2e9f5ec61b24d532865 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:54:13 -0700 Subject: [PATCH 05/11] fix(proxy): drop spend counters on global and organization spend resets `/global/spend/reset` and `/organization/update` with `spend` wrote only the DB `.spend` column, so with the longer counter TTL a blocked key or organization stayed blocked until its counter expired. Both now drop the affected counters (every `spend:key:*` and `spend:team:*` counter through a bounded Redis scan, the one `spend:org:` counter) so the reset takes effect on the next request. `RedisCache.async_scan_iter` now returns decoded `str` keys and stops at `count` without a mutable accumulator, and `delete_cache_keys` takes an `Iterable[str]`, so callers no longer see `bytes` keys. The audit's integration cells for the TTL (16 counter scopes and surfaces, 7 settings edges, 4 chaos cells) are added under tests/integration/spend. The leaked-reservation chaos cell asserts the shipped behavior: a reservation left by a hard-killed worker stays on the counter for the counter TTL. --- litellm/caching/redis_cache.py | 32 +- .../organization_endpoints.py | 5 + litellm/proxy/proxy_server.py | 26 + .../spend_management_endpoints.py | 4 +- tests/integration/_support/spend_counters.py | 33 ++ tests/integration/_support/upstream.py | 91 +++- .../spend/test_spend_counter_redis_ttl.py | 488 ++++++++++++++++++ .../spend/test_spend_counter_ttl_chaos.py | 296 +++++++++++ .../spend/test_spend_counter_ttl_settings.py | 228 ++++++++ tests/unit/caching/test_redis_cache.py | 19 + .../common_utils/test_reset_budget_job.py | 32 ++ .../test_organization_endpoints.py | 36 ++ .../test_spend_management_endpoints.py | 54 ++ 13 files changed, 1315 insertions(+), 29 deletions(-) create mode 100644 tests/integration/_support/spend_counters.py create mode 100644 tests/integration/spend/test_spend_counter_redis_ttl.py create mode 100644 tests/integration/spend/test_spend_counter_ttl_chaos.py create mode 100644 tests/integration/spend/test_spend_counter_ttl_settings.py diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3a847044e3b..108cc29b509 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,7 +18,7 @@ import json import logging import threading import time -from collections.abc import Awaitable, Callable, Iterator, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -498,6 +498,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +async def _first_keys(scan: AsyncIterator[object], count: int) -> AsyncIterator[str]: + taken: Final = itertools.count(1) + async for key in scan: + yield key.decode() if isinstance(key, bytes) else str(key) + if next(taken) >= count: + return + + class RedisCircuitBreakerOpenError(Exception): pass @@ -1007,10 +1015,9 @@ class RedisCache(BaseCache): return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) @_redis_circuit_breaker_guard - async def async_scan_iter(self, pattern: str, count: int = 100) -> list: + async def async_scan_iter(self, pattern: str, count: int = 100) -> list[str]: start_time: Final = time.time() try: - keys: Final = [] _redis_client: Final = self.init_async_client() if not hasattr(_redis_client, "scan_iter"): verbose_logger.debug( @@ -1018,11 +1025,13 @@ class RedisCache(BaseCache): ) return [] - pattern = self.check_and_fix_namespace(key=pattern) - async for key in _redis_client.scan_iter(match=pattern + "*", count=count): - keys.append(key) - if len(keys) >= count: - break + namespaced_pattern: Final = self.check_and_fix_namespace(key=pattern) + keys: Final = [ + key + async for key in _first_keys( + _redis_client.scan_iter(match=namespaced_pattern + "*", count=count), count + ) + ] ## LOGGING ## end_time = time.time() @@ -1904,11 +1913,10 @@ class RedisCache(BaseCache): raise e @_redis_circuit_breaker_guard - async def delete_cache_keys(self, keys): + async def delete_cache_keys(self, keys: Iterable[str]) -> None: _redis_client: Final = self._async_commands() - keys = [self.check_and_fix_namespace(key=key) for key in keys] - # keys is a list, unpack it so it gets passed as individual elements to delete - await _redis_client.delete(*keys) + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in keys] + await _redis_client.delete(*namespaced_keys) def client_list(self) -> list: client_list: Final[list] = self.redis_client.client_list() diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index b8ffec639c0..8928832626e 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -816,6 +816,11 @@ async def update_organization( include={"members": True, "teams": True, "litellm_budget_table": True}, ) + if data.spend is not None: + from litellm.proxy.proxy_server import invalidate_spend_counter + + await invalidate_spend_counter(counter_key=f"spend:org:{data.organization_id}") + return response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6c09cbeccf4..2e63584f60e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -51,6 +51,7 @@ import websockets import websockets.exceptions from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config from pydantic.fields import FieldInfo, PydanticUndefined +from redis.exceptions import RedisError from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._uuid import uuid @@ -3948,6 +3949,31 @@ async def invalidate_spend_counter(counter_key: str) -> None: _invalidate_spend_counter: Final = invalidate_spend_counter +_IN_MEMORY_COUNTER_KEYS: Final = TypeAdapter(tuple[str, ...]) + + +_SPEND_COUNTER_STORE_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, ValidationError) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def invalidate_spend_counters_with_prefix(prefix: str) -> None: + for counter_key in _IN_MEMORY_COUNTER_KEYS.validate_python(tuple(spend_counter_cache.in_memory_cache.cache_dict)): + if counter_key.startswith(prefix): + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + return + try: + while scanned := await redis_cache.async_scan_iter(pattern=prefix, count=500): + await redis_cache.delete_cache_keys(scanned) + except _SPEND_COUNTER_STORE_FAILURES: + verbose_proxy_logger.warning( + "Unable to delete the spend counters %s* after the spend reset; budget checks use the stale counters until " + "they expire", + prefix, + exc_info=True, + ) + async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: if _defer_spend_counter_increments(pending): diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5482661c3de..8384e5e9177 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3684,7 +3684,7 @@ async def global_spend_reset(): 3. LiteLLM_TeamTable spend will be set = 0 """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import invalidate_spend_counters_with_prefix, prisma_client if prisma_client is None: raise ProxyException( @@ -3696,6 +3696,8 @@ async def global_spend_reset(): await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + await invalidate_spend_counters_with_prefix(prefix="spend:key:") + await invalidate_spend_counters_with_prefix(prefix="spend:team:") return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/tests/integration/_support/spend_counters.py b/tests/integration/_support/spend_counters.py new file mode 100644 index 00000000000..a8c8b738df5 --- /dev/null +++ b/tests/integration/_support/spend_counters.py @@ -0,0 +1,33 @@ +import os +from collections.abc import Iterator +from contextlib import contextmanager +from hashlib import sha256 +from typing import Final + +from redis import Redis + +SPEND_COUNTER_TTL_SECONDS: Final = 3600 +TTL_SLACK_SECONDS: Final = 600 + + +@contextmanager +def rig_cache() -> Iterator[Redis]: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + yield cache + + +def key_counter(key: str) -> str: + return f"spend:key:{sha256(key.encode()).hexdigest()}" + + +def counter_ttl(cache: Redis, counter_key: str) -> int: + return int(str(cache.ttl(counter_key))) + + +def counter_value(cache: Redis, counter_key: str) -> float | None: + raw: Final = cache.get(counter_key) + return None if raw is None else float(str(raw, "utf-8") if isinstance(raw, bytes) else str(raw)) + + +def outlives(ttl: int, expected: int = SPEND_COUNTER_TTL_SECONDS) -> bool: + return expected - TTL_SLACK_SECONDS < ttl <= expected diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 6b36d56a757..831a4543a55 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -6,6 +6,7 @@ import base64 import json import os import struct +import time import uuid import zlib from collections import deque @@ -60,6 +61,46 @@ INTERNAL_FIELDS: Final = frozenset( ) +RESPONSES_USAGE: Final = {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40} + + +def responses_body(model: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{uuid.uuid4().hex[:24]}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": model, + "output": [ + { + "id": f"msg_{uuid.uuid4().hex[:24]}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Hello.", "annotations": []}], + } + ], + "usage": dict(RESPONSES_USAGE), + } + + +async def responses_stream(completed: Mapping[str, JsonValue]) -> AsyncIterator[str]: + item: Final = cast(dict[str, JsonValue], cast(list[JsonValue], completed["output"])[0]) + events: Final = ( + ("response.created", {"response": {**completed, "status": "in_progress", "output": [], "usage": None}}), + ("response.output_item.added", {"output_index": 0, "item": {**item, "status": "in_progress", "content": []}}), + ( + "response.output_text.delta", + {"output_index": 0, "content_index": 0, "item_id": item["id"], "delta": "Hello."}, + ), + ("response.output_item.done", {"output_index": 0, "item": item}), + ("response.completed", {"response": dict(completed)}), + ) + for sequence_number, (event_type, payload) in enumerate(events): + data: Final = json.dumps({"type": event_type, "sequence_number": sequence_number, **payload}) + yield f"event: {event_type}\ndata: {data}\n\n" + + def error_type(status: int) -> str: if status == 429: return "rate_limit_error" @@ -183,24 +224,41 @@ class Provider: for message in messages ): return JSONResponse({"error": {"message": "Invalid selected message contract"}}, status_code=400) - script: Final = self.scripts.get(str(body["model"])) - if script is not None: - if not script: - return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500) - status: Final = script.popleft() - if status != 200: - return JSONResponse( - { - "error": { - "message": "Controlled provider failure", - "type": error_type(status), - "code": str(status), - } - }, - status_code=status, - ) + failure: Final = self._scripted_failure(str(body["model"])) + if failure is not None: + return failure return await chat_completions(request) + def _scripted_failure(self, model: str) -> Response | None: + script: Final = self.scripts.get(model) + if script is None: + return None + if not script: + return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500) + status: Final = script.popleft() + if status == 200: + return None + return JSONResponse( + {"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}}, + status_code=status, + ) + + async def responses(self, request: Request) -> Response: + body: Final = JSON_OBJECT.validate_json(await request.body()) + self.observations.put( + Observation(request.url.path, request.headers.get("authorization", ""), body, request.method) + ) + model: Final = body.get("model") + if not isinstance(model, str) or "input" not in body: + return JSONResponse({"error": {"message": "model and input are required"}}, status_code=400) + failure: Final = self._scripted_failure(model) + if failure is not None: + return failure + completed: Final = responses_body(model) + if body.get("stream") is True: + return StreamingResponse(responses_stream(completed), media_type="text/event-stream") + return JSONResponse(completed) + async def vector_store_search(self, request: Request) -> Response: body: Final = JSON_OBJECT.validate_json(await request.body()) self.observations.put( @@ -560,6 +618,7 @@ class Provider: Route("/_cost_map", self.cost_map, methods=["GET"]), Route("/_oauth/token", self.oauth_token, methods=["POST"]), Route("/v1/chat/completions", self.chat, methods=["POST"]), + Route("/v1/responses", self.responses, methods=["POST"]), Route("/v1/completions", completions, methods=["POST"]), Route("/v1/embeddings", embeddings, methods=["POST"]), Route("/v1/moderations", moderations, methods=["POST"]), diff --git a/tests/integration/spend/test_spend_counter_redis_ttl.py b/tests/integration/spend/test_spend_counter_redis_ttl.py new file mode 100644 index 00000000000..ac4e338993b --- /dev/null +++ b/tests/integration/spend/test_spend_counter_redis_ttl.py @@ -0,0 +1,488 @@ +import asyncio +import json +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from hashlib import sha256 +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.spend_counters import ( + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import JsonValue +from redis import Redis + +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} + + +@dataclass(frozen=True, slots=True) +class Entities: + model: str + key: str + team: str + user: str + organization: str + end_user: str + tag: str + + def counters(self) -> Mapping[str, str]: + return { + "key": key_counter(self.key), + "team": f"spend:team:{self.team}", + "team_member": f"spend:team_member:{self.user}:{self.team}", + "user": f"spend:user:{self.user}", + "end_user": f"spend:end_user:{self.end_user}", + "tag": f"spend:tag:{self.tag}", + "org": f"spend:org:{self.organization}", + } + + +def _entities(gateway: Gateway, scenario: Scenario) -> Entities: + model: Final = scenario.model(**RATES) + models: Final = [model] + organization: Final = scenario.organization(max_budget=50.0) + team: Final = scenario.team(models=models, max_budget=20.0, organization_id=organization) + user: Final = scenario.user(max_budget=20.0) + added: Final = gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0}, + ) + assert added.status_code == 200, added.text + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + key: Final = scenario.key(team_id=team, user_id=user, models=models, max_budget=5.0, tags=[tag]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + return Entities(model, key, team, user, organization, end_user, tag) + + +def _base(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _chat_openai_sync(gateway: Gateway, entities: Entities, text: str) -> str: + with openai.OpenAI(base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60) as client: + response: Final = client.chat.completions.create( + model=entities.model, messages=[{"role": "user", "content": text}], user=entities.end_user + ) + assert response.usage is not None and response.usage.total_tokens == 40, response + return response.id + + +def _chat_openai_async_stream(gateway: Gateway, entities: Entities, text: str) -> str: + async def run() -> str: + async with openai.AsyncOpenAI( + base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60 + ) as client: + stream: Final = await client.chat.completions.create( + model=entities.model, + messages=[{"role": "user", "content": text}], + user=entities.end_user, + stream=True, + stream_options={"include_usage": True}, + ) + ids: Final = [chunk.id async for chunk in stream] + assert len(set(ids)) == 1, ids + return ids[0] + + return asyncio.run(run()) + + +def _chat_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str: + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": entities.model, + "messages": [{"role": "user", "content": text}], + "user": entities.end_user, + "stream": True, + "stream_options": {"include_usage": True}, + }, + headers={"Authorization": f"Bearer {entities.key}"}, + ) as response: + assert response.status_code == 200, response.read() + frames: Final = [line.removeprefix("data: ") for line in response.iter_lines() if line.startswith("data: ")] + ids: Final = {json.loads(frame)["id"] for frame in frames if frame != "[DONE]"} + assert len(ids) == 1, frames + return ids.pop() + + +def _messages_anthropic_sync(gateway: Gateway, entities: Entities, text: str) -> str: + with anthropic.Anthropic(base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60) as client: + message: Final = client.messages.create( + model=entities.model, + max_tokens=32, + messages=[{"role": "user", "content": text}], + metadata={"user_id": entities.end_user}, + ) + assert message.usage.output_tokens > 0, message + return message.id + + +def _messages_anthropic_async_stream(gateway: Gateway, entities: Entities, text: str) -> str: + async def run() -> str: + async with anthropic.AsyncAnthropic( + base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60 + ) as client: + async with client.messages.stream( + model=entities.model, + max_tokens=32, + messages=[{"role": "user", "content": text}], + metadata={"user_id": entities.end_user}, + ) as stream: + async for _ in stream.text_stream: + pass + final: Final = await stream.get_final_message() + assert final.usage.output_tokens > 0, final + return final.id + + return asyncio.run(run()) + + +def _messages_httpx(gateway: Gateway, entities: Entities, text: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": entities.model, + "max_tokens": 32, + "messages": [{"role": "user", "content": text}], + "metadata": {"user_id": entities.end_user}, + }, + key=entities.key, + headers=ANTHROPIC_HEADERS, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + +def _responses_httpx(gateway: Gateway, entities: Entities, text: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": entities.model, "input": text, "user": entities.end_user}, + key=entities.key, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + +def _responses_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str: + with gateway.client.stream( + "POST", + "/v1/responses", + json={"model": entities.model, "input": text, "user": entities.end_user, "stream": True}, + headers={"Authorization": f"Bearer {entities.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = [line for line in response.iter_lines() if line.startswith("data: ")] + events: Final = [json.loads(line.removeprefix("data: ")) for line in lines if line != "data: [DONE]"] + completed: Final = [event for event in events if event.get("type") == "response.completed"] + assert len(completed) == 1, [event.get("type") for event in events] + return str(completed[0]["response"]["id"]) + + +@dataclass(frozen=True, slots=True) +class Caller: + call: Callable[[Gateway, Entities, str], str] + upstream_path: str + exact_cost: float | None + + +CALLERS: Final[Mapping[str, Caller]] = { + "chat-openai-sync": Caller(_chat_openai_sync, "/v1/chat/completions", ONE_CALL_COST), + "chat-openai-async-stream": Caller(_chat_openai_async_stream, "/v1/chat/completions", ONE_CALL_COST), + "chat-httpx-stream": Caller(_chat_httpx_stream, "/v1/chat/completions", ONE_CALL_COST), + "messages-anthropic-sync": Caller(_messages_anthropic_sync, "/v1/responses", ONE_CALL_COST), + "messages-anthropic-async-stream": Caller(_messages_anthropic_async_stream, "/v1/responses", ONE_CALL_COST), + "messages-httpx": Caller(_messages_httpx, "/v1/responses", ONE_CALL_COST), + "responses-httpx": Caller(_responses_httpx, "/v1/responses", ONE_CALL_COST), + "responses-httpx-stream": Caller(_responses_httpx_stream, "/v1/responses", ONE_CALL_COST), +} + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (sha256(key.encode()).hexdigest(),), + ) + + +def _one_spend_row(key: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) + return rows[0] + + +def _observed(cache: Redis, counters: Mapping[str, str]) -> Mapping[str, tuple[float | None, int]]: + return {name: (counter_value(cache, counter), counter_ttl(cache, counter)) for name, counter in counters.items()} + + +def _assert_every_counter_outlives( + cache: Redis, counters: Mapping[str, str], spend: float +) -> Mapping[str, tuple[float | None, int]]: + observed: Final = eventually( + lambda: _observed(cache, counters), + lambda states: all( + value is not None and value == pytest.approx(spend, abs=1e-9) for value, _ in states.values() + ), + seconds=20, + ) + for name, (_, ttl) in observed.items(): + assert outlives(ttl), (name, observed) + return observed + + +@pytest.mark.parametrize("caller", sorted(CALLERS)) +def test_every_counter_scope_outlives_the_batch_write(gateway: Gateway, caller: str) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + rig_cache() as cache, + ): + entities: Final = _entities(gateway, scenario) + text: Final = f"counter ttl {caller} {uuid.uuid4().hex}" + upstream.get("/__observations").raise_for_status() + response_id: Final = CALLERS[caller].call(gateway, entities, text) + received: Final = upstream.get("/__observations").json()["requests"] + paths: Final = [request["path"] for request in received if text in json.dumps(request["body"])] + assert len(paths) == 1 and paths[0].endswith(CALLERS[caller].upstream_path), received + row: Final = _one_spend_row(entities.key) + spend: Final = float(str(row["spend"])) + assert spend == pytest.approx(CALLERS[caller].exact_cost or spend) and spend > 0, row + assert response_id, row + _assert_every_counter_outlives(cache, entities.counters(), spend) + + +def test_a_response_cache_hit_keeps_the_counter_and_its_ttl(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + rig_cache() as cache, + ): + entities: Final = _entities(gateway, scenario) + body: Final = { + "model": entities.model, + "messages": [{"role": "user", "content": f"cache twin {uuid.uuid4().hex}"}], + "user": entities.end_user, + } + first: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key) + assert first.status_code == 200, first.text + row: Final = _one_spend_row(entities.key) + spend: Final = float(str(row["spend"])) + before: Final = _assert_every_counter_outlives(cache, entities.counters(), spend) + upstream.get("/__observations").raise_for_status() + second: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key) + assert second.status_code == 200, second.text + assert second.json()["id"] == first.json()["id"], (first.text, second.text) + assert upstream.get("/__observations").json()["requests"] == [] + rows: Final = eventually(lambda: _spend_rows(entities.key), lambda values: len(values) == 2, seconds=70) + assert str(rows[1]["cache_hit"]) == "True" and float(str(rows[1]["spend"])) == 0.0, rows + after: Final = _assert_every_counter_outlives(cache, entities.counters(), spend) + assert {name: value for name, (value, _) in after.items()} == pytest.approx( + {name: value for name, (value, _) in before.items()}, abs=1e-9 + ), (before, after) + + +def _exhausted_key(gateway: Gateway, scenario: Scenario) -> tuple[str, str]: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + served: Final = gateway.chat(model, key=key, text=f"exhaust {uuid.uuid4().hex}") + assert served["usage"]["total_tokens"] == 40, served + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + return model, key + + +def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _assert_denied(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert response.json()["error"]["type"] == "budget_exceeded", response.text + + +def test_a_stale_low_counter_is_repaired_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key( + models=[model], budget_limits=[{"max_budget": ONE_CALL_COST, "budget_duration": "1d"}] + ) + counter: Final = f"{key_counter(key)}:window:1d" + assert gateway.chat(model, key=key, text=f"exhaust window {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + assert counter_value(cache, counter) == pytest.approx(ONE_CALL_COST) + + served: Final[list[int]] = [] + + def stale_then_request() -> int: + assert cache.set(counter, "0.01") + code: Final = _chat(gateway, model, key).status_code + served.append(code) + return code + + assert eventually(stale_then_request, lambda code: code == 422, seconds=30) == 422 + assert set(served) == {200, 422}, served + repaired: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + floor: Final = repaired[0] + assert floor is not None and ONE_CALL_COST - 1e-9 <= floor <= served.count(200) * ONE_CALL_COST + 1e-9, ( + repaired, + served, + ) + assert outlives(repaired[1]), repaired + + +def test_a_lost_counter_is_reseeded_from_the_db_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key = _exhausted_key(gateway, scenario) + counter: Final = key_counter(key) + assert cache.delete(counter) == 1 + _assert_denied(_chat(gateway, model, key)) + reseeded: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and state[0] >= ONE_CALL_COST and outlives(state[1]), + seconds=15, + ) + assert reseeded[0] == pytest.approx(ONE_CALL_COST), reseeded + + +def test_key_update_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert gateway.chat(model, key=key, text=f"before update {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + updated: Final = gateway.request("POST", "/key/update", {"key": key, "spend": 1.25}) + assert updated.status_code == 200, updated.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (1.25, True), updated.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value == pytest.approx(1.25 + ONE_CALL_COST), + seconds=15, + ) + assert outlives(counter_ttl(cache, counter)) + + +def test_key_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key = _exhausted_key(gateway, scenario) + counter: Final = key_counter(key) + _assert_denied(_chat(gateway, model, key)) + reset: Final = gateway.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert reset.json()["key_hash"] == sha256(key.encode()).hexdigest(), reset.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + + +def test_team_member_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + team: Final = scenario.team(models=[model]) + user: Final = scenario.user() + added: Final = gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": ONE_CALL_COST}, + ) + assert added.status_code == 200, added.text + key: Final = scenario.key(team_id=team, user_id=user, models=[model]) + counter: Final = f"spend:team_member:{user}:{team}" + assert gateway.chat(model, key=key, text=f"member {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', (team, user) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + _assert_denied(_chat(gateway, model, key)) + reset: Final = gateway.request("POST", f"/team/{team}/member/{user}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert reset.json()["previous_spend"] == pytest.approx(ONE_CALL_COST), reset.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + + +def test_malformed_spend_writes_are_rejected_and_leave_the_counter_alone(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + bystander: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert gateway.chat(model, key=key, text=f"before bad writes {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + before: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert before[0] == pytest.approx(ONE_CALL_COST) and outlives(before[1]), before + rejected: Final = tuple( + gateway.request(*attempt) + for attempt in ( + ("POST", "/key/update", {"key": key, "spend": "a" * 5000}), + ("POST", "/key/update", {"key": key, "spend": [1.0]}), + ("POST", "/key/update", {"key": key, "spend": {"spend": 1.0}}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": "abc"}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": True}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": [0.0]}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": ONE_CALL_COST * 10}), + ) + ) + for response in rejected: + assert response.status_code in (400, 422), response.text + assert "error" in response.json() or "detail" in response.json(), response.text + assert (counter_value(cache, counter), counter_ttl(cache, counter) <= before[1]) == (before[0], True) + assert outlives(counter_ttl(cache, counter)) + served: Final = _chat(gateway, model, bystander) + assert served.status_code == 200, served.text + + +def _landed_chat(gateway: Gateway, model: str, key: str, index: int) -> str: + served: Final = gateway.chat(model, key=key, text=f"repeat {index} {uuid.uuid4().hex}") + eventually(lambda: _spend_rows(key), lambda rows: len(rows) == index + 1, seconds=70) + return str(served["id"]) + + +def test_repeated_requests_log_once_each_and_the_counter_matches_their_sum(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + ids: Final = tuple(_landed_chat(gateway, model, key, index) for index in range(3)) + assert len(set(ids)) == 3, ids + rows: Final = _spend_rows(key) + assert len(rows) == 3, rows + assert sorted(str(row["request_id"]) for row in rows) == sorted(ids), rows + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx(3 * ONE_CALL_COST), rows + eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value == pytest.approx(total), + seconds=15, + ) + assert outlives(counter_ttl(cache, counter)) diff --git a/tests/integration/spend/test_spend_counter_ttl_chaos.py b/tests/integration/spend/test_spend_counter_ttl_chaos.py new file mode 100644 index 00000000000..b22f268b595 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_ttl_chaos.py @@ -0,0 +1,296 @@ +import os +import signal +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, group_members, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.spend_counters import ( + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import JsonValue +from redis import Redis + +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} +BURST: Final = 24 +UNBOUNDED_BUDGET: Final = 100000.0 +RESERVATION_ESTIMATE_CEILING: Final = 40.0 +SLOW_UPSTREAM_SECONDS: Final = 3.0 +LEAK_KEY_BUDGET: Final = 20.0 +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +RESPAWN_SECONDS: Final = graceful_stop_seconds() + 60 + + +def _chat(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"burst chat {index} {uuid.uuid4().hex}"}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + + +def _messages(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 32, + "messages": [{"role": "user", "content": f"burst messages {index} {uuid.uuid4().hex}"}], + "stream": stream, + }, + key=key, + headers=ANTHROPIC_HEADERS, + ) + + +def _responses(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": f"burst responses {index} {uuid.uuid4().hex}", "stream": stream}, + key=key, + ) + + +SURFACES: Final[tuple[Callable[[Gateway, str, str, int, bool], httpx.Response], ...]] = ( + _chat, + _messages, + _responses, +) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, status FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +def _successes(key: str, count: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"], + lambda values: len(values) >= count, + seconds=90, + ) + assert len(rows) == count, rows + assert len({str(row["request_id"]) for row in rows}) == count, rows + return rows + + +def _counter_matches(cache: Redis, counter: str, total: float) -> tuple[float | None, int]: + state: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda values: values[0] is not None and values[0] == pytest.approx(total, abs=1e-9), + seconds=20, + ) + assert outlives(state[1]), state + return state + + +def test_a_burst_with_upstream_failures_lands_each_success_once_with_a_long_ttl(gateway: Gateway) -> None: + with ( + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + gateway.scenario() as scenario, + rig_cache() as cache, + ): + provider_model: Final = f"burst-{uuid.uuid4().hex}" + model: Final = scenario.model(model=f"openai/{provider_model}", **RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + statuses: Final = [500] + [200, 200, 200] * (BURST // 4 + 2) + + def remove_script() -> None: + response: Final = upstream.delete(f"/__scripts/{provider_model}") + assert response.status_code in (200, 404), response.text + + scenario.cleanups.callback(remove_script) + configured: Final = upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}) + assert configured.status_code == 200, configured.text + + def attempt(index: int) -> httpx.Response: + stream: Final = index % 2 == 1 + return SURFACES[index % len(SURFACES)](gateway, model, key, index, stream) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + responses: Final = tuple(pool.map(attempt, range(BURST))) + codes: Final = [response.status_code for response in responses] + assert all(code in (200, 500) for code in codes), codes + succeeded: Final = codes.count(200) + assert 0 < succeeded < BURST, codes + rows: Final = _successes(key, succeeded) + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx(succeeded * 0.06), rows + _counter_matches(cache, key_counter(key), total) + + +def _workers(owned: OwnedProxy) -> tuple[int, ...]: + return tuple( + member.pid + for member in group_members(owned.process.pid) + if member.pid != owned.process.pid + and member.is_running() + and member.status() != psutil.STATUS_ZOMBIE + and any("spawn_main" in part for part in member.cmdline()) + ) + + +def _attempt(gateway: Gateway, model: str, key: str, index: int) -> int: + try: + return _chat(gateway, model, key, index, False).status_code + except httpx.HTTPError: + return 0 + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_killed_worker_does_not_shorten_or_lose_the_counter(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + rig_cache() as cache, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + before: Final = _workers(owned) + assert len(before) == 2, before + with ThreadPoolExecutor(max_workers=12) as pool: + first: Final = pool.submit(_attempt, owned.gateway, model, key, 0) + assert first.result(timeout=60) == 200 + in_flight: Final = [pool.submit(_attempt, owned.gateway, model, key, index) for index in range(1, 12)] + os.kill(before[0], signal.SIGKILL) + survivor: Final = _attempt(owned.gateway, model, key, 99) + codes: Final = [future.result(timeout=120) for future in in_flight] + assert survivor == 200, survivor + assert all(code in (0, 200) for code in codes), codes + replaced: Final = eventually( + lambda: _workers(owned), + lambda pids: len(pids) == 2 and before[0] not in pids and before[1] in pids, + seconds=RESPAWN_SECONDS, + ) + assert replaced + after: Final = [_attempt(owned.gateway, model, key, index) for index in range(12, 18)] + assert after == [200] * 6, after + answered: Final = 1 + 1 + codes.count(200) + 6 + landed: Final = eventually( + lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"], + lambda rows: len(rows) >= 1 + 1 + 6, + seconds=90, + ) + assert len({str(row["request_id"]) for row in landed}) == len(landed) <= answered, landed + ceiling: Final = answered * 0.06 + len(in_flight) * RESERVATION_ESTIMATE_CEILING + state: Final = eventually( + lambda: (counter_value(cache, key_counter(key)), counter_ttl(cache, key_counter(key))), + lambda values: values[0] is not None and 8 * 0.06 - 1e-9 <= values[0] <= ceiling, + seconds=20, + ) + assert outlives(state[1]), state + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_reservation_leaked_by_a_killed_worker_stays_for_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + rig_cache() as cache, + ): + slow: Final = scenario.model(model="openai/slow-endpoint", **RATES) + key: Final = scenario.key(models=[slow], max_budget=LEAK_KEY_BUDGET) + counter: Final = key_counter(key) + before: Final = _workers(owned) + assert len(before) == 2, before + with ThreadPoolExecutor(max_workers=1) as pool: + in_flight: Final = pool.submit(_attempt, owned.gateway, slow, key, 0) + reserved: Final = eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value >= LEAK_KEY_BUDGET - 1e-9, + seconds=SLOW_UPSTREAM_SECONDS, + ) + for pid in before: + os.kill(pid, signal.SIGKILL) + interrupted: Final = in_flight.result(timeout=60) + assert interrupted != 200, interrupted + leaked: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert leaked[0] == pytest.approx(reserved), (reserved, leaked) + respawned: Final = eventually( + lambda: _workers(owned), + lambda pids: len(pids) == 2 and not set(pids) & set(before), + seconds=RESPAWN_SECONDS, + ) + assert respawned + blocked: Final = eventually( + lambda: _attempt(owned.gateway, slow, key, 1), lambda code: code == 422, seconds=RESPAWN_SECONDS + ) + assert blocked == 422, (blocked, leaked) + held: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert held[0] == pytest.approx(reserved), (reserved, held) + assert outlives(held[1]), held + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_redis_outage_mid_burst_keeps_serving_and_rebuilds_the_counter_long(gateway: Gateway, tmp_path: Path) -> None: + with owned_redis(tmp_path) as redis: + overrides: Final[Mapping[str, str]] = { + "REDIS_HOST": redis.host, + "REDIS_PORT": str(redis.port), + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", + } + with ( + owned_proxy_process(gateway, tmp_path, overrides) as owned, + owned.gateway.scenario() as scenario, + Redis(host=redis.host, port=redis.port) as cache, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + counter: Final = key_counter(key) + + def burst(start: int) -> list[int]: + with ThreadPoolExecutor(max_workers=6) as pool: + return list( + pool.map(lambda index: _attempt(owned.gateway, model, key, index), range(start, start + 6)) + ) + + assert burst(0) == [200] * 6 + _counter_matches(cache, counter, 6 * 0.06) + redis.stop() + during: Final = burst(6) + assert during == [200] * 6, during + redis.start() + assert cache.ping() + probes: Final[list[int]] = [] + + def probe_until_the_counter_is_back() -> float | None: + probes.append(_attempt(owned.gateway, model, key, 100 + len(probes))) + return counter_value(cache, counter) + + eventually(probe_until_the_counter_is_back, lambda value: value is not None, seconds=30) + assert set(probes) == {200}, probes + after: Final = burst(12) + assert after == [200] * 6, after + rows: Final = _successes(key, 18 + len(probes)) + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx((18 + len(probes)) * 0.06), rows + rebuilt: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and 7 * 0.06 - 1e-9 <= state[0] <= total + 1e-9, + seconds=20, + ) + assert outlives(rebuilt[1]), rebuilt diff --git a/tests/integration/spend/test_spend_counter_ttl_settings.py b/tests/integration/spend/test_spend_counter_ttl_settings.py new file mode 100644 index 00000000000..ec9d725bf5c --- /dev/null +++ b/tests/integration/spend/test_spend_counter_ttl_settings.py @@ -0,0 +1,228 @@ +import json +import os +import time +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import psycopg +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.spend_counters import counter_ttl, counter_value, key_counter, outlives, rig_cache +from pydantic import JsonValue + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +SLOW_UPSTREAM_SECONDS: Final = 3.0 +SHORT_REDIS_TTL: Final = 4 +LONG_REDIS_TTL: Final = 86400 +ENV_REDIS_TTL: Final = "7200" +LONG_BATCH_WRITE_SECONDS: Final = 900 +STRICT_KEY_BUDGET: Final = 50.0 +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +FAST_RESET_TICK: Final = MappingProxyType( + {"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "2"} +) + + +def _config( + directory: Path, + *, + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), + general_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + configured: Final = { + **shared, + "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings}, + "general_settings": {**object_value(shared["general_settings"]), **general_settings}, + } + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(configured, sort_keys=False)) + return path + + +def _chat(gateway: Gateway, model: str, key: str) -> str: + served: Final = gateway.chat(model, key=key, text=f"settings {uuid.uuid4().hex}") + assert served["usage"]["total_tokens"] == 40, served + return str(served["id"]) + + +def _counter_state(counter: str) -> tuple[float | None, int]: + with rig_cache() as cache: + return counter_value(cache, counter), counter_ttl(cache, counter) + + +def _settled(counter: str, spend: float, expected_ttl: int) -> tuple[float | None, int]: + return eventually( + lambda: _counter_state(counter), + lambda state: state[0] is not None and state[0] == pytest.approx(spend) and outlives(state[1], expected_ttl), + seconds=15, + ) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_null_default_redis_ttl_falls_back_to_the_hour_floor(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": None}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + assert _chat(owned, model, key) + assert _settled(key_counter(key), ONE_CALL_COST, 3600) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_short_default_redis_ttl_never_shortens_the_counter_and_the_lease_refresh_keeps_it_long( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": SHORT_REDIS_TTL}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(model="openai/slow-endpoint", **RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + with ThreadPoolExecutor(max_workers=1) as pool: + slow: Final = pool.submit(lambda: (_chat(owned, model, key), time.monotonic())[1]) + eventually(lambda: cache.exists(counter), lambda present: present == 1, seconds=SLOW_UPSTREAM_SECONDS) + assert cache.expire(counter, 1000) + refreshed_at: Final = eventually( + lambda: (time.monotonic(), counter_ttl(cache, counter)), + lambda state: outlives(state[1]), + seconds=SLOW_UPSTREAM_SECONDS + 10, + )[0] + finished_at: Final = slow.result(timeout=60) + assert refreshed_at < finished_at, (refreshed_at, finished_at) + assert _settled(counter, ONE_CALL_COST, 3600) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_long_default_redis_ttl_wins_on_every_write_path(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": LONG_REDIS_TTL}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + assert _settled(counter, ONE_CALL_COST, LONG_REDIS_TTL) + updated: Final = owned.request("POST", "/key/update", {"key": key, "spend": 2.5}) + assert updated.status_code == 200, updated.text + assert _settled(counter, 2.5, LONG_REDIS_TTL) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_default_redis_ttl_given_as_an_environment_string_is_read_as_seconds( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": "os.environ/SPEND_COUNTER_REDIS_TTL"}) + overrides: Final = {"SPEND_COUNTER_REDIS_TTL": ENV_REDIS_TTL} + with owned_proxy(gateway, tmp_path, overrides, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + assert _settled(counter, ONE_CALL_COST, int(ENV_REDIS_TTL)) + reset: Final = owned.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert _settled(counter, 0.0, int(ENV_REDIS_TTL)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_long_batch_write_interval_lifts_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, general_settings={"proxy_batch_write_at": LONG_BATCH_WRITE_SECONDS}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + assert _chat(owned, model, key) + assert _settled(key_counter(key), ONE_CALL_COST, 10 * LONG_BATCH_WRITE_SECONDS) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_strict_enforcement_repairs_a_stale_low_key_counter_with_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, general_settings={"fail_closed_budget_enforcement": True}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=STRICT_KEY_BUDGET) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + attempts: Final[list[str]] = [] + + def stale_then_request() -> tuple[float | None, int]: + assert cache.set(counter, "0.01") + attempts.append(_chat(owned, model, key)) + return eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and state[0] < STRICT_KEY_BUDGET / 2, + seconds=15, + ) + + repaired: Final = eventually( + stale_then_request, + lambda state: state[0] is not None and state[0] > 0.01 + ONE_CALL_COST + 1e-6, + seconds=30, + ) + lifted: Final = repaired[0] + assert ( + lifted is not None and 2 * ONE_CALL_COST - 1e-9 <= lifted <= (len(attempts) + 1) * ONE_CALL_COST + 1e-9 + ), ( + repaired, + attempts, + ) + assert outlives(repaired[1]), repaired + + +def _window_reset_at(key: str) -> datetime: + rows: Final = read_rows( + 'SELECT budget_limits FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + assert len(rows) == 1, rows + raw: Final = rows[0]["budget_limits"] + windows: Final = json.loads(raw) if isinstance(raw, str) else raw + assert isinstance(windows, list) and len(windows) == 1, raw + return datetime.fromisoformat(str(object_value(windows[0])["reset_at"]).replace("Z", "+00:00")) + + +def _make_window_due(key: str) -> None: + due: Final = (datetime.now(UTC) - timedelta(days=1)).replace(tzinfo=None).isoformat() + windows: Final = [{"max_budget": 1.0, "budget_duration": "1d", "reset_at": due}] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + 'UPDATE "LiteLLM_VerificationToken" SET budget_limits = %s::jsonb WHERE token = %s', + (json.dumps(windows), sha256(key.encode()).hexdigest()), + ) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_reset_budget_window_keeps_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, FAST_RESET_TICK) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], budget_limits=[{"max_budget": 1.0, "budget_duration": "1d"}]) + window_counter: Final = f"{key_counter(key)}:window:1d" + assert _chat(owned, model, key) + assert _settled(window_counter, ONE_CALL_COST, 3600) + _make_window_due(key) + rolled: Final = eventually( + lambda: (_window_reset_at(key), counter_value(cache, window_counter), counter_ttl(cache, window_counter)), + lambda state: state[0].replace(tzinfo=None) > datetime.now(UTC).replace(tzinfo=None), + seconds=45, + ) + assert rolled[1] is not None and rolled[1] < ONE_CALL_COST, rolled + assert outlives(rolled[2]), rolled diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 57c32ace509..296bf678284 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -402,6 +402,25 @@ async def test_async_scan_iter_namespaces_pattern(namespace, expected_match, mon assert captured["match"] == expected_match +@pytest.mark.asyncio +async def test_async_scan_iter_decodes_bytes_keys_and_stops_at_count(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + def scan_iter(match, count): + async def gen(): + for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"): + yield key + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3) + assert keys == ["spend:key:a", "spend:key:b", "spend:key:c"] + + @pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) def test_increment_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index c17aed49014..afb5dce44b7 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -4704,3 +4704,35 @@ async def test_reset_budget_for_litellm_team_members_called(): assert len(team_member_writes) == 1 assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"] assert team_member_writes[0]["data"] == {"spend": 0} + + +@pytest.mark.asyncio +async def test_reset_expired_window_writes_the_redis_counter_with_the_spend_counter_ttl(monkeypatch): + """The window reset is a counter write like every other, so it must carry the spend counter TTL, not the 60s cache default.""" + from datetime import datetime + + from litellm.proxy.common_utils.reset_budget_job import BudgetResetSettings, Litellm_EntityType + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_redis_ttl = lambda: 4321 + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + spend_counter_cache = MagicMock() + spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + window = {"budget_duration": "1d", "max_budget": 1.0, "reset_at": "2026-01-01T00:00:00Z"} + + reset = await ResetBudgetJob._reset_expired_window( + window=window, + counter_key="spend:key:abc:window:1d", + spend_counter_cache=spend_counter_cache, + now=datetime(2026, 1, 2), + reset_settings=BudgetResetSettings(), + prisma_client=AsyncMock(), + entity_type=Litellm_EntityType.KEY, + entity_id="abc", + ) + + assert reset is True + spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( + key="spend:key:abc:window:1d", value=0.0, ttl=4321 + ) + spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(key="spend:key:abc:window:1d", value=0.0) diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index a98f1344279..92b2bf44a9f 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1573,3 +1573,39 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys) assert all(cache.get_cache(key=cache_key) == {"retained": True} for cache_key in kept_cache_keys) assert jwt_table.rows == (kept_row,) + + +def _org_spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + from litellm.caching.in_memory_cache import InMemoryCache + + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache = InMemoryCache() + spend_counter_cache.in_memory_cache.set_cache(key="spend:org:org-1", value=4.2) + spend_counter_cache.redis_cache.async_delete_cache = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + return spend_counter_cache + + +@pytest.mark.asyncio +async def test_legacy_update_with_spend_drops_the_org_spend_counter(monkeypatch): + """A direct spend write must drop the cross-pod counter, or the org stays at the stale value until the counter expires.""" + spend_counter_cache = _org_spend_counter_cache(monkeypatch) + + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "spend": 0.0}, existing_budget_id="budget-1" + ) + + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:org:org-1") + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") is None + + +@pytest.mark.asyncio +async def test_legacy_update_without_spend_keeps_the_org_spend_counter(monkeypatch): + spend_counter_cache = _org_spend_counter_cache(monkeypatch) + + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "organization_alias": "renamed"}, existing_budget_id="budget-1" + ) + + spend_counter_cache.redis_cache.async_delete_cache.assert_not_awaited() + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 4.2 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 25180038e2a..69058542ffd 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8059,3 +8059,57 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [ ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') ] + + +def _spend_counter_cache_for_reset(monkeypatch: pytest.MonkeyPatch, *, scan_results: object) -> MagicMock: + from litellm.caching.in_memory_cache import InMemoryCache + + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + mock_prisma_client.db.litellm_teamtable.update_many = AsyncMock() + monkeypatch.setattr(ps, "prisma_client", mock_prisma_client) + spend_counter_cache: Final = MagicMock() + spend_counter_cache.in_memory_cache = InMemoryCache() + spend_counter_cache.in_memory_cache.set_cache(key="spend:key:aaa", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t1", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend:user:u1", value=0.5) + spend_counter_cache.redis_cache.async_scan_iter = AsyncMock(side_effect=scan_results) + spend_counter_cache.redis_cache.delete_cache_keys = AsyncMock() + monkeypatch.setattr(ps, "spend_counter_cache", spend_counter_cache) + return spend_counter_cache + + +@pytest.mark.asyncio +async def test_global_spend_reset_drops_the_key_and_team_spend_counters(monkeypatch: pytest.MonkeyPatch) -> None: + spend_counter_cache: Final = _spend_counter_cache_for_reset( + monkeypatch, scan_results=[["spend:key:aaa", "spend:key:bbb"], [], ["spend:team:t1"], []] + ) + + response: Final = await spend_management_endpoints.global_spend_reset() + + assert response["status"] == "success" + assert [c.kwargs["pattern"] for c in spend_counter_cache.redis_cache.async_scan_iter.await_args_list] == [ + "spend:key:", + "spend:key:", + "spend:team:", + "spend:team:", + ] + assert [c.args for c in spend_counter_cache.redis_cache.delete_cache_keys.await_args_list] == [ + (["spend:key:aaa", "spend:key:bbb"],), + (["spend:team:t1"],), + ] + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") is None + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:team:t1") is None + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:user:u1") == 0.5 + + +@pytest.mark.asyncio +async def test_global_spend_reset_still_resets_the_rows_when_redis_is_down(monkeypatch: pytest.MonkeyPatch) -> None: + spend_counter_cache: Final = _spend_counter_cache_for_reset(monkeypatch, scan_results=ConnectionError("redis down")) + + response: Final = await spend_management_endpoints.global_spend_reset() + + assert response["status"] == "success" + ps.prisma_client.db.litellm_teamtable.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) + spend_counter_cache.redis_cache.delete_cache_keys.assert_not_awaited() + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") is None From 9e07878643400ee949f450691da17caaf58d37f8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 16:11:34 -0700 Subject: [PATCH 06/11] fix(proxy): set the reset spend counters and their floors on every pod A global or organization spend reset deleted the Redis counters, but every pod keeps a spend_db_floor marker for five seconds after it last checked the DB spend, and the next budget check on such a pod repaired the reseeded zero counter back to the old total with the hour-long TTL. Follow the key reset precedent instead: SET each counter and its floor marker to the new value in Redis and in memory, then broadcast both over the auth cache invalidation channel so every pod picks them up. The global reset lists the key and team counters with one unbounded SCAN, writes them in pipelines of 1000, and broadcasts 500 keys per message (old pods still read the single-key wire format). Window counters stay untouched because the reset does not change their DB source. --- litellm/caching/redis_cache.py | 8 +- .../auth_cache_invalidation_pubsub.py | 98 +++++++++++++------ .../organization_endpoints.py | 6 +- litellm/proxy/proxy_server.py | 71 +++++++++++--- .../spend_management_endpoints.py | 6 +- tests/unit/caching/test_redis_cache.py | 19 ++++ .../test_auth_cache_invalidation_pubsub.py | 51 ++++++++++ .../test_organization_endpoints.py | 30 ++++-- .../test_spend_management_endpoints.py | 70 +++++++++---- 9 files changed, 280 insertions(+), 79 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 108cc29b509..f179318e751 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -498,11 +498,11 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) -async def _first_keys(scan: AsyncIterator[object], count: int) -> AsyncIterator[str]: +async def _first_keys(scan: AsyncIterator[object], count: int | None) -> AsyncIterator[str]: taken: Final = itertools.count(1) async for key in scan: yield key.decode() if isinstance(key, bytes) else str(key) - if next(taken) >= count: + if count is not None and next(taken) >= count: return @@ -1015,7 +1015,7 @@ class RedisCache(BaseCache): return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) @_redis_circuit_breaker_guard - async def async_scan_iter(self, pattern: str, count: int = 100) -> list[str]: + async def async_scan_iter(self, pattern: str, count: int = 100, *, unbounded: bool = False) -> list[str]: start_time: Final = time.time() try: _redis_client: Final = self.init_async_client() @@ -1029,7 +1029,7 @@ class RedisCache(BaseCache): keys: Final = [ key async for key in _first_keys( - _redis_client.scan_iter(match=namespaced_pattern + "*", count=count), count + _redis_client.scan_iter(match=namespaced_pattern + "*", count=count), None if unbounded else count ) ] diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 2d874f7b473..3727f269adc 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -1,9 +1,11 @@ import asyncio import json -from collections.abc import Sequence -from dataclasses import asdict, dataclass +from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Final +from pydantic import TypeAdapter, ValidationError + from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.config_sync_pubsub import ( # noqa: F401 # legacy module exports @@ -24,6 +26,9 @@ AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" _POLL_TIMEOUT_SECONDS: Final = 1.0 _MAX_PENDING_PUBLISHES: Final = 1024 _MAX_IN_FLIGHT_PUBLISHES: Final = 16 +_KEYS_PER_MESSAGE: Final = 500 +_MESSAGE_CACHE_KEYS: Final = TypeAdapter(tuple[str, ...]) +_MESSAGE_FIELDS: Final = TypeAdapter(Mapping[str, object]) _pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep background publishes alive _in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES) _BACKOFF_INITIAL_SECONDS: Final = 5.0 @@ -38,14 +43,21 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str: @dataclass(frozen=True, slots=True) class _CacheInvalidationMessage: - cache_key: str + cache_keys: tuple[str, ...] new_value: float | None = None ttl: float | None = None -def _cache_invalidation_message_json(cache_key: str, new_value: float | None = None, ttl: float | None = None) -> str: - message: Final = asdict(_CacheInvalidationMessage(cache_key=cache_key, new_value=new_value, ttl=ttl)) - return json.dumps({field: value for field, value in message.items() if value is not None}) +def _cache_invalidation_message_json( + cache_keys: Sequence[str], new_value: float | None = None, ttl: float | None = None +) -> str: + fields: Final = { + "cache_key": cache_keys[0] if len(cache_keys) == 1 else None, + "cache_keys": list(cache_keys) if len(cache_keys) > 1 else None, + "new_value": new_value, + "ttl": ttl, + } + return json.dumps({field: value for field, value in fields.items() if value is not None}) def _finite_number_or_none(value: object) -> float | None: @@ -54,22 +66,29 @@ def _finite_number_or_none(value: object) -> float | None: return float(value) +def _message_cache_keys(single: object, many: object) -> tuple[str, ...]: + if isinstance(single, str): + return (single,) + try: + return _MESSAGE_CACHE_KEYS.validate_python(many) + except ValidationError: + return () + + def _message_from_data(data: object) -> _CacheInvalidationMessage | None: if isinstance(data, bytes): data = data.decode("utf-8", errors="replace") # rebind-ok: normalizing the wire payload to str if not isinstance(data, str): return None try: - parsed: Final = json.loads(data) - except json.JSONDecodeError: + parsed: Final = _MESSAGE_FIELDS.validate_json(data) + except ValidationError: return None - if not isinstance(parsed, dict): - return None - cache_key: Final = parsed.get("cache_key") - if not isinstance(cache_key, str): + cache_keys: Final = _message_cache_keys(parsed.get("cache_key"), parsed.get("cache_keys")) + if len(cache_keys) == 0: return None return _CacheInvalidationMessage( - cache_key=cache_key, + cache_keys=cache_keys, new_value=_finite_number_or_none(parsed.get("new_value")), ttl=_finite_number_or_none(parsed.get("ttl")), ) @@ -84,6 +103,26 @@ async def _publish_to_redis(redis_cache: "RedisCache", cache_key: str, message: verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e) +def _queue_publish( + redis_cache: "RedisCache", cache_keys: tuple[str, ...], new_value: float | None, ttl: float | None +) -> None: + _pending_publishes.difference_update({task for task in _pending_publishes if task.done()}) + if len(_pending_publishes) >= _MAX_PENDING_PUBLISHES: + verbose_proxy_logger.warning( + "auth cache invalidation publish for %s dropped: %d publishes already waiting on redis; " + "other workers keep their cached copy until its TTL expires", + cache_keys[0], + len(_pending_publishes), + ) + return + task: Final = asyncio.create_task( + _publish_to_redis( + redis_cache, cache_keys[0], _cache_invalidation_message_json(cache_keys, new_value=new_value, ttl=ttl) + ) + ) + _pending_publishes.add(task) + + async def publish_auth_cache_invalidation( cache_key: str, new_value: float | None = None, ttl: float | None = None ) -> None: @@ -106,24 +145,17 @@ async def publish_auth_cache_invalidation( Redis connection at once; the rest wait in the task set, so a wedge cannot drain the shared connection pool. """ + await publish_auth_cache_invalidations(cache_keys=(cache_key,), new_value=new_value, ttl=ttl) + + +async def publish_auth_cache_invalidations( + cache_keys: Sequence[str], new_value: float | None = None, ttl: float | None = None +) -> None: redis_cache: Final = coordination_redis_cache() if redis_cache is None: return - _pending_publishes.difference_update({task for task in _pending_publishes if task.done()}) - if len(_pending_publishes) >= _MAX_PENDING_PUBLISHES: - verbose_proxy_logger.warning( - "auth cache invalidation publish for %s dropped: %d publishes already waiting on redis; " - "other workers keep their cached copy until its TTL expires", - cache_key, - len(_pending_publishes), - ) - return - task: Final = asyncio.create_task( - _publish_to_redis( - redis_cache, cache_key, _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl) - ) - ) - _pending_publishes.add(task) + for start in range(0, len(cache_keys), _KEYS_PER_MESSAGE): + _queue_publish(redis_cache, tuple(cache_keys[start : start + _KEYS_PER_MESSAGE]), new_value, ttl) await asyncio.sleep(0) @@ -217,11 +249,13 @@ class AuthCacheInvalidationSubscriber: return if parsed.new_value is not None: for additional_cache in self._additional_in_memory_caches: - additional_cache.set_cache(parsed.cache_key, parsed.new_value, ttl=parsed.ttl) + for cache_key in parsed.cache_keys: + additional_cache.set_cache(cache_key, parsed.new_value, ttl=parsed.ttl) return - self._user_api_key_cache.in_memory_cache_for(parsed.cache_key).delete_cache(parsed.cache_key) - for additional_cache in self._additional_in_memory_caches: - additional_cache.delete_cache(parsed.cache_key) + for cache_key in parsed.cache_keys: + self._user_api_key_cache.in_memory_cache_for(cache_key).delete_cache(cache_key) + for additional_cache in self._additional_in_memory_caches: + additional_cache.delete_cache(cache_key) @staticmethod async def _close_pubsub(pubsub: ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8928832626e..9ebefe95fa2 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -817,9 +817,11 @@ async def update_organization( ) if data.spend is not None: - from litellm.proxy.proxy_server import invalidate_spend_counter + from litellm.proxy.proxy_server import set_spend_counters_with_floor_and_broadcast - await invalidate_spend_counter(counter_key=f"spend:org:{data.organization_id}") + await set_spend_counters_with_floor_and_broadcast( + counter_keys=(f"spend:org:{data.organization_id}",), value=data.spend + ) return response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2e63584f60e..196e66bb1f2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -66,6 +66,7 @@ from litellm.constants import ( BASE_MCP_ROUTE, DAILY_TAG_SPEND_BATCH_MULTIPLIER, DEFAULT_MAX_RECURSE_DEPTH, + DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, @@ -418,6 +419,7 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l ) from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, + publish_auth_cache_invalidations, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body @@ -2916,6 +2918,7 @@ async def _drain_spend_event_producer_on_shutdown() -> None: # Bounds authoritative DB re-reads when enforcing a budget against a # stale-low spend counter: at most one DB read per counter per window. SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5 +SPEND_DB_FLOOR_MARKER_PREFIX: Final = "spend_db_floor:" _SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE: Final = 10 @@ -3954,25 +3957,69 @@ _IN_MEMORY_COUNTER_KEYS: Final = TypeAdapter(tuple[str, ...]) _SPEND_COUNTER_STORE_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, ValidationError) +_WINDOW_COUNTER_SEGMENT: Final = ":window:" -@with_service_target(SPEND_COUNTERS_TARGET) -async def invalidate_spend_counters_with_prefix(prefix: str) -> None: - for counter_key in _IN_MEMORY_COUNTER_KEYS.validate_python(tuple(spend_counter_cache.in_memory_cache.cache_dict)): - if counter_key.startswith(prefix): - spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) - redis_cache: Final = spend_counter_cache.redis_cache - if redis_cache is None: - return + +def _spend_counter_keys_in_memory(prefix: str) -> frozenset[str]: + floor_prefix: Final = f"{SPEND_DB_FLOOR_MARKER_PREFIX}{prefix}" + cached_keys: Final = _IN_MEMORY_COUNTER_KEYS.validate_python(tuple(spend_counter_cache.in_memory_cache.cache_dict)) + return frozenset( + key.removeprefix(SPEND_DB_FLOOR_MARKER_PREFIX) for key in cached_keys if key.startswith((prefix, floor_prefix)) + ) + + +async def _spend_counter_keys_in_redis(redis_cache: RedisCache, prefix: str) -> frozenset[str]: + namespace_prefix: Final = f"{redis_cache.namespace}:" if redis_cache.namespace else "" try: - while scanned := await redis_cache.async_scan_iter(pattern=prefix, count=500): - await redis_cache.delete_cache_keys(scanned) + scanned: Final = await redis_cache.async_scan_iter(pattern=prefix, count=500, unbounded=True) except _SPEND_COUNTER_STORE_FAILURES: verbose_proxy_logger.warning( - "Unable to delete the spend counters %s* after the spend reset; budget checks use the stale counters until " - "they expire", + "Unable to list the spend counters %s* after the spend reset; budget checks on other pods use the stale " + "counters until they expire", prefix, exc_info=True, ) + return frozenset() + return frozenset(key.removeprefix(namespace_prefix) for key in scanned) + + +async def _write_spend_counters_to_redis(redis_cache: RedisCache, keys: Sequence[str], value: float) -> None: + ttl: Final = spend_counter_redis_ttl() + for start in range(0, len(keys), DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE): + await redis_cache.async_set_cache_pipeline_with_ttls( + [(key, value, ttl) for key in keys[start : start + DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE]] + ) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def set_spend_counters_with_floor_and_broadcast(counter_keys: Collection[str], value: float) -> None: + """A spend reset SETs each counter and its spend_db_floor marker on every pod, never + deletes them: a pod that cached the pre-reset floor within SPEND_DB_FLOOR_CACHE_TTL_SECONDS + would otherwise repair the reseeded counter right back to the old total.""" + keys: Final = tuple(sorted(counter_keys)) + if len(keys) == 0: + return + floor_keys: Final = tuple(f"{SPEND_DB_FLOOR_MARKER_PREFIX}{key}" for key in keys) + for key, floor_key in zip(keys, floor_keys, strict=True): + spend_counter_cache.in_memory_cache.set_cache(key=key, value=value, ttl=60) + spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is not None: + await _write_spend_counters_to_redis(redis_cache, keys, value) + await publish_auth_cache_invalidations(keys, new_value=value, ttl=60) + await publish_auth_cache_invalidations(floor_keys, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def reset_spend_counters_with_prefix(prefix: str) -> None: + redis_cache: Final = spend_counter_cache.redis_cache + in_redis: Final[frozenset[str]] = ( + frozenset() if redis_cache is None else await _spend_counter_keys_in_redis(redis_cache, prefix) + ) + counters: Final = frozenset( + key for key in _spend_counter_keys_in_memory(prefix) | in_redis if _WINDOW_COUNTER_SEGMENT not in key + ) + await set_spend_counters_with_floor_and_broadcast(counters, value=0.0) async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 8384e5e9177..de287452fa2 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3684,7 +3684,7 @@ async def global_spend_reset(): 3. LiteLLM_TeamTable spend will be set = 0 """ - from litellm.proxy.proxy_server import invalidate_spend_counters_with_prefix, prisma_client + from litellm.proxy.proxy_server import prisma_client, reset_spend_counters_with_prefix if prisma_client is None: raise ProxyException( @@ -3696,8 +3696,8 @@ async def global_spend_reset(): await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) - await invalidate_spend_counters_with_prefix(prefix="spend:key:") - await invalidate_spend_counters_with_prefix(prefix="spend:team:") + await reset_spend_counters_with_prefix(prefix="spend:key:") + await reset_spend_counters_with_prefix(prefix="spend:team:") return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 296bf678284..0f3cfdda21e 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -421,6 +421,25 @@ async def test_async_scan_iter_decodes_bytes_keys_and_stops_at_count(monkeypatch assert keys == ["spend:key:a", "spend:key:b", "spend:key:c"] +@pytest.mark.asyncio +async def test_async_scan_iter_unbounded_returns_every_matching_key(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + def scan_iter(match, count): + async def gen(): + for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"): + yield key + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3, unbounded=True) + assert keys == ["spend:key:a", "spend:key:b", "spend:key:c", "spend:key:d"] + + @pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) def test_increment_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 34d9741bcc0..9d3f458e2f5 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -16,6 +16,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, evict_and_broadcast, publish_auth_cache_invalidation, + publish_auth_cache_invalidations, ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -226,6 +227,56 @@ async def test_subscriber_deletes_additional_in_memory_cache_entry_on_message() assert spend_counter_in_memory_cache.get_cache("spend:team_member:u-1:t-1") is None +@pytest.mark.asyncio +async def test_publish_many_sends_one_message_per_five_hundred_keys() -> None: + client = _RecordingRedisClient() + keys = tuple(f"spend:key:{index}" for index in range(1001)) + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_FakeRedisCache(client=client), + ): + await publish_auth_cache_invalidations(cache_keys=keys, new_value=0.0, ttl=5) + await asyncio.gather(*pubsub_module._pending_publishes) + + assert [channel for channel, _ in client.published] == [AUTH_CACHE_INVALIDATION_CHANNEL] * 3 + assert [json.loads(message) for _, message in client.published] == [ + {"cache_keys": list(keys[:500]), "new_value": 0.0, "ttl": 5}, + {"cache_keys": list(keys[500:1000]), "new_value": 0.0, "ttl": 5}, + {"cache_key": keys[1000], "new_value": 0.0, "ttl": 5}, + ] + + +@pytest.mark.asyncio +async def test_subscriber_sets_every_key_of_a_multi_key_message_in_additional_caches() -> None: + cache = UserApiKeyCache() + spend_counter_in_memory_cache = InMemoryCache() + spend_counter_in_memory_cache.set_cache("spend:key:a", 999.0) + spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:key:a", 999.0) + message = { + "type": "message", + "data": json.dumps({"cache_keys": ["spend:key:a", "spend:key:b"], "new_value": 0.0, "ttl": 60}).encode(), + } + + pubsub = _QueuePubSub(initial_messages=[message]) + subscriber = AuthCacheInvalidationSubscriber( + redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), + user_api_key_cache=cache, + additional_in_memory_caches=(spend_counter_in_memory_cache,), + ) + subscriber.start() + try: + for _ in range(200): + if spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0: + break + await asyncio.sleep(0.01) + finally: + await subscriber.stop() + + assert spend_counter_in_memory_cache.get_cache("spend:key:a") == 0.0 + assert spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0 + assert spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:key:a") == 999.0 + + @pytest.mark.asyncio async def test_subscriber_ignores_malformed_messages() -> None: cache = UserApiKeyCache() diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index 92b2bf44a9f..c4fa4ff3286 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -2,7 +2,7 @@ import asyncio import json from types import MappingProxyType, SimpleNamespace from typing import Final, Mapping, Optional, cast -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest from fastapi import HTTPException @@ -647,6 +647,7 @@ async def test_organization_member_add_budget_omission_and_null_leave_budget_uns from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add user = LiteLLM_UserTable(user_id="user-1", user_role="internal_user") + async def create_membership(data): return LiteLLM_OrganizationMembershipTable( user_id="user-1", @@ -1581,22 +1582,35 @@ def _org_spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> MagicMock: spend_counter_cache = MagicMock() spend_counter_cache.in_memory_cache = InMemoryCache() spend_counter_cache.in_memory_cache.set_cache(key="spend:org:org-1", value=4.2) - spend_counter_cache.redis_cache.async_delete_cache = AsyncMock() + spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:org:org-1", value=4.2) + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + monkeypatch.setattr("litellm.proxy.proxy_server.publish_auth_cache_invalidations", AsyncMock()) return spend_counter_cache @pytest.mark.asyncio -async def test_legacy_update_with_spend_drops_the_org_spend_counter(monkeypatch): - """A direct spend write must drop the cross-pod counter, or the org stays at the stale value until the counter expires.""" +async def test_legacy_update_with_spend_sets_the_org_counter_and_its_floor_on_every_pod(monkeypatch): + """A direct spend write must SET the cross-pod counter and its floor marker: a pod that cached the old floor + would otherwise repair a deleted-and-reseeded counter back up, and the org stays at the stale value until + the counter expires.""" + from litellm.proxy import proxy_server + spend_counter_cache = _org_spend_counter_cache(monkeypatch) await _run_legacy_update_organization( - monkeypatch, body={"organization_id": "org-1", "spend": 0.0}, existing_budget_id="budget-1" + monkeypatch, body={"organization_id": "org-1", "spend": 1.5}, existing_budget_id="budget-1" ) - spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:org:org-1") - assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") is None + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_awaited_once_with( + [("spend:org:org-1", 1.5, proxy_server.spend_counter_redis_ttl())] + ) + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 1.5 + assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:org:org-1") == 1.5 + assert proxy_server.publish_auth_cache_invalidations.await_args_list == [ + call(("spend:org:org-1",), new_value=1.5, ttl=60), + call(("spend_db_floor:spend:org:org-1",), new_value=1.5, ttl=proxy_server.SPEND_DB_FLOOR_CACHE_TTL_SECONDS), + ] @pytest.mark.asyncio @@ -1607,5 +1621,5 @@ async def test_legacy_update_without_spend_keeps_the_org_spend_counter(monkeypat monkeypatch, body={"organization_id": "org-1", "organization_alias": "renamed"}, existing_budget_id="budget-1" ) - spend_counter_cache.redis_cache.async_delete_cache.assert_not_awaited() + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_not_awaited() assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 4.2 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 69058542ffd..11ae740f987 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -7,7 +7,7 @@ import json import re import sqlite3 from datetime import timezone -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, call, patch import prisma import pytest @@ -8071,45 +8071,79 @@ def _spend_counter_cache_for_reset(monkeypatch: pytest.MonkeyPatch, *, scan_resu spend_counter_cache: Final = MagicMock() spend_counter_cache.in_memory_cache = InMemoryCache() spend_counter_cache.in_memory_cache.set_cache(key="spend:key:aaa", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:key:aaa", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:key:floored", value=0.9) + spend_counter_cache.in_memory_cache.set_cache(key="spend:key:aaa:window:1d", value=0.3) spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t1", value=0.5) spend_counter_cache.in_memory_cache.set_cache(key="spend:user:u1", value=0.5) + spend_counter_cache.redis_cache.namespace = "ns" spend_counter_cache.redis_cache.async_scan_iter = AsyncMock(side_effect=scan_results) - spend_counter_cache.redis_cache.delete_cache_keys = AsyncMock() + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock() monkeypatch.setattr(ps, "spend_counter_cache", spend_counter_cache) + monkeypatch.setattr(ps, "publish_auth_cache_invalidations", AsyncMock()) return spend_counter_cache +_RESET_KEY_COUNTERS: Final = ("spend:key:aaa", "spend:key:bbb", "spend:key:floored") + + @pytest.mark.asyncio -async def test_global_spend_reset_drops_the_key_and_team_spend_counters(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_global_spend_reset_sets_every_key_and_team_counter_and_its_floor_to_zero( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Deleting the counters is not enough: a pod that cached the pre-reset floor within the last 5s would repair + the reseeded counter back to the old total for the whole counter TTL, so the reset SETs counter and floor + everywhere and leaves budget-window counters, which the reset does not touch in the DB, alone.""" spend_counter_cache: Final = _spend_counter_cache_for_reset( - monkeypatch, scan_results=[["spend:key:aaa", "spend:key:bbb"], [], ["spend:team:t1"], []] + monkeypatch, + scan_results=[["ns:spend:key:aaa", "ns:spend:key:bbb", "ns:spend:key:bbb:window:1d"], ["ns:spend:team:t1"]], ) + ttl: Final = ps.spend_counter_redis_ttl() response: Final = await spend_management_endpoints.global_spend_reset() assert response["status"] == "success" - assert [c.kwargs["pattern"] for c in spend_counter_cache.redis_cache.async_scan_iter.await_args_list] == [ - "spend:key:", - "spend:key:", - "spend:team:", - "spend:team:", + assert [c.kwargs for c in spend_counter_cache.redis_cache.async_scan_iter.await_args_list] == [ + {"pattern": "spend:key:", "count": 500, "unbounded": True}, + {"pattern": "spend:team:", "count": 500, "unbounded": True}, ] - assert [c.args for c in spend_counter_cache.redis_cache.delete_cache_keys.await_args_list] == [ - (["spend:key:aaa", "spend:key:bbb"],), - (["spend:team:t1"],), + assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [ + ([(counter, 0.0, ttl) for counter in _RESET_KEY_COUNTERS],), + ([("spend:team:t1", 0.0, ttl)],), + ] + in_memory: Final = spend_counter_cache.in_memory_cache + assert [in_memory.get_cache(key=counter) for counter in (*_RESET_KEY_COUNTERS, "spend:team:t1")] == [0.0] * 4 + assert [ + in_memory.get_cache(key=f"spend_db_floor:{counter}") for counter in (*_RESET_KEY_COUNTERS, "spend:team:t1") + ] == [0.0] * 4 + assert in_memory.get_cache(key="spend:key:aaa:window:1d") == 0.3 + assert in_memory.get_cache(key="spend:user:u1") == 0.5 + assert ps.publish_auth_cache_invalidations.await_args_list == [ + call(_RESET_KEY_COUNTERS, new_value=0.0, ttl=60), + call( + tuple(f"spend_db_floor:{counter}" for counter in _RESET_KEY_COUNTERS), + new_value=0.0, + ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ), + call(("spend:team:t1",), new_value=0.0, ttl=60), + call(("spend_db_floor:spend:team:t1",), new_value=0.0, ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS), ] - assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") is None - assert spend_counter_cache.in_memory_cache.get_cache(key="spend:team:t1") is None - assert spend_counter_cache.in_memory_cache.get_cache(key="spend:user:u1") == 0.5 @pytest.mark.asyncio -async def test_global_spend_reset_still_resets_the_rows_when_redis_is_down(monkeypatch: pytest.MonkeyPatch) -> None: +async def test_global_spend_reset_still_resets_the_rows_and_this_pod_when_redis_is_down( + monkeypatch: pytest.MonkeyPatch, +) -> None: spend_counter_cache: Final = _spend_counter_cache_for_reset(monkeypatch, scan_results=ConnectionError("redis down")) + ttl: Final = ps.spend_counter_redis_ttl() response: Final = await spend_management_endpoints.global_spend_reset() assert response["status"] == "success" ps.prisma_client.db.litellm_teamtable.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) - spend_counter_cache.redis_cache.delete_cache_keys.assert_not_awaited() - assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") is None + assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [ + ([("spend:key:aaa", 0.0, ttl), ("spend:key:floored", 0.0, ttl)],), + ([("spend:team:t1", 0.0, ttl)],), + ] + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") == 0.0 + assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:key:floored") == 0.0 From 244cae37d457b519a88fd0e0fb0b9cca73a74de4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:29:12 -0700 Subject: [PATCH 07/11] fix(proxy): fence a spend reset against stale floor repairs and report a lost Redis write A pod that read the DB floor within the last 5s could repair the counter a reset had just zeroed back to the old total for the whole counter TTL, since the floor broadcast reaches it after the counter SET. The reset now SETs a hash-tagged reset fence next to each counter in the same pipeline, the repair's Lua leaves a fenced counter alone, and the floors are broadcast before the counters. The reset also answered 200 when the Redis pipeline was lost, leaving every other pod on the old counters for up to an hour. The pipeline write returns the swallowed error as a value, the reset helpers carry it (a failed SCAN too), and /global/spend/reset and /organization/update with spend answer 500 with the DB write kept and a retry instruction. A value-carrying broadcast now lands only in the spend counter cache: written into every additional cache, a reset of thousands of counters evicted the cached BYOK credentials on every pod. A reset told window counters apart by the ":window:" substring, so a team whose id merely contains it kept its primary counter. The reset now uses the reseed's check, a validated trailing duration, which is the same rule that keeps a window counter out of the DB lookup --- litellm/caching/redis_cache.py | 48 +- .../auth_cache_invalidation_pubsub.py | 15 +- litellm/proxy/db/spend_counter_reseed.py | 4 +- .../organization_endpoints.py | 8 +- litellm/proxy/proxy_server.py | 98 ++- .../spend_management_endpoints.py | 11 +- .../test_batch_redis_requests_prefetch.py | 87 +++ .../spend/test_spend_counter_reset_fence.py | 26 + .../spend/test_spend_counter_reset_paths.py | 611 ++++++++++++++++++ tests/unit/caching/test_redis_cache.py | 59 ++ ...st_auth_checks_object_access_and_lookup.py | 1 + .../test_auth_cache_invalidation_pubsub.py | 11 +- .../test_organization_endpoints.py | 35 +- .../proxy/proxy_server/test_spend_counters.py | 12 +- .../test_spend_management_endpoints.py | 104 ++- 15 files changed, 1043 insertions(+), 87 deletions(-) create mode 100644 tests/integration/caching/test_batch_redis_requests_prefetch.py create mode 100644 tests/integration/spend/test_spend_counter_reset_fence.py create mode 100644 tests/integration/spend/test_spend_counter_reset_paths.py diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index f179318e751..4dc9e52c120 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1332,10 +1332,15 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard - async def async_set_cache_pipeline_with_ttls(self, cache_list: Sequence[tuple[str, object, float | None]]) -> None: - """One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL.""" + async def async_set_cache_pipeline_with_ttls( + self, cache_list: Sequence[tuple[str, object, float | None]] + ) -> Exception | None: + """One round trip for writes whose TTLs differ; a ``None`` TTL falls back to the default TTL. + + Returns the Redis error when the writes were lost, so a caller that must not report + success on a lost write can tell. An open circuit breaker still raises.""" if len(cache_list) == 0: - return + return None commands: Final = tuple( (self.check_and_fix_namespace(key=cache_key), json.dumps(cache_value), self.get_ttl(ttl=ttl)) for cache_key, cache_value, ttl in cache_list @@ -1346,16 +1351,6 @@ class RedisCache(BaseCache): for cache_key, json_cache_value, ttl in commands: pipe.set(name=cache_key, value=json_cache_value, ex=None if ttl is None else timedelta(seconds=ttl)) await pipe.execute() - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.REDIS, - duration=time.time() - start_time, - call_type="async_set_cache_pipeline_with_ttls", - caller=_get_call_stack_info(), - start_time=start_time, - end_time=time.time(), - ) - ) except Exception as e: asyncio.create_task( self.service_logger_obj.async_service_failure_hook( @@ -1375,6 +1370,19 @@ class RedisCache(BaseCache): e, ) _record_swallowed_redis_failure(self._circuit_breaker, e) + return e + else: + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=time.time() - start_time, + call_type="async_set_cache_pipeline_with_ttls", + caller=_get_call_stack_info(), + start_time=start_time, + end_time=time.time(), + ) + ) + return None async def _set_cache_sadd_helper( self, @@ -1551,12 +1559,18 @@ class RedisCache(BaseCache): ) raise e + def reset_fence_key(self, key: str) -> str: + """Hash-tagged with the namespaced ``key`` so a cluster keeps the fence and the key on one slot.""" + return f"{{{self.check_and_fix_namespace(key=key)}}}:reset_fence" + @_redis_circuit_breaker_guard async def async_set_max( self, key: str, value: float, ttl: int | None = None, + *, + fence_key: str | None = None, ) -> float | None: """Atomically set ``key`` to ``value`` only when ``value`` is greater than the stored value (or the key is unset), refreshing the TTL. @@ -1566,11 +1580,17 @@ class RedisCache(BaseCache): concurrent increment that has already pushed the counter higher. The GET/compare/SET runs in a single Lua call, so it is also atomic across racing callers and pods. Returns the resulting value. + + While ``fence_key`` exists the write is held and the stored value returned + as is: a spend reset sets one (``reset_fence_key``) so a repair still + carrying the pre-reset total cannot raise the counter the reset zeroed. """ _redis_client: Final = self.init_async_client() _used_ttl: Final = self.get_ttl(ttl=ttl) key = self.check_and_fix_namespace(key=key) + keys: Final = (key,) if fence_key is None else (key, self.check_and_fix_namespace(key=fence_key)) lua: Final = ( + "if KEYS[2] ~= nil and redis.call('EXISTS', KEYS[2]) == 1 then return redis.call('GET', KEYS[1]) end " "local cur = redis.call('GET', KEYS[1]) " "if cur == false or tonumber(cur) < tonumber(ARGV[1]) then " "redis.call('SET', KEYS[1], ARGV[1]) " @@ -1580,7 +1600,7 @@ class RedisCache(BaseCache): ) result = cast( "str | bytes | int | float | None", - await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))), + await _redis_client.eval(lua, len(keys), *keys, str(value), str(int(_used_ttl or 0))), ) if result is None: return None diff --git a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py index 3727f269adc..b6350024a7a 100644 --- a/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py +++ b/litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py @@ -182,17 +182,23 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us class AuthCacheInvalidationSubscriber: - __slots__ = ("_additional_in_memory_caches", "_redis_cache", "_task", "_user_api_key_cache") + """``counter_cache`` is the one cache a value-carrying message (a spend reset) is written into; + a delete reaches every additional cache. Writing a reset's thousands of counters into every + additional cache would evict what the others hold, the BYOK credentials among them.""" + + __slots__ = ("_additional_in_memory_caches", "_counter_cache", "_redis_cache", "_task", "_user_api_key_cache") def __init__( self, redis_cache: "RedisCache", user_api_key_cache: "UserApiKeyCache", additional_in_memory_caches: Sequence["InMemoryCache"] = (), + counter_cache: "InMemoryCache | None" = None, ) -> None: self._redis_cache = redis_cache self._user_api_key_cache = user_api_key_cache self._additional_in_memory_caches = tuple(additional_in_memory_caches) + self._counter_cache = counter_cache self._task: asyncio.Task[None] | None = None def start(self) -> None: @@ -248,9 +254,10 @@ class AuthCacheInvalidationSubscriber: if parsed is None: return if parsed.new_value is not None: - for additional_cache in self._additional_in_memory_caches: - for cache_key in parsed.cache_keys: - additional_cache.set_cache(cache_key, parsed.new_value, ttl=parsed.ttl) + if self._counter_cache is None: + return + for cache_key in parsed.cache_keys: + self._counter_cache.set_cache(cache_key, parsed.new_value, ttl=parsed.ttl) return for cache_key in parsed.cache_keys: self._user_api_key_cache.in_memory_cache_for(cache_key).delete_cache(cache_key) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index ab7ec7d3b1d..57d84d473b0 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -137,7 +137,7 @@ class SpendCounterReseed: # Per-window key/team counters share prefixes with primary counters # but don't correspond to a DB row. Do not reject arbitrary entity IDs # or tag names that merely contain ":window:". - if SpendCounterReseed._is_key_or_team_window_counter(counter_key): + if SpendCounterReseed.is_key_or_team_window_counter(counter_key): return None try: row: Final = await bounded_db_lookup( @@ -199,7 +199,7 @@ class SpendCounterReseed: return float(row.spend or 0.0) @staticmethod - def _is_key_or_team_window_counter(counter_key: str) -> bool: + def is_key_or_team_window_counter(counter_key: str) -> bool: for prefix in ("spend:key:", "spend:team:"): if not counter_key.startswith(prefix): continue diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 9ebefe95fa2..21378a7a097 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -817,11 +817,15 @@ async def update_organization( ) if data.spend is not None: - from litellm.proxy.proxy_server import set_spend_counters_with_floor_and_broadcast + from litellm.proxy.proxy_server import ( + raise_if_spend_counters_unwritten, + set_spend_counters_with_floor_and_broadcast, + ) - await set_spend_counters_with_floor_and_broadcast( + lost: Final = await set_spend_counters_with_floor_and_broadcast( counter_keys=(f"spend:org:{data.organization_id}",), value=data.spend ) + raise_if_spend_counters_unwritten(lost, saved="The organization spend was saved") return response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 196e66bb1f2..6742db8279f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2918,6 +2918,7 @@ async def _drain_spend_event_producer_on_shutdown() -> None: # Bounds authoritative DB re-reads when enforcing a budget against a # stale-low spend counter: at most one DB read per counter per window. SPEND_DB_FLOOR_CACHE_TTL_SECONDS: Final = 5 +SPEND_RESET_FENCE_TTL_SECONDS: Final = 2 * SPEND_DB_FLOOR_CACHE_TTL_SECONDS SPEND_DB_FLOOR_MARKER_PREFIX: Final = "spend_db_floor:" _SPEND_COUNTER_TTL_BATCH_WRITE_MULTIPLE: Final = 10 @@ -3067,7 +3068,10 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None forget_spend_counter(counter_key) try: repaired: Final = await spend_counter_cache.redis_cache.async_set_max( - key=counter_key, value=db_spend, ttl=spend_counter_redis_ttl() + key=counter_key, + value=db_spend, + ttl=spend_counter_redis_ttl(), + fence_key=spend_counter_cache.redis_cache.reset_fence_key(counter_key), ) except Exception: verbose_proxy_logger.debug( @@ -3957,8 +3961,6 @@ _IN_MEMORY_COUNTER_KEYS: Final = TypeAdapter(tuple[str, ...]) _SPEND_COUNTER_STORE_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, ValidationError) -_WINDOW_COUNTER_SEGMENT: Final = ":window:" - def _spend_counter_keys_in_memory(prefix: str) -> frozenset[str]: floor_prefix: Final = f"{SPEND_DB_FLOOR_MARKER_PREFIX}{prefix}" @@ -3968,58 +3970,99 @@ def _spend_counter_keys_in_memory(prefix: str) -> frozenset[str]: ) -async def _spend_counter_keys_in_redis(redis_cache: RedisCache, prefix: str) -> frozenset[str]: +async def _spend_counter_keys_in_redis(redis_cache: RedisCache, prefix: str) -> frozenset[str] | Exception: namespace_prefix: Final = f"{redis_cache.namespace}:" if redis_cache.namespace else "" try: scanned: Final = await redis_cache.async_scan_iter(pattern=prefix, count=500, unbounded=True) - except _SPEND_COUNTER_STORE_FAILURES: - verbose_proxy_logger.warning( - "Unable to list the spend counters %s* after the spend reset; budget checks on other pods use the stale " - "counters until they expire", - prefix, - exc_info=True, - ) - return frozenset() + except _SPEND_COUNTER_STORE_FAILURES as e: + return e return frozenset(key.removeprefix(namespace_prefix) for key in scanned) -async def _write_spend_counters_to_redis(redis_cache: RedisCache, keys: Sequence[str], value: float) -> None: +_KEYS_PER_RESET_PIPELINE: Final = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE // 2 + + +async def _write_reset_pipeline( + redis_cache: RedisCache, batch: Sequence[str], value: float, ttl: int +) -> Exception | None: + """The fences go first, so a repair that lands between a fence and its counter already finds the fence.""" + fences: Final = tuple((redis_cache.reset_fence_key(key), 1, SPEND_RESET_FENCE_TTL_SECONDS) for key in batch) + try: + return await redis_cache.async_set_cache_pipeline_with_ttls((*fences, *((key, value, ttl) for key in batch))) + except RedisCircuitBreakerOpenError as e: + return e + + +async def _write_spend_counters_to_redis( + redis_cache: RedisCache, keys: Sequence[str], value: float +) -> Exception | None: ttl: Final = spend_counter_redis_ttl() - for start in range(0, len(keys), DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE): - await redis_cache.async_set_cache_pipeline_with_ttls( - [(key, value, ttl) for key in keys[start : start + DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE]] - ) + for start in range(0, len(keys), _KEYS_PER_RESET_PIPELINE): + if ( + lost := await _write_reset_pipeline(redis_cache, keys[start : start + _KEYS_PER_RESET_PIPELINE], value, ttl) + ) is not None: + return lost + return None @with_service_target(SPEND_COUNTERS_TARGET) -async def set_spend_counters_with_floor_and_broadcast(counter_keys: Collection[str], value: float) -> None: +async def set_spend_counters_with_floor_and_broadcast(counter_keys: Collection[str], value: float) -> Exception | None: """A spend reset SETs each counter and its spend_db_floor marker on every pod, never deletes them: a pod that cached the pre-reset floor within SPEND_DB_FLOOR_CACHE_TTL_SECONDS - would otherwise repair the reseeded counter right back to the old total.""" + would otherwise repair the reseeded counter right back to the old total. The floors are + broadcast before the counters so a subscriber never holds the new counter under the old + floor, and a reset fence in Redis holds any repair that still races the broadcast. + + Returns the Redis error when the counters could not be written there: this pod and the + broadcast still carry the new value, but the other pods fall back to the stale Redis + counters once their copies expire, so the caller must not report success.""" keys: Final = tuple(sorted(counter_keys)) if len(keys) == 0: - return + return None floor_keys: Final = tuple(f"{SPEND_DB_FLOOR_MARKER_PREFIX}{key}" for key in keys) for key, floor_key in zip(keys, floor_keys, strict=True): spend_counter_cache.in_memory_cache.set_cache(key=key, value=value, ttl=60) spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) redis_cache: Final = spend_counter_cache.redis_cache - if redis_cache is not None: - await _write_spend_counters_to_redis(redis_cache, keys, value) - await publish_auth_cache_invalidations(keys, new_value=value, ttl=60) + lost: Final = None if redis_cache is None else await _write_spend_counters_to_redis(redis_cache, keys, value) await publish_auth_cache_invalidations(floor_keys, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + await publish_auth_cache_invalidations(keys, new_value=value, ttl=60) + return lost @with_service_target(SPEND_COUNTERS_TARGET) -async def reset_spend_counters_with_prefix(prefix: str) -> None: +async def reset_spend_counters_with_prefix(prefix: str) -> Exception | None: + """Resets every counter under ``prefix`` this pod or Redis knows; returns the Redis error + when the counters there could not be listed or written, with this pod already reset.""" redis_cache: Final = spend_counter_cache.redis_cache - in_redis: Final[frozenset[str]] = ( + in_redis: Final[frozenset[str] | Exception] = ( frozenset() if redis_cache is None else await _spend_counter_keys_in_redis(redis_cache, prefix) ) + listed: Final[frozenset[str]] = frozenset() if isinstance(in_redis, Exception) else in_redis counters: Final = frozenset( - key for key in _spend_counter_keys_in_memory(prefix) | in_redis if _WINDOW_COUNTER_SEGMENT not in key + key + for key in _spend_counter_keys_in_memory(prefix) | listed + if not SpendCounterReseed.is_key_or_team_window_counter(key) + ) + lost: Final = await set_spend_counters_with_floor_and_broadcast(counters, value=0.0) + return in_redis if isinstance(in_redis, Exception) else lost + + +def raise_if_spend_counters_unwritten(lost: Exception | None, saved: str) -> None: + """Maps a lost Redis counter write to the public error: the database write stands, so the + caller retries the same request once Redis accepts writes rather than undoing anything.""" + if lost is None: + return + raise ProxyException( + message=( + f"{saved}, but the cached spend counters could not be reset in Redis ({lost}). Budget checks on other " + "instances keep the previous spend until those counters expire, so retry this request once Redis " + "accepts writes." + ), + type=ProxyErrorTypes.internal_server_error, + param="None", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) - await set_spend_counters_with_floor_and_broadcast(counters, value=0.0) async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: @@ -8543,6 +8586,7 @@ class ProxyConfig: redis_cache=redis_cache, user_api_key_cache=user_api_key_cache, additional_in_memory_caches=(spend_counter_cache.in_memory_cache, byok_credential_cache), + counter_cache=spend_counter_cache.in_memory_cache, ) self.auth_cache_invalidation_subscriber = subscriber subscriber.start() diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index de287452fa2..392d574f82b 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3684,7 +3684,11 @@ async def global_spend_reset(): 3. LiteLLM_TeamTable spend will be set = 0 """ - from litellm.proxy.proxy_server import prisma_client, reset_spend_counters_with_prefix + from litellm.proxy.proxy_server import ( + prisma_client, + raise_if_spend_counters_unwritten, + reset_spend_counters_with_prefix, + ) if prisma_client is None: raise ProxyException( @@ -3696,8 +3700,9 @@ async def global_spend_reset(): await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) - await reset_spend_counters_with_prefix(prefix="spend:key:") - await reset_spend_counters_with_prefix(prefix="spend:team:") + keys_lost: Final = await reset_spend_counters_with_prefix(prefix="spend:key:") + teams_lost: Final = await reset_spend_counters_with_prefix(prefix="spend:team:") + raise_if_spend_counters_unwritten(keys_lost if keys_lost is not None else teams_lost, saved="Spend was reset") return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/tests/integration/caching/test_batch_redis_requests_prefetch.py b/tests/integration/caching/test_batch_redis_requests_prefetch.py new file mode 100644 index 00000000000..e2d7fa4e0e4 --- /dev/null +++ b/tests/integration/caching/test_batch_redis_requests_prefetch.py @@ -0,0 +1,87 @@ +"""The ``batch_redis_requests`` callback prefetches a key's response cache entries +into memory with one Redis scan, so the next identical request is served without a +Redis round trip for the entry.""" + +import os +import threading +import uuid +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import OwnedRedis, owned_redis +from redis import Redis +from redis.exceptions import TimeoutError as RedisTimeoutError + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) + + +def _config_with_batch_redis_requests(directory: Path) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + litellm_settings: Final = {**object_value(shared["litellm_settings"]), "callbacks": ["batch_redis_requests"]} + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False)) + return path + + +@contextmanager +def _monitored(redis: OwnedRedis) -> Iterator[list[str]]: + commands: Final[list[str]] = [] + stopped: Final = threading.Event() + with Redis(host=redis.host, port=redis.port, socket_timeout=0.3) as client, client.monitor() as monitor: + + def pump() -> None: + while not stopped.is_set(): + try: + commands.append(str(monitor.next_command()["command"])) + except RedisTimeoutError: + continue + + thread: Final = threading.Thread(target=pump, daemon=True) + thread.start() + try: + yield commands + finally: + stopped.set() + thread.join(timeout=5) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_the_prefetched_response_cache_entry_is_served_from_memory_on_the_next_request( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_redis(tmp_path) as redis: + overrides: Final = {"REDIS_HOST": redis.host, "REDIS_PORT": str(redis.port)} + with ( + owned_proxy_process( + gateway, tmp_path, overrides, config=_config_with_batch_redis_requests(tmp_path) + ) as owned, + owned.gateway.scenario() as scenario, + Redis(host=redis.host, port=redis.port) as cache, + httpx.Client(base_url=gateway.upstream_url, timeout=30) as upstream, + ): + model: Final = scenario.model() + key: Final = scenario.key(models=[model]) + body: Final = {"model": model, "messages": [{"role": "user", "content": f"prefetch {uuid.uuid4().hex}"}]} + namespace: Final = f"litellm:{key}:acompletion" + first: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert first.status_code == 200, first.text + entries: Final = eventually(lambda: cache.keys(f"{namespace}:*"), lambda keys: len(keys) == 1, seconds=10) + upstream.get("/__observations").raise_for_status() + with _monitored(redis) as commands: + second: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert second.status_code == 200, second.text + assert second.json()["id"] == first.json()["id"], (first.text, second.text) + assert upstream.get("/__observations").json()["requests"] == [] + fetched: Final = [command for command in commands if command.startswith(f"GET {namespace}:")] + assert fetched == [], (entries, commands) + third: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) + assert third.status_code == 200 and third.json()["id"] == first.json()["id"], third.text + assert upstream.get("/__observations").json()["requests"] == [] diff --git a/tests/integration/spend/test_spend_counter_reset_fence.py b/tests/integration/spend/test_spend_counter_reset_fence.py new file mode 100644 index 00000000000..4fb321d48b8 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_reset_fence.py @@ -0,0 +1,26 @@ +"""The reset fence: a spend reset SETs it in the same pipeline as the zeroed counter, and the stale-floor +repair's Lua script leaves a fenced counter alone, so a pod still holding the pre-reset floor cannot raise +the counter the reset zeroed for the whole counter TTL.""" + +import os +import uuid +from typing import Final + +from integration._support.spend_counters import rig_cache + +from litellm.caching.redis_cache import RedisCache + + +async def test_set_max_holds_while_the_reset_fence_exists() -> None: + cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace="fence") + counter: Final = f"spend:key:{uuid.uuid4().hex}" + fence: Final = cache.reset_fence_key(counter) + assert await cache.async_set_cache_pipeline_with_ttls(((fence, 1, 10), (counter, 0.0, 3600))) is None + + assert await cache.async_set_max(key=counter, value=12.0, ttl=3600, fence_key=fence) == 0.0 + + with rig_cache() as raw: + assert raw.delete(cache.check_and_fix_namespace(key=fence)) == 1 + assert await cache.async_set_max(key=counter, value=12.0, ttl=3600, fence_key=fence) == 12.0 + with rig_cache() as raw: + raw.delete(cache.check_and_fix_namespace(key=counter)) diff --git a/tests/integration/spend/test_spend_counter_reset_paths.py b/tests/integration/spend/test_spend_counter_reset_paths.py new file mode 100644 index 00000000000..b2c6b5c1282 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_reset_paths.py @@ -0,0 +1,611 @@ +"""The two spend reset paths that touch the Redis spend counters. + +``POST /global/spend/reset`` zeroes every key and team spend in the database. The +counters that enforce those budgets must follow it on every pod, or a key the +operator just reset stays blocked until the counter's TTL runs out. ``PATCH +/organization/update`` with ``spend`` sets one organization's counter to the given +value. Both write a reset fence next to each zeroed counter so a pod still holding +the pre-reset floor cannot repair the counter back up, and both answer 500 when the +counters could not be written to Redis, so the operator retries once it is back. +Every cell drives a real proxy: the cells that reset every key or stop Redis boot +their own proxy on a scratch database and an owned Redis. +""" + +import asyncio +import os +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import OwnedRedis, owned_redis +from integration._support.spend_counters import ( + SPEND_COUNTER_TTL_SECONDS, + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import BaseModel, ConfigDict +from redis import Redis + +from litellm.caching.redis_cache import RedisCache + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" +REDIS_SETTINGS: Final = frozenset({"cache", "cache_params", "enable_redis_auth_cache"}) +SEEDED_COUNTERS: Final = 1200 +SETTLE_SECONDS: Final = 20 +BROADCAST_SECONDS: Final = 10 +RACE_WINDOW_SECONDS: Final = 2 +DATABASE_SECONDS: Final = 70 +FENCE_SECONDS: Final = 10 +RECOVERY_SECONDS: Final = 30 +LOST_COUNTERS: Final = "but the cached spend counters could not be reset in Redis" +RETRY_HINT: Final = "retry this request once Redis accepts writes." +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) + + +class _BroadcastMessage(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + cache_key: str | None = None + cache_keys: tuple[str, ...] = () + new_value: float | None = None + ttl: float | None = None + + def names(self) -> tuple[str, ...]: + return (self.cache_key,) if self.cache_key is not None else self.cache_keys + + +@dataclass(frozen=True, slots=True) +class Broadcast: + key: str + value: float | None + ttl: float | None + + +class _Broadcasts: + def __init__(self, cache: Redis) -> None: + self._pubsub: Final = cache.pubsub() + self._pubsub.subscribe(INVALIDATION_CHANNEL) + confirmation: Final = self._pubsub.get_message(timeout=5) + assert confirmation is not None and confirmation["type"] == "subscribe", confirmation + self._condition: Final = threading.Condition() + self._received: Final[list[Broadcast]] = [] + self._stopped: Final = threading.Event() + self._thread: Final = threading.Thread(target=self._pump, daemon=True) + self._thread.start() + + def _pump(self) -> None: + while not self._stopped.is_set(): + self._record(self._pubsub.get_message(timeout=0.2)) + + def _record(self, message: Mapping[str, object] | None) -> None: + if message is None or message["type"] != "message": + return + data: Final = message["data"] + assert isinstance(data, bytes | str), message + parsed: Final = _BroadcastMessage.model_validate_json(data) + with self._condition: + self._received.extend(Broadcast(key, parsed.new_value, parsed.ttl) for key in parsed.names()) + self._condition.notify_all() + + def wait_for(self, satisfied: Callable[[Sequence[Broadcast]], bool], seconds: float) -> bool: + with self._condition: + return self._condition.wait_for(lambda: satisfied(tuple(self._received)), timeout=seconds) + + def received(self) -> tuple[Broadcast, ...]: + with self._condition: + return tuple(self._received) + + def close(self) -> None: + self._stopped.set() + self._thread.join(timeout=5) + self._pubsub.close() + + +@contextmanager +def _broadcasts(cache: Redis) -> Iterator[_Broadcasts]: + broadcasts: Final = _Broadcasts(cache) + try: + yield broadcasts + finally: + broadcasts.close() + + +def _zeroed(counter: str) -> Callable[[Sequence[Broadcast]], bool]: + return lambda received: any(broadcast.key == counter and broadcast.value == 0.0 for broadcast in received) + + +def _zeroed_with_its_floor(counter: str) -> Callable[[Sequence[Broadcast]], bool]: + return lambda received: ( + sum(1 for broadcast in received if broadcast.key.endswith(counter) and broadcast.value == 0.0) >= 2 + ) + + +@dataclass(frozen=True, slots=True) +class Isolated: + gateway: Gateway + cache: Redis + redis: OwnedRedis + database_url: str + log: Path + + +@contextmanager +def _isolated(gateway: Gateway, directory: Path, *, workers: int = 1) -> Iterator[Isolated]: + with scratch_database() as database_url, owned_redis(directory) as redis: + overrides: Final[Mapping[str, str]] = { + "DATABASE_URL": database_url, + "REDIS_HOST": redis.host, + "REDIS_PORT": str(redis.port), + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", + } + with ( + owned_proxy_process( + gateway, directory, overrides, workers=workers, remove_environment=("DATABASE_URL_READ_REPLICA",) + ) as owned, + Redis(host=redis.host, port=redis.port) as cache, + ): + yield Isolated(owned.gateway, cache, redis, database_url, owned.log) + + +def _config_without_redis(directory: Path) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + litellm_settings: Final = { + name: value for name, value in object_value(shared["litellm_settings"]).items() if name not in REDIS_SETTINGS + } + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False)) + return path + + +@contextmanager +def _isolated_without_redis(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, str]]: + with ( + scratch_database() as database_url, + owned_proxy_process( + gateway, + directory, + {"DATABASE_URL": database_url}, + config=_config_without_redis(directory), + remove_environment=("REDIS_HOST", "REDIS_PORT", "DATABASE_URL_READ_REPLICA"), + ) as owned, + ): + yield owned.gateway, database_url + + +def _one_spend(rows: Sequence[Mapping[str, object]]) -> float: + assert len(rows) == 1, rows + return float(str(rows[0]["spend"])) + + +def _key_spend(key: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', + (sha256(key.encode()).hexdigest(),), + database_url=database_url, + ) + ) + + +def _team_spend(team: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,), database_url=database_url) + ) + + +def _user_spend(user: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=%s', (user,), database_url=database_url) + ) + + +def _organization_spend(organization: str, database_url: str | None = None) -> float: + return _one_spend( + read_rows( + 'SELECT spend FROM "LiteLLM_OrganizationTable" WHERE organization_id=%s', + (organization,), + database_url=database_url, + ) + ) + + +def _landed(read: Callable[[], float], least: float) -> float: + return eventually(read, lambda spend: spend >= least - 1e-9, seconds=DATABASE_SECONDS) + + +def _chat(client: httpx.Client, model: str, key: str, *, end_user: str | None = None) -> httpx.Response: + return client.post( + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": f"reset path {uuid.uuid4().hex}"}], + **({} if end_user is None else {"user": end_user}), + }, + headers={"Authorization": f"Bearer {key}"}, + ) + + +def _assert_served(response: httpx.Response) -> str: + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["usage"]["total_tokens"] == 40, body + return str(body["id"]) + + +def _assert_denied(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert response.json()["error"]["type"] == "budget_exceeded", response.text + + +def _assert_counters_unwritten(response: httpx.Response, saved: str) -> None: + assert response.status_code == 500, response.text + error: Final = response.json()["error"] + assert error["type"] == "internal_server_error", response.text + message: Final = str(error["message"]) + assert message.startswith(f"{saved}, {LOST_COUNTERS}"), message + assert message.endswith(RETRY_HINT), message + + +def _accepted_once_redis_is_back(request: Callable[[], httpx.Response]) -> None: + def attempt() -> int: + response: Final = request() + assert response.status_code in {200, 500}, response.text + return response.status_code + + eventually(attempt, lambda code: code == 200, seconds=RECOVERY_SECONDS) + + +def _fence(counter: str) -> str: + return f"{{{counter}}}:reset_fence" + + +async def _repair_as_a_pod_would(redis: OwnedRedis, counter: str, value: float) -> float | None: + cache: Final = RedisCache(host=redis.host, port=redis.port) + assert cache.reset_fence_key(counter) == _fence(counter) + return await cache.async_set_max( + key=counter, value=value, ttl=SPEND_COUNTER_TTL_SECONDS, fence_key=cache.reset_fence_key(counter) + ) + + +def _state(cache: Redis, counter: str) -> tuple[float | None, int]: + return counter_value(cache, counter), counter_ttl(cache, counter) + + +def _settled(cache: Redis, counter: str, spend: float) -> int: + _, ttl = eventually( + lambda: _state(cache, counter), + lambda state: state[0] is not None and state[0] == pytest.approx(spend, abs=1e-9), + seconds=SETTLE_SECONDS, + ) + return ttl + + +def _values(cache: Redis, counters: Sequence[str]) -> tuple[float | None, ...]: + return tuple(None if raw is None else float(raw) for raw in cache.mget(counters)) + + +def _ttls(cache: Redis, counters: Sequence[str]) -> tuple[int, ...]: + with cache.pipeline() as pipeline: + for counter in counters: + pipeline.ttl(counter) + return tuple(int(ttl) for ttl in pipeline.execute()) + + +def _exhaust(gateway: Gateway, cache: Redis, database_url: str, model: str, key: str) -> str: + counter: Final = key_counter(key) + _assert_served(_chat(gateway.client, model, key)) + _settled(cache, counter, ONE_CALL_COST) + _landed(lambda: _key_spend(key, database_url), ONE_CALL_COST) + _assert_denied(_chat(gateway.client, model, key)) + return counter + + +def _global_reset(gateway: Gateway) -> None: + reset: Final = gateway.request("POST", "/global/spend/reset") + assert reset.status_code == 200, reset.text + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_admits_a_blocked_key_on_the_next_request(gateway: Gateway, tmp_path: Path) -> None: + with ( + _isolated(gateway, tmp_path, workers=2) as isolated, + isolated.gateway.scenario() as scenario, + httpx.Client(base_url=str(isolated.gateway.client.base_url), timeout=60) as other_connection, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + _global_reset(isolated.gateway) + assert outlives(_settled(isolated.cache, counter, 0.0)) + assert _key_spend(key, isolated.database_url) == 0.0 + _assert_served(_chat(other_connection, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(other_connection, model, key)) + _landed(lambda: _key_spend(key, isolated.database_url), ONE_CALL_COST) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_admits_a_key_blocked_by_its_team_budget(gateway: Gateway, tmp_path: Path) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + team: Final = scenario.team(models=[model], max_budget=ONE_CALL_COST) + key: Final = scenario.key(team_id=team, models=[model]) + counter: Final = f"spend:team:{team}" + _assert_served(_chat(isolated.gateway.client, model, key)) + _settled(isolated.cache, counter, ONE_CALL_COST) + _landed(lambda: _team_spend(team, isolated.database_url), ONE_CALL_COST) + _assert_denied(_chat(isolated.gateway.client, model, key)) + _global_reset(isolated.gateway) + assert outlives(_settled(isolated.cache, counter, 0.0)) + assert _team_spend(team, isolated.database_url) == 0.0 + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_leaves_the_other_scopes_and_a_non_master_caller_alone( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + organization: Final = scenario.organization(max_budget=50.0) + team: Final = scenario.team(models=[model], max_budget=20.0, organization_id=organization) + user: Final = scenario.user(max_budget=20.0) + added: Final = isolated.gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0}, + ) + assert added.status_code == 200, added.text + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + key: Final = scenario.key(team_id=team, user_id=user, models=[model], max_budget=5.0, tags=[tag]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + reset_counters: Final = (key_counter(key), f"spend:team:{team}") + kept_counters: Final = ( + f"spend:team_member:{user}:{team}", + f"spend:user:{user}", + f"spend:end_user:{end_user}", + f"spend:tag:{tag}", + f"spend:org:{organization}", + ) + _assert_served(_chat(isolated.gateway.client, model, key, end_user=end_user)) + for counter in (*reset_counters, *kept_counters): + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)), counter + _landed(lambda: _user_spend(user, isolated.database_url), ONE_CALL_COST) + _landed(lambda: _organization_spend(organization, isolated.database_url), ONE_CALL_COST) + + forbidden: Final = isolated.gateway.request("POST", "/global/spend/reset", key=key) + assert forbidden.status_code == 401, forbidden.text + assert forbidden.json()["error"]["type"] == "auth_error", forbidden.text + assert _values(isolated.cache, (*reset_counters, *kept_counters)) == pytest.approx( + (ONE_CALL_COST,) * 7, abs=1e-9 + ) + + _global_reset(isolated.gateway) + for counter in reset_counters: + assert outlives(_settled(isolated.cache, counter, 0.0)), counter + assert _values(isolated.cache, kept_counters) == pytest.approx((ONE_CALL_COST,) * 5, abs=1e-9) + assert all(outlives(ttl) for ttl in _ttls(isolated.cache, kept_counters)), _ttls(isolated.cache, kept_counters) + assert _key_spend(key, isolated.database_url) == 0.0 + assert _team_spend(team, isolated.database_url) == 0.0 + assert _user_spend(user, isolated.database_url) >= ONE_CALL_COST - 1e-9 + assert _organization_spend(organization, isolated.database_url) >= ONE_CALL_COST - 1e-9 + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_zeroes_every_key_counter_past_one_scan_page_and_broadcasts_them( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated: + stem: Final = f"spend:key:audit-{uuid.uuid4().hex}" + counters: Final = tuple(f"{stem}-{index}" for index in range(SEEDED_COUNTERS)) + windows: Final = tuple(f"{stem}-w{index}:window:1d" for index in range(3)) + with isolated.cache.pipeline() as pipeline: + for name in (*counters, *windows): + pipeline.set(name, "1.0", ex=3600) + pipeline.execute() + with _broadcasts(isolated.cache) as broadcasts: + _global_reset(isolated.gateway) + eventually( + lambda: _values(isolated.cache, counters), + lambda values: all(value == 0.0 for value in values), + seconds=SETTLE_SECONDS, + ) + ttls: Final = _ttls(isolated.cache, counters) + assert all(outlives(ttl) for ttl in ttls), sorted(ttls)[:5] + assert _values(isolated.cache, windows) == (1.0, 1.0, 1.0) + assert broadcasts.wait_for( + lambda received: {broadcast.key for broadcast in received if broadcast.value == 0.0} >= set(counters), + BROADCAST_SECONDS, + ), len(broadcasts.received()) + assert not any(broadcast.key in windows for broadcast in broadcasts.received()) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_worker_that_cached_the_pre_reset_floor_admits_the_next_request_after_a_global_reset( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _isolated(gateway, tmp_path, workers=2) as isolated, + isolated.gateway.scenario() as scenario, + httpx.Client(base_url=str(isolated.gateway.client.base_url), timeout=60) as pinned, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + with _broadcasts(isolated.cache) as broadcasts: + assert isolated.cache.set(counter, "0.0", ex=3600) + _assert_denied(_chat(pinned, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _global_reset(isolated.gateway) + broadcasts.wait_for(_zeroed_with_its_floor(counter), RACE_WINDOW_SECONDS) + assert pinned.get("/health/liveliness").status_code == 200 + assert pinned.get("/health/liveliness").status_code == 200 + _assert_served(_chat(pinned, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(pinned, model, key)) + assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received() + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_fences_the_counter_against_a_stale_floor_repair(gateway: Gateway, tmp_path: Path) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + _global_reset(isolated.gateway) + assert outlives(_settled(isolated.cache, counter, 0.0)) + assert isolated.cache.get(_fence(counter)) == b"1" + assert 0 < counter_ttl(isolated.cache, _fence(counter)) <= FENCE_SECONDS + assert asyncio.run(_repair_as_a_pod_would(isolated.redis, counter, ONE_CALL_COST)) == 0.0 + value, ttl = _state(isolated.cache, counter) + assert value == 0.0 and outlives(ttl), (value, ttl) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + eventually(lambda: isolated.cache.exists(_fence(counter)), lambda found: found == 0, seconds=FENCE_SECONDS + 5) + late_total: Final = 2 * ONE_CALL_COST + assert asyncio.run(_repair_as_a_pod_would(isolated.redis, counter, late_total)) == pytest.approx(late_total) + value_after, ttl_after = _state(isolated.cache, counter) + assert value_after == pytest.approx(late_total) and outlives(ttl_after), (value_after, ttl_after) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_while_redis_is_down_answers_500_and_the_retry_resets_the_counters( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + isolated.redis.stop() + _assert_counters_unwritten(isolated.gateway.request("POST", "/global/spend/reset"), "Spend was reset") + assert _key_spend(key, isolated.database_url) == 0.0 + _assert_served(_chat(isolated.gateway.client, model, key)) + _landed(lambda: _key_spend(key, isolated.database_url), ONE_CALL_COST) + isolated.redis.start() + assert isolated.cache.ping() + _accepted_once_redis_is_back(lambda: isolated.gateway.request("POST", "/global/spend/reset")) + assert _key_spend(key, isolated.database_url) == 0.0 + assert outlives(_settled(isolated.cache, counter, 0.0)) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_global_spend_reset_without_redis_admits_the_next_request(gateway: Gateway, tmp_path: Path) -> None: + with _isolated_without_redis(gateway, tmp_path) as (owned, database_url), owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + _assert_served(_chat(owned.client, model, key)) + _landed(lambda: _key_spend(key, database_url), ONE_CALL_COST) + _assert_denied(_chat(owned.client, model, key)) + _global_reset(owned) + assert _key_spend(key, database_url) == 0.0 + _assert_served(_chat(owned.client, model, key)) + _landed(lambda: _key_spend(key, database_url), ONE_CALL_COST) + _assert_denied(_chat(owned.client, model, key)) + + +def _organization_entities( + gateway: Gateway, scenario: Scenario, database_url: str | None = None +) -> tuple[str, str, str]: + model: Final = scenario.model(**RATES) + organization: Final = scenario.organization(max_budget=ONE_CALL_COST) + team: Final = scenario.team(models=[model], organization_id=organization) + key: Final = scenario.key(team_id=team, models=[model]) + counter: Final = f"spend:org:{organization}" + _assert_served(_chat(gateway.client, model, key)) + _landed(lambda: _organization_spend(organization, database_url), ONE_CALL_COST) + _assert_denied(_chat(gateway.client, model, key)) + return model, key, counter + + +def _organization_update(gateway: Gateway, counter: str, **fields: object) -> httpx.Response: + return gateway.request( + "PATCH", "/organization/update", {"organization_id": counter.removeprefix("spend:org:"), **fields} + ) + + +def _update_organization(gateway: Gateway, counter: str, **fields: object) -> None: + updated: Final = _organization_update(gateway, counter, **fields) + assert updated.status_code == 200, updated.text + + +def test_an_organization_spend_set_to_zero_admits_a_key_blocked_by_the_organization_budget(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key, counter = _organization_entities(gateway, scenario) + _settled(cache, counter, ONE_CALL_COST) + with _broadcasts(cache) as broadcasts: + _update_organization(gateway, counter, spend=0.0) + assert outlives(_settled(cache, counter, 0.0)) + assert _organization_spend(counter.removeprefix("spend:org:")) == 0.0 + assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received() + _assert_served(_chat(gateway.client, model, key)) + assert outlives(_settled(cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(gateway.client, model, key)) + + +def test_an_organization_update_without_spend_keeps_the_counter_and_the_block(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key, counter = _organization_entities(gateway, scenario) + _settled(cache, counter, ONE_CALL_COST) + _update_organization(gateway, counter, metadata={"note": f"audit {uuid.uuid4().hex}"}) + value, ttl = _state(cache, counter) + assert value == pytest.approx(ONE_CALL_COST, abs=1e-9) and outlives(ttl), (value, ttl) + _assert_denied(_chat(gateway.client, model, key)) + + +def test_an_organization_spend_set_above_the_budget_keeps_the_block_until_set_to_zero(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key, counter = _organization_entities(gateway, scenario) + _settled(cache, counter, ONE_CALL_COST) + _update_organization(gateway, counter, spend=1.5) + assert outlives(_settled(cache, counter, 1.5)) + _assert_denied(_chat(gateway.client, model, key)) + _update_organization(gateway, counter, spend=0.0) + assert outlives(_settled(cache, counter, 0.0)) + _assert_served(_chat(gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_an_organization_spend_set_while_redis_is_down_answers_500_and_the_retry_resets_the_counter( + gateway: Gateway, tmp_path: Path +) -> None: + with _isolated(gateway, tmp_path) as isolated, isolated.gateway.scenario() as scenario: + model, key, counter = _organization_entities(isolated.gateway, scenario, isolated.database_url) + organization: Final = counter.removeprefix("spend:org:") + _settled(isolated.cache, counter, ONE_CALL_COST) + isolated.redis.stop() + _assert_counters_unwritten( + _organization_update(isolated.gateway, counter, spend=0.0), "The organization spend was saved" + ) + assert _organization_spend(organization, isolated.database_url) == 0.0 + _assert_served(_chat(isolated.gateway.client, model, key)) + isolated.redis.start() + assert isolated.cache.ping() + _accepted_once_redis_is_back(lambda: _organization_update(isolated.gateway, counter, spend=0.0)) + assert _organization_spend(organization, isolated.database_url) == 0.0 + assert outlives(_settled(isolated.cache, counter, 0.0)) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 0f3cfdda21e..0823e755df1 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1492,6 +1492,65 @@ async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypat ] +class _LostPipeline(_SetRecordingPipeline): + async def execute(self) -> list[bool]: + from redis.exceptions import ConnectionError as RedisConnectionError + + raise RedisConnectionError("redis down") + + +@pytest.mark.asyncio +async def test_async_set_cache_pipeline_with_ttls_returns_the_error_of_a_lost_write(monkeypatch, redis_no_ping): + """A spend reset must not report success on a write Redis never applied, so the lost write comes back + as a value while every other caller keeps ignoring it.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + client = MagicMock() + client.pipeline = MagicMock(side_effect=(_LostPipeline(), _SetRecordingPipeline())) + + with patch.object(redis_cache, "init_async_client", return_value=client): + lost = await redis_cache.async_set_cache_pipeline_with_ttls((("k", 1, 60),)) + kept = await redis_cache.async_set_cache_pipeline_with_ttls((("k", 1, 60),)) + + assert isinstance(lost, RedisConnectionError) + assert kept is None + + +@pytest.mark.parametrize("namespace", [None, "ns"]) +def test_reset_fence_key_shares_the_cluster_slot_of_its_counter(namespace, monkeypatch, redis_no_ping): + """The repair checks the fence and the counter in one Lua call, which a cluster only runs on one slot.""" + from redis.crc import key_slot + + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace=namespace) + counter = redis_cache.check_and_fix_namespace(key="spend:key:abc") + fence = redis_cache.check_and_fix_namespace(key=redis_cache.reset_fence_key("spend:key:abc")) + + assert fence != counter + assert key_slot(fence.encode()) == key_slot(counter.encode()) + + +@pytest.mark.asyncio +async def test_async_set_max_declares_the_reset_fence_as_a_second_key(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache(namespace="ns") + client = MagicMock() + client.eval = AsyncMock(return_value=b"3.0") + + with patch.object(redis_cache, "init_async_client", return_value=client): + held = await redis_cache.async_set_max( + key="spend:key:abc", value=12.0, ttl=60, fence_key=redis_cache.reset_fence_key("spend:key:abc") + ) + unfenced = await redis_cache.async_set_max(key="spend:key:abc", value=12.0, ttl=60) + + assert (held, unfenced) == (3.0, 3.0) + fenced_call, plain_call = client.eval.await_args_list + assert fenced_call.args[1:4] == (2, "ns:spend:key:abc", "ns:{ns:spend:key:abc}:reset_fence") + assert plain_call.args[1:3] == (1, "ns:spend:key:abc") + + class _ListPipeline: def __init__(self, rows: list[str]) -> None: self.rows = rows diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index de8ff5c0545..c9d7f66733d 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -10163,6 +10163,7 @@ async def test_invalidate_team_member_spend_state_broadcasts_the_spend_counter_t redis_cache=_FakeRedisCache(), user_api_key_cache=remote_user_api_key_cache, additional_in_memory_caches=(remote_spend_counter_in_memory_cache,), + counter_cache=remote_spend_counter_in_memory_cache, ) for cache_key in ("spend:team_member:user-1:team-1", "spend_db_floor:spend:team_member:user-1:team-1"): remote_subscriber._apply_message( # pyright: ignore[reportPrivateUsage] # exercising the real cross-worker message handler, not a public API diff --git a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py index 9d3f458e2f5..a190ec9bf3c 100644 --- a/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py +++ b/tests/unit/proxy/common_utils/test_auth_cache_invalidation_pubsub.py @@ -247,11 +247,15 @@ async def test_publish_many_sends_one_message_per_five_hundred_keys() -> None: @pytest.mark.asyncio -async def test_subscriber_sets_every_key_of_a_multi_key_message_in_additional_caches() -> None: +async def test_subscriber_sets_every_key_of_a_multi_key_message_in_the_counter_cache_only() -> None: + """A global reset broadcasts thousands of counters; written into every additional cache they would evict + the BYOK credentials cached next to them.""" cache = UserApiKeyCache() spend_counter_in_memory_cache = InMemoryCache() spend_counter_in_memory_cache.set_cache("spend:key:a", 999.0) spend_counter_in_memory_cache.set_cache("spend_db_floor:spend:key:a", 999.0) + byok_cache = InMemoryCache() + byok_cache.set_cache("byok:cred-1", {"api_key": "sk-cred"}) message = { "type": "message", "data": json.dumps({"cache_keys": ["spend:key:a", "spend:key:b"], "new_value": 0.0, "ttl": 60}).encode(), @@ -261,7 +265,8 @@ async def test_subscriber_sets_every_key_of_a_multi_key_message_in_additional_ca subscriber = AuthCacheInvalidationSubscriber( redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])), user_api_key_cache=cache, - additional_in_memory_caches=(spend_counter_in_memory_cache,), + additional_in_memory_caches=(spend_counter_in_memory_cache, byok_cache), + counter_cache=spend_counter_in_memory_cache, ) subscriber.start() try: @@ -275,6 +280,8 @@ async def test_subscriber_sets_every_key_of_a_multi_key_message_in_additional_ca assert spend_counter_in_memory_cache.get_cache("spend:key:a") == 0.0 assert spend_counter_in_memory_cache.get_cache("spend:key:b") == 0.0 assert spend_counter_in_memory_cache.get_cache("spend_db_floor:spend:key:a") == 999.0 + assert byok_cache.get_cache("byok:cred-1") == {"api_key": "sk-cred"} + assert byok_cache.get_cache("spend:key:a") is None @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index c4fa4ff3286..b57949c5dec 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1576,14 +1576,15 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert jwt_table.rows == (kept_row,) -def _org_spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> MagicMock: +def _org_spend_counter_cache(monkeypatch: pytest.MonkeyPatch, *, lost_write: Exception | None = None) -> MagicMock: from litellm.caching.in_memory_cache import InMemoryCache spend_counter_cache = MagicMock() spend_counter_cache.in_memory_cache = InMemoryCache() spend_counter_cache.in_memory_cache.set_cache(key="spend:org:org-1", value=4.2) spend_counter_cache.in_memory_cache.set_cache(key="spend_db_floor:spend:org:org-1", value=4.2) - spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock() + spend_counter_cache.redis_cache.reset_fence_key = lambda key: f"{{{key}}}:reset_fence" + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock(return_value=lost_write) monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) monkeypatch.setattr("litellm.proxy.proxy_server.publish_auth_cache_invalidations", AsyncMock()) return spend_counter_cache @@ -1603,16 +1604,42 @@ async def test_legacy_update_with_spend_sets_the_org_counter_and_its_floor_on_ev ) spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.assert_awaited_once_with( - [("spend:org:org-1", 1.5, proxy_server.spend_counter_redis_ttl())] + ( + ("{spend:org:org-1}:reset_fence", 1, proxy_server.SPEND_RESET_FENCE_TTL_SECONDS), + ("spend:org:org-1", 1.5, proxy_server.spend_counter_redis_ttl()), + ) ) assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 1.5 assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:org:org-1") == 1.5 assert proxy_server.publish_auth_cache_invalidations.await_args_list == [ - call(("spend:org:org-1",), new_value=1.5, ttl=60), call(("spend_db_floor:spend:org:org-1",), new_value=1.5, ttl=proxy_server.SPEND_DB_FLOOR_CACHE_TTL_SECONDS), + call(("spend:org:org-1",), new_value=1.5, ttl=60), ] +@pytest.mark.asyncio +async def test_legacy_update_with_spend_answers_500_when_the_redis_counter_write_is_lost(monkeypatch): + """The row is saved either way, but a lost counter write leaves the other pods enforcing the old spend + for up to the counter TTL, so the admin hears about it and retries instead of reading success.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + from litellm.proxy import proxy_server + from litellm.proxy._types import ProxyException + + spend_counter_cache = _org_spend_counter_cache(monkeypatch, lost_write=RedisConnectionError("redis down")) + + with pytest.raises(ProxyException) as raised: + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "spend": 1.5}, existing_budget_id="budget-1" + ) + + assert raised.value.code == "500" + assert raised.value.message.startswith("The organization spend was saved, but the cached spend counters") + proxy_server.prisma_client.db.litellm_organizationtable.update.assert_awaited_once() + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 1.5 + assert proxy_server.publish_auth_cache_invalidations.await_count == 2 + + @pytest.mark.asyncio async def test_legacy_update_without_spend_keeps_the_org_spend_counter(monkeypatch): spend_counter_cache = _org_spend_counter_cache(monkeypatch) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 64c730e4f81..836896c9bea 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -57,6 +57,7 @@ def _make_spend_counter_cache( cache.redis_cache.async_delete_cache = AsyncMock() cache.redis_cache.async_set_cache = AsyncMock() cache.redis_cache.async_set_max = AsyncMock() + cache.redis_cache.reset_fence_key = lambda key: f"{{{key}}}:reset_fence" cache.redis_cache.async_increment_pipeline = AsyncMock(return_value=None) cache.redis_cache.get_ttl = MagicMock(return_value=None) else: @@ -135,7 +136,7 @@ async def test_get_current_spend_floors_stale_low_counter_against_db(monkeypatch # monotonic set-max so other workers read the corrected total, and a # concurrent increment cannot be clobbered fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:key:abc", value=12.0, ttl=ps.spend_counter_redis_ttl() + key="spend:key:abc", value=12.0, ttl=ps.spend_counter_redis_ttl(), fence_key="{spend:key:abc}:reset_fence" ) @@ -285,7 +286,10 @@ async def test_get_current_spend_end_user_floor_repairs_a_stale_low_counter(monk assert result == 12.0 fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key="spend:end_user:customer-42", value=12.0, ttl=ps.spend_counter_redis_ttl() + key="spend:end_user:customer-42", + value=12.0, + ttl=ps.spend_counter_redis_ttl(), + fence_key="{spend:end_user:customer-42}:reset_fence", ) @@ -331,7 +335,7 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch): assert result == 15.0 assert wfsl.await_count == 1 fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0, ttl=ps.spend_counter_redis_ttl() + key=counter_key, value=15.0, ttl=ps.spend_counter_redis_ttl(), fence_key=f"{{{counter_key}}}:reset_fence" ) @@ -374,7 +378,7 @@ async def test_get_current_spend_floors_window_against_maintained_row(monkeypatc assert result == 15.0 fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited() fake_cache.redis_cache.async_set_max.assert_awaited_once_with( - key=counter_key, value=15.0, ttl=ps.spend_counter_redis_ttl() + key=counter_key, value=15.0, ttl=ps.spend_counter_redis_ttl(), fence_key=f"{{{counter_key}}}:reset_fence" ) 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 11ae740f987..26fbe6780c4 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8061,7 +8061,9 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader ] -def _spend_counter_cache_for_reset(monkeypatch: pytest.MonkeyPatch, *, scan_results: object) -> MagicMock: +def _spend_counter_cache_for_reset( + monkeypatch: pytest.MonkeyPatch, *, scan_results: object, lost_write: Exception | None = None +) -> MagicMock: from litellm.caching.in_memory_cache import InMemoryCache mock_prisma_client: Final = MagicMock() @@ -8077,14 +8079,36 @@ def _spend_counter_cache_for_reset(monkeypatch: pytest.MonkeyPatch, *, scan_resu spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t1", value=0.5) spend_counter_cache.in_memory_cache.set_cache(key="spend:user:u1", value=0.5) spend_counter_cache.redis_cache.namespace = "ns" + spend_counter_cache.redis_cache.reset_fence_key = lambda key: f"{{ns:{key}}}:reset_fence" spend_counter_cache.redis_cache.async_scan_iter = AsyncMock(side_effect=scan_results) - spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock() + spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls = AsyncMock(return_value=lost_write) monkeypatch.setattr(ps, "spend_counter_cache", spend_counter_cache) monkeypatch.setattr(ps, "publish_auth_cache_invalidations", AsyncMock()) return spend_counter_cache _RESET_KEY_COUNTERS: Final = ("spend:key:aaa", "spend:key:bbb", "spend:key:floored") +_RESET_TEAM_COUNTERS: Final = ("spend:team:sales:window:east", "spend:team:t1") + + +def _reset_pipeline(counters: tuple[str, ...]) -> tuple[tuple[str, float, int], ...]: + """Every fence is SET before any counter, so a repair racing the pipeline already finds its fence.""" + return ( + *((f"{{ns:{counter}}}:reset_fence", 1, ps.SPEND_RESET_FENCE_TTL_SECONDS) for counter in counters), + *((counter, 0.0, ps.spend_counter_redis_ttl()) for counter in counters), + ) + + +def _reset_broadcasts(counters: tuple[str, ...]) -> tuple[object, object]: + """The floors go out first: a subscriber that applied the new counter under the old floor would repair it.""" + return ( + call( + tuple(f"spend_db_floor:{counter}" for counter in counters), + new_value=0.0, + ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS, + ), + call(counters, new_value=0.0, ttl=60), + ) @pytest.mark.asyncio @@ -8093,12 +8117,16 @@ async def test_global_spend_reset_sets_every_key_and_team_counter_and_its_floor_ ) -> None: """Deleting the counters is not enough: a pod that cached the pre-reset floor within the last 5s would repair the reseeded counter back to the old total for the whole counter TTL, so the reset SETs counter and floor - everywhere and leaves budget-window counters, which the reset does not touch in the DB, alone.""" + everywhere, fences each counter against that repair, and leaves budget-window counters, which the reset + does not touch in the DB, alone. A window counter is told by its trailing duration, so a team whose id + merely contains ``:window:`` is still reset.""" spend_counter_cache: Final = _spend_counter_cache_for_reset( monkeypatch, - scan_results=[["ns:spend:key:aaa", "ns:spend:key:bbb", "ns:spend:key:bbb:window:1d"], ["ns:spend:team:t1"]], + scan_results=[ + ["ns:spend:key:aaa", "ns:spend:key:bbb", "ns:spend:key:bbb:window:1d"], + ["ns:spend:team:t1", "ns:spend:team:t1:window:1d", "ns:spend:team:sales:window:east"], + ], ) - ttl: Final = ps.spend_counter_redis_ttl() response: Final = await spend_management_endpoints.global_spend_reset() @@ -8108,42 +8136,68 @@ async def test_global_spend_reset_sets_every_key_and_team_counter_and_its_floor_ {"pattern": "spend:team:", "count": 500, "unbounded": True}, ] assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [ - ([(counter, 0.0, ttl) for counter in _RESET_KEY_COUNTERS],), - ([("spend:team:t1", 0.0, ttl)],), + (_reset_pipeline(_RESET_KEY_COUNTERS),), + (_reset_pipeline(_RESET_TEAM_COUNTERS),), ] in_memory: Final = spend_counter_cache.in_memory_cache - assert [in_memory.get_cache(key=counter) for counter in (*_RESET_KEY_COUNTERS, "spend:team:t1")] == [0.0] * 4 - assert [ - in_memory.get_cache(key=f"spend_db_floor:{counter}") for counter in (*_RESET_KEY_COUNTERS, "spend:team:t1") - ] == [0.0] * 4 + reset_counters: Final = (*_RESET_KEY_COUNTERS, *_RESET_TEAM_COUNTERS) + assert [in_memory.get_cache(key=counter) for counter in reset_counters] == [0.0] * 5 + assert [in_memory.get_cache(key=f"spend_db_floor:{counter}") for counter in reset_counters] == [0.0] * 5 assert in_memory.get_cache(key="spend:key:aaa:window:1d") == 0.3 + assert in_memory.get_cache(key="spend:team:t1:window:1d") is None assert in_memory.get_cache(key="spend:user:u1") == 0.5 assert ps.publish_auth_cache_invalidations.await_args_list == [ - call(_RESET_KEY_COUNTERS, new_value=0.0, ttl=60), - call( - tuple(f"spend_db_floor:{counter}" for counter in _RESET_KEY_COUNTERS), - new_value=0.0, - ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS, - ), - call(("spend:team:t1",), new_value=0.0, ttl=60), - call(("spend_db_floor:spend:team:t1",), new_value=0.0, ttl=ps.SPEND_DB_FLOOR_CACHE_TTL_SECONDS), + *_reset_broadcasts(_RESET_KEY_COUNTERS), + *_reset_broadcasts(_RESET_TEAM_COUNTERS), ] @pytest.mark.asyncio -async def test_global_spend_reset_still_resets_the_rows_and_this_pod_when_redis_is_down( +async def test_global_spend_reset_answers_500_after_resetting_the_rows_and_this_pod_when_redis_cannot_be_listed( monkeypatch: pytest.MonkeyPatch, ) -> None: + """The counters Redis holds for the other pods stay at the old spend for the whole counter TTL when the + scan fails, so the reset says so instead of reporting success, after resetting what this pod knows.""" spend_counter_cache: Final = _spend_counter_cache_for_reset(monkeypatch, scan_results=ConnectionError("redis down")) - ttl: Final = ps.spend_counter_redis_ttl() - response: Final = await spend_management_endpoints.global_spend_reset() + with pytest.raises(ProxyException) as raised: + await spend_management_endpoints.global_spend_reset() - assert response["status"] == "success" + assert raised.value.code == "500" + assert "retry this request once Redis accepts writes" in raised.value.message + ps.prisma_client.db.litellm_verificationtoken.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) ps.prisma_client.db.litellm_teamtable.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) assert [c.args for c in spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_args_list] == [ - ([("spend:key:aaa", 0.0, ttl), ("spend:key:floored", 0.0, ttl)],), - ([("spend:team:t1", 0.0, ttl)],), + (_reset_pipeline(("spend:key:aaa", "spend:key:floored")),), + (_reset_pipeline(("spend:team:t1",)),), ] assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") == 0.0 assert spend_counter_cache.in_memory_cache.get_cache(key="spend_db_floor:spend:key:floored") == 0.0 + + +@pytest.mark.asyncio +async def test_global_spend_reset_answers_500_when_the_redis_counter_write_is_lost( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A lost pipeline leaves the other pods on the old counters for up to the counter TTL, so the reset still + resets the teams, this pod, and the broadcast, then reports the loss so the admin retries.""" + from redis.exceptions import ConnectionError as RedisConnectionError + + spend_counter_cache: Final = _spend_counter_cache_for_reset( + monkeypatch, + scan_results=[["ns:spend:key:aaa"], ["ns:spend:team:t1"]], + lost_write=RedisConnectionError("redis down"), + ) + + with pytest.raises(ProxyException) as raised: + await spend_management_endpoints.global_spend_reset() + + assert raised.value.code == "500" + assert raised.value.message.startswith("Spend was reset, but the cached spend counters could not be reset in Redis") + assert "redis down" in raised.value.message + assert spend_counter_cache.redis_cache.async_set_cache_pipeline_with_ttls.await_count == 2 + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:team:t1") == 0.0 + assert ps.publish_auth_cache_invalidations.await_args_list == [ + *_reset_broadcasts(("spend:key:aaa", "spend:key:floored")), + *_reset_broadcasts(("spend:team:t1",)), + ] From 4e2d6ab15f7a9b81e211bb0c54b1f5828d693740 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:01:51 -0700 Subject: [PATCH 08/11] test(proxy): give the spend counter test fakes the reset fence hook CountingRedis skipped RedisCache.__init__ and _ExpiringRedisCache never had reset_fence_key, so the repair's broad except swallowed the AttributeError and the counter stayed stale in five tests --- tests/unit/proxy/spend_tracking/test_spend_counter_batch.py | 1 + tests/unit/proxy/test_budget_reservation.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py index bdb430e1945..d7844e58273 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py +++ b/tests/unit/proxy/spend_tracking/test_spend_counter_batch.py @@ -41,6 +41,7 @@ class CountingRedis(RedisCache): self.store: dict[str, object] = dict(store or {}) self.fail = fail self.commands: list[str] = [] + self.namespace = None async def async_get_cache(self, key: str, **kwargs: object) -> object: if self.fail: diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 9fe30f4c696..d62659212d0 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -2366,6 +2366,9 @@ class _ExpiringRedisCache: self.expires_at[key] = self._expiry(kwargs.get("ttl")) return self.store[key] + def reset_fence_key(self, key: str) -> str: + return f"{{{key}}}:reset_fence" + async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: self._evict_expired(key) self.store[key] = max(self.store.get(key, float("-inf")), float(value)) From 07612dcc09631a143ac7e90133d82520223da0a5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 19:23:10 -0700 Subject: [PATCH 09/11] fix(proxy): keep at least one key per spend reset pipeline Each counter in a reset pipeline rides with its fence, so a pipeline takes half of DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE keys. A batch size of 1 made that step 0, and every reset then raised in range() instead of writing the counters --- litellm/proxy/proxy_server.py | 6 +++++- tests/unit/proxy/proxy_server/test_spend_counters.py | 8 ++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6742db8279f..a9c2335e2ab 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3979,7 +3979,11 @@ async def _spend_counter_keys_in_redis(redis_cache: RedisCache, prefix: str) -> return frozenset(key.removeprefix(namespace_prefix) for key in scanned) -_KEYS_PER_RESET_PIPELINE: Final = DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE // 2 +def _keys_per_reset_pipeline(max_batch_size: int) -> int: + return max(1, max_batch_size // 2) + + +_KEYS_PER_RESET_PIPELINE: Final = _keys_per_reset_pipeline(DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE) async def _write_reset_pipeline( diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 836896c9bea..b1db27c166e 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -1612,3 +1612,11 @@ async def test_update_cache_user_cache_failure_invalid_state_is_swallowed(monkey ) assert result is None + + +@pytest.mark.parametrize(("max_batch_size", "keys_per_pipeline"), [(1, 1), (2, 1), (3, 1), (1000, 500)]) +def test_a_reset_pipeline_takes_half_the_redis_batch_size_and_never_less_than_one_key( + max_batch_size: int, keys_per_pipeline: int +) -> None: + """Each key rides with its fence, so half the batch size; a step of zero would make every reset raise.""" + assert ps._keys_per_reset_pipeline(max_batch_size) == keys_per_pipeline From f710aea00140f6b5459b6a2c6e0a30f82a3ec2dc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 10:30:33 -0700 Subject: [PATCH 10/11] test(proxy): run the spend reset floor race across two pods and keep the prefetch cell as a control --- .../test_batch_redis_requests_prefetch.py | 3 +- .../spend/test_spend_counter_reset_paths.py | 100 ++++++++++++++---- 2 files changed, 80 insertions(+), 23 deletions(-) diff --git a/tests/integration/caching/test_batch_redis_requests_prefetch.py b/tests/integration/caching/test_batch_redis_requests_prefetch.py index e2d7fa4e0e4..5177f2d30f0 100644 --- a/tests/integration/caching/test_batch_redis_requests_prefetch.py +++ b/tests/integration/caching/test_batch_redis_requests_prefetch.py @@ -7,6 +7,7 @@ import threading import uuid from collections.abc import Iterator from contextlib import contextmanager +from hashlib import sha256 from pathlib import Path from typing import Final @@ -70,7 +71,7 @@ def test_the_prefetched_response_cache_entry_is_served_from_memory_on_the_next_r model: Final = scenario.model() key: Final = scenario.key(models=[model]) body: Final = {"model": model, "messages": [{"role": "user", "content": f"prefetch {uuid.uuid4().hex}"}]} - namespace: Final = f"litellm:{key}:acompletion" + namespace: Final = f"litellm:{sha256(key.encode()).hexdigest()}:acompletion" first: Final = owned.gateway.request("POST", "/v1/chat/completions", body, key=key) assert first.status_code == 200, first.text entries: Final = eventually(lambda: cache.keys(f"{namespace}:*"), lambda keys: len(keys) == 1, seconds=10) diff --git a/tests/integration/spend/test_spend_counter_reset_paths.py b/tests/integration/spend/test_spend_counter_reset_paths.py index b2c6b5c1282..ef853e79667 100644 --- a/tests/integration/spend/test_spend_counter_reset_paths.py +++ b/tests/integration/spend/test_spend_counter_reset_paths.py @@ -46,7 +46,8 @@ ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).re RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} ONE_CALL_COST: Final = 0.06 INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation" -REDIS_SETTINGS: Final = frozenset({"cache", "cache_params", "enable_redis_auth_cache"}) +REDIS_AUTH_CACHE_SETTING: Final = frozenset({"enable_redis_auth_cache"}) +REDIS_SETTINGS: Final = frozenset({"cache", "cache_params", *REDIS_AUTH_CACHE_SETTING}) SEEDED_COUNTERS: Final = 1200 SETTLE_SECONDS: Final = 20 BROADCAST_SECONDS: Final = 10 @@ -146,28 +147,57 @@ class Isolated: log: Path +def _pod_environment(database_url: str, redis: OwnedRedis) -> Mapping[str, str]: + return { + "DATABASE_URL": database_url, + "REDIS_HOST": redis.host, + "REDIS_PORT": str(redis.port), + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", + } + + @contextmanager -def _isolated(gateway: Gateway, directory: Path, *, workers: int = 1) -> Iterator[Isolated]: +def _isolated( + gateway: Gateway, + directory: Path, + *, + workers: int = 1, + settings: tuple[tuple[str, str], ...] = (), + config: Path | None = None, +) -> Iterator[Isolated]: with scratch_database() as database_url, owned_redis(directory) as redis: - overrides: Final[Mapping[str, str]] = { - "DATABASE_URL": database_url, - "REDIS_HOST": redis.host, - "REDIS_PORT": str(redis.port), - "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", - } + overrides: Final[Mapping[str, str]] = {**_pod_environment(database_url, redis), **dict(settings)} with ( owned_proxy_process( - gateway, directory, overrides, workers=workers, remove_environment=("DATABASE_URL_READ_REPLICA",) + gateway, + directory, + overrides, + workers=workers, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), ) as owned, Redis(host=redis.host, port=redis.port) as cache, ): yield Isolated(owned.gateway, cache, redis, database_url, owned.log) -def _config_without_redis(directory: Path) -> Path: +@contextmanager +def _second_pod(gateway: Gateway, directory: Path, isolated: Isolated, config: Path) -> Iterator[Gateway]: + with owned_proxy_process( + gateway, + directory, + _pod_environment(isolated.database_url, isolated.redis), + config=config, + database_setup=(), + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned: + yield owned.gateway + + +def _config_without(directory: Path, dropped: frozenset[str]) -> Path: shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) litellm_settings: Final = { - name: value for name, value in object_value(shared["litellm_settings"]).items() if name not in REDIS_SETTINGS + name: value for name, value in object_value(shared["litellm_settings"]).items() if name not in dropped } path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" path.write_text(yaml.safe_dump({**shared, "litellm_settings": litellm_settings}, sort_keys=False)) @@ -182,7 +212,7 @@ def _isolated_without_redis(gateway: Gateway, directory: Path) -> Iterator[tuple gateway, directory, {"DATABASE_URL": database_url}, - config=_config_without_redis(directory), + config=_config_without(directory, REDIS_SETTINGS), remove_environment=("REDIS_HOST", "REDIS_PORT", "DATABASE_URL_READ_REPLICA"), ) as owned, ): @@ -416,7 +446,8 @@ def test_a_global_spend_reset_zeroes_every_key_counter_past_one_scan_page_and_br ) -> None: with _isolated(gateway, tmp_path) as isolated: stem: Final = f"spend:key:audit-{uuid.uuid4().hex}" - counters: Final = tuple(f"{stem}-{index}" for index in range(SEEDED_COUNTERS)) + team_named_like_a_window: Final = f"spend:team:audit-{uuid.uuid4().hex}:window:prod" + counters: Final = (*(f"{stem}-{index}" for index in range(SEEDED_COUNTERS)), team_named_like_a_window) windows: Final = tuple(f"{stem}-w{index}:window:1d" for index in range(3)) with isolated.cache.pipeline() as pipeline: for name in (*counters, *windows): @@ -440,28 +471,53 @@ def test_a_global_spend_reset_zeroes_every_key_counter_past_one_scan_page_and_br @pytest.mark.timeout(CELL_TIMEOUT) -def test_a_worker_that_cached_the_pre_reset_floor_admits_the_next_request_after_a_global_reset( +def test_a_global_spend_reset_with_a_redis_batch_size_of_one_resets_and_fences_every_counter( gateway: Gateway, tmp_path: Path ) -> None: with ( - _isolated(gateway, tmp_path, workers=2) as isolated, + _isolated(gateway, tmp_path, settings=(("DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE", "1"),)) as isolated, + isolated.gateway.scenario() as scenario, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) + seeded: Final = tuple(f"spend:key:audit-{uuid.uuid4().hex}" for _ in range(2)) + with isolated.cache.pipeline() as pipeline: + for name in seeded: + pipeline.set(name, "1.0", ex=3600) + pipeline.execute() + counters: Final = (counter, *seeded) + _global_reset(isolated.gateway) + assert isolated.cache.exists(*(_fence(name) for name in counters)) == len(counters) + assert _values(isolated.cache, counters) == (0.0, 0.0, 0.0) + assert all(outlives(ttl) for ttl in _ttls(isolated.cache, counters)), _ttls(isolated.cache, counters) + _assert_served(_chat(isolated.gateway.client, model, key)) + assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(isolated.gateway.client, model, key)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_pod_that_cached_the_pre_reset_floor_admits_the_next_request_after_a_reset_on_another_pod( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + _isolated(gateway, tmp_path, config=_config_without(tmp_path, REDIS_AUTH_CACHE_SETTING)) as isolated, + _second_pod(gateway, tmp_path, isolated, _config_without(tmp_path, REDIS_AUTH_CACHE_SETTING)) as other_pod, isolated.gateway.scenario() as scenario, - httpx.Client(base_url=str(isolated.gateway.client.base_url), timeout=60) as pinned, ): model: Final = scenario.model(**RATES) key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) counter: Final = _exhaust(isolated.gateway, isolated.cache, isolated.database_url, model, key) with _broadcasts(isolated.cache) as broadcasts: assert isolated.cache.set(counter, "0.0", ex=3600) - _assert_denied(_chat(pinned, model, key)) - assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) + _assert_denied(_chat(other_pod.client, model, key)) + _settled(isolated.cache, counter, ONE_CALL_COST) _global_reset(isolated.gateway) broadcasts.wait_for(_zeroed_with_its_floor(counter), RACE_WINDOW_SECONDS) - assert pinned.get("/health/liveliness").status_code == 200 - assert pinned.get("/health/liveliness").status_code == 200 - _assert_served(_chat(pinned, model, key)) + assert other_pod.client.get("/health/liveliness").status_code == 200 + _assert_served(_chat(other_pod.client, model, key)) assert outlives(_settled(isolated.cache, counter, ONE_CALL_COST)) - _assert_denied(_chat(pinned, model, key)) + _assert_denied(_chat(other_pod.client, model, key)) assert broadcasts.wait_for(_zeroed(counter), BROADCAST_SECONDS), broadcasts.received() From ce18fc6634b45e24cf6f433314bd0ebb6d8df55f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:29:50 -0700 Subject: [PATCH 11/11] test(integration): kill spend counter chaos workers through the owned proxy helpers --- .../spend/test_spend_counter_ttl_chaos.py | 26 +++++++------------ 1 file changed, 10 insertions(+), 16 deletions(-) diff --git a/tests/integration/spend/test_spend_counter_ttl_chaos.py b/tests/integration/spend/test_spend_counter_ttl_chaos.py index b22f268b595..8c75a7d8482 100644 --- a/tests/integration/spend/test_spend_counter_ttl_chaos.py +++ b/tests/integration/spend/test_spend_counter_ttl_chaos.py @@ -1,5 +1,3 @@ -import os -import signal import uuid from collections.abc import Callable, Mapping from concurrent.futures import ThreadPoolExecutor @@ -8,11 +6,10 @@ from pathlib import Path from typing import Final import httpx -import psutil import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows -from integration._support.process import OwnedProxy, graceful_stop_seconds, group_members, owned_proxy_process +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.redis_process import owned_redis from integration._support.spend_counters import ( counter_ttl, @@ -144,14 +141,7 @@ def test_a_burst_with_upstream_failures_lands_each_success_once_with_a_long_ttl( def _workers(owned: OwnedProxy) -> tuple[int, ...]: - return tuple( - member.pid - for member in group_members(owned.process.pid) - if member.pid != owned.process.pid - and member.is_running() - and member.status() != psutil.STATUS_ZOMBIE - and any("spawn_main" in part for part in member.cmdline()) - ) + return tuple(sorted(owned.worker_pids())) def _attempt(gateway: Gateway, model: str, key: str, index: int) -> int: @@ -176,8 +166,9 @@ def test_a_killed_worker_does_not_shorten_or_lose_the_counter(gateway: Gateway, first: Final = pool.submit(_attempt, owned.gateway, model, key, 0) assert first.result(timeout=60) == 200 in_flight: Final = [pool.submit(_attempt, owned.gateway, model, key, index) for index in range(1, 12)] - os.kill(before[0], signal.SIGKILL) - survivor: Final = _attempt(owned.gateway, model, key, 99) + with owned.replacing_worker(before[0]) as kill: + kill() + survivor: Final = _attempt(owned.gateway, model, key, 99) codes: Final = [future.result(timeout=120) for future in in_flight] assert survivor == 200, survivor assert all(code in (0, 200) for code in codes), codes @@ -224,8 +215,11 @@ def test_a_reservation_leaked_by_a_killed_worker_stays_for_the_counter_ttl(gatew lambda value: value is not None and value >= LEAK_KEY_BUDGET - 1e-9, seconds=SLOW_UPSTREAM_SECONDS, ) - for pid in before: - os.kill(pid, signal.SIGKILL) + with owned.replacing_worker(before[0]) as first, owned.replacing_worker(before[1]) as second: + first.freeze() + second.freeze() + first() + second() interrupted: Final = in_flight.result(timeout=60) assert interrupted != 200, interrupted leaked: Final = (counter_value(cache, counter), counter_ttl(cache, counter))