Tag query fix (#25094)

* feat(tag-spend): implement separate scheduler job for daily tag spend updates

* fix(docker): add g++ to build dependencies in Dockerfile

* initial test cases. TODO: check scheduler init and test cases in proxy_server related to it

* resolved QPS issue when redis transaction buffer is enabled

* resolving circular import error flagged by greptile
This commit is contained in:
Harish 2026-04-04 09:51:46 -07:00 committed by GitHub
parent a5322c6efc
commit f4e69f48e2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 298 additions and 60 deletions

View file

@ -15,6 +15,7 @@ USER root
# Install build dependencies in one layer
RUN apt-get update && apt-get install -y --no-install-recommends \
gcc \
g++ \
python3-dev \
libssl-dev \
pkg-config \

View file

@ -1341,6 +1341,7 @@ LITELLM_UI_SESSION_DURATION = os.getenv("LITELLM_UI_SESSION_DURATION", "24h")
########################### DB CRON JOB NAMES ###########################
DB_SPEND_UPDATE_JOB_NAME = "db_spend_update_job"
DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME = "db_daily_tag_spend_update_job"
PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME = "prometheus_emit_budget_metrics"
CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME = "cloudzero_export_usage_data"
CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int(
@ -1393,6 +1394,10 @@ APSCHEDULER_REPLACE_EXISTING = os.getenv(
"1",
] # always replace existing jobs
# The number of tag entries are higher than number of user, team entries. This leads to a higher QPS.
# This will run tag spcific tasks at a later time to smooth QPS
DAILY_TAG_SPEND_BATCH_MULTIPLIER = 2.3
DEFAULT_HEALTH_CHECK_INTERVAL = int(
os.getenv("DEFAULT_HEALTH_CHECK_INTERVAL", 300)
) # 5 minutes

View file

@ -28,7 +28,7 @@ from typing import (
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache, RedisCache
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME,DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
@ -797,7 +797,6 @@ class DBSpendUpdateWriter:
daily_org_spend_update_queue=self.daily_org_spend_update_queue,
daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue,
daily_agent_spend_update_queue=self.daily_agent_spend_update_queue,
daily_tag_spend_update_queue=self.daily_tag_spend_update_queue,
)
# Only commit from redis to db if this pod is the leader
@ -814,7 +813,6 @@ class DBSpendUpdateWriter:
daily_org_spend_update_transactions,
daily_end_user_spend_update_transactions,
daily_agent_spend_update_transactions,
daily_tag_spend_update_transactions,
) = (
await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
)
@ -890,13 +888,6 @@ class DBSpendUpdateWriter:
daily_spend_transactions=daily_org_spend_update_transactions,
)
if daily_tag_spend_update_transactions is not None:
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
if daily_end_user_spend_update_transactions is not None:
await DBSpendUpdateWriter.update_daily_end_user_spend(
n_retry_times=n_retry_times,
@ -991,19 +982,7 @@ class DBSpendUpdateWriter:
daily_spend_transactions=daily_org_spend_update_transactions,
)
################## Daily Tag Spend Update Transactions ##################
# Aggregate all in memory daily tag spend transactions and commit to db
daily_tag_spend_update_transactions = cast(
Dict[str, DailyTagSpendTransaction],
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
# NOTE: Daily tag spend is committed by a separate scheduler job.
################## Daily End-User Spend Update Transactions ##################
# Aggregate all in memory daily end-user spend transactions and commit to db
@ -1032,10 +1011,75 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_agent_spend_update_transactions,
)
################## Tool Registry Upserts ##################
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
async def _commit_daily_tag_spend_to_db(
self,
prisma_client: PrismaClient,
n_retry_times: int,
proxy_logging_obj: ProxyLogging,
):
"""
Commit only tag spend updates to database.
This is called by a separate scheduler job at a longer interval.
"""
daily_tag_spend_update_transactions = cast(
Dict[str, DailyTagSpendTransaction],
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
if daily_tag_spend_update_transactions:
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
async def _commit_daily_tag_spend_to_db_with_redis(
self,
prisma_client: PrismaClient,
n_retry_times: int,
proxy_logging_obj: ProxyLogging,
):
"""
Commit daily tag spend updates using Redis buffering.
This lets the dedicated daily tag scheduler drain both in-memory and
Redis-backed tag transactions.
"""
await self.redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(
daily_tag_spend_update_queue=self.daily_tag_spend_update_queue,
)
if await self.pod_lock_manager.acquire_lock(
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
):
verbose_proxy_logger.debug("acquired lock for daily tag spend updates")
try:
daily_tag_spend_update_transactions = await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
if daily_tag_spend_update_transactions:
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to commit daily tag spend updates from Redis to DB. "
"Data already popped from Redis may be lost. Error: %s\n%s",
str(e),
traceback.format_exc(),
)
finally:
await self.pod_lock_manager.release_lock(
cronjob_id=DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME,
)
async def _flush_tool_discovery_queue(
self,
prisma_client: PrismaClient,

View file

@ -131,7 +131,6 @@ class RedisUpdateBuffer:
daily_org_spend_update_queue: DailySpendUpdateQueue,
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
daily_agent_spend_update_queue: DailySpendUpdateQueue,
daily_tag_spend_update_queue: DailySpendUpdateQueue,
):
"""
Stores the in-memory spend updates to Redis
@ -202,9 +201,6 @@ class RedisUpdateBuffer:
daily_agent_spend_update_transactions = (
await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
daily_tag_spend_update_transactions = (
await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
verbose_proxy_logger.debug(
"ALL DB SPEND UPDATE TRANSACTIONS: %s", db_spend_update_transactions
@ -245,11 +241,6 @@ class RedisUpdateBuffer:
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
),
(
daily_tag_spend_update_transactions,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
),
]
rpush_list: List[RedisPipelineRpushOperation] = []
@ -376,22 +367,20 @@ class RedisUpdateBuffer:
Optional[Dict[str, DailyOrganizationSpendTransaction]],
Optional[Dict[str, DailyEndUserSpendTransaction]],
Optional[Dict[str, DailyAgentSpendTransaction]],
Optional[Dict[str, DailyTagSpendTransaction]],
]:
"""
Drains all 7 Redis buffer queues in a single pipeline round-trip.
Drains the main 6 Redis buffer queues in a single pipeline round-trip.
Returns a 7-tuple of parsed results in this order:
Returns a 6-tuple of parsed results in this order:
0: DBSpendUpdateTransactions
1: daily user spend
2: daily team spend
3: daily org spend
4: daily end-user spend
5: daily agent spend
6: daily tag spend
"""
if self.redis_cache is None:
return None, None, None, None, None, None, None
return None, None, None, None, None, None
lpop_list: List[RedisPipelineLpopOperation] = [
RedisPipelineLpopOperation(
@ -417,16 +406,12 @@ class RedisUpdateBuffer:
key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
]
raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
# Pad with None if pipeline returned fewer results than expected
while len(raw_results) < 7:
while len(raw_results) < 6:
raw_results.append(None)
# Slot 0: DBSpendUpdateTransactions
@ -436,9 +421,9 @@ class RedisUpdateBuffer:
if len(parsed) > 0:
db_spend = self._combine_list_of_transactions(parsed)
# Slots 1-6: daily spend categories
# Slots 1-5: daily spend categories
daily_results: List[Optional[Dict[str, Any]]] = []
for slot in range(1, 7):
for slot in range(1, 6):
if raw_results[slot] is None:
daily_results.append(None)
else:
@ -457,7 +442,22 @@ class RedisUpdateBuffer:
),
cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]),
cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]),
cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]),
)
async def store_in_memory_daily_tag_spend_updates_in_redis(
self,
daily_tag_spend_update_queue: DailySpendUpdateQueue,
) -> None:
"""
Flush in-memory daily tag spend updates and append them to Redis.
"""
daily_tag_spend_update_transactions = (
await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
await self._store_transactions_in_redis(
transactions=daily_tag_spend_update_transactions,
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
)
async def get_all_daily_spend_update_transactions_from_redis_buffer(

View file

@ -54,6 +54,7 @@ from litellm.constants import (
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
LITELLM_UI_ALLOW_HEADERS,
LITELLM_UI_SESSION_DURATION,
DAILY_TAG_SPEND_BATCH_MULTIPLIER
)
from litellm.litellm_core_utils.litellm_logging import (
_init_custom_logger_compatible_class,
@ -6218,6 +6219,25 @@ class ProxyStartupEvent:
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
### 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 = 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",
seconds=tag_spend_update_interval,
args=[prisma_client, proxy_logging_obj],
id="update_daily_tag_spend_job",
replace_existing=True,
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
)
verbose_proxy_logger.info(
f"Tag spend update job scheduled at {tag_spend_update_interval}s interval "
f"({tag_spend_update_interval / batch_writing_interval:.1f}x main job interval)"
)
### MONITOR SPEND LOGS QUEUE (queue-size-based job) ###
if general_settings.get("disable_spend_logs", False) is False:
from litellm.proxy.utils import _monitor_spend_logs_queue

View file

@ -4827,6 +4827,9 @@ async def update_spend( # noqa: PLR0915
Triggered every minute.
NOTE: This job now skips tag spend updates, which are handled by a separate
scheduler job (update_daily_tag_spend) at a longer interval to reduce contention.
Requires:
user_id_list: dict,
keys_list: list,
@ -4859,6 +4862,46 @@ async def update_spend( # noqa: PLR0915
)
async def update_daily_tag_spend(
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
):
"""
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.
This is called by a dedicated scheduler job and does NOT process:
- Regular spend updates (user, key, team, org)
- End-user spend
- Agent spend
- Spend logs
Only processes tag spend transactions from the daily_tag_spend_update_queue.
Args:
prisma_client: PrismaClient instance
proxy_logging_obj: ProxyLogging instance for error handling
"""
n_retry_times = 3
try:
if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis():
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis(
prisma_client=prisma_client,
n_retry_times=n_retry_times,
proxy_logging_obj=proxy_logging_obj,
)
else:
await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db(
prisma_client=prisma_client,
n_retry_times=n_retry_times,
proxy_logging_obj=proxy_logging_obj,
)
except Exception as e:
verbose_proxy_logger.error(f"Error updating daily tag spend: {e}")
async def update_spend_logs_job(
prisma_client: PrismaClient,
db_writer_client: Optional[AsyncHTTPHandler],

View file

@ -0,0 +1,134 @@
from typing import Dict
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.proxy.utils import update_daily_tag_spend
from litellm.proxy._types import DailyTagSpendTransaction
import httpx
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
@pytest.mark.asyncio
async def test_update_daily_tag_spend_delegates_to_tag_commit_writer():
prisma_client = MagicMock()
proxy_logging_obj = MagicMock()
redis_update_buffer = MagicMock()
redis_update_buffer._should_commit_spend_updates_to_redis.return_value = False
proxy_logging_obj.db_spend_update_writer = MagicMock()
proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock()
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock()
await update_daily_tag_spend(
prisma_client,
proxy_logging_obj,
)
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_awaited_once_with(
prisma_client=prisma_client,
n_retry_times=3,
proxy_logging_obj=proxy_logging_obj,
)
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_daily_tag_spend_logs_error_and_does_not_raise():
prisma_client = MagicMock()
proxy_logging_obj = MagicMock()
redis_update_buffer = MagicMock()
redis_update_buffer._should_commit_spend_updates_to_redis.return_value = False
proxy_logging_obj.db_spend_update_writer = MagicMock()
proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock(
side_effect=ValueError("boom")
)
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock()
with patch("litellm.proxy.utils.verbose_proxy_logger.error") as error_logger:
await update_daily_tag_spend(
prisma_client,
proxy_logging_obj,
)
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_awaited_once()
error_logger.assert_called_once()
@pytest.mark.asyncio
async def test_update_daily_tag_spend_uses_redis_writer_when_enabled():
prisma_client = MagicMock()
proxy_logging_obj = MagicMock()
redis_update_buffer = MagicMock()
redis_update_buffer._should_commit_spend_updates_to_redis.return_value = True
proxy_logging_obj.db_spend_update_writer = MagicMock()
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db = AsyncMock()
proxy_logging_obj.db_spend_update_writer.redis_update_buffer = redis_update_buffer
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis = AsyncMock()
await update_daily_tag_spend(
prisma_client,
proxy_logging_obj,
)
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis.assert_awaited_once_with(
prisma_client=prisma_client,
n_retry_times=3,
proxy_logging_obj=proxy_logging_obj,
)
proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db.assert_not_awaited()
@pytest.mark.asyncio
async def test_daily_tag_spend_retries_then_succeeds():
prisma_client = MagicMock()
proxy_logging_obj = MagicMock()
mock_batcher = MagicMock()
mock_table = MagicMock()
mock_batcher.litellm_dailytagspend = mock_table
# Fail entering batch context 3 times with retryable DB errors, then succeed.
prisma_client.db.batch_.return_value.__aenter__ = AsyncMock(
side_effect=[
httpx.ConnectError("x"),
httpx.ConnectError("x"),
httpx.ConnectError("x"),
mock_batcher,
]
)
daily_spend_transactions: Dict[str, DailyTagSpendTransaction] = {
"k": {
"tag": "prod-tag",
"date": "2026-04-03",
"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,
"spend": 0.01,
"api_requests": 1,
"successful_requests": 1,
"failed_requests": 0,
"request_id": None,
}
}
with patch("asyncio.sleep", new_callable=AsyncMock) as sleep_mock, patch(
"random.uniform", return_value=0
):
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=3,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_transactions,
)
assert prisma_client.db.batch_.return_value.__aenter__.await_count == 4
assert sleep_mock.await_count == 3
mock_table.upsert.assert_called_once()

View file

@ -35,7 +35,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
"""
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2])
# Create mock queues - only 3 of 7 have data
# Create mock queues - only 3 of 6 have data
spend_update_queue = AsyncMock()
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
return_value={"key_list_transactions": {"key1": 1.0}}
@ -67,11 +67,6 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
return_value={}
)
daily_tag_queue = AsyncMock()
daily_tag_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
return_value={}
)
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
spend_update_queue=spend_update_queue,
daily_spend_update_queue=daily_spend_queue,
@ -79,7 +74,6 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
daily_org_spend_update_queue=daily_org_queue,
daily_end_user_spend_update_queue=daily_end_user_queue,
daily_agent_spend_update_queue=daily_agent_queue,
daily_tag_spend_update_queue=daily_tag_queue,
)
# Should be called exactly once (pipeline)
@ -117,7 +111,6 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(
daily_org_spend_update_queue=empty_daily_queue,
daily_end_user_spend_update_queue=empty_daily_queue,
daily_agent_spend_update_queue=empty_daily_queue,
daily_tag_spend_update_queue=empty_daily_queue,
)
mock_redis_cache.async_rpush_pipeline.assert_not_called()
@ -131,7 +124,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
and aggregates results from async_lpop_pipeline.
"""
# Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories
# Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories
db_spend_json = json.dumps(
{
"key_list_transactions": {"key1": 1.0, "key2": 2.0},
@ -154,14 +147,13 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
None, # slot 3: daily org (empty)
None, # slot 4: daily end-user (empty)
None, # slot 5: daily agent (empty)
None, # slot 6: daily tag (empty)
]
)
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert len(result) == 7
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result
assert len(result) == 6
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent = result
# Verify db spend was parsed correctly
assert db_spend is not None
@ -181,7 +173,6 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
assert daily_org is None
assert daily_end_user is None
assert daily_agent is None
assert daily_tag is None
# Verify pipeline was called once with correct keys
mock_redis_cache.async_lpop_pipeline.assert_called_once()
@ -192,7 +183,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
"""When redis_cache is None, should return all Nones"""
buffer = RedisUpdateBuffer(redis_cache=None)
result = await buffer.get_all_transactions_from_redis_buffer_pipeline()
assert result == (None, None, None, None, None, None, None)
assert result == (None, None, None, None, None, None)
def test_validate_redis_transaction_buffer_raises_without_redis():