From 74e3248ecaf6527882702f95a903c1817d7fae83 Mon Sep 17 00:00:00 2001 From: Nitish Agarwal <1592163+nitishagar@users.noreply.github.com> Date: Mon, 15 Jun 2026 20:19:04 +0530 Subject: [PATCH] fix(proxy): drain all daily-spend batches per flush cycle (#30281) --- litellm/proxy/db/db_spend_update_writer.py | 350 +++++++++--------- .../proxy/db/test_db_spend_update_writer.py | 48 +++ 2 files changed, 231 insertions(+), 167 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e7f14df5294d..e2c15f1c57fc 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1613,199 +1613,215 @@ async def _update_daily_spend( start_time = time.time() try: - for i in range(n_retry_times + 1): - try: - # Sort the transactions to minimize the probability of deadlocks by reducing the chance of concurrent - # trasactions locking the same rows/ranges in different orders. - transactions_to_process = dict( - sorted( - daily_spend_transactions.items(), - # Normally to avoid deadlocks we would sort by the index, but since we have sprinkled indexes - # on our schema like we're discount Salt Bae, we just sort by all fields that have an index, - # in an ad-hoc (but hopefully sensible) order of indexes. The actual ordering matters less than - # ensuring that all concurrent transactions sort in the same order. - # We could in theory use the dict key, as it contains basically the same fields, but this is more - # robust to future changes in the key format. - # If _update_daily_spend ever gets the ability to write to multiple tables at once, the sorting - # should sort by the table first. - key=lambda x: ( - x[1].get("date") or "", - x[1].get(entity_id_field) or "", - x[1].get("api_key") or "", - x[1].get("model") or "", - x[1].get("custom_llm_provider") or "", - ), - )[:BATCH_SIZE] - ) - - if len(transactions_to_process) == 0: - verbose_proxy_logger.debug( - f"No new transactions to process for daily {entity_type} spend update" + while daily_spend_transactions: + for i in range(n_retry_times + 1): + try: + # Sort the transactions to minimize the probability of deadlocks by reducing the chance of concurrent + # trasactions locking the same rows/ranges in different orders. + transactions_to_process = dict( + sorted( + daily_spend_transactions.items(), + # Normally to avoid deadlocks we would sort by the index, but since we have sprinkled indexes + # on our schema like we're discount Salt Bae, we just sort by all fields that have an index, + # in an ad-hoc (but hopefully sensible) order of indexes. The actual ordering matters less than + # ensuring that all concurrent transactions sort in the same order. + # We could in theory use the dict key, as it contains basically the same fields, but this is more + # robust to future changes in the key format. + # If _update_daily_spend ever gets the ability to write to multiple tables at once, the sorting + # should sort by the table first. + key=lambda x: ( + x[1].get("date") or "", + x[1].get(entity_id_field) or "", + x[1].get("api_key") or "", + x[1].get("model") or "", + x[1].get("custom_llm_provider") or "", + ), + )[:BATCH_SIZE] ) - break - try: - async with prisma_client.db.batch_() as batcher: - for _, transaction in transactions_to_process.items(): - entity_id = transaction.get(entity_id_field) + if len(transactions_to_process) == 0: + verbose_proxy_logger.debug( + f"No new transactions to process for daily {entity_type} spend update" + ) + return + + try: + async with prisma_client.db.batch_() as batcher: + for _, transaction in transactions_to_process.items(): + entity_id = transaction.get(entity_id_field) + + # Construct the where clause dynamically + where_clause = { + unique_constraint_name: { + entity_id_field: entity_id, + "date": transaction["date"], + "api_key": transaction["api_key"], + "model": transaction["model"], + "custom_llm_provider": transaction.get( + "custom_llm_provider" + ) + or "", + "mcp_namespaced_tool_name": transaction.get( + "mcp_namespaced_tool_name" + ) + or "", + "endpoint": transaction.get("endpoint") + or "", + } + } - # Construct the where clause dynamically - where_clause = { - unique_constraint_name: { + # Get the table dynamically + table = getattr(batcher, table_name) + + # Common data structure for both create and update + common_data = { entity_id_field: entity_id, "date": transaction["date"], "api_key": transaction["api_key"], - "model": transaction["model"], - "custom_llm_provider": transaction.get( - "custom_llm_provider" - ) - or "", + "model": transaction.get("model"), + "model_group": transaction.get("model_group"), "mcp_namespaced_tool_name": transaction.get( "mcp_namespaced_tool_name" ) or "", + "custom_llm_provider": transaction.get( + "custom_llm_provider" + ), "endpoint": transaction.get("endpoint") or "", + "prompt_tokens": transaction["prompt_tokens"], + "completion_tokens": transaction[ + "completion_tokens" + ], + "spend": transaction["spend"], + "api_requests": transaction["api_requests"], + "successful_requests": transaction[ + "successful_requests" + ], + "failed_requests": transaction[ + "failed_requests" + ], } - } - - # Get the table dynamically - table = getattr(batcher, table_name) - - # Common data structure for both create and update - common_data = { - entity_id_field: entity_id, - "date": transaction["date"], - "api_key": transaction["api_key"], - "model": transaction.get("model"), - "model_group": transaction.get("model_group"), - "mcp_namespaced_tool_name": transaction.get( - "mcp_namespaced_tool_name" - ) - or "", - "custom_llm_provider": transaction.get( - "custom_llm_provider" - ), - "endpoint": transaction.get("endpoint") or "", - "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction[ - "completion_tokens" - ], - "spend": transaction["spend"], - "api_requests": transaction["api_requests"], - "successful_requests": transaction[ - "successful_requests" - ], - "failed_requests": transaction["failed_requests"], - } - - # Add cache-related fields if they exist - if "cache_read_input_tokens" in transaction: - common_data["cache_read_input_tokens"] = ( - transaction.get("cache_read_input_tokens", 0) - ) - if "cache_creation_input_tokens" in transaction: - common_data["cache_creation_input_tokens"] = ( - transaction.get( - "cache_creation_input_tokens", 0 - ) - ) - - if entity_type == "tag" and "request_id" in transaction: - common_data["request_id"] = transaction.get( - "request_id" - ) - # Create update data structure - update_data = { - "prompt_tokens": { - "increment": transaction["prompt_tokens"] - }, - "completion_tokens": { - "increment": transaction["completion_tokens"] - }, - "spend": {"increment": transaction["spend"]}, - "api_requests": { - "increment": transaction["api_requests"] - }, - "successful_requests": { - "increment": transaction["successful_requests"] - }, - "failed_requests": { - "increment": transaction["failed_requests"] - }, - } + # Add cache-related fields if they exist + if "cache_read_input_tokens" in transaction: + common_data["cache_read_input_tokens"] = ( + transaction.get( + "cache_read_input_tokens", 0 + ) + ) + if "cache_creation_input_tokens" in transaction: + common_data["cache_creation_input_tokens"] = ( + transaction.get( + "cache_creation_input_tokens", 0 + ) + ) - # Add cache-related fields to update if they exist - if "cache_read_input_tokens" in transaction: - update_data["cache_read_input_tokens"] = { - "increment": transaction.get( - "cache_read_input_tokens", 0 + if ( + entity_type == "tag" + and "request_id" in transaction + ): + common_data["request_id"] = transaction.get( + "request_id" ) + + # Create update data structure + update_data = { + "prompt_tokens": { + "increment": transaction["prompt_tokens"] + }, + "completion_tokens": { + "increment": transaction[ + "completion_tokens" + ] + }, + "spend": {"increment": transaction["spend"]}, + "api_requests": { + "increment": transaction["api_requests"] + }, + "successful_requests": { + "increment": transaction[ + "successful_requests" + ] + }, + "failed_requests": { + "increment": transaction["failed_requests"] + }, } - if "cache_creation_input_tokens" in transaction: - update_data["cache_creation_input_tokens"] = { - "increment": transaction.get( - "cache_creation_input_tokens", 0 + + # Add cache-related fields to update if they exist + if "cache_read_input_tokens" in transaction: + update_data["cache_read_input_tokens"] = { + "increment": transaction.get( + "cache_read_input_tokens", 0 + ) + } + if "cache_creation_input_tokens" in transaction: + update_data["cache_creation_input_tokens"] = { + "increment": transaction.get( + "cache_creation_input_tokens", 0 + ) + } + + if ( + entity_type == "tag" + and "request_id" in transaction + ): + update_data["request_id"] = transaction.get( + "request_id" ) - } - if entity_type == "tag" and "request_id" in transaction: - update_data["request_id"] = transaction.get( - "request_id" + # Add endpoint to update_data so existing rows get their endpoint field updated + update_data["endpoint"] = ( + transaction.get("endpoint") or "" ) - # Add endpoint to update_data so existing rows get their endpoint field updated - update_data["endpoint"] = ( - transaction.get("endpoint") or "" - ) + table.upsert( + where=where_clause, + data={ + "create": common_data, + "update": update_data, + }, + ) + except Exception as batch_error: + # Log detailed error information for debugging batch upsert failures + # This helps diagnose issues like unique constraint violations + spend_log_error( + "Daily %s spend batch upsert failed. " + "Table: %s, Constraint: %s, Batch size: %d, Error: %s", + entity_type, + table_name, + unique_constraint_name, + len(transactions_to_process), + str(batch_error), + exc=batch_error, + ) + raise - table.upsert( - where=where_clause, - data={ - "create": common_data, - "update": update_data, - }, - ) - except Exception as batch_error: - # Log detailed error information for debugging batch upsert failures - # This helps diagnose issues like unique constraint violations - spend_log_error( - "Daily %s spend batch upsert failed. " - "Table: %s, Constraint: %s, Batch size: %d, Error: %s", - entity_type, - table_name, - unique_constraint_name, - len(transactions_to_process), - str(batch_error), - exc=batch_error, + verbose_proxy_logger.debug( + f"Processed {len(transactions_to_process)} daily {entity_type} transactions in {time.time() - start_time:.2f}s" ) - raise - verbose_proxy_logger.debug( - f"Processed {len(transactions_to_process)} daily {entity_type} transactions in {time.time() - start_time:.2f}s" - ) + # Remove processed transactions + for key in transactions_to_process.keys(): + daily_spend_transactions.pop(key, None) - # Remove processed transactions - for key in transactions_to_process.keys(): - daily_spend_transactions.pop(key, None) - - break + break - except DB_CONNECTION_ERROR_TYPES as e: - if i >= n_retry_times: - _raise_failed_update_spend_exception( - e=e, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, + except DB_CONNECTION_ERROR_TYPES as e: + if i >= n_retry_times: + _raise_failed_update_spend_exception( + e=e, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + await asyncio.sleep( + # Sleep a random amount to avoid retrying and deadlocking again: when two transactions deadlock they are + # cancelled basically at the same time, so if they wait the same time they will also retry at the same time + # and thus they are more likely to deadlock again. + # Instead, we sleep a random amount so that they retry at slightly different times, lowering the chance of + # repeated deadlocks, and therefore of exceeding the retry limit. + random.uniform(2**i, 2 ** (i + 1)) ) - await asyncio.sleep( - # Sleep a random amount to avoid retrying and deadlocking again: when two transactions deadlock they are - # cancelled basically at the same time, so if they wait the same time they will also retry at the same time - # and thus they are more likely to deadlock again. - # Instead, we sleep a random amount so that they retry at slightly different times, lowering the chance of - # repeated deadlocks, and therefore of exceeding the retry limit. - random.uniform(2**i, 2 ** (i + 1)) - ) except Exception as e: if "transactions_to_process" in locals(): diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 79e6494eab0e..04c93f48ca98 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -239,6 +239,54 @@ async def test_update_daily_spend_sorting(): mock_table.upsert.assert_has_calls(upsert_calls) +@pytest.mark.asyncio +async def test_update_daily_spend_drains_all_batches_over_batch_size(): + """ + Regression for #30281: >BATCH_SIZE (100) unique entities in one flush must all + be written and the in-memory dict fully drained within a single call. Pre-fix, + only the first 100 sorted items were upserted then the method returned, silently + dropping the remaining entities. + """ + mock_prisma_client = MagicMock() + mock_batcher = MagicMock() + mock_table = MagicMock() + mock_prisma_client.db.batch_.return_value.__aenter__.return_value = mock_batcher + mock_batcher.litellm_dailyuserspend = mock_table + + num_entities = 250 + daily_spend_transactions = { + f"test_key_{i}": { + "user_id": f"user{i:04d}", + "date": "2024-01-01", + "api_key": "test-api-key", + "model": "gpt-4", + "custom_llm_provider": "openai", + "prompt_tokens": 10, + "completion_tokens": 20, + "spend": 0.1, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + for i in range(num_entities) + } + + await DBSpendUpdateWriter._update_daily_spend( + n_retry_times=1, + prisma_client=mock_prisma_client, + proxy_logging_obj=MagicMock(), + daily_spend_transactions=daily_spend_transactions, + entity_type="user", + entity_id_field="user_id", + table_name="litellm_dailyuserspend", + unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint", + ) + + assert mock_table.upsert.call_count == num_entities + assert mock_prisma_client.db.batch_.call_count == 3 + assert daily_spend_transactions == {} + + @pytest.mark.asyncio async def test_update_daily_spend_tag_with_request_id(): """