diff --git a/litellm/constants.py b/litellm/constants.py index a2fd373a61..c8248f548a 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -20,9 +20,14 @@ DEFAULT_IMAGE_HEIGHT = 300 DEFAULT_MAX_TOKENS = 256 # used when providers need a default MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = 1024 # 1MB = 1024KB SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD = 1000 # Minimum number of requests to consider "reasonable traffic". Used for single-deployment cooldown logic. + +########### v2 Architecture constants for managing writing updates to the database ########### REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT = 100 +MAX_SIZE_IN_MEMORY_QUEUE = 10000 +MAX_IN_MEMORY_QUEUE_FLUSH_COUNT = 1000 +############################################################################################### MINIMUM_PROMPT_CACHE_TOKEN_COUNT = ( 1024 # minimum number of tokens to cache a prompt by Anthropic ) diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py index 2bf2393127..873164c047 100644 --- a/litellm/proxy/db/db_transaction_queue/base_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -10,6 +10,7 @@ from litellm._service_logger import ServiceLogging service_logger_obj = ( ServiceLogging() ) # used for tracking metrics for In memory buffer, redis buffer, pod lock manager +from litellm.constants import MAX_IN_MEMORY_QUEUE_FLUSH_COUNT, MAX_SIZE_IN_MEMORY_QUEUE class BaseUpdateQueue: @@ -17,6 +18,7 @@ class BaseUpdateQueue: def __init__(self): self.update_queue = asyncio.Queue() + self.MAX_SIZE_IN_MEMORY_QUEUE = MAX_SIZE_IN_MEMORY_QUEUE async def add_update(self, update): """Enqueue an update.""" @@ -30,6 +32,12 @@ class BaseUpdateQueue: """Get all updates from the queue.""" updates = [] while not self.update_queue.empty(): + # Circuit breaker to ensure we're not stuck dequeuing updates. Protect CPU utilization + if len(updates) >= MAX_IN_MEMORY_QUEUE_FLUSH_COUNT: + verbose_proxy_logger.warning( + "Max in memory queue flush count reached, stopping flush" + ) + break updates.append(await self.update_queue.get()) return updates diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index afae431370..c61d24d50e 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -1,4 +1,5 @@ import asyncio +from copy import deepcopy from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger @@ -56,6 +57,29 @@ class DailySpendUpdateQueue(BaseUpdateQueue): Dict[str, DailyUserSpendTransaction] ] = asyncio.Queue() + async def add_update(self, update: Dict[str, DailyUserSpendTransaction]): + """Enqueue an update.""" + verbose_proxy_logger.debug("Adding update to queue: %s", update) + await self.update_queue.put(update) + if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE: + verbose_proxy_logger.warning( + "Spend update queue is full. Aggregating all entries in queue to concatenate entries." + ) + await self.aggregate_queue_updates() + + async def aggregate_queue_updates(self): + """ + Combine all updates in the queue into a single update. + This is used to reduce the size of the in-memory queue. + """ + updates: List[ + Dict[str, DailyUserSpendTransaction] + ] = await self.flush_all_updates_from_in_memory_queue() + aggregated_updates = self.get_aggregated_daily_spend_update_transactions( + updates + ) + await self.update_queue.put(aggregated_updates) + async def flush_and_get_aggregated_daily_spend_update_transactions( self, ) -> Dict[str, DailyUserSpendTransaction]: @@ -95,7 +119,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue): ] daily_transaction["failed_requests"] += payload["failed_requests"] else: - aggregated_daily_spend_update_transactions[_key] = payload + aggregated_daily_spend_update_transactions[_key] = deepcopy(payload) return aggregated_daily_spend_update_transactions async def _emit_new_item_added_to_queue_event( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 60e9379751..98a9e5088a 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -1,5 +1,5 @@ import asyncio -from typing import List, Optional +from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -31,6 +31,98 @@ class SpendUpdateQueue(BaseUpdateQueue): verbose_proxy_logger.debug("Aggregating updates by entity type: %s", updates) return self.get_aggregated_db_spend_update_transactions(updates) + async def add_update(self, update: SpendUpdateQueueItem): + """Enqueue an update to the spend update queue""" + verbose_proxy_logger.debug("Adding update to queue: %s", update) + await self.update_queue.put(update) + + # if the queue is full, aggregate the updates + if self.update_queue.qsize() >= self.MAX_SIZE_IN_MEMORY_QUEUE: + verbose_proxy_logger.warning( + "Spend update queue is full. Aggregating all entries in queue to concatenate entries." + ) + await self.aggregate_queue_updates() + + async def aggregate_queue_updates(self): + """Concatenate all updates in the queue to reduce the size of in-memory queue""" + updates: List[ + SpendUpdateQueueItem + ] = await self.flush_all_updates_from_in_memory_queue() + aggregated_updates = self._get_aggregated_spend_update_queue_item(updates) + for update in aggregated_updates: + await self.update_queue.put(update) + return + + def _get_aggregated_spend_update_queue_item( + self, updates: List[SpendUpdateQueueItem] + ) -> List[SpendUpdateQueueItem]: + """ + This is used to reduce the size of the in-memory queue by aggregating updates by entity type + id + + + Aggregate updates by entity type + id + + eg. + + ``` + [ + { + "entity_type": "user", + "entity_id": "123", + "response_cost": 100 + }, + { + "entity_type": "user", + "entity_id": "123", + "response_cost": 200 + } + ] + + ``` + + becomes + + ``` + + [ + { + "entity_type": "user", + "entity_id": "123", + "response_cost": 300 + } + ] + + ``` + """ + verbose_proxy_logger.debug( + "Aggregating spend updates, current queue size: %s", + self.update_queue.qsize(), + ) + aggregated_spend_updates: List[SpendUpdateQueueItem] = [] + + _in_memory_map: Dict[str, SpendUpdateQueueItem] = {} + """ + Used for combining several updates into a single update + Key=entity_type:entity_id + Value=SpendUpdateQueueItem + """ + for update in updates: + _key = f"{update.get('entity_type')}:{update.get('entity_id')}" + if _key not in _in_memory_map: + _in_memory_map[_key] = update + else: + current_cost = _in_memory_map[_key].get("response_cost", 0) or 0 + update_cost = update.get("response_cost", 0) or 0 + _in_memory_map[_key]["response_cost"] = current_cost + update_cost + + for _key, update in _in_memory_map.items(): + aggregated_spend_updates.append(update) + + verbose_proxy_logger.debug( + "Aggregated spend updates: %s", aggregated_spend_updates + ) + return aggregated_spend_updates + def get_aggregated_db_spend_update_transactions( self, updates: List[SpendUpdateQueueItem] ) -> DBSpendUpdateTransactions: diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_base_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_base_update_queue.py new file mode 100644 index 0000000000..e1d4cb0541 --- /dev/null +++ b/tests/litellm/proxy/db/db_transaction_queue/test_base_update_queue.py @@ -0,0 +1,41 @@ +import asyncio +import json +import os +import sys + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.constants import MAX_IN_MEMORY_QUEUE_FLUSH_COUNT +from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue + + +@pytest.mark.asyncio +async def test_queue_flush_limit(): + """ + Test to ensure we don't dequeue more than MAX_IN_MEMORY_QUEUE_FLUSH_COUNT items. + """ + # Arrange + queue = BaseUpdateQueue() + # Add more items than the max flush count + items_to_add = MAX_IN_MEMORY_QUEUE_FLUSH_COUNT + 100 + + for i in range(items_to_add): + await queue.add_update(f"test_update_{i}") + + # Act + flushed_updates = await queue.flush_all_updates_from_in_memory_queue() + + # Assert + assert ( + len(flushed_updates) == MAX_IN_MEMORY_QUEUE_FLUSH_COUNT + ), f"Expected {MAX_IN_MEMORY_QUEUE_FLUSH_COUNT} items, but got {len(flushed_updates)}" + + # Verify remaining items are still in queue + assert ( + queue.update_queue.qsize() == 100 + ), "Expected 100 items to remain in the queue" diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py index 228dfe64b1..b92f2919ed 100644 --- a/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py +++ b/tests/litellm/proxy/db/db_transaction_queue/test_daily_spend_update_queue.py @@ -6,6 +6,11 @@ import sys import pytest from fastapi.testclient import TestClient +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path +import litellm +from litellm.constants import MAX_SIZE_IN_MEMORY_QUEUE from litellm.proxy._types import ( DailyUserSpendTransaction, Litellm_EntityType, @@ -16,10 +21,6 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - @pytest.fixture def daily_spend_update_queue(): @@ -254,7 +255,7 @@ async def test_flush_and_get_aggregated_daily_spend_update_transactions( await daily_spend_update_queue.add_update({test_key: test_transaction1}) await daily_spend_update_queue.add_update({test_key: test_transaction2}) - # Test full workflow + # Flush and get aggregated transactions result = ( await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) @@ -262,3 +263,170 @@ async def test_flush_and_get_aggregated_daily_spend_update_transactions( assert len(result) == 1 assert test_key in result assert result[test_key] == expected_transaction + + +@pytest.mark.asyncio +async def test_queue_max_size_triggers_aggregation( + monkeypatch, daily_spend_update_queue +): + """Test that reaching MAX_SIZE_IN_MEMORY_QUEUE triggers aggregation""" + # Override MAX_SIZE_IN_MEMORY_QUEUE for testing + litellm._turn_on_debug() + monkeypatch.setattr(daily_spend_update_queue, "MAX_SIZE_IN_MEMORY_QUEUE", 6) + + test_key = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction = { + "spend": 1.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + # Add 6 identical updates (exceeding the max size of 5) + for i in range(6): + await daily_spend_update_queue.add_update({test_key: test_transaction}) + + # Queue should have aggregated to a single item + assert daily_spend_update_queue.update_queue.qsize() == 1 + + # Verify the aggregated values + result = ( + await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert result[test_key]["spend"] == 6.0 + assert result[test_key]["prompt_tokens"] == 600 + assert result[test_key]["completion_tokens"] == 300 + assert result[test_key]["api_requests"] == 6 + assert result[test_key]["successful_requests"] == 6 + assert result[test_key]["failed_requests"] == 0 + + +@pytest.mark.asyncio +async def test_aggregate_queue_updates_accuracy(daily_spend_update_queue): + """Test that queue aggregation correctly combines metrics by transaction key""" + # Add multiple updates for different transaction keys + test_key1 = "user1_2023-01-01_key123_gpt-4_openai" + test_transaction1 = { + "spend": 10.0, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + test_key2 = "user1_2023-01-01_key123_gpt-4_openai" # Same key + test_transaction2 = { + "spend": 5.0, + "prompt_tokens": 200, + "completion_tokens": 30, + "api_requests": 1, + "successful_requests": 0, + "failed_requests": 1, + } + + test_key3 = "user2_2023-01-01_key456_gpt-3.5-turbo_openai" # Different key + test_transaction3 = { + "spend": 3.0, + "prompt_tokens": 150, + "completion_tokens": 25, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + # Add updates directly to the queue + await daily_spend_update_queue.update_queue.put({test_key1: test_transaction1}) + await daily_spend_update_queue.update_queue.put({test_key2: test_transaction2}) + await daily_spend_update_queue.update_queue.put({test_key3: test_transaction3}) + + # Force aggregation + await daily_spend_update_queue.aggregate_queue_updates() + updates = await daily_spend_update_queue.flush_all_updates_from_in_memory_queue() + print("AGGREGATED UPDATES", json.dumps(updates, indent=4)) + daily_spend_update_transactions = updates[0] + + # Should have 2 keys after aggregation (test_key1/test_key2 combined, and test_key3) + assert len(daily_spend_update_transactions) == 2 + + # Check aggregated values for test_key1 (which is the same as test_key2) + assert daily_spend_update_transactions[test_key1]["spend"] == 15.0 + assert daily_spend_update_transactions[test_key1]["prompt_tokens"] == 300 + assert daily_spend_update_transactions[test_key1]["completion_tokens"] == 80 + assert daily_spend_update_transactions[test_key1]["api_requests"] == 2 + assert daily_spend_update_transactions[test_key1]["successful_requests"] == 1 + assert daily_spend_update_transactions[test_key1]["failed_requests"] == 1 + + # Check values for test_key3 remain the same + assert daily_spend_update_transactions[test_key3]["spend"] == 3.0 + assert daily_spend_update_transactions[test_key3]["prompt_tokens"] == 150 + assert daily_spend_update_transactions[test_key3]["completion_tokens"] == 25 + assert daily_spend_update_transactions[test_key3]["api_requests"] == 1 + assert daily_spend_update_transactions[test_key3]["successful_requests"] == 1 + assert daily_spend_update_transactions[test_key3]["failed_requests"] == 0 + + +@pytest.mark.asyncio +async def test_queue_size_reduction_with_large_volume( + monkeypatch, daily_spend_update_queue +): + """Test that queue size is actually reduced when dealing with many items""" + # Set a smaller MAX_SIZE for testing + monkeypatch.setattr(daily_spend_update_queue, "MAX_SIZE_IN_MEMORY_QUEUE", 10) + + # Create transaction templates + user1_key = "user1_2023-01-01_key123_gpt-4_openai" + user1_transaction = { + "spend": 0.5, + "prompt_tokens": 100, + "completion_tokens": 50, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + user2_key = "user2_2023-01-01_key456_gpt-3.5-turbo_openai" + user2_transaction = { + "spend": 1.0, + "prompt_tokens": 200, + "completion_tokens": 30, + "api_requests": 1, + "successful_requests": 1, + "failed_requests": 0, + } + + # Add 30 updates (200 for user1, 10 for user2) + for i in range(200): + await daily_spend_update_queue.add_update({user1_key: user1_transaction}) + + # At this point, aggregation should have happened at least once + # Queue size should be much less than 10 + assert daily_spend_update_queue.update_queue.qsize() <= 10 + + for i in range(100): + await daily_spend_update_queue.add_update({user2_key: user2_transaction}) + + # Queue should have at most 10 items after all this activity + assert daily_spend_update_queue.update_queue.qsize() <= 10 + + # Verify total costs are correct + result = ( + await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + print("RESULT", json.dumps(result, indent=4)) + + assert result[user1_key]["spend"] == 200 * 0.5 # 10.0 + assert result[user1_key]["prompt_tokens"] == 200 * 100 # 2000 + assert result[user1_key]["completion_tokens"] == 200 * 50 # 1000 + assert result[user1_key]["api_requests"] == 200 + assert result[user1_key]["successful_requests"] == 200 + assert result[user1_key]["failed_requests"] == 0 + + assert result[user2_key]["spend"] == 100 * 1.0 # 10.0 + assert result[user2_key]["prompt_tokens"] == 100 * 200 # 2000 + assert result[user2_key]["completion_tokens"] == 100 * 30 # 300 + assert result[user2_key]["api_requests"] == 100 + assert result[user2_key]["successful_requests"] == 100 + assert result[user2_key]["failed_requests"] == 0 diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py index 98d3b4e4c7..9993b25dfd 100644 --- a/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py +++ b/tests/litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py @@ -6,6 +6,7 @@ import sys import pytest from fastapi.testclient import TestClient +from litellm.constants import MAX_SIZE_IN_MEMORY_QUEUE from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue @@ -150,3 +151,115 @@ async def test_missing_entity_type(spend_queue): # Should ignore updates without entity type assert all(len(transactions) == 0 for transactions in aggregated.values()) + + +@pytest.mark.asyncio +async def test_queue_max_size_triggers_aggregation(monkeypatch, spend_queue): + """Test that reaching MAX_SIZE_IN_MEMORY_QUEUE triggers aggregation""" + # Override MAX_SIZE_IN_MEMORY_QUEUE for testing + monkeypatch.setattr(spend_queue, "MAX_SIZE_IN_MEMORY_QUEUE", 6) + + # Add 6 updates for the same user (exceeding the max size) + for i in range(6): + update: SpendUpdateQueueItem = { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user123", + "response_cost": 1.0, + } + await spend_queue.add_update(update) + + # Queue should have been aggregated, resulting in a single entry + assert spend_queue.update_queue.qsize() == 1 + + # Verify the aggregated cost is correct + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + assert aggregated["user_list_transactions"]["user123"] == 6.0 + + +@pytest.mark.asyncio +async def test_aggregate_queue_updates_accuracy(spend_queue): + """Test that queue aggregation correctly combines costs by entity type and ID""" + # Add multiple updates for different entities + updates = [ + { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user1", + "response_cost": 1.5, + }, + { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user1", + "response_cost": 2.5, + }, + { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user2", + "response_cost": 3.0, + }, + { + "entity_type": Litellm_EntityType.TEAM, + "entity_id": "team1", + "response_cost": 5.0, + }, + ] + + for update in updates: + await spend_queue.update_queue.put(update) + + # Force aggregation + await spend_queue.aggregate_queue_updates() + + # Queue size should now be 3 (user1, user2, team1) + assert spend_queue.update_queue.qsize() == 3 + + # Flush and verify aggregated values + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + print("aggregated values", aggregated) + + assert aggregated["user_list_transactions"]["user1"] == 4.0 # 1.5 + 2.5 + assert aggregated["user_list_transactions"]["user2"] == 3.0 + assert aggregated["team_list_transactions"]["team1"] == 5.0 + + +@pytest.mark.asyncio +async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue): + """Test that queue size is actually reduced when dealing with many items""" + # Set a smaller MAX_SIZE for testing + monkeypatch.setattr(spend_queue, "MAX_SIZE_IN_MEMORY_QUEUE", 10) + + # Add 30 updates (200 for user1, 10 for key1) + for i in range(200): + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.USER, + "entity_id": "user1", + "response_cost": 0.5, + } + ) + + # At this point, aggregation should have happened at least once + # Queue size should be much less than 20 + assert spend_queue.update_queue.qsize() <= 10 + + for i in range(300): + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.KEY, + "entity_id": "key1", + "response_cost": 1.0, + } + ) + + # Queue should have at most 2 items after all this activity + assert spend_queue.update_queue.qsize() <= 10 + + # Verify total costs are correct + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + assert aggregated["user_list_transactions"]["user1"] == 200 * 0.5 + assert aggregated["key_list_transactions"]["key1"] == 300 * 1.0