diff --git a/litellm/accounting/metrics.py b/litellm/accounting/metrics.py index d4650813549..ee4b8b1c2b5 100644 --- a/litellm/accounting/metrics.py +++ b/litellm/accounting/metrics.py @@ -1,7 +1,9 @@ +import time +from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Final -from litellm.repositories.accounting_delivery_repository import DeliveryMetrics +from litellm.repositories.accounting_delivery_repository import DeliveryBacklog, DeliveryStorage if TYPE_CHECKING: from prometheus_client import Gauge, Histogram @@ -44,12 +46,32 @@ def _instruments() -> _Instruments | None: _INSTRUMENTS: Final = _instruments() -def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None: +def record_backlog(backlog: DeliveryBacklog, seconds: float) -> None: if _INSTRUMENTS is None: return - _INSTRUMENTS.pending.set(metrics.pending_count) - _INSTRUMENTS.parked.set(metrics.parked_count) - _INSTRUMENTS.oldest.set(max(0.0, metrics.oldest_pending_seconds)) + _INSTRUMENTS.pending.set(backlog.pending_count) + _INSTRUMENTS.parked.set(backlog.parked_count) + _INSTRUMENTS.oldest.set(max(0.0, backlog.oldest_pending_seconds)) _INSTRUMENTS.batch_time.observe(seconds) - _INSTRUMENTS.batch_count.set(metrics.batch_count) - _INSTRUMENTS.table_bytes.set(metrics.table_bytes) + + +def record_storage(storage: DeliveryStorage) -> None: + if _INSTRUMENTS is None: + return + _INSTRUMENTS.batch_count.set(storage.batch_count) + _INSTRUMENTS.table_bytes.set(storage.table_bytes) + + +class Cadence: + __slots__ = ("_clock", "_last") + + def __init__(self, clock: Callable[[], float] = time.monotonic) -> None: + self._clock = clock + self._last: float | None = None + + def due(self, interval_seconds: float) -> bool: + now: Final = self._clock() + if self._last is not None and now - self._last < interval_seconds: + return False + self._last = now + return True diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index be18d8ac6b1..0f7e80123c2 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -10,7 +10,7 @@ from typing import ( from litellm._logging import verbose_proxy_logger from litellm.accounting.application import apply_deliveries -from litellm.accounting.metrics import record_deliveries +from litellm.accounting.metrics import record_backlog, record_storage from litellm.accounting.projection import Contribution from litellm.repositories.accounting_delivery_repository import ( AccountingDeliveryRepository, @@ -96,24 +96,37 @@ async def flush_accounting_deliveries( or time.monotonic() >= cleanup_deadline ): break - metrics: Final = await repository.metrics() - record_deliveries(metrics, time.monotonic() - started) - if metrics.batch_count >= int( - os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000") - ) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")): - await _alert( - proxy_logging_obj, - message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.", - ) - if metrics.oldest_pending_seconds > 300: - await _alert( - proxy_logging_obj, - message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.", - ) + await _report(writer, repository, proxy_logging_obj, time.monotonic() - started) finally: if acquired: await writer.pod_lock_manager.release_lock(cronjob_id=_JOB) +async def _report( + writer: "DBSpendUpdateWriter", + repository: AccountingDeliveryRepository, + proxy_logging_obj: _AccountingAlerts, + seconds: float, +) -> None: + backlog: Final = await repository.backlog() + record_backlog(backlog, seconds) + interval: Final = float(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_CHECK_INTERVAL_SECONDS", "900")) + if writer.accounting_storage_checks.due(interval): + storage: Final = await repository.storage() + record_storage(storage) + if storage.batch_count >= int( + os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000") + ) or storage.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")): + await _alert( + proxy_logging_obj, + message=f"Accounting delivery storage requires review. Retained batches={storage.batch_count}, table bytes={storage.table_bytes}.", + ) + if backlog.oldest_pending_seconds > 300 and writer.accounting_delay_alerts.due(interval): + await _alert( + proxy_logging_obj, + message=f"Accounting deliveries are delayed. Pending={backlog.pending_count}, oldest={backlog.oldest_pending_seconds:.0f}s.", + ) + + async def _alert(alerts: _AccountingAlerts, message: str) -> None: await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ff57596232f..4a1871ff1f5 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -28,6 +28,7 @@ from typing_extensions import LiteralString, ReadOnly, TypedDict import litellm from litellm._internal_context import service_target, with_service_target from litellm._logging import verbose_proxy_logger +from litellm.accounting.metrics import Cadence from litellm.caching import RedisCache from litellm.constants import ( BILLED_RETRY_ATTEMPT_METADATA_KEY, @@ -682,8 +683,11 @@ class DBSpendUpdateWriter: def __init__( self, redis_cache: RedisCache | None = None, + accounting_clock: Callable[[], float] = time.monotonic, ): self.redis_cache = redis_cache + self.accounting_storage_checks = Cadence(accounting_clock) + self.accounting_delay_alerts = Cadence(accounting_clock) self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_lock_manager = PodLockManager() self.spend_update_queue = SpendUpdateQueue() diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index d5bbe891089..624c2b8851a 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -60,12 +60,17 @@ class DeliveryBatchResult: error: str | None = None -class DeliveryMetrics(BaseModel): +class DeliveryBacklog(BaseModel): model_config = ConfigDict(frozen=True) pending_count: int parked_count: int oldest_pending_seconds: float + + +class DeliveryStorage(BaseModel): + model_config = ConfigDict(frozen=True) + batch_count: int table_bytes: int @@ -88,6 +93,14 @@ SET attempts = attempts + 1, last_error = $3, state = $4::"AccountingDeliveryState" WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer RETURNING delivery_id, attempts""" +_BACKLOG: Final = """SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, + COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, + COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) + AS oldest_pending_seconds +FROM "LiteLLM_SpendLogDelivery" +WHERE state IN ('pending', 'parked')""" +_STORAGE: Final = """SELECT (SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE is_batch) AS batch_count, + pg_total_relation_size('"LiteLLM_SpendLogDelivery"') AS table_bytes""" _UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery" SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload' WHERE delivery_id = ANY($1::text[]) AND state = 'pending'""" @@ -114,17 +127,13 @@ class AccountingDeliveryRepository: def __init__(self, database: LedgerDatabase) -> None: self._database = database - async def metrics(self) -> DeliveryMetrics: + async def backlog(self) -> DeliveryBacklog: async with self._database.tx(timeout=_TIMEOUT) as transaction: - rows: Final = await transaction.query_raw( - "SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, " - "COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, " - "COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) " - "AS oldest_pending_seconds, COUNT(*) FILTER (WHERE is_batch) AS batch_count, " - "pg_total_relation_size('\"LiteLLM_SpendLogDelivery\"') AS table_bytes " - 'FROM "LiteLLM_SpendLogDelivery"' - ) - return DeliveryMetrics.model_validate(rows[0]) + return DeliveryBacklog.model_validate((await transaction.query_raw(_BACKLOG))[0]) + + async def storage(self) -> DeliveryStorage: + async with self._database.tx(timeout=_TIMEOUT) as transaction: + return DeliveryStorage.model_validate((await transaction.query_raw(_STORAGE))[0]) async def cleanup_applied(self, *, retention_days: int = 1, batch_size: int = 1000) -> int: if retention_days < 1 or batch_size < 1: diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py index d1227e85565..2b0be742cd4 100644 --- a/tests/integration/database/test_spendlogs_delivery.py +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -373,15 +373,45 @@ async def test_receipt_cleanup_keeps_pending_and_parked_rows_and_reports_lag( ) repository = AccountingDeliveryRepository(delivery_client) assert await repository.cleanup_applied() == 1 - metrics = await repository.metrics() - assert metrics.pending_count == 1 - assert metrics.parked_count == 1 - assert metrics.oldest_pending_seconds >= 40 * 86400 + backlog = await repository.backlog() + assert backlog.pending_count == 1 + assert backlog.parked_count == 1 + assert backlog.oldest_pending_seconds >= 40 * 86400 assert delivery_database.execute( 'SELECT request_id FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' ).fetchall() == [("parked",), ("pending",)] +@pytest.mark.asyncio +async def test_backlog_reads_only_unapplied_deliveries_however_many_receipts_are_retained( + delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma +) -> None: + from litellm.repositories.accounting_delivery_repository import _BACKLOG + + delivery_database.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'applied-' || n, 'applied-' || n, jsonb_build_object('payload', repeat(md5(n::text), 32)), " + "n % 10 = 0, 'applied', CURRENT_TIMESTAMP FROM generate_series(1, 20000) AS n" + ) + _insert(delivery_database, tuple(_row(state) for state in ("pending", "parked"))) + delivery_database.execute( + 'UPDATE "LiteLLM_SpendLogDelivery" SET state = request_id::"AccountingDeliveryState" ' + "WHERE request_id IN ('pending', 'parked')" + ) + delivery_database.execute('ANALYZE "LiteLLM_SpendLogDelivery"') + explained = delivery_database.execute(f"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {_BACKLOG}").fetchone() + assert explained is not None + plan = explained[0][0]["Plan"] + table_pages = delivery_database.execute( + "SELECT relpages FROM pg_class WHERE oid = '\"LiteLLM_SpendLogDelivery\"'::regclass" + ).fetchone() + assert table_pages is not None and table_pages[0] > 1000 + assert plan["Shared Hit Blocks"] + plan["Shared Read Blocks"] < table_pages[0] / 20 + + backlog = await AccountingDeliveryRepository(delivery_client).backlog() + assert (backlog.pending_count, backlog.parked_count) == (1, 1) + + @pytest.mark.parametrize("delivery_database", ["public", "ledger_tenant"], indirect=True) @pytest.mark.parametrize("state", [None, "pending", "parked", "applied"]) def test_delivery_migration_replay_preserves_ledger_receipts(delivery_database, state): @@ -467,9 +497,9 @@ async def test_retention_keeps_batch_receipts_indefinitely(delivery_database, de assert await repository.cleanup_applied() == 1 assert await repository.cleanup_applied() == 0 assert delivery_database.execute('SELECT request_id FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("batch",)] - metrics = await repository.metrics() - assert metrics.batch_count == 1 - assert metrics.table_bytes > 0 + storage = await repository.storage() + assert storage.batch_count == 1 + assert storage.table_bytes > 0 @pytest.mark.parametrize("disabled", [False, True]) diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index 60f3c074cca..ef87712b788 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -127,6 +127,66 @@ async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_on await client.db.disconnect() +@pytest.mark.asyncio +async def test_worker_checks_storage_and_repeats_delay_alerts_once_per_interval( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + class _Alerts: + def __init__(self) -> None: + self.messages: list[str] = [] # mutable-ok: alert sink the worker appends to + + async def alerting_handler(self, message: str, level: str, alert_type: object) -> None: + self.messages.append(message) + + class _Clock: + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1") + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" ' + "(delivery_id, request_id, accounting_payload, is_batch, attempts, created_at, next_attempt_at) " + "VALUES ('stuck', 'stuck', '{}'::jsonb, false, 1, CURRENT_TIMESTAMP - INTERVAL '1 hour', " + "CURRENT_TIMESTAMP + INTERVAL '1 hour')" + ) + client: Final = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + clock: Final = _Clock() + writer: Final = DBSpendUpdateWriter(accounting_clock=clock) + alerts: Final = _Alerts() + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_CHECK_INTERVAL_SECONDS", "100") + + async def flush_at(moment: float) -> tuple[str, ...]: + clock.now = moment + before: Final = len(alerts.messages) + await flush_accounting_deliveries(writer, client, alerts) + return tuple(message.split(",")[0] for message in alerts.messages[before:]) + + await client.db.connect() + try: + first: Final = await flush_at(0.0) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch, state) ' + "VALUES ('batch', 'batch', '{}'::jsonb, true, 'applied')" + ) + within: Final = (await flush_at(10.0), await flush_at(99.0)) + after: Final = await flush_at(100.0) + finally: + await client.db.disconnect() + + delayed: Final = "Accounting deliveries are delayed. Pending=1" + assert first == ("Accounting delivery storage requires review. Retained batches=0", delayed) + assert within == ((), ()) + assert after == ("Accounting delivery storage requires review. Retained batches=1", delayed) + + @pytest.mark.parametrize("ledger", [False, True]) @pytest.mark.parametrize("exclude_messages", [False, True]) @pytest.mark.asyncio diff --git a/tests/unit/accounting/test_metrics.py b/tests/unit/accounting/test_metrics.py index 054ceaba967..d05d860bd66 100644 --- a/tests/unit/accounting/test_metrics.py +++ b/tests/unit/accounting/test_metrics.py @@ -5,12 +5,11 @@ import pytest from prometheus_client.metrics import MetricWrapperBase from litellm.accounting import metrics -from litellm.accounting.metrics import record_deliveries -from litellm.repositories.accounting_delivery_repository import DeliveryMetrics +from litellm.accounting.metrics import Cadence, record_backlog, record_storage +from litellm.repositories.accounting_delivery_repository import DeliveryBacklog, DeliveryStorage -_DELIVERY_METRICS = DeliveryMetrics( - pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192 -) +_BACKLOG = DeliveryBacklog(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0) +_STORAGE = DeliveryStorage(batch_count=7, table_bytes=8192) def _sample(metric: MetricWrapperBase, suffix: str = "") -> float: @@ -23,7 +22,8 @@ def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None: assert instruments is not None batches_before = _sample(instruments.batch_time, "_count") - record_deliveries(_DELIVERY_METRICS, 0.5) + record_backlog(_BACKLOG, 0.5) + record_storage(_STORAGE) assert ( _sample(instruments.pending), @@ -46,4 +46,17 @@ def test_delivery_metrics_are_skipped_when_prometheus_client_is_not_installed( spec.loader.exec_module(module) - assert module.record_deliveries(_DELIVERY_METRICS, 0.5) is None + assert module.record_backlog(_BACKLOG, 0.5) is None + assert module.record_storage(_STORAGE) is None + + +def test_cadence_is_due_once_per_interval() -> None: + now = [100.0] + cadence = Cadence(lambda: now[0]) + + observed = [] + for moment in (100.0, 101.0, 999.9, 1000.0, 1000.0, 1899.0, 1900.0): + now[0] = moment + observed.append(cadence.due(900)) + + assert observed == [True, False, False, True, False, False, True] diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index cb903fb134a..25c62034f56 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -25,6 +25,7 @@ from fakeredis import TcpFakeServer from fastapi import HTTPException from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.route_llm_request import ( @@ -2545,11 +2546,38 @@ def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, obj } +async def _settle_router_accounting() -> None: + try: + await GLOBAL_LOGGING_WORKER.flush() + finally: + await GLOBAL_LOGGING_WORKER.stop() + while pending := {task for task in asyncio.all_tasks() - {asyncio.current_task()} if not _is_periodic_sync(task)}: + _, still_running = await asyncio.wait(pending, timeout=5) + assert not still_running, still_running + + +def _is_periodic_sync(task: asyncio.Task[object]) -> bool: + return task.get_coro().__qualname__.endswith(".periodic_sync_in_memory_spend_with_redis") + + async def _user_config_request_outcome( redis_port: int | None, routing_settings: Mapping[str, object], deployment_limits: Mapping[str, int], logging_obj: LiteLLMLoggingObj | None, +) -> str: + outcome: Final = await _unsettled_user_config_request_outcome( + redis_port, routing_settings, deployment_limits, logging_obj + ) + await _settle_router_accounting() + return outcome + + +async def _unsettled_user_config_request_outcome( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + logging_obj: LiteLLMLoggingObj | None, ) -> str: data: Final = { "model": "rate-limited", @@ -2566,8 +2594,6 @@ async def _user_config_request_outcome( await (await route_request(data, None, None, "acompletion")) except (litellm.RateLimitError, RouterRateLimitError): return "429" - finally: - await asyncio.sleep(0.3) return "200" @@ -2606,7 +2632,7 @@ async def _proxy_router_outcomes( except (litellm.RateLimitError, RouterRateLimitError): return "429" finally: - await asyncio.sleep(0.3) + await _settle_router_accounting() return "200" try: @@ -2700,16 +2726,21 @@ async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_l tpm_for_three_requests: Final = await _tpm_for_three_requests() def request() -> Coroutine[None, None, str]: - return _user_config_request_outcome( + return _unsettled_user_config_request_outcome( fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"tpm": tpm_for_three_requests}, _proxy_logging_obj() if arrives_with_logging_obj else None, ) + async def settled(*requests: Coroutine[None, None, str]) -> list[str]: + outcomes: Final = await asyncio.gather(*requests) + await _settle_router_accounting() + return outcomes + await _wait_out_a_rate_limit_minute_rollover() - concurrent_outcomes: Final = await asyncio.gather(request(), request()) - later_outcomes: Final = [await request() for _ in range(2)] + concurrent_outcomes: Final = await settled(request(), request()) + later_outcomes: Final = [*(await settled(request())), *(await settled(request()))] assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] @@ -2717,12 +2748,13 @@ async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_l @pytest.mark.asyncio async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: def request() -> Coroutine[None, None, str]: - return _user_config_request_outcome( + return _unsettled_user_config_request_outcome( fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj() ) await _wait_out_a_rate_limit_minute_rollover() outcomes: Final = await asyncio.gather(*(request() for _ in range(4))) + await _settle_router_accounting() assert outcomes.count("200") <= 2, outcomes assert "429" in outcomes, outcomes @@ -2747,13 +2779,14 @@ async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_ async def outcomes() -> list[str]: first_instance: Final = [await request() for _ in range(2)] user_config_routers.clear() + await _settle_router_accounting() second_instance: Final = [await request() for _ in range(3)] return [*first_instance, *second_instance] await _wait_out_a_rate_limit_minute_rollover() alone: Final = await outcomes() user_config_routers.clear() - await asyncio.sleep(0.3) + await _settle_router_accounting() redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall() proxy_router: Final = litellm.Router( routing_strategy=routing_strategy, @@ -2781,7 +2814,7 @@ async def test_user_config_requests_stay_under_the_proxy_provider_budget( ) try: await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}]) - await asyncio.sleep(0.3) + await _settle_router_accounting() with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj()) finally: