From d7f19bbfe3db824c0635abfd644b16b1e474cb59 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 30 May 2025 22:34:39 -0700 Subject: [PATCH] [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 | 39 ++++++--- litellm/constants.py | 2 +- .../specialty_caches/dynamic_logging_cache.py | 48 ++++++++++- .../test_dynamic_logging_cache.py | 81 +++++++++++++++++++ 4 files changed, 158 insertions(+), 12 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index e9c3f7ba44..a3ee881344 100644 --- a/litellm/caching/in_memory_cache.py +++ b/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]: """ diff --git a/litellm/constants.py b/litellm/constants.py index da9778922b..d36ad0f228 100644 --- a/litellm/constants.py +++ b/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 ############### diff --git a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py b/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py index 704803c78b..c2acc708bb 100644 --- a/litellm/litellm_core_utils/specialty_caches/dynamic_logging_cache.py +++ b/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) diff --git a/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py b/tests/test_litellm/litellm_core_utils/specialty_caches/test_dynamic_logging_cache.py new file mode 100644 index 0000000000..85fcc2700b --- /dev/null +++ b/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()