mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
fix(proxy): park requeued spend logs in Redis so they survive a pod restart during a DB outage
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
a93bfdc749
commit
3a0cabacf8
10 changed files with 599 additions and 14 deletions
|
|
@ -1999,6 +1999,51 @@ class RedisCache(BaseCache):
|
|||
log_redis_failure(verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH: - Got exception from REDIS", e)
|
||||
raise e
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_rpush_and_trim(
|
||||
self,
|
||||
key: str,
|
||||
values: Sequence[str | bytes | int | float],
|
||||
max_len: int,
|
||||
) -> int:
|
||||
"""Append values and keep only the newest ``max_len`` entries in one MULTI/EXEC.
|
||||
|
||||
Returns the list length right after the push, so callers can tell how many
|
||||
of the oldest entries the trim dropped.
|
||||
"""
|
||||
_redis_client: Final = self._async_commands()
|
||||
namespaced_key: Final = self.check_and_fix_namespace(key=key)
|
||||
start_time: Final = time.time()
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=True) as pipe:
|
||||
pipe.rpush(namespaced_key, *values)
|
||||
pipe.ltrim(namespaced_key, -max_len, -1)
|
||||
results: Final = await pipe.execute()
|
||||
for r in results:
|
||||
if isinstance(r, Exception):
|
||||
raise r
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=time.time() - start_time,
|
||||
call_type=f"async_rpush_and_trim <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
return int(results[0])
|
||||
except Exception as e:
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=time.time() - start_time,
|
||||
error=e,
|
||||
call_type=f"async_rpush_and_trim <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
log_redis_failure(
|
||||
verbose_logger, logging.ERROR, "LiteLLM Redis Cache RPUSH+LTRIM: - Got exception from REDIS", e
|
||||
)
|
||||
raise e
|
||||
|
||||
async def _pipeline_rpush_helper(
|
||||
self,
|
||||
pipe: pipeline,
|
||||
|
|
|
|||
|
|
@ -370,6 +370,9 @@ REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_up
|
|||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
|
||||
MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
|
||||
REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer"
|
||||
REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000
|
||||
REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT: Final = 1000
|
||||
# Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth
|
||||
LITELLM_ASYNCIO_QUEUE_MAXSIZE: Final = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000))
|
||||
TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60))
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ This is to prevent deadlocks and improve reliability
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from functools import reduce
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
|
||||
|
||||
|
|
@ -22,6 +23,8 @@ from litellm.constants import (
|
|||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_SPEND_LOGS_BUFFER_KEY,
|
||||
REDIS_SPEND_LOGS_BUFFER_MAX_ROWS,
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
|
||||
)
|
||||
|
|
@ -48,6 +51,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
|||
WindowSpendUpdateQueue,
|
||||
to_wire_payload,
|
||||
)
|
||||
from litellm.proxy.db.spend_log_batching import SpendLogRow
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineLpopOperation,
|
||||
|
|
@ -93,6 +97,19 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
|
|||
_ValueT = TypeVar("_ValueT")
|
||||
|
||||
|
||||
def _spend_log_json_default(value: object) -> str:
|
||||
return value.isoformat() if isinstance(value, datetime) else str(value)
|
||||
|
||||
|
||||
def _encode_spend_log_row(row: SpendLogRow) -> str:
|
||||
return json.dumps(row, default=_spend_log_json_default)
|
||||
|
||||
|
||||
def _decode_spend_log_row(encoded: str) -> dict[str, object] | None:
|
||||
decoded: Final = json.loads(encoded)
|
||||
return decoded if isinstance(decoded, dict) else None
|
||||
|
||||
|
||||
def _accumulated_spend(totals: Mapping[str, float], entities: Mapping[str, float]) -> dict[str, float]:
|
||||
return {**totals, **{entity_id: totals.get(entity_id, 0) + amount for entity_id, amount in entities.items()}}
|
||||
|
||||
|
|
@ -526,6 +543,49 @@ class RedisUpdateBuffer:
|
|||
str(e),
|
||||
)
|
||||
|
||||
async def store_spend_logs_in_redis(
|
||||
self,
|
||||
rows: Sequence[SpendLogRow],
|
||||
max_rows: int = REDIS_SPEND_LOGS_BUFFER_MAX_ROWS,
|
||||
) -> bool:
|
||||
"""Park spend-log rows in Redis so they outlive this pod, dropping the oldest past ``max_rows``."""
|
||||
if self.redis_cache is None or len(rows) == 0 or not self._should_commit_spend_updates_to_redis():
|
||||
return False
|
||||
try:
|
||||
buffer_size: Final = await self.redis_cache.async_rpush_and_trim(
|
||||
key=REDIS_SPEND_LOGS_BUFFER_KEY,
|
||||
values=[_encode_spend_log_row(row) for row in rows],
|
||||
max_len=max_rows,
|
||||
)
|
||||
overflow: Final = buffer_size - max_rows
|
||||
if overflow > 0:
|
||||
verbose_proxy_logger.error(
|
||||
"Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs",
|
||||
max_rows,
|
||||
overflow,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # the caller falls back to the in-memory queue on any Redis fault
|
||||
verbose_proxy_logger.error(
|
||||
"Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), str(e)
|
||||
)
|
||||
return False
|
||||
verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows))
|
||||
return True
|
||||
|
||||
async def get_spend_logs_from_redis_buffer(self, limit: int) -> tuple[dict[str, object], ...]:
|
||||
"""Atomically take up to ``limit`` parked spend-log rows out of Redis."""
|
||||
if self.redis_cache is None or not self._should_commit_spend_updates_to_redis():
|
||||
return ()
|
||||
popped: Final[str | list[str] | None] = await self.redis_cache.async_lpop(
|
||||
key=REDIS_SPEND_LOGS_BUFFER_KEY,
|
||||
count=limit,
|
||||
)
|
||||
if popped is None:
|
||||
return ()
|
||||
encoded_rows: Final = popped if isinstance(popped, list) else [popped]
|
||||
decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in encoded_rows)
|
||||
return tuple(row for row in decoded_rows if row is not None)
|
||||
|
||||
@staticmethod
|
||||
def _number_of_transactions_to_store_in_redis(
|
||||
db_spend_update_transactions: DBSpendUpdateTransactions,
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.constants import (
|
|||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
MAX_TEAM_LIST_LIMIT,
|
||||
REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT,
|
||||
SPEND_LOG_QUEUE_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
|
||||
|
|
@ -4167,6 +4168,7 @@ class PrismaClient:
|
|||
spend_log_flush_requested: "asyncio.Event | None" = None
|
||||
spend_log_queue_bytes: ClassVar[int] = 0
|
||||
spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None
|
||||
spend_log_write_lock = asyncio.Lock()
|
||||
tool_usage_transactions: list["ToolUsageTransaction"] = []
|
||||
_tool_usage_transactions_lock = asyncio.Lock()
|
||||
autorouter_turn_transactions: ClassVar[
|
||||
|
|
@ -7062,7 +7064,7 @@ class ProxyUpdateSpend:
|
|||
except Exception as e:
|
||||
if not _is_transient_spend_log_write_error(e):
|
||||
if PrismaDBExceptionHandler.is_prisma_error(e):
|
||||
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
||||
await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process)
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s",
|
||||
len(logs_to_process),
|
||||
|
|
@ -7077,7 +7079,7 @@ class ProxyUpdateSpend:
|
|||
str(e),
|
||||
)
|
||||
if i >= n_retry_times:
|
||||
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
||||
await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process)
|
||||
raise
|
||||
await asyncio.sleep(2**i)
|
||||
except Exception as e:
|
||||
|
|
@ -7127,6 +7129,7 @@ async def update_spend(
|
|||
)
|
||||
|
||||
### UPDATE SPEND LOGS ###
|
||||
await recover_parked_spend_logs(prisma_client, proxy_logging_obj)
|
||||
# Check queue size with lock protection
|
||||
queue_size: Final = await _total_queued_spend_transactions(prisma_client)
|
||||
verbose_proxy_logger.debug("Spend Logs transactions: %s", queue_size)
|
||||
|
|
@ -7144,6 +7147,51 @@ async def update_spend(
|
|||
)
|
||||
|
||||
|
||||
async def _park_spend_logs_in_redis(proxy_logging_obj: ProxyLogging, rows: Sequence[Mapping[str, object]]) -> bool:
|
||||
try:
|
||||
return await proxy_logging_obj.db_spend_update_writer.redis_update_buffer.store_spend_logs_in_redis(rows)
|
||||
except Exception as e: # noqa: BLE001 # a Redis fault falls back to the in-memory queue, never loses the rows
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - could not park spend logs in Redis, keeping them in memory: %s", e
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def requeue_spend_logs(
|
||||
prisma_client: PrismaClient,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
rows: Sequence[Mapping[str, object]],
|
||||
) -> None:
|
||||
"""Park rows from a failed or cancelled write in Redis, falling back to the head of the in-memory queue."""
|
||||
if await _park_spend_logs_in_redis(proxy_logging_obj, rows):
|
||||
return
|
||||
await enqueue_spend_logs(prisma_client, rows, at_head=True)
|
||||
|
||||
|
||||
async def recover_parked_spend_logs(
|
||||
prisma_client: PrismaClient,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
limit: int = REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT,
|
||||
) -> int:
|
||||
"""Move spend-log rows parked in Redis back to the head of the in-memory queue for the next write."""
|
||||
try:
|
||||
rows: Final = (
|
||||
await proxy_logging_obj.db_spend_update_writer.redis_update_buffer.get_spend_logs_from_redis_buffer(limit)
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # Redis being down must not stop the regular in-memory flush
|
||||
verbose_proxy_logger.warning("Spend tracking - could not read parked spend logs from Redis: %s", e)
|
||||
return 0
|
||||
if len(rows) == 0:
|
||||
return 0
|
||||
try:
|
||||
await enqueue_spend_logs(prisma_client, rows, at_head=True)
|
||||
except BaseException:
|
||||
await _park_spend_logs_in_redis(proxy_logging_obj, rows)
|
||||
raise
|
||||
verbose_proxy_logger.info("Spend tracking - recovered %d parked spend log rows from Redis", len(rows))
|
||||
return len(rows)
|
||||
|
||||
|
||||
async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int:
|
||||
"""Pending entries across every request-time spend queue, sized under each queue's
|
||||
lock. Every drain trigger reads this one owner, so a queue added later joins the
|
||||
|
|
@ -7215,14 +7263,19 @@ async def update_spend_logs_job(
|
|||
This job is triggered based on queue size rather than time.
|
||||
Pops the batch once, writes spend logs, then runs guardrail usage tracking.
|
||||
"""
|
||||
n_retry_times: Final = 3
|
||||
MAX_LOGS_PER_INTERVAL: Final = 10000
|
||||
|
||||
# Atomically pop batch from queue. The tool usage queue counts toward the
|
||||
# emptiness check: a spend-log write failure aborts a run before the tool
|
||||
# drain below, and those entries must not strand once the spend queue drains.
|
||||
if await _total_queued_spend_transactions(prisma_client) == 0:
|
||||
return
|
||||
async with prisma_client.spend_log_write_lock:
|
||||
await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj)
|
||||
|
||||
|
||||
async def _run_spend_logs_job(
|
||||
prisma_client: PrismaClient,
|
||||
db_writer_client: AsyncHTTPHandler | None,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
n_retry_times: Final = 3
|
||||
MAX_LOGS_PER_INTERVAL: Final = 10000
|
||||
|
||||
logs_to_process: Final = await dequeue_spend_logs(prisma_client, MAX_LOGS_PER_INTERVAL)
|
||||
|
||||
|
|
@ -7235,7 +7288,7 @@ async def update_spend_logs_job(
|
|||
logs_to_process=logs_to_process,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
|
||||
await requeue_spend_logs(prisma_client, proxy_logging_obj, logs_to_process)
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - spend log write cancelled, requeued %d rows for the next flush",
|
||||
len(logs_to_process),
|
||||
|
|
@ -7321,14 +7374,22 @@ async def drain_spend_logs_queue(
|
|||
await monitor_task
|
||||
prisma_client.spend_logs_queue_monitor_task = None # rebind-ok: the client owns its monitor handle
|
||||
|
||||
async with prisma_client.spend_log_write_lock:
|
||||
try:
|
||||
await _drain_spend_logs_queue_to_db(prisma_client, db_writer_client, proxy_logging_obj)
|
||||
finally:
|
||||
await _park_remaining_spend_logs(prisma_client, proxy_logging_obj)
|
||||
|
||||
|
||||
async def _drain_spend_logs_queue_to_db(
|
||||
prisma_client: PrismaClient,
|
||||
db_writer_client: "AsyncHTTPHandler | None",
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
for _ in range(MAX_SPEND_LOG_DRAIN_ITERATIONS):
|
||||
if await _total_queued_spend_transactions(prisma_client) == 0:
|
||||
return
|
||||
await update_spend_logs_job(
|
||||
prisma_client=prisma_client,
|
||||
db_writer_client=db_writer_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj)
|
||||
|
||||
remaining: Final = await _total_queued_spend_transactions(prisma_client)
|
||||
if remaining > 0:
|
||||
|
|
@ -7339,6 +7400,17 @@ async def drain_spend_logs_queue(
|
|||
)
|
||||
|
||||
|
||||
async def _park_remaining_spend_logs(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None:
|
||||
rows: Final = await dequeue_spend_logs(prisma_client, sys.maxsize)
|
||||
if len(rows) == 0 or await _park_spend_logs_in_redis(proxy_logging_obj, rows):
|
||||
return
|
||||
await enqueue_spend_logs(prisma_client, rows, at_head=True)
|
||||
spend_log_error(
|
||||
"Spend tracking - %d spend log rows could not be written or parked in Redis and will be lost on exit",
|
||||
len(rows),
|
||||
)
|
||||
|
||||
|
||||
async def _monitor_spend_logs_queue(
|
||||
prisma_client: PrismaClient,
|
||||
db_writer_client: AsyncHTTPHandler | None,
|
||||
|
|
@ -7372,6 +7444,7 @@ async def _monitor_spend_logs_queue(
|
|||
|
||||
while True:
|
||||
try:
|
||||
await recover_parked_spend_logs(prisma_client, proxy_logging_obj)
|
||||
# Check queue sizes with lock protection; the tool usage queue keeps
|
||||
# the monitor firing when a prior failed run left it nonempty.
|
||||
queue_size = await _total_queued_spend_transactions(prisma_client)
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ class MockPrismaClient:
|
|||
import asyncio
|
||||
|
||||
self._spend_log_transactions_lock = asyncio.Lock()
|
||||
self.spend_log_write_lock = asyncio.Lock()
|
||||
self._tool_usage_transactions_lock = asyncio.Lock()
|
||||
self._autorouter_turn_transactions_lock = asyncio.Lock()
|
||||
|
||||
|
|
|
|||
|
|
@ -1502,3 +1502,51 @@ async def test_async_set_cache_pipeline_with_ttls_keeps_each_entry_ttl(monkeypat
|
|||
("ns:u1", '{"user_id": "u1"}', timedelta(seconds=7)),
|
||||
("ns:org_id:o1", '{"a": 1}', timedelta(seconds=300)),
|
||||
]
|
||||
|
||||
|
||||
class _ListPipeline:
|
||||
def __init__(self, rows: list[str]) -> None:
|
||||
self.rows = rows
|
||||
self.queued: list[tuple[str, ...]] = []
|
||||
|
||||
async def __aenter__(self) -> "_ListPipeline":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
return None
|
||||
|
||||
def rpush(self, key: str, *values: str) -> None:
|
||||
self.queued.append(("rpush", key, *values))
|
||||
|
||||
def ltrim(self, key: str, start: int, end: int) -> None:
|
||||
self.queued.append(("ltrim", key, str(start), str(end)))
|
||||
|
||||
async def execute(self) -> list[object]:
|
||||
results: list[object] = []
|
||||
for op in self.queued:
|
||||
if op[0] == "rpush":
|
||||
self.rows.extend(op[2:])
|
||||
results.append(len(self.rows))
|
||||
else:
|
||||
start, end = int(op[2]), int(op[3])
|
||||
del self.rows[: max(len(self.rows) + start, 0) if start < 0 else start]
|
||||
results.append(True)
|
||||
return results
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_and_trim_runs_push_and_trim_in_one_transaction(monkeypatch, redis_no_ping):
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache(namespace="ns")
|
||||
rows = ["a", "b"]
|
||||
pipe = _ListPipeline(rows)
|
||||
client = MagicMock()
|
||||
client.pipeline = MagicMock(return_value=pipe)
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=client):
|
||||
pushed_len = await redis_cache.async_rpush_and_trim(key="buf", values=["c", "d"], max_len=3)
|
||||
|
||||
client.pipeline.assert_called_once_with(transaction=True)
|
||||
assert pushed_len == 4
|
||||
assert rows == ["b", "c", "d"]
|
||||
assert [op[:2] for op in pipe.queued] == [("rpush", "ns:buf"), ("ltrim", "ns:buf")]
|
||||
|
|
|
|||
|
|
@ -651,3 +651,53 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu
|
|||
restored = await window_queue.flush_and_get_aggregated_window_spend_transactions()
|
||||
assert [payload["spend"] for payload in restored] == [4.0]
|
||||
assert [payload["entity_id"] for payload in restored] == ["team-1"]
|
||||
|
||||
|
||||
class _ListRedis:
|
||||
def __init__(self) -> None:
|
||||
self.rows: list[str] = []
|
||||
|
||||
async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int:
|
||||
self.rows.extend(values)
|
||||
pushed_len = len(self.rows)
|
||||
del self.rows[:-max_len]
|
||||
return pushed_len
|
||||
|
||||
async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> list[str] | None:
|
||||
if not self.rows:
|
||||
return None
|
||||
popped = self.rows[:count]
|
||||
del self.rows[:count]
|
||||
return popped
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap():
|
||||
redis = _ListRedis()
|
||||
buffer = RedisUpdateBuffer(redis_cache=redis)
|
||||
buffer._should_commit_spend_updates_to_redis = MagicMock(return_value=True)
|
||||
|
||||
assert await buffer.store_spend_logs_in_redis([{"request_id": "old"}, {"request_id": "mid"}], max_rows=2) is True
|
||||
assert await buffer.store_spend_logs_in_redis([{"request_id": "new"}], max_rows=2) is True
|
||||
|
||||
parked = await buffer.get_spend_logs_from_redis_buffer(limit=10)
|
||||
assert [row["request_id"] for row in parked] == ["mid", "new"]
|
||||
assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_spend_logs_in_redis_reports_failure_without_redis():
|
||||
buffer = RedisUpdateBuffer(redis_cache=None)
|
||||
|
||||
assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False
|
||||
assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_enabled():
|
||||
redis = _ListRedis()
|
||||
buffer = RedisUpdateBuffer(redis_cache=redis)
|
||||
buffer._should_commit_spend_updates_to_redis = MagicMock(return_value=False)
|
||||
|
||||
assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False
|
||||
assert redis.rows == []
|
||||
|
|
|
|||
|
|
@ -130,6 +130,7 @@ def mock_prisma_client() -> MagicMock:
|
|||
client.spend_log_transactions = []
|
||||
client._spend_log_transactions_lock = asyncio.Lock()
|
||||
client.spend_logs_queue_monitor_task = None
|
||||
client.spend_log_write_lock = asyncio.Lock()
|
||||
client.tool_usage_transactions = []
|
||||
client._tool_usage_transactions_lock = asyncio.Lock()
|
||||
client.jsonify_object = lambda data: dict(data)
|
||||
|
|
@ -313,6 +314,54 @@ def make_spend_log_row() -> Callable[..., Dict[str, Any]]:
|
|||
return _make
|
||||
|
||||
|
||||
class FakeRedisList:
|
||||
def __init__(self) -> None:
|
||||
self.items: dict[str, list[str]] = {}
|
||||
self.down = False
|
||||
|
||||
def _check_up(self) -> None:
|
||||
if self.down:
|
||||
raise ConnectionError("redis unreachable")
|
||||
|
||||
async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int:
|
||||
self._check_up()
|
||||
stored = self.items.setdefault(key, [])
|
||||
stored.extend(str(v) for v in values)
|
||||
pushed_len = len(stored)
|
||||
del stored[:-max_len]
|
||||
return pushed_len
|
||||
|
||||
async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> str | list[str] | None:
|
||||
self._check_up()
|
||||
stored = self.items.get(key, [])
|
||||
if not stored:
|
||||
return None
|
||||
if count is None:
|
||||
return stored.pop(0)
|
||||
popped = stored[:count]
|
||||
del stored[:count]
|
||||
return popped
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_redis() -> FakeRedisList:
|
||||
return FakeRedisList()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proxy_logging_with_redis(fake_redis: FakeRedisList) -> MagicMock:
|
||||
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
|
||||
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
proxy_logging.db_spend_update_writer = MagicMock()
|
||||
proxy_logging.db_spend_update_writer.db_update_spend_transaction_handler = AsyncMock()
|
||||
buffer = RedisUpdateBuffer(redis_cache=fake_redis)
|
||||
buffer._should_commit_spend_updates_to_redis = MagicMock(return_value=True)
|
||||
proxy_logging.db_spend_update_writer.redis_update_buffer = buffer
|
||||
return proxy_logging
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SentMessage:
|
||||
from_addr: Optional[str]
|
||||
|
|
|
|||
|
|
@ -883,3 +883,37 @@ def test_disable_spend_updates_error_when_general_settings_unavailable(
|
|||
monkeypatch.delattr(proxy_server_mod, "general_settings", raising=False)
|
||||
with pytest.raises(ImportError):
|
||||
ProxyUpdateSpend.disable_spend_updates()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_parks_failed_batch_in_redis_with_wire_safe_datetimes(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any
|
||||
) -> None:
|
||||
"""Regression: a batch the DB rejected used to go back to process memory only. With Redis
|
||||
wired in it must be parked there, and datetimes must come back as ISO strings the DB write
|
||||
accepts, since the row is replayed by a process that never saw the original objects.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from prisma.errors import TableNotFoundError
|
||||
|
||||
started = datetime(2026, 9, 19, 20, 0, 5, 123000, tzinfo=timezone.utc)
|
||||
err = TableNotFoundError(
|
||||
{"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}}
|
||||
)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err)
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
|
||||
with pytest.raises(TableNotFoundError):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=2,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
logs_to_process=[make_spend_log_row(request_id="a", startTime=started)],
|
||||
)
|
||||
|
||||
buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer
|
||||
parked = await buffer.get_spend_logs_from_redis_buffer(limit=10)
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
assert [(row["request_id"], row["startTime"]) for row in parked] == [("a", started.isoformat())]
|
||||
|
|
|
|||
|
|
@ -11,17 +11,20 @@ Symbols pinned here:
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from contextlib import suppress
|
||||
from typing import Any, Dict, Final, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
|
||||
from litellm.proxy.utils import (
|
||||
MAX_SPEND_LOG_DRAIN_ITERATIONS,
|
||||
_monitor_spend_logs_queue,
|
||||
_raise_failed_update_spend_exception,
|
||||
drain_spend_logs_queue,
|
||||
recover_parked_spend_logs,
|
||||
update_daily_tag_spend,
|
||||
update_spend,
|
||||
update_spend_logs_job,
|
||||
|
|
@ -719,3 +722,222 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None:
|
|||
|
||||
with pytest.raises(ValueError, match="specific"):
|
||||
asyncio.run(_runner())
|
||||
|
||||
|
||||
def _table_gone_error() -> Exception:
|
||||
from prisma.errors import TableNotFoundError
|
||||
|
||||
return TableNotFoundError(
|
||||
{"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}}
|
||||
)
|
||||
|
||||
|
||||
def _parked_request_ids(fake_redis: Any) -> list[str]:
|
||||
return [json.loads(row)["request_id"] for row in fake_redis.items.get(REDIS_SPEND_LOGS_BUFFER_KEY, [])]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_spend_logs_queue_parks_unwritable_rows_in_redis_on_shutdown(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any
|
||||
) -> None:
|
||||
from prisma.errors import TableNotFoundError
|
||||
|
||||
mock_prisma_client.spend_log_transactions = [
|
||||
make_spend_log_row(request_id="r1"),
|
||||
make_spend_log_row(request_id="r2"),
|
||||
]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_table_gone_error())
|
||||
|
||||
with pytest.raises(TableNotFoundError):
|
||||
await drain_spend_logs_queue(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
assert sorted(_parked_request_ids(fake_redis)) == ["r1", "r2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_spend_logs_queue_waits_for_an_in_flight_write_before_parking(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any
|
||||
) -> None:
|
||||
db_outage_seen: Final = asyncio.Event()
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="in-flight")]
|
||||
|
||||
async def _fail_once_shutdown_starts(*args: Any, **kwargs: Any) -> None:
|
||||
await db_outage_seen.wait()
|
||||
raise _table_gone_error()
|
||||
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_fail_once_shutdown_starts)
|
||||
scheduler_write: Final = asyncio.ensure_future(
|
||||
update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
|
||||
async def _release_after_shutdown_started() -> None:
|
||||
await asyncio.sleep(0.05)
|
||||
db_outage_seen.set()
|
||||
|
||||
release: Final = asyncio.ensure_future(_release_after_shutdown_started())
|
||||
await drain_spend_logs_queue(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
|
||||
assert _parked_request_ids(fake_redis) == ["in-flight"]
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
await release
|
||||
with suppress(Exception):
|
||||
await scheduler_write
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_spend_logs_queue_parks_rows_left_after_max_passes(
|
||||
mock_prisma_client: Any,
|
||||
make_spend_log_row: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
proxy_logging_with_redis: MagicMock,
|
||||
fake_redis: Any,
|
||||
) -> None:
|
||||
import litellm.proxy.db.spend_log_tool_index as tool_mod
|
||||
import litellm.proxy.guardrails.usage_tracking as guard_mod
|
||||
|
||||
monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False)
|
||||
monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False)
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r0")]
|
||||
|
||||
async def _write_and_refill(*args: Any, **kwargs: Any) -> None:
|
||||
mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="late"))
|
||||
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_write_and_refill)
|
||||
|
||||
await drain_spend_logs_queue(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
assert _parked_request_ids(fake_redis) == ["late"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_spend_logs_queue_keeps_rows_in_memory_when_redis_is_down(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any
|
||||
) -> None:
|
||||
from prisma.errors import TableNotFoundError
|
||||
|
||||
fake_redis.down = True
|
||||
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_table_gone_error())
|
||||
|
||||
with pytest.raises(TableNotFoundError):
|
||||
await drain_spend_logs_queue(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
|
||||
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r1"]
|
||||
assert fake_redis.items == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_writes_rows_parked_in_redis_by_a_previous_pod(
|
||||
mock_prisma_client: Any,
|
||||
make_spend_log_row: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
proxy_logging_with_redis: MagicMock,
|
||||
fake_redis: Any,
|
||||
) -> None:
|
||||
import litellm.proxy.db.spend_log_tool_index as tool_mod
|
||||
import litellm.proxy.guardrails.usage_tracking as guard_mod
|
||||
|
||||
monkeypatch.setattr(guard_mod, "process_spend_logs_guardrail_usage", AsyncMock(), raising=False)
|
||||
monkeypatch.setattr(tool_mod, "flush_tool_usage_transactions", AsyncMock(), raising=False)
|
||||
buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer
|
||||
assert await buffer.store_spend_logs_in_redis([make_spend_log_row(request_id="parked")]) is True
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock()
|
||||
|
||||
await update_spend(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
|
||||
written = mock_prisma_client.db.litellm_spendlogs.create_many.await_args.kwargs["data"]
|
||||
assert [row["request_id"] for row in written] == ["parked"]
|
||||
assert _parked_request_ids(fake_redis) == []
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_parked_spend_logs_re_parks_rows_when_the_enqueue_is_cancelled(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any
|
||||
) -> None:
|
||||
buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer
|
||||
assert await buffer.store_spend_logs_in_redis([make_spend_log_row(request_id="parked")]) is True
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
await mock_prisma_client._spend_log_transactions_lock.acquire()
|
||||
recovery: Final = asyncio.ensure_future(
|
||||
recover_parked_spend_logs(prisma_client=mock_prisma_client, proxy_logging_obj=proxy_logging_with_redis)
|
||||
)
|
||||
await asyncio.sleep(0.01)
|
||||
assert _parked_request_ids(fake_redis) == []
|
||||
|
||||
recovery.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await recovery
|
||||
mock_prisma_client._spend_log_transactions_lock.release()
|
||||
|
||||
assert _parked_request_ids(fake_redis) == ["parked"]
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_monitor_spend_logs_queue_pulls_parked_rows_before_each_flush(
|
||||
mock_prisma_client: Any,
|
||||
make_spend_log_row: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
proxy_logging_with_redis: MagicMock,
|
||||
) -> None:
|
||||
import litellm.constants as constants_mod
|
||||
import litellm.proxy.utils as utils_mod
|
||||
|
||||
monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 0.0, raising=False)
|
||||
buffer = proxy_logging_with_redis.db_spend_update_writer.redis_update_buffer
|
||||
assert await buffer.store_spend_logs_in_redis([make_spend_log_row(request_id="parked")]) is True
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
seen: list[list[str]] = []
|
||||
polls = {"n": 0}
|
||||
|
||||
async def _fake_job(*args: Any, **kwargs: Any) -> None:
|
||||
seen.append([row["request_id"] for row in mock_prisma_client.spend_log_transactions])
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
async def _poll(*args: Any, **kwargs: Any) -> bool:
|
||||
polls["n"] += 1
|
||||
if polls["n"] >= 3:
|
||||
raise asyncio.CancelledError()
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job)
|
||||
monkeypatch.setattr(utils_mod, "_wait_for_spend_log_flush_request", _poll)
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await _monitor_spend_logs_queue(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging_with_redis,
|
||||
)
|
||||
|
||||
assert seen == [["parked"]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue