mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): flush the daily tag spend queue on shutdown
The shutdown flush committed every in-memory spend queue except daily tag spend, which only update_daily_tag_spend_job drains. That job runs at DAILY_TAG_SPEND_BATCH_MULTIPLIER times the batch interval, so every worker exit (deploy, SIGTERM, max_requests_before_restart) dropped the tag rows buffered since its last run and LiteLLM_DailyTagSpend undercounted LiteLLM_SpendLogs. Shutdown now runs the tag job once more, in both flush passes around the logging drain, before Prisma disconnects
This commit is contained in:
parent
dddbaa803d
commit
018ad5592b
4 changed files with 105 additions and 7 deletions
|
|
@ -1900,7 +1900,7 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# NOTE: Daily tag spend is committed by a separate scheduler job.
|
||||
# NOTE: Daily tag spend is committed by update_daily_tag_spend, on its own interval and at shutdown.
|
||||
|
||||
################## Daily End-User Spend Update Transactions ##################
|
||||
# Aggregate all in memory daily end-user spend transactions and commit to db
|
||||
|
|
|
|||
|
|
@ -953,6 +953,7 @@ from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports
|
|||
migrate_passwords_to_scrypt_async,
|
||||
model_dump_with_preserved_fields,
|
||||
prefetch_config_params,
|
||||
update_daily_tag_spend,
|
||||
update_spend,
|
||||
)
|
||||
from litellm.proxy.video_endpoints.endpoints import router as video_router
|
||||
|
|
@ -1237,6 +1238,7 @@ async def flush_spend_counters_on_shutdown() -> None:
|
|||
)
|
||||
except Exception as e: # noqa: BLE001 # shutdown must continue even if the commit fails
|
||||
verbose_proxy_logger.exception("Error flushing spend counters on shutdown: %s", e)
|
||||
await update_daily_tag_spend(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
|
||||
|
||||
|
||||
async def _flush_spend_logs_queue_on_shutdown() -> None:
|
||||
|
|
@ -11312,8 +11314,6 @@ class ProxyStartupEvent:
|
|||
### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ###
|
||||
## Reduces QPS as there are more tags for a single request
|
||||
tag_spend_update_interval: Final = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER)
|
||||
from litellm.proxy.utils import update_daily_tag_spend
|
||||
|
||||
scheduler.add_job(
|
||||
update_daily_tag_spend,
|
||||
"interval",
|
||||
|
|
|
|||
|
|
@ -8328,9 +8328,10 @@ async def update_daily_tag_spend(
|
|||
Separate scheduler job to commit daily tag spend updates.
|
||||
|
||||
Runs at a longer interval (2.3x default) than the main update_spend job
|
||||
to reduce query contention for DailyTagSpend table.
|
||||
to reduce query contention for DailyTagSpend table. The proxy shutdown
|
||||
flush also calls it once, so the tags buffered since the last tick are kept.
|
||||
|
||||
This is called by a dedicated scheduler job and does NOT process:
|
||||
It does NOT process:
|
||||
- Regular spend updates (user, key, team, org)
|
||||
- End-user spend
|
||||
- Agent spend
|
||||
|
|
|
|||
|
|
@ -29,10 +29,13 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
|
||||
import pytest
|
||||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
|
||||
from litellm.proxy._types import BaseDailySpendTransaction, DailyTagSpendTransaction
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
from litellm.proxy.proxy_server import (
|
||||
ProxyStartupEvent,
|
||||
initialize_shared_aiohttp_session,
|
||||
|
|
@ -54,6 +57,8 @@ from litellm.proxy.management_helpers.audit_logs import track_audit_task
|
|||
|
||||
from .conftest import normalize
|
||||
|
||||
_TAG_BATCH_ADAPTER: Final = TypeAdapter(dict[str, DailyTagSpendTransaction])
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# cleanup_router_config_variables
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -412,6 +417,96 @@ async def test_flush_spend_counters_on_shutdown_logs_and_swallows_commit_errors(
|
|||
assert "Error flushing spend counters on shutdown: db gone" in caplog.text
|
||||
|
||||
|
||||
def _daily_tag_spend(tag: str) -> dict[str, BaseDailySpendTransaction]:
|
||||
return {
|
||||
f"{tag}_2026-10-26_key-1_gpt-4o_openai": DailyTagSpendTransaction(
|
||||
tag=tag,
|
||||
date="2026-10-26",
|
||||
api_key="key-1",
|
||||
model="gpt-4o",
|
||||
model_group=None,
|
||||
custom_llm_provider="openai",
|
||||
mcp_namespaced_tool_name="",
|
||||
endpoint="",
|
||||
prompt_tokens=10,
|
||||
completion_tokens=5,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
compression_saved_tokens=0,
|
||||
compression_savings_spend=0.0,
|
||||
prompt_caching_savings_spend=0.0,
|
||||
gateway_injected_caching_savings_spend=0.0,
|
||||
spend=0.01,
|
||||
api_requests=1,
|
||||
successful_requests=1,
|
||||
failed_requests=0,
|
||||
request_id=None,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_spend_counters_on_shutdown_commits_the_daily_tag_spend_buffered_since_the_last_tag_job_run(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""The tag job runs every batch interval x DAILY_TAG_SPEND_BATCH_MULTIPLIER, so a worker that exits
|
||||
between runs must still write the tag rows it buffered or every tag rollup undercounts LiteLLM_SpendLogs."""
|
||||
writer: Final = DBSpendUpdateWriter()
|
||||
await writer.daily_tag_spend_update_queue.add_update(_daily_tag_spend("prod-tag"))
|
||||
fake_prisma: Final = MagicMock()
|
||||
execute_raw: Final = AsyncMock(return_value=1)
|
||||
fake_prisma.db.execute_raw = execute_raw
|
||||
fake_prisma.db.tx.return_value.__aenter__.return_value.execute_raw = execute_raw
|
||||
monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False)
|
||||
monkeypatch.setattr(ps, "general_settings", {}, raising=False)
|
||||
monkeypatch.setattr(ps.proxy_logging_obj, "db_spend_update_writer", writer)
|
||||
monkeypatch.setattr(writer, "db_update_spend_transaction_handler", AsyncMock())
|
||||
|
||||
await ps.flush_spend_counters_on_shutdown()
|
||||
|
||||
tag_upserts: Final = [
|
||||
call.args for call in execute_raw.await_args_list if '"LiteLLM_DailyTagSpend"' in str(call.args[0])
|
||||
]
|
||||
observed: Final = {
|
||||
"tag_upserts": len(tag_upserts),
|
||||
"tag_written": "prod-tag" in tag_upserts[0][1:] if tag_upserts else False,
|
||||
"left_in_memory": writer.daily_tag_spend_update_queue.update_queue.qsize(),
|
||||
}
|
||||
assert observed == {"tag_upserts": 1, "tag_written": True, "left_in_memory": 0}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_spend_counters_on_shutdown_hands_daily_tag_spend_to_redis_when_another_pod_holds_the_lock(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""With the Redis transaction buffer, a pod that is not the tag-commit leader keeps its tag spend by pushing
|
||||
it to the shared buffer on shutdown, where the leader's next tag job commits it."""
|
||||
async_rpush: Final = AsyncMock(return_value=1)
|
||||
redis_cache: Final = MagicMock()
|
||||
redis_cache.async_rpush = async_rpush
|
||||
writer: Final = DBSpendUpdateWriter(redis_cache=redis_cache)
|
||||
await writer.daily_tag_spend_update_queue.add_update(_daily_tag_spend("prod-tag"))
|
||||
monkeypatch.setattr(ps, "prisma_client", MagicMock(), raising=False)
|
||||
monkeypatch.setattr(ps, "general_settings", {"use_redis_transaction_buffer": True}, raising=False)
|
||||
monkeypatch.setattr(ps.proxy_logging_obj, "db_spend_update_writer", writer)
|
||||
monkeypatch.setattr(writer, "db_update_spend_transaction_handler", AsyncMock())
|
||||
monkeypatch.setattr(writer.pod_lock_manager, "acquire_lock", AsyncMock(return_value=False))
|
||||
|
||||
await ps.flush_spend_counters_on_shutdown()
|
||||
|
||||
pushed_batches: Final = [
|
||||
_TAG_BATCH_ADAPTER.validate_json(value)
|
||||
for call in async_rpush.await_args_list
|
||||
if call.kwargs["key"] == REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
|
||||
for value in call.kwargs["values"]
|
||||
]
|
||||
observed: Final = {
|
||||
"pushed_tags": [txn["tag"] for batch in pushed_batches for txn in batch.values()],
|
||||
"left_in_memory": writer.daily_tag_spend_update_queue.update_queue.qsize(),
|
||||
}
|
||||
assert observed == {"pushed_tags": ["prod-tag"], "left_in_memory": 0}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _initialize_shared_aiohttp_session
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -1008,7 +1103,9 @@ async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_tas
|
|||
|
||||
task_may_finish: Final = asyncio.Event()
|
||||
task_finished: Final = asyncio.Event()
|
||||
flushes: Final[list[tuple[str, bool]]] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion
|
||||
flushes: Final[
|
||||
list[tuple[str, bool]]
|
||||
] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion
|
||||
|
||||
async def slow_logging_task() -> None:
|
||||
await task_may_finish.wait()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue