fix(proxy): add a process's cached spend to a cold spend counter once

Overlapping cold starts of one spend counter inside one proxy process each added
the process's whole cached spend, so N concurrent requests left N times the spend
and refused a customer still under budget. Requests that read the counter cold
before this process's own add landed now join that add, or re-read the counter
and add again only when it is gone. Other processes still add their own cached
spend as before.

Refs #47062

litellm/proxy/db/spend_counter_reseed.py

@@ -13,8 +13,10 @@ memory in long-lived deployments.
 """
 
 import asyncio
+import functools
 from collections import OrderedDict
-from collections.abc import Mapping
+from collections.abc import Awaitable, Callable, Mapping
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from types import MappingProxyType
 from typing import TYPE_CHECKING, ClassVar, Final, Optional
@@ -27,8 +29,10 @@ from litellm.proxy._types import Litellm_EntityType
 from litellm.proxy.db.db_lookup_gate import bounded_db_lookup, db_lookup_gate
 from litellm.proxy.spend_tracking.spend_counter_batch import (
     SPEND_COUNTERS_TARGET,
+    forget_spend_counter,
     read_batched_spend_counter,
     record_spend_counter_value,
+    spend_counter_tick,
 )
 from litellm.repositories.organization_repository import OrganizationRepository
 from litellm.repositories.project_repository import ProjectRepository
@@ -68,6 +72,30 @@ _WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
 )
 
 
+@dataclass(slots=True)
+class _CachedSpendSeed:
+    """This process adding its cached spend to one cold counter: the add in flight, and the tick the last one landed."""
+
+    landed_at: int = 0
+    in_flight: asyncio.Task[float | None] | None = None
+
+    def start(self, add: Callable[[], Awaitable[float | None]]) -> asyncio.Task[float | None]:
+        in_flight: Final = asyncio.ensure_future(add())
+        in_flight.add_done_callback(self._settle)
+        self.in_flight = in_flight
+        return in_flight
+
+    def _settle(self, done: asyncio.Future[float | None]) -> None:
+        self.in_flight = None
+        if not done.cancelled() and done.exception() is None:
+            self.landed_at = spend_counter_tick()
+
+
+@functools.lru_cache(maxsize=SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE)
+def _cached_spend_seed(spend_counter_cache: "DualCache", counter_key: str) -> _CachedSpendSeed:
+    return _CachedSpendSeed()
+
+
 def _as_utc(value: datetime) -> datetime:
     return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
 
@@ -112,6 +140,27 @@ class SpendCounterReseed:
                 SpendCounterReseed._locks.popitem(last=False)
             return lock
 
+    @staticmethod
+    async def add_cached_spend_once(
+        spend_counter_cache: "DualCache",
+        counter_key: str,
+        read_at: int,
+        add: Callable[[], Awaitable[float | None]],
+        is_warm: Callable[[], Awaitable[bool]],
+    ) -> None:
+        """Run ``add``, this process adding its cached spend to a cold counter, once for all of this process's
+        requests that read the counter cold (at or after tick ``read_at``) before that add landed. A request joins the
+        add in flight and gets its result or error. When the add already landed, the request re-reads the counter and
+        adds again only if the counter is gone. Each other process still adds its own cached spend."""
+        seed: Final = _cached_spend_seed(spend_counter_cache, counter_key)
+        if seed.in_flight is None and seed.landed_at > read_at:
+            forget_spend_counter(counter_key)
+            if await is_warm():
+                return
+        added: Final = await asyncio.shield(seed.in_flight or seed.start(add))
+        if added is not None:
+            record_spend_counter_value(counter_key, added)
+
     @staticmethod
     @with_service_target(SPEND_COUNTERS_TARGET)
     async def increment_in_memory(spend_counter_cache: "DualCache", counter_key: str, increment: float) -> float | None:

litellm/proxy/proxy_server.py

@@ -936,6 +936,7 @@ from litellm.proxy.spend_tracking.spend_counter_batch import (
     read_batched_spend_counter,
     record_spend_counter_value,
     spend_counter_batch_scope,
+    spend_counter_read_tick,
 )
 from litellm.proxy.spend_tracking.spend_management_endpoints import (
     router as spend_management_router,
@@ -4085,6 +4086,7 @@ async def ensure_spend_counter_initialized(
     counter_key: str,
     source_cache_key: str | list[str],
 ) -> None:
+    read_at: Final = spend_counter_read_tick()
     is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key)
     if is_warm is False:
         # Shares the per-counter lock with get_current_spend.
@@ -4098,7 +4100,13 @@ async def ensure_spend_counter_initialized(
             # DB unavailable - fall back to in-process cache (may be stale).
             base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key)
             if base_spend > 0:
-                await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend)
+                await SpendCounterReseed.add_cached_spend_once(
+                    spend_counter_cache=spend_counter_cache,
+                    counter_key=counter_key,
+                    read_at=read_at,
+                    add=lambda: increment_spend_counter_cache(counter_key=counter_key, increment=base_spend),
+                    is_warm=lambda: _is_spend_counter_cache_warm(counter_key=counter_key),
+                )
 
 
 _ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized

litellm/proxy/spend_tracking/spend_counter_batch.py

@@ -1,6 +1,7 @@
 """One Redis MGET per phase (admission, reservation, post-call) for the spend counters it reads, not one GET each."""
 
 import asyncio
+import itertools
 from collections.abc import Iterator, Mapping, Sequence
 from contextvars import ContextVar, Token
 from dataclasses import dataclass
@@ -22,6 +23,12 @@ from litellm.proxy.common_utils.user_api_key_cache import (
 _CounterValues: Final = TypeAdapter(dict[str, float | None])
 _NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
 SPEND_COUNTERS_TARGET: Final = "spend_counters"
+_ticks: Final = itertools.count(1)
+
+
+def spend_counter_tick() -> int:
+    """The next number in a per-process sequence that orders spend counter reads and cold seeds."""
+    return next(_ticks)
 
 
 @dataclass(frozen=True, slots=True)
@@ -39,10 +46,21 @@ class SpendCounterBatch:
     hook declares whatever is bound but unread, so whoever flushes first (the auth object prefetch, usually)
     carries the spend counters in the same round trip."""
 
