From df8cb3eb06943bc32183c86d9d1ba3e3a7669a00 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 17 Feb 2026 17:48:29 +0000 Subject: [PATCH] fix: memory leak investigation - close HTTP clients on cache eviction, fix Sagemaker NoneType error Three fixes addressing the v1.81.12.rc.1 memory leak report: 1. InMemoryCache: Add on_evict callback to _remove_key() so cached objects can be cleaned up deterministically when evicted (TTL expiry or size limit). 2. LLMClientCache: Use on_evict to close HTTP clients when they are evicted from the client cache, preventing connection pool leaks that accumulate under high load (previously relied on unreliable __del__/GC). 3. Sagemaker handler: Fix AttributeError 'NoneType' object has no attribute 'get' when exception.response is explicitly None. Also remove unreachable dead code in make_async_call. https://claude.ai/code/session_01H8Eja2TiMCqi6R4nXknyVq --- litellm/caching/in_memory_cache.py | 16 +- litellm/caching/llm_caching_handler.py | 32 ++ litellm/llms/sagemaker/completion/handler.py | 18 +- .../test_memory_leak_investigation.py | 305 ++++++++++++++++++ 4 files changed, 356 insertions(+), 15 deletions(-) create mode 100644 tests/test_litellm/test_memory_leak_investigation.py diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 5239fa1f4b0..907bbb1b5d5 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -12,7 +12,7 @@ import json import sys import time import heapq -from typing import TYPE_CHECKING, Any, List, Optional +from typing import TYPE_CHECKING, Any, Callable, List, Optional if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -32,9 +32,12 @@ class InMemoryCache(BaseCache): int ] = 600, # default ttl is 10 minutes. At maximum litellm rate limiting logic requires objects to be in memory for 1 minute max_size_per_item: Optional[int] = 1024, # 1MB = 1024KB + on_evict: Optional[Callable[[Any], None]] = None, ): """ max_size_in_memory [int]: Maximum number of items in cache. done to prevent memory leaks. Use 200 items as a default + on_evict [Callable]: Optional callback called with evicted value when an item is removed from cache. + Used by LLMClientCache to close HTTP clients when they expire. """ self.max_size_in_memory = ( max_size_in_memory if max_size_in_memory is not None else 200 @@ -43,6 +46,7 @@ class InMemoryCache(BaseCache): self.max_size_per_item = ( max_size_per_item or MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB ) # 1MB = 1024KB + self.on_evict = on_evict # in-memory cache self.cache_dict: dict = {} @@ -97,10 +101,16 @@ class InMemoryCache(BaseCache): def _remove_key(self, key: str) -> None: """ - Remove a key from both cache_dict and ttl_dict + Remove a key from both cache_dict and ttl_dict. + Calls on_evict callback if set, to allow resource cleanup (e.g. closing HTTP clients). """ - self.cache_dict.pop(key, None) + evicted_value = self.cache_dict.pop(key, None) self.ttl_dict.pop(key, None) + if evicted_value is not None and self.on_evict is not None: + try: + self.on_evict(evicted_value) + except Exception: + pass def evict_cache(self): """ diff --git a/litellm/caching/llm_caching_handler.py b/litellm/caching/llm_caching_handler.py index 16eb824f4c9..6c810f77a07 100644 --- a/litellm/caching/llm_caching_handler.py +++ b/litellm/caching/llm_caching_handler.py @@ -4,10 +4,42 @@ Add the event loop to the cache key, to prevent event loop closed errors. import asyncio +from litellm._logging import verbose_logger + from .in_memory_cache import InMemoryCache +def _close_http_client_on_evict(client) -> None: + """ + Callback invoked when an HTTP client is evicted from the LLMClientCache. + + Ensures the underlying connection pool is closed deterministically instead + of relying on __del__ / garbage collection, which is unreliable for async + resources and can cause connection pool leaks under high load. + """ + close_fn = getattr(client, "close", None) + if close_fn is None: + return + + if asyncio.iscoroutinefunction(close_fn): + # Schedule async close on the running event loop if available + try: + loop = asyncio.get_running_loop() + loop.create_task(close_fn()) + except RuntimeError: + # No running loop – best-effort sync fallback + pass + else: + try: + close_fn() + except Exception: + pass + + class LLMClientCache(InMemoryCache): + def __init__(self, **kwargs): + super().__init__(on_evict=_close_http_client_on_evict, **kwargs) + def update_cache_key_with_event_loop(self, key): """ Add the event loop to the cache key, to prevent event loop closed errors. diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 2a30dc5ef38..013c7df03d0 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -330,13 +330,14 @@ class SagemakerLLM(BaseAWSLLM): raise e except Exception as e: verbose_logger.error("Sagemaker error %s", str(e)) + _response = getattr(e, "response", None) or {} status_code = ( - getattr(e, "response", {}) + _response .get("ResponseMetadata", {}) .get("HTTPStatusCode", 500) ) error_message = ( - getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) + _response.get("Error", {}).get("Message", str(e)) ) if "Inference Component Name header is required" in error_message: error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`" @@ -386,14 +387,6 @@ class SagemakerLLM(BaseAWSLLM): return completion_stream - # LOGGING - logging_obj.post_call( - input=[], - api_key="", - original_response="first stream response received", - additional_args={"complete_input_dict": data}, - ) - except httpx.HTTPStatusError as err: error_code = err.response.status_code raise SagemakerError(status_code=error_code, message=err.response.text) @@ -653,13 +646,14 @@ class SagemakerLLM(BaseAWSLLM): CustomAttributes="accept_eula=true", ) except Exception as e: + _response = getattr(e, "response", None) or {} status_code = ( - getattr(e, "response", {}) + _response .get("ResponseMetadata", {}) .get("HTTPStatusCode", 500) ) error_message = ( - getattr(e, "response", {}).get("Error", {}).get("Message", str(e)) + _response.get("Error", {}).get("Message", str(e)) ) raise SagemakerError(status_code=status_code, message=error_message) diff --git a/tests/test_litellm/test_memory_leak_investigation.py b/tests/test_litellm/test_memory_leak_investigation.py new file mode 100644 index 00000000000..771273b08a8 --- /dev/null +++ b/tests/test_litellm/test_memory_leak_investigation.py @@ -0,0 +1,305 @@ +""" +Memory Leak Investigation Tests for v1.81.12.rc.1 + +These tests reproduce and verify fixes for memory leak patterns identified +in the investigation of the v1.81.12.rc.1 memory leak report. + +Key issues investigated: +1. HTTP client cache eviction without closing connections +2. Sagemaker error handler NoneType crash (AttributeError: 'NoneType' object has no attribute 'get') +3. InMemoryCache eviction callback support +""" + +import asyncio +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +class TestInMemoryCacheEvictionCallback: + """Test that InMemoryCache properly calls on_evict when items are removed.""" + + def test_eviction_callback_called_on_ttl_expiry(self): + """When a cached item expires via TTL, the on_evict callback should fire.""" + from litellm.caching.in_memory_cache import InMemoryCache + + evicted_values = [] + + def on_evict(value): + evicted_values.append(value) + + cache = InMemoryCache( + max_size_in_memory=10, + default_ttl=1, # 1 second TTL + on_evict=on_evict, + ) + + cache.set_cache("key1", "value1", ttl=0.01) # expires almost immediately + time.sleep(0.05) + + # Trigger eviction by setting another key (evict_cache is called on set) + cache.set_cache("key2", "value2") + + # Force eviction check + cache.evict_cache() + + assert "value1" in evicted_values, "on_evict should be called when item expires" + + def test_eviction_callback_called_on_size_limit(self): + """When cache exceeds max_size_in_memory, eviction callback should fire.""" + from litellm.caching.in_memory_cache import InMemoryCache + + evicted_values = [] + + def on_evict(value): + evicted_values.append(value) + + cache = InMemoryCache( + max_size_in_memory=2, + default_ttl=600, + on_evict=on_evict, + ) + + cache.set_cache("key1", "value1") + cache.set_cache("key2", "value2") + cache.set_cache("key3", "value3") # Should trigger eviction of oldest + + assert len(evicted_values) >= 1, "on_evict should be called when cache is full" + + def test_eviction_callback_not_called_when_not_set(self): + """When on_evict is None, no callback should be attempted.""" + from litellm.caching.in_memory_cache import InMemoryCache + + cache = InMemoryCache( + max_size_in_memory=1, + default_ttl=600, + on_evict=None, + ) + + # This should not raise an error + cache.set_cache("key1", "value1") + cache.set_cache("key2", "value2") # Triggers eviction + + def test_eviction_callback_exception_does_not_propagate(self): + """If on_evict raises, it should not break cache operations.""" + from litellm.caching.in_memory_cache import InMemoryCache + + def bad_callback(value): + raise RuntimeError("callback error") + + cache = InMemoryCache( + max_size_in_memory=1, + default_ttl=600, + on_evict=bad_callback, + ) + + cache.set_cache("key1", "value1") + # Should not raise even though callback throws + cache.set_cache("key2", "value2") + + def test_delete_cache_calls_eviction_callback(self): + """Explicit delete_cache should also trigger on_evict.""" + from litellm.caching.in_memory_cache import InMemoryCache + + evicted_values = [] + + def on_evict(value): + evicted_values.append(value) + + cache = InMemoryCache( + max_size_in_memory=10, + default_ttl=600, + on_evict=on_evict, + ) + + cache.set_cache("key1", "value1") + cache.delete_cache("key1") + + assert "value1" in evicted_values + + +class TestLLMClientCacheCleanup: + """Test that LLMClientCache closes HTTP clients on eviction.""" + + def test_sync_client_closed_on_eviction(self): + """Sync HTTP clients should have close() called when evicted.""" + from litellm.caching.llm_caching_handler import LLMClientCache + + mock_client = MagicMock() + mock_client.close = MagicMock() + + cache = LLMClientCache(max_size_in_memory=1, default_ttl=600) + cache.set_cache("client1", mock_client) + cache.set_cache("client2", MagicMock()) # Triggers eviction of client1 + + mock_client.close.assert_called_once() + + def test_async_client_close_scheduled_on_eviction(self): + """Async HTTP clients should have close() scheduled when evicted.""" + from litellm.caching.llm_caching_handler import LLMClientCache + + mock_client = MagicMock() + mock_close = AsyncMock() + mock_client.close = mock_close + + cache = LLMClientCache(max_size_in_memory=1, default_ttl=600) + + async def run(): + cache.set_cache("client1", mock_client) + cache.set_cache("client2", MagicMock()) # Triggers eviction + # Give the scheduled task a chance to run + await asyncio.sleep(0.01) + + asyncio.get_event_loop().run_until_complete(run()) + + def test_client_without_close_does_not_crash(self): + """Objects without a close() method should not cause errors on eviction.""" + from litellm.caching.llm_caching_handler import LLMClientCache + + cache = LLMClientCache(max_size_in_memory=1, default_ttl=600) + cache.set_cache("key1", "just_a_string") + cache.set_cache("key2", "another_string") # Should not raise + + +class TestSagemakerNoneTypeError: + """Test fix for AttributeError: 'NoneType' object has no attribute 'get' + in Sagemaker error handling (handler.py line 658).""" + + def test_sagemaker_error_with_none_response(self): + """When exception.response is explicitly None, error handling should not crash.""" + from litellm.llms.sagemaker.common_utils import SagemakerError + + class FakeException(Exception): + response = None # Explicitly None, not missing + + # Simulate the fixed error handling logic + e = FakeException("timeout error") + _response = getattr(e, "response", None) or {} + status_code = ( + _response.get("ResponseMetadata", {}).get("HTTPStatusCode", 500) + ) + error_message = _response.get("Error", {}).get("Message", str(e)) + + assert status_code == 500 + assert error_message == "timeout error" + + def test_sagemaker_error_with_valid_response(self): + """When exception has a proper response dict, it should extract correctly.""" + + class FakeException(Exception): + response = { + "ResponseMetadata": {"HTTPStatusCode": 429}, + "Error": {"Message": "Rate limit exceeded"}, + } + + e = FakeException("rate limit") + _response = getattr(e, "response", None) or {} + status_code = ( + _response.get("ResponseMetadata", {}).get("HTTPStatusCode", 500) + ) + error_message = _response.get("Error", {}).get("Message", str(e)) + + assert status_code == 429 + assert error_message == "Rate limit exceeded" + + def test_sagemaker_error_with_missing_response(self): + """When exception has no response attribute, defaults should be used.""" + + e = Exception("generic error") + _response = getattr(e, "response", None) or {} + status_code = ( + _response.get("ResponseMetadata", {}).get("HTTPStatusCode", 500) + ) + error_message = _response.get("Error", {}).get("Message", str(e)) + + assert status_code == 500 + assert error_message == "generic error" + + +class TestHTTPClientCacheKeyGrowth: + """Test that HTTP client cache keys don't grow unboundedly.""" + + def test_cache_bounded_by_max_size(self): + """Cache should never exceed max_size_in_memory items.""" + from litellm.caching.in_memory_cache import InMemoryCache + + cache = InMemoryCache(max_size_in_memory=5, default_ttl=600) + + for i in range(20): + cache.set_cache(f"key_{i}", f"value_{i}") + + assert ( + len(cache.cache_dict) <= 5 + ), f"Cache should be bounded to 5 items, got {len(cache.cache_dict)}" + + def test_llm_client_cache_bounded(self): + """LLMClientCache should respect max_size_in_memory.""" + from litellm.caching.llm_caching_handler import LLMClientCache + + cache = LLMClientCache(max_size_in_memory=3, default_ttl=600) + + for i in range(10): + cache.set_cache(f"client_{i}", MagicMock()) + + assert ( + len(cache.cache_dict) <= 3 + ), f"LLMClientCache should be bounded to 3 items, got {len(cache.cache_dict)}" + + +class TestRouterCooldownIsolation: + """Test that router cooldown system isolates failures per deployment.""" + + def test_cooldown_uses_deployment_id_not_model_group(self): + """Cooldown keys should be per-deployment, not per-model-group.""" + from litellm.router_utils.cooldown_cache import CooldownCache + + cache = CooldownCache( + default_cooldown_time=60, + cache=MagicMock(), + ) + + # Verify cooldown key format uses model_id (deployment-specific) + key = cache.get_cooldown_cache_key("deployment_abc123") + assert "deployment_abc123" in key + assert "deployment:" in key + + +class TestMemoryLeakReproduction: + """ + Simulated reproduction of the memory leak pattern. + + The primary leak sources in v1.81.12.rc.1 were: + 1. APScheduler jitter causing normalize() memory explosion + 2. Unbounded asyncio.Queue instances for spend tracking + 3. HTTP clients not being closed when evicted from cache + """ + + def test_http_client_eviction_closes_resources(self): + """Simulate high-volume HTTP client creation and verify cleanup.""" + from litellm.caching.llm_caching_handler import LLMClientCache + + close_count = 0 + + def make_mock_client(): + nonlocal close_count + client = MagicMock() + + def on_close(): + nonlocal close_count + close_count += 1 + + client.close = on_close + return client + + cache = LLMClientCache(max_size_in_memory=5, default_ttl=600) + + # Simulate creating many clients (like different providers under load) + for i in range(50): + cache.set_cache(f"client_{i}", make_mock_client()) + + # At least 45 clients should have been closed (50 - 5 remaining) + assert close_count >= 44, ( + f"Expected at least 44 clients to be closed on eviction, " + f"but only {close_count} were closed" + )