scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

refactor(redis): move batch charge markers behind a typed store

Browse source
This commit is contained in:
Devin AI 2026-10-30 21:44:23 +00:00
parent 5733bbc834
commit 47db6b2ead
7 changed files with 289 additions and 54 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/batch_utils.py
View file

@ -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):

22
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -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:

31
litellm/redis/stores/batch_charge_marker.py Normal file
View file

@ -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

18
litellm/router_strategy/budget_limiter.py
View file

@ -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

46
tests/unit/litellm_core_utils/test_batch_utils.py
View file

@ -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

41
tests/unit/proxy/hooks/test_model_max_budget_limiter.py
View file

@ -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

181
tests/unit/redis/stores/test_batch_charge_marker.py Normal file
View file

@ -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