mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
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🔑*` and `spend:team:*` counter through a bounded Redis scan, the one
`spend:org:<id>` 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.
This commit is contained in:
mateo-berri
parent
21aa23caae
commit
2cc0c79d07
13 changed files with 1315 additions and 29 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
33
tests/integration/_support/spend_counters.py
Normal file
33
tests/integration/_support/spend_counters.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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"]),
|
||||
|
|
|
|||
488
tests/integration/spend/test_spend_counter_redis_ttl.py
Normal file
488
tests/integration/spend/test_spend_counter_redis_ttl.py
Normal file
|
|
@ -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))
|
||||
296
tests/integration/spend/test_spend_counter_ttl_chaos.py
Normal file
296
tests/integration/spend/test_spend_counter_ttl_chaos.py
Normal file
|
|
@ -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
|
||||
228
tests/integration/spend/test_spend_counter_ttl_settings.py
Normal file
228
tests/integration/spend/test_spend_counter_ttl_settings.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue