From 018ad5592b79a63d427b9f4c83fb38d3d248bf20 Mon Sep 17 00:00:00 2001 From: yehuda-alt <226970295+yehuda-alt@users.noreply.github.com> Date: Mon, 26 Oct 2026 10:13:33 +0000 Subject: [PATCH 1/3] 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 --- litellm/proxy/db/db_spend_update_writer.py | 2 +- litellm/proxy/proxy_server.py | 4 +- litellm/proxy/utils.py | 5 +- .../unit/proxy/proxy_server/test_lifecycle.py | 101 +++++++++++++++++- 4 files changed, 105 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8ac2597c99b..bbc1c913ad3 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index df948b69a67..3a1b60c84e1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ce7e6afde24..8ed163df36f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 11a44a0ec16..bb6fc272e94 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -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() From a19c2c2a624ea22428833a47775801f84c1f86af Mon Sep 17 00:00:00 2001 From: yehuda-alt <226970295+yehuda-alt@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:12:10 +0000 Subject: [PATCH 2/3] test(proxy): cover the shutdown tag spend flush end to end and without doubles The regression now lives beside the other shutdown flush contracts in tests/integration/spend/test_shutdown_flush.py, against a real proxy, Postgres and a scripted upstream. The unit test keeps only the Redis handoff a non-leader pod makes, on a real RedisCache over fakeredis instead of patching the writer and the pod lock --- .../integration/spend/test_shutdown_flush.py | 37 ++++++++ .../unit/proxy/proxy_server/test_lifecycle.py | 92 ++++++++----------- 2 files changed, 74 insertions(+), 55 deletions(-) diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index e96511511d9..b4a7f2c72fd 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -13,6 +13,7 @@ import httpx import psycopg import pytest import yaml +from pydantic import JsonValue from integration._support.client import Gateway, delete_key_if_present, eventually, string_value from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process @@ -25,6 +26,7 @@ BATCH_DRAINED_LOG_LINE: Final = f"flushed {REQUESTS_WHILE_BLOCKED} daily spend u MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new" COMMIT_DELAY_SECONDS: Final = 15 BURST_REQUESTS: Final = 30 +TAGGED_REQUESTS: Final = 5 def _api_requests(table: str, column: str, identity: str) -> int: @@ -168,6 +170,14 @@ class _Shutdown: self.owned.process.wait(timeout=120) +def _config_with_batch_write_interval(tmp_path: Path, seconds: int) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"]["proxy_batch_write_at"] = seconds + path: Final = tmp_path / f"batch-write-{seconds}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["general_settings"]["database_connection_pool_limit"] = pool_limit @@ -310,3 +320,30 @@ def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_ex shutdown.terminate_once(lambda: _committing_daily_user_spend() >= 1, lambda: None) finally: _drop_slow_commit(suffix) + + +def test_daily_tag_spend_buffered_since_the_last_tag_job_run_is_written_by_the_shutdown_flush( + gateway: Gateway, tmp_path: Path +) -> None: + tag: Final = f"integration-shutdown-tag-{uuid.uuid4().hex}" + with gateway.scenario() as scenario, wire_server(_provider) as wire: + model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) + with owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": os.environ["DATABASE_URL"]}, + config=_config_with_batch_write_interval(tmp_path, 600), + remove_environment=("DATABASE_URL_READ_REPLICA",), + ) as owned: + for _ in range(TAGGED_REQUESTS): + body: dict[str, JsonValue] = { + "model": model, + "messages": [{"role": "user", "content": f"tag {uuid.uuid4().hex}"}], + } + response = owned.gateway.request("POST", "/v1/chat/completions", body, headers={"x-litellm-tags": tag}) + assert response.status_code == 200, response.text + owned.begin_shutdown() + owned.process.wait(timeout=120) + assert _api_requests("LiteLLM_DailyTagSpend", "tag", tag) == TAGGED_REQUESTS, ( + "the tag job runs every batch interval x 2.3, so only the shutdown flush can have written these rows" + ) diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index bb6fc272e94..5a20b0c3171 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -29,13 +29,18 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from apscheduler.schedulers.asyncio import AsyncIOScheduler +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection from pydantic import BaseModel, TypeAdapter +from redis.asyncio import BlockingConnectionPool 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.caching.redis_cache import RedisCache +from litellm.constants import DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME, 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.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.proxy_server import ( ProxyStartupEvent, initialize_shared_aiohttp_session, @@ -58,6 +63,7 @@ from litellm.proxy.management_helpers.audit_logs import track_audit_task from .conftest import normalize _TAG_BATCH_ADAPTER: Final = TypeAdapter(dict[str, DailyTagSpendTransaction]) +_REDIS_LIST_ADAPTER: Final = TypeAdapter(list[bytes]) # --------------------------------------------------------------------------- # cleanup_router_config_variables @@ -445,66 +451,42 @@ def _daily_tag_spend(tag: str) -> dict[str, BaseDailySpendTransaction]: } -@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)) + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url="redis://spend-buffer.example:6379/0") + writer: Final = DBSpendUpdateWriter(redis_cache=redis) + writer.pod_lock_manager.redis_cache = redis + sync_client.set(PodLockManager.get_redis_lock_key(DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME), "other-pod") + 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) - await ps.flush_spend_counters_on_shutdown() + 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} + buffered_tags: Final = [ + txn["tag"] + for value in _REDIS_LIST_ADAPTER.validate_python( + sync_client.lrange(REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, 0, -1) + ) + for txn in _TAG_BATCH_ADAPTER.validate_json(value).values() + ] + assert buffered_tags == ["prod-tag"], ( + "a pod that is not the tag-commit leader must leave its buffered tag spend in Redis for the leader" + ) + finally: + await pool.disconnect() + sync_client.close() # --------------------------------------------------------------------------- From 380a2601ab6d94cd7273e797dfe93a8341af76ab Mon Sep 17 00:00:00 2001 From: yehuda-alt <226970295+yehuda-alt@users.noreply.github.com> Date: Wed, 28 Oct 2026 09:08:33 +0000 Subject: [PATCH 3/3] refactor(proxy): drop the tag spend note from the spend flush --- litellm/proxy/db/db_spend_update_writer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index bbc1c913ad3..16dbe2b4131 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1900,8 +1900,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) - # 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 await self._flush_daily_spend_queue(