mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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:
parent
a5322c6efc
commit
f4e69f48e2
8 changed files with 298 additions and 60 deletions
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
134
tests/proxy_unit_tests/test_update_daily_tag_spend.py
Normal file
134
tests/proxy_unit_tests/test_update_daily_tag_spend.py
Normal 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()
|
||||
|
|
@ -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():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue