Merge a19c2c2a62 into e0d1e64591

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

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",

litellm/proxy/utils.py

@@ -8416,9 +8416,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

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"
+    )

tests/unit/proxy/proxy_server/test_lifecycle.py

@@ -29,10 +29,18 @@ from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
-from pydantic import BaseModel
+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.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,
@@ -54,6 +62,9 @@ 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
 # ---------------------------------------------------------------------------
@@ -412,6 +423,72 @@ 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_hands_daily_tag_spend_to_redis_when_another_pod_holds_the_lock(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    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()
+
+            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()
+
+
 # ---------------------------------------------------------------------------
 # _initialize_shared_aiohttp_session
 # ---------------------------------------------------------------------------
@@ -1008,7 +1085,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()