From 47db6b2eadcbd0c83503b15a233b0f5b7071821e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:44:23 +0000 Subject: [PATCH] refactor(redis): move batch charge markers behind a typed store --- litellm/litellm_core_utils/batch_utils.py | 4 - .../proxy/hooks/model_max_budget_limiter.py | 22 ++- litellm/redis/stores/batch_charge_marker.py | 31 +++ litellm/router_strategy/budget_limiter.py | 18 +- .../litellm_core_utils/test_batch_utils.py | 46 ++--- .../hooks/test_model_max_budget_limiter.py | 41 ++-- .../redis/stores/test_batch_charge_marker.py | 181 ++++++++++++++++++ 7 files changed, 289 insertions(+), 54 deletions(-) create mode 100644 litellm/redis/stores/batch_charge_marker.py create mode 100644 tests/unit/redis/stores/test_batch_charge_marker.py diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 5b94d4b4e5b..de5ffebbf3a 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/litellm/litellm_core_utils/batch_utils.py @@ -289,10 +289,6 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES -def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str: - return f"{spend_key}:batch:{batch_id}" - - def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None: """A finished batch reports its whole cost on every poll, so its id is charged once per counter.""" if response_cost <= 0 or not is_batch_retrieve_call_type(call_type): diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 1c5d6e05f09..f373b7205d5 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -5,7 +5,7 @@ import json import time from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -19,7 +19,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once +from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.bedrock.common_utils import get_bedrock_base_model, get_bedrock_cross_region_inference_regions @@ -30,6 +30,12 @@ from litellm.proxy.hooks.model_budget_shared_windows import ( SharedWindow, SharedWindows, ) +from litellm.redis.stores.batch_charge_marker import ( + BatchChargeMarkerStore, + TransportBatchChargeMarkerStore, + charge_marker_key, +) +from litellm.redis.transport.contract import RedisFailure from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, LlmProviders, StandardLoggingPayload @@ -1363,11 +1369,15 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: - marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) - polls: Final = await self.dual_cache.async_increment_cache( - key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True + local_polls: Final = await self.dual_cache.in_memory_cache.async_increment( + key=charge_marker_key(spend_key, batch_id), value=1, ttl=ttl_seconds ) - return polls == 1 + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None: + return local_polls == 1 + store: Final[BatchChargeMarkerStore] = TransportBatchChargeMarkerStore(redis_cache.transport()) + claimed: Final = await store.claim_batch(spend_key, batch_id, ttl=timedelta(seconds=ttl_seconds)) + return local_polls == 1 if isinstance(claimed, RedisFailure) else claimed def _local_windows(dual_cache: DualCache) -> LocalTTLCache: diff --git a/litellm/redis/stores/batch_charge_marker.py b/litellm/redis/stores/batch_charge_marker.py new file mode 100644 index 00000000000..1e6db0059a2 --- /dev/null +++ b/litellm/redis/stores/batch_charge_marker.py @@ -0,0 +1,31 @@ +from dataclasses import dataclass +from datetime import timedelta +from typing import Final, Protocol + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + + +def charge_marker_key(spend_key: str, batch_id: str) -> str: + return f"{spend_key}:batch:{batch_id}" + + +class BatchChargeMarkerStore(Protocol): + async def claim_batch(self, spend_key: str, batch_id: str, *, ttl: timedelta) -> bool | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportBatchChargeMarkerStore: + transport: RedisTransport + + async def claim_batch(self, spend_key: str, batch_id: str, *, ttl: timedelta) -> bool | RedisFailure: + key: Final = charge_marker_key(spend_key, batch_id) + batch: Final = self.transport.batch(name="batch_charge_marker") + polled: Final = batch.incrbyfloat(key, 1) + expired: Final = batch.expire_seconds(key, ttl) + polls: Final = await polled + expiry: Final = await expired + if isinstance(polls, RedisFailure): + return polls + if isinstance(expiry, RedisFailure): + return expiry + return polls == 1 diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 1d502dd8d95..423785ff82b 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -36,11 +36,17 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once +from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from +from litellm.redis.stores.batch_charge_marker import ( + BatchChargeMarkerStore, + TransportBatchChargeMarkerStore, + charge_marker_key, +) +from litellm.redis.transport.contract import RedisFailure, raise_redis_failure from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs from litellm.router_utils.cooldown_callbacks import ( get_prometheus_logger_from_callbacks, @@ -640,13 +646,15 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost) async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: - marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is not None: - polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True) - return polls == 1 + store: Final[BatchChargeMarkerStore] = TransportBatchChargeMarkerStore(redis_cache.transport()) + claimed: Final = await store.claim_batch(spend_key, batch_id, ttl=timedelta(seconds=ttl_seconds)) + if isinstance(claimed, RedisFailure): + raise_redis_failure(claimed) + return claimed local_polls: Final = await self.dual_cache.in_memory_cache.async_increment( - key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True + key=charge_marker_key(spend_key, batch_id), value=1, ttl=ttl_seconds, refresh_ttl=True ) return local_polls == 1 diff --git a/tests/unit/litellm_core_utils/test_batch_utils.py b/tests/unit/litellm_core_utils/test_batch_utils.py index cd37d635c43..9dc18596d5b 100644 --- a/tests/unit/litellm_core_utils/test_batch_utils.py +++ b/tests/unit/litellm_core_utils/test_batch_utils.py @@ -2,7 +2,7 @@ import asyncio import importlib import json from io import BytesIO -from typing import Dict, Final, List +from typing import Final import pytest @@ -18,7 +18,7 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco # Fixtures @pytest.fixture -def sample_jsonl_data() -> List[Dict]: +def sample_jsonl_data() -> list[dict]: """Fixture providing sample JSONL data""" return [ { @@ -101,17 +101,9 @@ def test_file_like_object(sample_file_like): def test_router_metadata_variable_name(): """Test that the variable name is correct""" assert get_router_metadata_variable_name(function_name="completion") == "metadata" - assert ( - get_router_metadata_variable_name(function_name="batch") == "litellm_metadata" - ) - assert ( - get_router_metadata_variable_name(function_name="acreate_file") - == "litellm_metadata" - ) - assert ( - get_router_metadata_variable_name(function_name="aget_file") - == "litellm_metadata" - ) + assert get_router_metadata_variable_name(function_name="batch") == "litellm_metadata" + assert get_router_metadata_variable_name(function_name="acreate_file") == "litellm_metadata" + assert get_router_metadata_variable_name(function_name="aget_file") == "litellm_metadata" def test_non_json_input(): @@ -384,8 +376,13 @@ def setup_and_teardown(): @pytest.mark.parametrize( ("call_type", "cost", "eligible"), - (("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False), - ("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)), + ( + ("aretrieve_batch", 0.25, True), + ("retrieve_batch", 0.25, True), + ("acompletion", 0.25, False), + ("aretrieve_batch", 0, False), + ("aretrieve_batch", -0.25, False), + ), ) def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None: from openai.types import Batch @@ -393,8 +390,13 @@ def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, from litellm.litellm_core_utils.batch_utils import batch_id_to_charge_once batch: Final = Batch( - id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions", - input_file_id="file-budget", object="batch", status="completed", + id="batch-budget", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-budget", + object="batch", + status="completed", ) assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None) assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None @@ -402,10 +404,10 @@ def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, def test_batch_charge_markers_keep_counters_and_batches_independent() -> None: from litellm.caching.in_memory_cache import InMemoryCache - from litellm.litellm_core_utils.batch_utils import batch_charged_once_marker_key + from litellm.redis.stores.batch_charge_marker import charge_marker_key cache: Final = InMemoryCache() - cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1) - assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1 - assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None - assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None + cache.set_cache(charge_marker_key("provider-spend", "first-batch"), 1) + assert cache.get_cache(charge_marker_key("provider-spend", "first-batch")) == 1 + assert cache.get_cache(charge_marker_key("deployment-spend", "first-batch")) is None + assert cache.get_cache(charge_marker_key("provider-spend", "second-batch")) is None diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index f5a2ccf7b76..1e0b5cce281 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -7,11 +7,14 @@ from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType, SimpleNamespace from typing import Final +from unittest.mock import patch +import fakeredis import pytest import litellm from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction from litellm.proxy.hooks.model_budget_shared_windows import DatabaseSharedWindows, SharedWindows @@ -27,6 +30,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.router import Router from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import LiteLLMBatch, Usage +from tests._support.fake_redis import fake_redis_cache KEY_HASH: Final = "key-hash-batch" USER_ID: Final = "user-batch" @@ -261,10 +265,8 @@ class _SharedRedisDouble: return [await self.async_increment(op["key"], op["increment_value"], ttl=op["ttl"]) for op in increment_list] -def _worker(redis: _SharedRedisDouble) -> PROXY_VirtualKeyModelMaxBudgetLimiter: - return PROXY_VirtualKeyModelMaxBudgetLimiter( - dual_cache=DualCache(redis_cache=redis) # pyright: ignore[reportArgumentType] # duck-typed Redis double - ) +def _worker(redis: RedisCache, clock: Callable[[], float] = time.time) -> PROXY_VirtualKeyModelMaxBudgetLimiter: + return PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=redis), clock=clock) async def _drain_redis_pushes() -> None: @@ -299,7 +301,7 @@ async def test_a_second_batch_and_chat_requests_still_charge_the_budget(): @pytest.mark.asyncio async def test_two_workers_polling_the_same_finished_batch_at_once_charge_it_once(): - redis: Final = _SharedRedisDouble() + redis: Final = fake_redis_cache(fakeredis.FakeServer()) worker_a: Final = _worker(redis) worker_b: Final = _worker(redis) finished: Final = _batch("batch_first", "completed") @@ -314,16 +316,18 @@ async def test_two_workers_polling_the_same_finished_batch_at_once_charge_it_onc @pytest.mark.asyncio async def test_a_batch_polled_within_every_budget_window_is_never_charged_again(): clock: Final = _Clock() - limiter: Final = _worker(_SharedRedisDouble(now=clock.now)) + limiter: Final = _worker(fake_redis_cache(fakeredis.FakeServer()), clock=clock.now) finished: Final = _batch("batch_first", "completed") - await _poll(limiter, finished, BATCH_COST) - clock.advance(12 * 3600) - await _poll(limiter, finished, BATCH_COST) - clock.advance(18 * 3600) - await _poll(limiter, finished, BATCH_COST) - - assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + clock.advance(WINDOW_EPOCH) + with patch("time.time", clock.now): + await _poll(limiter, finished, BATCH_COST) + assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + clock.advance(12 * 3600) + await _poll(limiter, finished, BATCH_COST) + clock.advance(18 * 3600) + await _poll(limiter, finished, BATCH_COST) + assert _local_spend(limiter, KEY_SPEND_KEY) == 0.0 OPUS_GROUP: Final = "claude-opus" @@ -660,9 +664,10 @@ async def test_spend_charged_while_the_database_is_unreachable_reaches_the_share shared_windows=database.replica(), ) assert usage["claude-opus*"]["current_spend"] == round(CHAT_COST * 3, 4) - assert usage["claude-opus*"]["budget_reset_at"] == datetime.fromtimestamp( - WINDOW_EPOCH + 24 * 3600, tz=timezone.utc - ).isoformat() + assert ( + usage["claude-opus*"]["budget_reset_at"] + == datetime.fromtimestamp(WINDOW_EPOCH + 24 * 3600, tz=timezone.utc).isoformat() + ) @pytest.mark.asyncio @@ -715,7 +720,9 @@ def test_an_access_group_named_after_a_served_model_never_widens_that_models_ent resolved: Final = resolve_model_budget( model="gpt-4o-mini", model_max_budget=model_max_budget, - model_access_groups=lambda model, team_id: model_access_groups_of(router=router, model=model, team_id=team_id), + model_access_groups=lambda model, team_id: model_access_groups_of( + router=router, model=model, team_id=team_id + ), ) return None if resolved is None else resolved.budget_model diff --git a/tests/unit/redis/stores/test_batch_charge_marker.py b/tests/unit/redis/stores/test_batch_charge_marker.py new file mode 100644 index 00000000000..b9148440dbd --- /dev/null +++ b/tests/unit/redis/stores/test_batch_charge_marker.py @@ -0,0 +1,181 @@ +from collections.abc import Iterator +from datetime import timedelta +from typing import Final, Protocol, cast +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import ConnectionError as RedisConnectionError +from redis.exceptions import ResponseError + +from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter +from litellm.redis.stores.batch_charge_marker import TransportBatchChargeMarkerStore +from litellm.redis.transport.contract import RedisFailure +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting, budget_window_seconds +from litellm.types.utils import LiteLLMBatch +from tests._support.fake_redis import fake_redis_cache + +NOW: Final = 1800000000.0 +PROVIDER_SPEND: Final = "provider_spend:openai:1d" +KEY_SPEND: Final = "virtual_key_spend:key-hash-batch:batch-qa-primary:1d" +BATCH_COST: Final = 0.25 + + +class LegacyIncrement(Protocol): + async def async_increment( + self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False + ) -> float: ... + + +@pytest.fixture(autouse=True) +def clock() -> Iterator[None]: + with patch("time.time", return_value=NOW): + yield + + +def _batch(batch_id: str) -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-batch", + object="batch", + status="completed", + ) + + +def _router_limiter(cache: RedisCache) -> RouterBudgetLimiting: + return RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=cache), + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + clock=lambda: NOW, + ) + + +async def _router_poll(limiter: RouterBudgetLimiting, batch_id: str) -> None: + kwargs: Final = { + "call_type": "aretrieve_batch", + "litellm_params": {"model": "openai/gpt-4o"}, + "standard_logging_object": { + "response_cost": BATCH_COST, + "model_id": "deployment-1", + "custom_llm_provider": "openai", + }, + } + await limiter.async_log_success_event(kwargs, _batch(batch_id), None, None) + + +def _model_limiter(cache: RedisCache) -> PROXY_VirtualKeyModelMaxBudgetLimiter: + return PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=cache)) + + +async def _model_poll(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, batch_id: str) -> None: + budget: Final = {"batch-qa-primary": {"budget_limit": 100, "time_period": "1d"}} + event: Final = { + "call_type": "aretrieve_batch", + "standard_logging_object": { + "call_type": "aretrieve_batch", + "response_cost": BATCH_COST, + "model": "openai/gpt-5.4-mini", + "model_group": "batch-qa-primary", + "metadata": {"user_api_key_hash": "key-hash-batch"}, + }, + "litellm_params": {"metadata": {"user_api_key_model_max_budget": budget}}, + } + await limiter.async_log_success_event(event, response_obj=_batch(batch_id), start_time=None, end_time=None) + + +async def _shared_spend(cache: RedisCache, key: str, *limiters: RouterBudgetLimiting) -> float: + _: Final = [ + await limiter._push_in_memory_increments_to_redis() # pyright: ignore[reportPrivateUsage] # flushes queued spend the way the periodic sync does + for limiter in limiters + ] + spend: Final = await cache.transport().get(key) + assert not isinstance(spend, RedisFailure) + return 0.0 if spend is None else float(spend) + + +def _local_spend(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, key: str) -> float: + return limiter.dual_cache.in_memory_cache.get_cache(key=key) or 0.0 + + +@pytest.mark.asyncio +async def test_router_marker_bytes_ttl_and_namespace_are_shared_by_workers(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server, namespace="tenant") + first: Final = _router_limiter(cache) + second: Final = _router_limiter(cache) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + marker: Final = "tenant:" + PROVIDER_SPEND + ":batch:batch-shared" + await _router_poll(first, "batch-shared") + assert await raw.get(marker) == b"1" + assert await raw.pttl(marker) == budget_window_seconds(NOW, "1d") * 1000 + await _router_poll(second, "batch-shared") + assert await raw.get(marker) == b"2" + assert await _shared_spend(cache, PROVIDER_SPEND, first, second) == BATCH_COST + await _router_poll(second, "batch-other") + assert await _shared_spend(cache, PROVIDER_SPEND, second) == 2 * BATCH_COST + + +@pytest.mark.asyncio +async def test_old_and_new_pods_see_each_others_markers(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + legacy: Final = cast(LegacyIncrement, cache) # cast-ok: rolling-deploy writer on the same keyspace + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await legacy.async_increment(PROVIDER_SPEND + ":batch:batch-old", 1, ttl=86400, refresh_ttl=True) + limiter: Final = _router_limiter(cache) + await _router_poll(limiter, "batch-old") + assert await _shared_spend(cache, PROVIDER_SPEND, limiter) == 0.0 + await _router_poll(limiter, "batch-new") + assert await _shared_spend(cache, PROVIDER_SPEND, limiter) == BATCH_COST + assert await legacy.async_increment(PROVIDER_SPEND + ":batch:batch-new", 1, ttl=86400, refresh_ttl=True) == 2.0 + assert await raw.get(PROVIDER_SPEND + ":batch:batch-new") == b"2" + + +@pytest.mark.asyncio +async def test_router_marker_failures_still_raise_and_charge_nothing(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + limiter: Final = _router_limiter(cache) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(PROVIDER_SPEND + ":batch:batch-bad", b"not-a-number") + with pytest.raises(ResponseError): + await _router_poll(limiter, "batch-bad") + server.connected = False + with pytest.raises(RedisConnectionError): + await _router_poll(limiter, "batch-down") + server.connected = True + assert await _shared_spend(cache, PROVIDER_SPEND, limiter) == 0.0 + + +@pytest.mark.asyncio +async def test_model_budget_marker_is_shared_and_falls_back_to_the_local_marker(server: fakeredis.FakeServer) -> None: + cache: Final = fake_redis_cache(server) + first: Final = _model_limiter(cache) + second: Final = _model_limiter(cache) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await _model_poll(first, "batch-shared") + await _model_poll(second, "batch-shared") + assert _local_spend(first, KEY_SPEND) + _local_spend(second, KEY_SPEND) == BATCH_COST + assert await raw.get(KEY_SPEND + ":batch:batch-shared") == b"2" + assert await raw.pttl(KEY_SPEND + ":batch:batch-shared") == 86400000 + await raw.set(KEY_SPEND + ":batch:batch-bad", b"not-a-number") + await _model_poll(second, "batch-bad") + await _model_poll(second, "batch-bad") + assert _local_spend(first, KEY_SPEND) + _local_spend(second, KEY_SPEND) == 2 * BATCH_COST + assert await raw.get(KEY_SPEND + ":batch:batch-bad") == b"not-a-number" + + +@pytest.mark.asyncio +async def test_marker_outage_is_a_transport_failure_and_model_budgets_fall_back_locally() -> None: + server: Final = fakeredis.FakeServer() + cache: Final = fake_redis_cache(server) + store: Final = TransportBatchChargeMarkerStore(cache.transport()) + server.connected = False + outage: Final = await store.claim_batch(KEY_SPEND, "batch-down", ttl=timedelta(days=1)) + assert isinstance(outage, RedisFailure) and outage.kind == "transport" + limiter: Final = _model_limiter(cache) + await _model_poll(limiter, "batch-down") + await _model_poll(limiter, "batch-down") + assert _local_spend(limiter, KEY_SPEND) == BATCH_COST