-    __slots__ = ("_fetched", "_inflight", "_keys", "_loaded", "_lock", "_open", "_redis_cache", "_request_batch")
+    __slots__ = (
+        "_fetched",
+        "_inflight",
+        "_keys",
+        "_loaded",
+        "_lock",
+        "_open",
+        "_redis_cache",
+        "_request_batch",
+        "opened_at",
+    )
 
     def __init__(self, redis_cache: RedisCache) -> None:
         self._redis_cache: Final = redis_cache
+        self.opened_at: Final = spend_counter_tick()
         self._lock: Final = asyncio.Lock()
         self._open = True
         self._keys: frozenset[str] = frozenset()
@@ -145,6 +163,12 @@ def active_spend_counter_batch() -> SpendCounterBatch | None:
     return _active_batch.get()
 
 
+def spend_counter_read_tick() -> int:
+    """A tick no later than this context's next spend counter read: the open batch's opening tick, else a new one."""
+    batch: Final = _active_batch.get()
+    return batch.opened_at if batch is not None and batch.is_open else spend_counter_tick()
+
+
 class spend_counter_batch_scope:
     """Reads inside the scope share one MGET for the keys bound here or by ``bind_*`` calls inside it.
     Opened inside a scope whose batch is still open, it binds into that batch so both phases share the MGET."""

tests/integration/spend/test_end_user_budget_enforcement.py

@@ -1,17 +1,20 @@
 import json
+import os
 import re
 import uuid
 from collections.abc import Callable, Mapping
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import Final
 
 import httpx
 import pytest
 from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
-from integration._support.database import read_rows
+from integration._support.database import read_rows, write_rows
 from integration._support.upstream import delete_scenario, register_scenario
 from integration.cost_calculation.cost_tracking_case import JsonResponse
 from pydantic import BaseModel, JsonValue
