ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #13852 from michal-otmianowski/async-s3-implementation

Browse Source 
Async s3 implementation
This commit is contained in:
Krish Dholakia 2025-08-21 23:19:51 -07:00 committed by GitHub
commit 59e43595b9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 233 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
litellm/caching/caching.py
View File

@ -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

8
litellm/caching/caching_handler.py
View File

@ -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(

38
litellm/caching/s3_cache.py
View File

@ -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

222
tests/test_litellm/caching/test_s3_cache.py
View File

@ -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/"
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()