mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(accounting): keep delivering spend when another worker's maintenance lock is orphaned (#50774)
* fix(accounting): keep delivering spend when another worker's maintenance lock is orphaned Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): run worker regression tests in CI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(accounting): inject fake Redis and colocate worker coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Mateo Wang
GitHub
parent
a56a27deb6
commit
ac7eb4e6bf
2 changed files with 130 additions and 32 deletions
|
|
@ -46,9 +46,6 @@ def _progressed(result: DeliveryBatchResult) -> bool:
|
|||
async def flush_accounting_deliveries(
|
||||
writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging"
|
||||
) -> None:
|
||||
acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False)
|
||||
if acquired is False:
|
||||
return
|
||||
applied: tuple[Contribution, ...] = () # rebind-ok: keep committed contributions for cache invalidation
|
||||
|
||||
async def apply(transaction: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None:
|
||||
|
|
@ -57,36 +54,39 @@ async def flush_accounting_deliveries(
|
|||
transaction, deliveries, prisma_client=prisma_client, writer=writer, proxy_logging_obj=proxy_logging_obj
|
||||
)
|
||||
|
||||
database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction
|
||||
LedgerDatabase, prisma_client
|
||||
)
|
||||
repository: Final = AccountingDeliveryRepository(database)
|
||||
started: Final = time.monotonic()
|
||||
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500"))
|
||||
|
||||
async def run_batch(*, retry_only: bool) -> DeliveryBatchResult:
|
||||
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
|
||||
if result.applied:
|
||||
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
|
||||
if result.parked:
|
||||
verbose_proxy_logger.error(
|
||||
"Accounting delivery parked %d requests after isolated failures", len(result.parked)
|
||||
)
|
||||
await _alert(
|
||||
proxy_logging_obj,
|
||||
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
|
||||
)
|
||||
return result
|
||||
|
||||
fresh_deadline: Final = time.monotonic() + 5
|
||||
for _ in range(batch_size):
|
||||
if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline:
|
||||
break
|
||||
retry_deadline: Final = time.monotonic() + 5
|
||||
for _ in range(batch_size):
|
||||
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
|
||||
break
|
||||
acquired: Final = await writer.pod_lock_manager.acquire_lock(cronjob_id=_JOB, allow_reentrant=False)
|
||||
if acquired is False:
|
||||
return
|
||||
try:
|
||||
database: Final = cast( # cast-ok: PrismaClient.tx owns writer routing and yields the generated transaction
|
||||
LedgerDatabase, prisma_client
|
||||
)
|
||||
repository: Final = AccountingDeliveryRepository(database)
|
||||
started: Final = time.monotonic()
|
||||
batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500"))
|
||||
|
||||
async def run_batch(*, retry_only: bool) -> DeliveryBatchResult:
|
||||
result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only)
|
||||
if result.applied:
|
||||
await writer.invalidate_accounting_entities(applied, proxy_logging_obj)
|
||||
if result.parked:
|
||||
verbose_proxy_logger.error(
|
||||
"Accounting delivery parked %d requests after isolated failures", len(result.parked)
|
||||
)
|
||||
await _alert(
|
||||
proxy_logging_obj,
|
||||
message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.",
|
||||
)
|
||||
return result
|
||||
|
||||
fresh_deadline: Final = time.monotonic() + 5
|
||||
for _ in range(batch_size):
|
||||
if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline:
|
||||
break
|
||||
retry_deadline: Final = time.monotonic() + 5
|
||||
for _ in range(batch_size):
|
||||
if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)):
|
||||
break
|
||||
retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1"))
|
||||
cleanup_deadline: Final = time.monotonic() + 5
|
||||
for _ in range(batch_size):
|
||||
|
|
|
|||
98
tests/unit/accounting/test_worker.py
Normal file
98
tests/unit/accounting/test_worker.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
import json
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fakeredis import FakeRedis as FakeSyncRedis
|
||||
from fakeredis import FakeServer
|
||||
from fakeredis.aioredis import FakeRedis
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm import in_memory_llm_clients_cache
|
||||
from litellm.accounting.worker import flush_accounting_deliveries
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
||||
_FIELDS: Final = TypeAdapter(dict[str, object])
|
||||
_INCREMENTS: Final = TypeAdapter(dict[str, dict[str, float]])
|
||||
_PAYLOAD: Final = json.dumps(
|
||||
{"request_id": "completed-create", "spend": 0.25, "metadata": {"accounting_entities": {"api_key": "creator"}}}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Ledger:
|
||||
pending: bool
|
||||
unavailable: bool
|
||||
spend: float = 0.0
|
||||
attempts: int = 0
|
||||
|
||||
async def query(self, query: str, *args: object) -> tuple[dict[str, object], ...]:
|
||||
if query.startswith("SELECT delivery_id, attempts"):
|
||||
if not self.pending or self.attempts or args[1] is True:
|
||||
return ()
|
||||
return ({"delivery_id": "delivery", "attempts": self.attempts, "accounting_payload": _PAYLOAD},)
|
||||
if "SET attempts = attempts + 1" in query:
|
||||
self.attempts += 1
|
||||
return ({"delivery_id": "delivery", "attempts": self.attempts},)
|
||||
if "pending_count" in query:
|
||||
return ({"pending_count": int(self.pending), "parked_count": 0, "oldest_pending_seconds": 0},)
|
||||
return ({"batch_count": 0, "table_bytes": 0},)
|
||||
|
||||
async def execute(self, query: str, *args: object) -> int:
|
||||
if "SET state = 'applied'" in query:
|
||||
self.pending = False
|
||||
return 0
|
||||
|
||||
def update_key(self, *, where: object, data: object) -> None:
|
||||
if self.unavailable:
|
||||
raise OSError("database connection lost before the update")
|
||||
assert _FIELDS.validate_python(where) == {"token": "creator"}
|
||||
self.spend += _INCREMENTS.validate_python({"spend": _FIELDS.validate_python(data)["spend"]})["spend"][
|
||||
"increment"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("locked", [False, True], ids=["without-redis", "orphaned-redis-lock"])
|
||||
@pytest.mark.parametrize("scenario", ["charge", "empty", "unavailable"])
|
||||
async def test_delivery_charges_the_creating_key_despite_an_orphaned_maintenance_lock(
|
||||
locked: bool, scenario: Literal["charge", "empty", "unavailable"]
|
||||
) -> None:
|
||||
ledger: Final = _Ledger(pending=scenario != "empty", unavailable=scenario == "unavailable")
|
||||
database: Final = MagicMock(spec=PrismaClient)
|
||||
transaction: Final = MagicMock()
|
||||
transaction.query_raw = AsyncMock(side_effect=ledger.query)
|
||||
transaction.execute_raw = AsyncMock(side_effect=ledger.execute)
|
||||
batcher: Final = MagicMock()
|
||||
batcher.litellm_verificationtoken.update_many.side_effect = ledger.update_key
|
||||
transaction.batch_.return_value.__aenter__ = AsyncMock(return_value=batcher)
|
||||
transaction.batch_.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
database.tx.return_value.__aenter__ = AsyncMock(return_value=transaction)
|
||||
database.tx.return_value.__aexit__ = AsyncMock(return_value=False)
|
||||
writer: Final = DBSpendUpdateWriter(accounting_clock=lambda: 0.0)
|
||||
logger: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache())
|
||||
server: Final = FakeServer()
|
||||
redis: Final = FakeRedis(server=server)
|
||||
sync_redis: Final = FakeSyncRedis(server=server)
|
||||
cache: Final = RedisCache(host="accounting-unit.example", connection_pool=sync_redis.connection_pool)
|
||||
cache_key: Final = cache._get_async_client_cache_key()
|
||||
in_memory_llm_clients_cache.set_cache(key=cache_key, value=redis)
|
||||
await redis.set("cronjob_lock:accounting_delivery", json.dumps("terminated-worker"), ex=60)
|
||||
try:
|
||||
if locked:
|
||||
writer.pod_lock_manager.redis_cache = cache
|
||||
await flush_accounting_deliveries(writer, database, logger)
|
||||
await flush_accounting_deliveries(writer, database, logger)
|
||||
|
||||
assert ledger.spend == (0.25 if scenario == "charge" else 0.0)
|
||||
assert ledger.pending is (scenario == "unavailable")
|
||||
assert ledger.attempts == (1 if scenario == "unavailable" else 0)
|
||||
assert await redis.get("cronjob_lock:accounting_delivery") == json.dumps("terminated-worker").encode()
|
||||
finally:
|
||||
in_memory_llm_clients_cache.delete_cache(key=cache_key)
|
||||
await redis.aclose()
|
||||
sync_redis.close()
|
||||
Loading…
Add table
Reference in a new issue