diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 0a8b6ef18d1..82fc37e0cb4 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -774,11 +774,9 @@ class Cache: """ Internal method to check if the cache type supports async get/set operations - Only S3 Cache Does NOT support async operations + All cache types now support async operations """ - if self.type and self.type == LiteLLMCacheType.S3: - return False return True diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 7580752c307..1dcc0f1fdb2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -599,7 +599,7 @@ class LLMCachingHandler: cached_result = await litellm.cache.async_get_cache( dynamic_cache_object=self.dual_cache, **new_kwargs ) - else: # for s3 caching. [NOT RECOMMENDED IN PROD - this will slow down responses since boto3 is sync] + else: # fallback for caches that don't support async cached_result = litellm.cache.get_cache( dynamic_cache_object=self.dual_cache, **new_kwargs ) @@ -806,12 +806,6 @@ class LLMCachingHandler: result, dynamic_cache_object=self.dual_cache, **new_kwargs ) ) - elif isinstance(litellm.cache.cache, S3Cache): - threading.Thread( - target=litellm.cache.add_cache, - args=(result,), - kwargs=new_kwargs, - ).start() else: asyncio.create_task( litellm.cache.async_add_cache( diff --git a/litellm/caching/s3_cache.py b/litellm/caching/s3_cache.py index 6095d547aab..929de6493cd 100644 --- a/litellm/caching/s3_cache.py +++ b/litellm/caching/s3_cache.py @@ -1,17 +1,17 @@ """ S3 Cache implementation -WARNING: DO NOT USE THIS IN PRODUCTION - This is not ASYNC Has 4 methods: - set_cache - get_cache - - async_set_cache - - async_get_cache + - async_set_cache (uses run_in_executor) + - async_get_cache (uses run_in_executor) """ import ast import asyncio import json +from functools import partial from typing import Optional from litellm._logging import print_verbose, verbose_logger @@ -72,7 +72,7 @@ class S3Cache(BaseCache): import datetime # Calculate expiration time - expiration_time = datetime.datetime.now() + ttl + expiration_time = datetime.datetime.now() + datetime.timedelta(seconds=ttl) # Upload the data to S3 with the calculated expiration time self.s3_client.put_object( @@ -102,7 +102,18 @@ class S3Cache(BaseCache): print_verbose(f"S3 Caching: set_cache() - Got exception from S3: {e}") async def async_set_cache(self, key, value, **kwargs): - self.set_cache(key=key, value=value, **kwargs) + """ + Asynchronously set cache using run_in_executor to avoid blocking the event loop. + Compatible with Python 3.8+. + """ + try: + verbose_logger.debug(f"LiteLLM ASYNC SET Cache - S3. Key={key}. Value={value}") + loop = asyncio.get_event_loop() + func = partial(self.set_cache, key, value, **kwargs) + await loop.run_in_executor(None, func) + except Exception as e: + # NON blocking - notify users S3 is throwing an exception + verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") def get_cache(self, key, **kwargs): import botocore @@ -142,13 +153,26 @@ class S3Cache(BaseCache): return None except Exception as e: - # NON blocking - notify users S3 is throwing an exception verbose_logger.error( f"S3 Caching: get_cache() - Got exception from S3: {e}" ) async def async_get_cache(self, key, **kwargs): - return self.get_cache(key=key, **kwargs) + """ + Asynchronously get cache using run_in_executor to avoid blocking the event loop. + Compatible with Python 3.8+. + """ + try: + verbose_logger.debug(f"Get ASYNC S3 Cache: key: {key}") + loop = asyncio.get_event_loop() + func = partial(self.get_cache, key, **kwargs) + result = await loop.run_in_executor(None, func) + return result + except Exception as e: + verbose_logger.error( + f"S3 Caching: async_get_cache() - Got exception from S3: {e}" + ) + return None def flush_cache(self): pass diff --git a/tests/test_litellm/caching/test_s3_cache_async.py b/tests/test_litellm/caching/test_s3_cache_async.py new file mode 100644 index 00000000000..4c18a5f2955 --- /dev/null +++ b/tests/test_litellm/caching/test_s3_cache_async.py @@ -0,0 +1,207 @@ +import os +import sys +from unittest.mock import MagicMock, patch +import json +import datetime +import asyncio + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from litellm.caching.s3_cache import S3Cache + + +@pytest.fixture +def mock_s3_dependencies(): + mock_s3_client = MagicMock() + + with patch("boto3.client", return_value=mock_s3_client): + yield {"s3_client": mock_s3_client} + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache(mock_s3_dependencies): + """Test async_set_cache functionality using run_in_executor""" + cache = S3Cache("test-bucket") + test_value = {"key": "value", "number": 42} + + await cache.async_set_cache("test_key", test_value) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + assert call_args[1]["Bucket"] == "test-bucket" + assert call_args[1]["Key"] == "test_key" + assert call_args[1]["Body"] == json.dumps(test_value) + assert call_args[1]["ContentType"] == "application/json" + assert call_args[1]["ContentLanguage"] == "en" + assert call_args[1]["ContentDisposition"] == 'inline; filename="test_key.json"' + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache_with_ttl(mock_s3_dependencies): + """Test async_set_cache with TTL functionality""" + cache = S3Cache("test-bucket") + test_value = {"key": "value"} + ttl = datetime.timedelta(seconds=3600) # 1 hour + + await cache.async_set_cache("test_key", test_value, ttl=ttl) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + assert "Expires" in call_args[1] + assert "CacheControl" in call_args[1] + assert "max-age=1:00:00" in call_args[1]["CacheControl"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache(mock_s3_dependencies): + """Test async_get_cache functionality using run_in_executor""" + cache = S3Cache("test-bucket") + + mock_response = { + "Body": MagicMock() + } + mock_response["Body"].read.return_value = b'{"key": "value", "number": 42}' + cache.s3_client.get_object.return_value = mock_response + + result = await cache.async_get_cache("test_key") + + cache.s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="test_key" + ) + + assert result == {"key": "value", "number": 42} + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache_not_found(mock_s3_dependencies): + """Test async_get_cache when key is not found""" + import botocore.exceptions + + cache = S3Cache("test-bucket") + + error_response = {"Error": {"Code": "NoSuchKey"}} + cache.s3_client.get_object.side_effect = botocore.exceptions.ClientError( + error_response, "GetObject" + ) + + result = await cache.async_get_cache("nonexistent_key") + + cache.s3_client.get_object.assert_called_once_with( + Bucket="test-bucket", + Key="nonexistent_key" + ) + assert result is None + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache_pipeline(mock_s3_dependencies): + """Test async_set_cache_pipeline functionality""" + cache = S3Cache("test-bucket") + + cache_list = [ + ("key1", {"data": "value1"}), + ("key2", {"data": "value2"}), + ("key3", {"data": "value3"}), + ] + + await cache.async_set_cache_pipeline(cache_list) + + # Should have called put_object 3 times + assert cache.s3_client.put_object.call_count == 3 + + # Verify each call + calls = cache.s3_client.put_object.call_args_list + for i, (key, value) in enumerate(cache_list): + call_args = calls[i][1] + assert call_args["Bucket"] == "test-bucket" + assert call_args["Key"] == key + assert call_args["Body"] == json.dumps(value) + + +@pytest.mark.asyncio +async def test_s3_cache_concurrent_async_operations(mock_s3_dependencies): + """Test concurrent async operations to ensure they don't block each other""" + cache = S3Cache("test-bucket") + + # Create multiple concurrent set operations + tasks = [] + for i in range(5): + key = f"concurrent_key_{i}" + value = {"id": i, "data": f"test_data_{i}"} + tasks.append(cache.async_set_cache(key, value)) + + # Execute all tasks concurrently + await asyncio.gather(*tasks) + + # Verify all operations were called + assert cache.s3_client.put_object.call_count == 5 + + # Verify each call had correct parameters + calls = cache.s3_client.put_object.call_args_list + for i, call in enumerate(calls): + call_args = call[1] + assert call_args["Bucket"] == "test-bucket" + assert f"concurrent_key_{i}" == call_args["Key"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_error_handling(mock_s3_dependencies): + """Test that async methods handle errors gracefully""" + cache = S3Cache("test-bucket") + + # Test async_set_cache error handling + cache.s3_client.put_object.side_effect = Exception("S3 Error") + + # Should not raise exception, just log it + await cache.async_set_cache("error_key", {"data": "value"}) + + # Test async_get_cache error handling + cache.s3_client.get_object.side_effect = Exception("S3 Error") + + result = await cache.async_get_cache("error_key") + assert result is None + + +@pytest.mark.asyncio +async def test_s3_cache_async_with_key_prefix(mock_s3_dependencies): + """Test async operations with s3_path prefix""" + cache = S3Cache("test-bucket", s3_path="cache/data") + test_value = {"key": "value"} + + await cache.async_set_cache("namespace:key", test_value) + + cache.s3_client.put_object.assert_called_once() + call_args = cache.s3_client.put_object.call_args + + # Should transform key with prefix and colon replacement + assert call_args[1]["Key"] == "cache/data/namespace/key" + + +def test_s3_cache_supports_async(): + """Test that S3Cache now supports async operations""" + from litellm.caching.caching import Cache, LiteLLMCacheType + + cache = Cache(type=LiteLLMCacheType.S3, s3_bucket_name="test-bucket") + + # Should now return True for async support + assert cache._supports_async() is True + + +@pytest.mark.asyncio +async def test_s3_cache_async_disconnect(mock_s3_dependencies): + """Test async disconnect method""" + cache = S3Cache("test-bucket") + + # Should not raise any exceptions + await cache.disconnect() + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"])