Merge branch 'main' into litellm_metrics_pod_lock_manager

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
 )

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
 

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(

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:

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"

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

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