[Bug]: Performance Fix Max langfuse clients reached: 20 is greater than 20 (#11285)

* fix: initializing langfuse clients

* fix: initializing langfuse clients

* tests: tests for langfuse cache

litellm/caching/in_memory_cache.py

@@ -84,6 +84,19 @@ class InMemoryCache(BaseCache):
         except Exception:
             return False
 
+    def _is_key_expired(self, key: str) -> bool:
+        """
+        Check if a specific key is expired
+        """
+        return key in self.ttl_dict and time.time() > self.ttl_dict[key]
+
+    def _remove_key(self, key: str) -> None:
+        """
+        Remove a key from both cache_dict and ttl_dict
+        """
+        self.cache_dict.pop(key, None)
+        self.ttl_dict.pop(key, None)
+
     def evict_cache(self):
         """
         Eviction policy:
@@ -97,9 +110,8 @@ class InMemoryCache(BaseCache):
 
         """
         for key in list(self.ttl_dict.keys()):
-            if time.time() > self.ttl_dict[key]:
-                self.cache_dict.pop(key, None)
-                self.ttl_dict.pop(key, None)
+            if self._is_key_expired(key):
+                self._remove_key(key)
 
                 # de-reference the removed item
                 # https://www.geeksforgeeks.org/diagnosing-and-fixing-memory-leaks-in-python/
@@ -153,13 +165,21 @@ class InMemoryCache(BaseCache):
         self.set_cache(key, init_value, ttl=ttl)
         return value
 
+    def evict_element_if_expired(self, key: str) -> bool:
+        """
+        Returns True if the element is expired and removed from the cache
+
+        Returns False if the element is not expired
+        """
+        if self._is_key_expired(key):
+            self._remove_key(key)
+            return True
+        return False
+
     def get_cache(self, key, **kwargs):
         if key in self.cache_dict:
-            if key in self.ttl_dict:
-                if time.time() > self.ttl_dict[key]:
-                    self.cache_dict.pop(key, None)
-                    self.ttl_dict.pop(key, None)
-                    return None
+            if self.evict_element_if_expired(key):
+                return None
             original_cached_response = self.cache_dict[key]
             try:
                 cached_response = json.loads(original_cached_response)
@@ -207,8 +227,7 @@ class InMemoryCache(BaseCache):
         pass
 
     def delete_cache(self, key):
-        self.cache_dict.pop(key, None)
-        self.ttl_dict.pop(key, None)
+        self._remove_key(key)
 
     async def async_get_ttl(self, key: str) -> Optional[int]:
         """

litellm/constants.py

@@ -154,7 +154,7 @@ FIREWORKS_AI_80_B = int(os.getenv("FIREWORKS_AI_80_B", 80))
 #### Logging callback constants ####
 REDACTED_BY_LITELM_STRING = "REDACTED_BY_LITELM"
 MAX_LANGFUSE_INITIALIZED_CLIENTS = int(
-    os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 20)
+    os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50)
 )
 
 ############### LLM Provider Constants ###############

litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py

@@ -1,10 +1,56 @@
+"""
+This is a cache for LangfuseLoggers.
+
+Langfuse Python SDK initializes a thread for each client. 
+
+This ensures we do 
+1. Proper cleanup of Langfuse initialized clients.
+2. Re-use created langfuse clients.
+"""
 import hashlib
 import json
 from typing import Any, Optional
 
+import litellm
+from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
+
 from ...caching import InMemoryCache
 
 
+class LangfuseInMemoryCache(InMemoryCache):
+    """
+    Ensures we do proper cleanup of Langfuse initialized clients.
+
+    Langfuse Python SDK initializes a thread for each client, we need to call Langfuse.shutdown() to properly cleanup.
+
+    This ensures we do proper cleanup of Langfuse initialized clients.
+    """
+
+    def _remove_key(self, key: str) -> None:
+        """
+        Override _remove_key in InMemoryCache to ensure we do proper cleanup of Langfuse initialized clients.
+
+        LangfuseLoggers consume threads when initalized, this shuts them down when they are expired
+
+        Relevant Issue: https://github.com/BerriAI/litellm/issues/11169
+        """
+        from litellm.integrations.langfuse.langfuse import LangFuseLogger
+
+        if isinstance(self.cache_dict[key], LangFuseLogger):
+            _created_langfuse_logger: LangFuseLogger = self.cache_dict[key]
+            #########################################################
+            # Clean up Langfuse initialized clients
+            #########################################################
+            litellm.initialized_langfuse_clients -= 1
+            _created_langfuse_logger.Langfuse.flush()
+            _created_langfuse_logger.Langfuse.shutdown()
+
+        #########################################################
+        # Call parent class to remove key from cache
+        #########################################################
+        return super()._remove_key(key)
+
+
 class DynamicLoggingCache:
     """
     Prevent memory leaks caused by initializing new logging clients on each request.
@@ -13,7 +59,7 @@ class DynamicLoggingCache:
     """
 
     def __init__(self) -> None:
-        self.cache = InMemoryCache()
+        self.cache = LangfuseInMemoryCache(default_ttl=_DEFAULT_TTL_FOR_HTTPX_CLIENTS)
 
     def get_cache_key(self, args: dict) -> str:
         args_str = json.dumps(args, sort_keys=True)

tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py

@@ -0,0 +1,81 @@
+import json
+import os
+import sys
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+import litellm
+from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
+    LangfuseInMemoryCache,
+)
+
+
+class TestLangfuseInMemoryCache:
+    """Simple tests to ensure langfuse client cleanup works correctly."""
+
+    def setup_method(self):
+        """Set up test fixtures before each test method."""
+        self.cache = LangfuseInMemoryCache(max_size_in_memory=2, default_ttl=1)
+
+    @patch("litellm.initialized_langfuse_clients", 5)
+    def test_langfuse_client_count_decrements_on_eviction(self):
+        """Test that langfuse client count decrements when elements get evicted from cache."""
+
+        # Create a mock LangFuseLogger class
+        class MockLangFuseLogger:
+            def __init__(self):
+                self.Langfuse = MagicMock()
+                self.Langfuse.flush = MagicMock()
+                self.Langfuse.shutdown = MagicMock()
+
+        mock_logger = MockLangFuseLogger()
+
+        # Patch the LangFuseLogger import to return our mock class
+        with patch(
+            "litellm.integrations.langfuse.langfuse.LangFuseLogger", MockLangFuseLogger
+        ):
+            # Add the mock logger to cache with expired TTL
+            self.cache.cache_dict["test_key"] = mock_logger
+            self.cache.ttl_dict["test_key"] = time.time() - 1  # Already expired
+
+            initial_count = litellm.initialized_langfuse_clients
+
+            # Trigger eviction
+            self.cache.evict_cache()
+
+            # Verify client count was decremented
+            assert litellm.initialized_langfuse_clients == initial_count - 1
+
+    @patch("litellm.initialized_langfuse_clients", 3)
+    def test_langfuse_client_shutdown_called_on_eviction(self):
+        """Test that langfuse client shutdown is called to close the thread."""
+
+        # Create a mock LangFuseLogger class
+        class MockLangFuseLogger:
+            def __init__(self):
+                self.Langfuse = MagicMock()
+                self.Langfuse.flush = MagicMock()
+                self.Langfuse.shutdown = MagicMock()
+
+        mock_logger = MockLangFuseLogger()
+
+        # Patch the LangFuseLogger import to return our mock class
+        with patch(
+            "litellm.integrations.langfuse.langfuse.LangFuseLogger", MockLangFuseLogger
+        ):
+            # Add the mock logger to cache
+            self.cache.cache_dict["test_key"] = mock_logger
+            self.cache.ttl_dict["test_key"] = time.time() + 100
+
+            # Remove the key (this should trigger cleanup)
+            self.cache._remove_key("test_key")
+
+            # Verify flush and shutdown were called
+            mock_logger.Langfuse.flush.assert_called_once()
+            mock_logger.Langfuse.shutdown.assert_called_once()