diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 0a8b6ef18d..82fc37e0cb 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 7580752c30..1dcc0f1fdb 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 6095d547aa..15f7a5c1e1 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( @@ -98,11 +98,20 @@ class S3Cache(BaseCache): ContentDisposition=f'inline; filename="{key}.json"', ) except Exception as e: - # NON blocking - notify users S3 is throwing an exception 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"Set ASYNC S3 Cache: 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: + verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}") def get_cache(self, key, **kwargs): import botocore @@ -142,13 +151,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.py b/tests/test_litellm/caching/test_s3_cache.py index a37f7f5706..dce3f7d585 100644 --- a/tests/test_litellm/caching/test_s3_cache.py +++ b/tests/test_litellm/caching/test_s3_cache.py @@ -3,6 +3,7 @@ import sys from unittest.mock import MagicMock, patch import json import datetime +import asyncio import pytest @@ -16,7 +17,7 @@ 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} @@ -25,12 +26,12 @@ def test_s3_cache_set_cache(mock_s3_dependencies): """Test basic set_cache functionality""" cache = S3Cache("test-bucket") test_value = {"key": "value", "number": 42} - + cache.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) @@ -43,7 +44,7 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): """Test set_cache with TTL functionality""" cache = S3Cache("test-bucket") test_value = {"key": "value"} - ttl = datetime.timedelta(seconds=3600) # 1 hour + ttl = 3600 # 1 hour in seconds cache.set_cache("test_key", test_value, ttl=ttl) @@ -52,44 +53,44 @@ def test_s3_cache_set_cache_with_ttl(mock_s3_dependencies): assert "Expires" in call_args[1] assert "CacheControl" in call_args[1] - assert "max-age=1:00:00" in call_args[1]["CacheControl"] + assert "max-age=3600" in call_args[1]["CacheControl"] def test_s3_cache_get_cache(mock_s3_dependencies): """Test basic get_cache functionality""" 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 = cache.get_cache("test_key") - + cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", + Bucket="test-bucket", Key="test_key" ) - + assert result == {"key": "value", "number": 42} def test_s3_cache_get_cache_not_found(mock_s3_dependencies): """Test 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 = cache.get_cache("nonexistent_key") - + cache.s3_client.get_object.assert_called_once_with( - Bucket="test-bucket", + Bucket="test-bucket", Key="nonexistent_key" ) assert result is None @@ -98,16 +99,16 @@ def test_s3_cache_get_cache_not_found(mock_s3_dependencies): def test_s3_key_transformation(): """Test the _to_s3_key method for key transformation""" cache = S3Cache("test-bucket") - + # Test basic key transformation (colon to slash) result = cache._to_s3_key("user:123:session:456") assert result == "user/123/session/456" - + # Test with s3_path prefix cache_with_prefix = S3Cache("test-bucket", s3_path="cache/data") result = cache_with_prefix._to_s3_key("namespace:key") assert result == "cache/data/namespace/key" - + # Test with s3_path that has trailing slash cache_with_slash = S3Cache("test-bucket", s3_path="cache/data/") result = cache_with_slash._to_s3_key("namespace:key") @@ -120,7 +121,186 @@ def test_s3_cache_initialization(): cache = S3Cache("test-bucket") assert cache.bucket_name == "test-bucket" assert cache.key_prefix == "" - + # Test with s3_path cache_with_path = S3Cache("test-bucket", s3_path="my/cache/path") - assert cache_with_path.key_prefix == "my/cache/path/" \ No newline at end of file + assert cache_with_path.key_prefix == "my/cache/path/" + + +# ============================================================================ +# ASYNC TESTS +# ============================================================================ + + +@pytest.mark.asyncio +async def test_s3_cache_async_set_cache(mock_s3_dependencies): + 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): + cache = S3Cache("test-bucket") + test_value = {"key": "value"} + ttl = 3600 # 1 hour in seconds + + 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=3600" in call_args[1]["CacheControl"] + + +@pytest.mark.asyncio +async def test_s3_cache_async_get_cache(mock_s3_dependencies): + 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()