+from redis import Redis
 
 PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key"
 CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint."
@@ -639,3 +642,133 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) ->
         assert "data:" not in denied.text, denied.text
         _assert_budget_refusal(denied, end_user, 0.05)
         assert _observed_requests(upstream) == []
+
+
+_COLD_BURST: Final = 40
+_COLD_SEED_SPEND: Final = 6.0
+_COLD_BUDGET: Final = 10.0
+_CHAT_COST: Final = 20 * 0.001 + 20 * 0.002
+
+
+def _serves_model(gateway: Gateway, model: str) -> bool:
+    response: Final = gateway.request("GET", "/model/info")
+    return response.status_code == 200 and any(
+        object_value(entry)["model_name"] == model for entry in object_value(response.json())["data"]
+    )
+
+
+def _counter_cache() -> Redis:
+    return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
+
+
+def _counter(cache: Redis, end_user: str) -> float | None:
+    raw: Final = cache.get(f"spend:end_user:{end_user}")
+    return None if raw is None else float(raw)
+
+
+def _lose_counter(cache: Redis, end_user: str) -> None:
+    cache.delete(f"spend:end_user:{end_user}")
+
+
+def _set_end_user_row_spend(end_user: str, spend: float) -> None:
+    write_rows('UPDATE "LiteLLM_EndUserTable" SET spend = %s WHERE user_id = %s', (str(spend), end_user))
+
+
+def _await_row_spend(end_user: str, spend: float) -> None:
+    eventually(
+        lambda: _end_user_row(end_user),
+        lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(spend),
+        seconds=70,
+    )
+
+
+def _refused_spend(denied: httpx.Response, end_user: str) -> float:
+    error: Final = _denial_error(denied)
+    assert error["type"] == "budget_exceeded", denied.text
+    matched: Final = _REFUSAL_MESSAGE.match(string_value(error["message"]))
+    assert matched is not None, denied.text
+    assert matched.group("user") == end_user
+    return float(matched.group("spend"))
+
+
+@pytest.mark.parametrize("spelling", ("user", "responses_user", "messages_metadata_user_id"))
+def test_overlapping_cold_starts_count_the_end_user_spend_once(gateway: Gateway, spelling: str) -> None:
+    chosen: Final = SPELLINGS[spelling]
+    with (
+        gateway.scenario() as scenario,
+        httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
+        _counter_cache() as cache,
+    ):
+        deployment: Final = _deployment(scenario, chosen)
+        key: Final = scenario.key(models=[deployment.model])
+        end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
+        _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
+        scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
+        assert _customer_info(gateway, end_user).spend == _COLD_SEED_SPEND
+        assert _counter(cache, end_user) is None
+        upstream.delete("/__observations").raise_for_status()
+
+        with ThreadPoolExecutor(max_workers=_COLD_BURST) as pool:
+            burst: Final = tuple(
+                pool.map(
+                    lambda index: _send(gateway, chosen, deployment, key, end_user, f"cold {index}"),
+                    range(_COLD_BURST),
+                )
+            )
+        refused: Final = [response.text for response in burst if response.status_code != 200]
+        assert refused == [], refused
+        follow_up: Final = _send(gateway, chosen, deployment, key, end_user, "cold follow-up")
+        _assert_served(chosen, follow_up)
+        assert len(_observed_requests(upstream)) == _COLD_BURST + 1
+
+        expected_spend: Final = _COLD_SEED_SPEND + (_COLD_BURST + 1) * _CHAT_COST
+        _await_row_spend(end_user, expected_spend)
+        assert _customer_info(gateway, end_user).spend == pytest.approx(expected_spend)
+        counted: Final = eventually(
+            lambda: _counter(cache, end_user),
+            lambda value: value is not None and value == pytest.approx(expected_spend),
+            seconds=30,
+            return_last_on_timeout=True,
+        )
+        assert counted == pytest.approx(expected_spend)
+
+
+def test_a_lost_counter_keeps_spend_this_proxy_has_not_written_yet(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario, _counter_cache() as cache:
+        model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.25)
+        key: Final = scenario.key(models=[model])
+        end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
+        call_cost: Final = 20 * 0.001 + 20 * 0.25
+        _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
+        scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
+        served: Final = _chat(gateway, model, key, end_user)
+        assert served.status_code == 200, served.text
+        _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost)
+        _set_end_user_row_spend(end_user, _COLD_SEED_SPEND)
+        _lose_counter(cache, end_user)
+
+        denied: Final = _chat(gateway, model, key, end_user)
+
+        assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + call_cost)
+
+
+def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: Gateway, peer: Gateway) -> None:
+    with gateway.scenario() as scenario, _counter_cache() as cache:
+        model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.149)
+        eventually(lambda: _serves_model(peer, model), bool, seconds=30)
+        key: Final = scenario.key(models=[model])
+        end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
+        call_cost: Final = 20 * 0.001 + 20 * 0.149
+        _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND)
+        scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
+        first: Final = _chat(peer, model, key, end_user)
+        assert first.status_code == 200, first.text
+        _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost)
+        _set_end_user_row_spend(end_user, _COLD_SEED_SPEND)
+        _lose_counter(cache, end_user)
+        second: Final = _chat(peer, model, key, end_user)
+        assert second.status_code == 200, second.text
+
+        denied: Final = _chat(gateway, model, key, end_user)
+
+        assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + 2 * call_cost)

tests/unit/proxy/proxy_server/test_spend_counter_cold_seed.py

@@ -0,0 +1,179 @@
+"""Overlapping cold starts of one spend counter inside one process add that process's cached spend once."""
+
+import asyncio
+from typing import Final, cast
+
+import pytest
+
+import litellm.proxy.proxy_server as ps
+from litellm.caching.dual_cache import DualCache
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.caching.redis_cache import RedisCache
+from litellm.proxy.spend_tracking.spend_counter_batch import active_spend_counter_batch, spend_counter_batch_scope
+
+SOURCE: Final = "end_user_id:customer-1"
+
+
+class _SharedRedis:
+    """The counter commands a cold start sends, yielding like a network call so overlapping requests interleave."""
+
+    def __init__(self) -> None:
+        self.values: Final[dict[str, float]] = {}  # mutable-ok: stands in for the Redis keyspace
+        self.failing = False
+
+    async def async_batch_get_cache(self, key_list: list[str], **_: object) -> dict[str, float | None]:
+        await asyncio.sleep(0)
+        return {key: self.values.get(key) for key in key_list}
+
+    async def async_get_cache(self, key: str, **_: object) -> float | None:
+        await asyncio.sleep(0)
+        return self.values.get(key)
+
+    async def async_increment(self, key: str, value: float, **_: object) -> float:
+        await asyncio.sleep(0)
+        if self.failing:
+            raise ConnectionError("redis down")
+        self.values[key] = self.values.get(key, 0.0) + value
+        return self.values[key]
+
+    async def async_delete_cache(self, key: str) -> None:
+        await asyncio.sleep(0)
+        if not self.failing:
+            self.values.pop(key, None)
+
+
+def _process(monkeypatch: pytest.MonkeyPatch, redis: _SharedRedis | None, cached_spend: float) -> DualCache:
+    counters: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, redis))
+    objects: Final = DualCache(in_memory_cache=InMemoryCache())
+    objects.in_memory_cache.set_cache(key=SOURCE, value={"spend": cached_spend})
+    monkeypatch.setattr(ps, "spend_counter_cache", counters)
+    monkeypatch.setattr(ps, "user_api_key_cache", objects)
+    monkeypatch.setattr(ps, "prisma_client", object())
+    return counters
+
+
+async def _cold_start(redis: _SharedRedis | None, counter: str) -> None:
+    with spend_counter_batch_scope(cast(RedisCache, redis) if redis is not None else None, frozenset((counter,))):
+        await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE)
+
+
+async def _read_cold(counter: str) -> None:
+    batch: Final = active_spend_counter_batch()
+    assert batch is not None
+    assert await batch.read(counter) == (None, True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("with_redis", (True, False))
+async def test_overlapping_cold_starts_in_one_process_add_its_cached_spend_once(
+    monkeypatch: pytest.MonkeyPatch, with_redis: bool
+) -> None:
+    counter: Final = f"spend:end_user:once-{with_redis}"
+    redis: Final = _SharedRedis() if with_redis else None
+    counters: Final = _process(monkeypatch, redis, cached_spend=6.0)
+
+    await asyncio.gather(*(_cold_start(redis, counter) for _ in range(5)))
+
+    seeded: Final = redis.values.get(counter) if redis is not None else counters.in_memory_cache.get_cache(key=counter)
+    assert seeded == 6.0
+
+
+@pytest.mark.asyncio
+async def test_a_request_that_read_the_counter_before_this_process_seeded_it_does_not_add_again(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    counter: Final = "spend:end_user:stale-reader"
+    redis: Final = _SharedRedis()
+    _process(monkeypatch, redis, cached_spend=6.0)
+    read: Final = asyncio.Event()
+
+    async def _stale_request() -> None:
+        with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))):
+            await _read_cold(counter)
+            read.set()
+            await seeded.wait()
+            await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE)
+
+    seeded: Final = asyncio.Event()
+
+    async def _first_request() -> None:
+        await read.wait()
+        await _cold_start(redis, counter)
+        seeded.set()
+
+    await asyncio.gather(_stale_request(), _first_request())
+
+    assert redis.values[counter] == 6.0
+
+
+@pytest.mark.asyncio
+async def test_a_counter_lost_again_after_this_process_seeded_it_is_seeded_again(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    counter: Final = "spend:end_user:lost-again"
+    redis: Final = _SharedRedis()
+    _process(monkeypatch, redis, cached_spend=6.0)
+    read: Final = asyncio.Event()
+    seeded: Final = asyncio.Event()
+
+    async def _stale_request() -> None:
+        with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))):
+            await _read_cold(counter)
+            read.set()
+            await seeded.wait()
+            del redis.values[counter]
+            await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE)
+
+    async def _first_request() -> None:
+        await read.wait()
+        await _cold_start(redis, counter)
+        seeded.set()
+
+    await asyncio.gather(_stale_request(), _first_request())
+
+    assert redis.values[counter] == 6.0
+
+
+@pytest.mark.asyncio
+async def test_another_process_seeding_after_this_request_read_the_counter_still_sums(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    counter: Final = "spend:end_user:two-processes"
+    redis: Final = _SharedRedis()
+    read: Final = asyncio.Event()
+    a_seeded: Final = asyncio.Event()
+
+    async def _process_b() -> None:
+        with spend_counter_batch_scope(cast(RedisCache, redis), frozenset((counter,))):
+            await _read_cold(counter)
+            read.set()
+            await a_seeded.wait()
+            _process(monkeypatch, redis, cached_spend=5.0)
+            await ps.ensure_spend_counter_initialized(counter_key=counter, source_cache_key=SOURCE)
+
+    async def _process_a() -> None:
+        await read.wait()
+        _process(monkeypatch, redis, cached_spend=3.0)
+        await _cold_start(redis, counter)
+        a_seeded.set()
+
+    await asyncio.gather(_process_b(), _process_a())
+
+    assert redis.values[counter] == 8.0
+
+
+@pytest.mark.asyncio
+async def test_a_failed_seed_fails_the_requests_that_joined_it_and_the_next_cold_start_seeds(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    counter: Final = "spend:end_user:failed-seed"
+    redis: Final = _SharedRedis()
+    _process(monkeypatch, redis, cached_spend=6.0)
+    redis.failing = True
+
+    outcomes: Final = await asyncio.gather(*(_cold_start(redis, counter) for _ in range(3)), return_exceptions=True)
+    redis.failing = False
+    await _cold_start(redis, counter)
+
+    assert [type(outcome) for outcome in outcomes] == [ConnectionError] * 3
+    assert redis.values[counter] == 6.0