[Perf] - Add Async + Batched S3 Logging (#11340)

* fix: add s3 v2 async

* fix: add s3 v2 async

* fix: add s3 v2 async

* test: s3 v2 logging

* fixes: s3 logging

* fixes: s3 logging use max upload batch size

* fixes: s3 logging tests

* fixes: s3 logging tests

* fixes: s3 logging tests
This commit is contained in:
Ishaan Jaff 2025-06-02 21:52:34 -07:00 • committed by GitHub
parent d408814978
commit 3db272b6d2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 536 additions and 3 deletions

View file

@ -1158,6 +1158,7 @@ jobs:
pip install "google-cloud-aiplatform==1.43.0"
pip install "mlflow==2.17.2"
pip install "anthropic==0.52.0"
pip install "blockbuster==1.5.24"
# Run pytest and generate JUnit XML report
- setup_litellm_enterprise_pip
- run:

View file

@ -407,6 +407,8 @@ router_settings:
| DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND | Default price per second for Replicate GPU. Default is 0.001400
| DEFAULT_REPLICATE_POLLING_DELAY_SECONDS | Default delay in seconds for Replicate polling. Default is 1
| DEFAULT_REPLICATE_POLLING_RETRIES | Default number of retries for Replicate polling. Default is 5
| DEFAULT_S3_BATCH_SIZE | Default batch size for S3 logging. Default is 512
| DEFAULT_S3_FLUSH_INTERVAL_SECONDS | Default flush interval for S3 logging. Default is 10
| DEFAULT_SLACK_ALERTING_THRESHOLD | Default threshold for Slack alerting. Default is 300
| DEFAULT_SOFT_BUDGET | Default soft budget for LiteLLM proxy keys. Default is 50.0
| DEFAULT_TRIM_RATIO | Default ratio of tokens to trim from prompt end. Default is 0.75

View file

@ -119,6 +119,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"resend_email",
"smtp_email",
"deepeval",
"s3_v2",
]
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
_known_custom_logger_compatible_callbacks: List = list(

View file

@ -4,6 +4,10 @@ from typing import List, Literal
ROUTER_MAX_FALLBACKS = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
DEFAULT_BATCH_SIZE = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
DEFAULT_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
DEFAULT_S3_FLUSH_INTERVAL_SECONDS = int(
os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)
)
DEFAULT_S3_BATCH_SIZE = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512))
DEFAULT_MAX_RETRIES = int(os.getenv("DEFAULT_MAX_RETRIES", 2))
DEFAULT_MAX_RECURSE_DEPTH = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100))
DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int(

View file

@ -0,0 +1,438 @@
"""
s3 Bucket Logging Integration
async_log_success_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3
NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to upload each element individually
"""
import asyncio
import json
from datetime import datetime
from typing import List, Optional, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS
from litellm.integrations.s3 import get_s3_object_key
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
from litellm.types.utils import StandardLoggingPayload
from .custom_batch_logger import CustomBatchLogger
class S3Logger(CustomBatchLogger, BaseAWSLLM):
def __init__(
self,
s3_bucket_name: Optional[str] = None,
s3_path: Optional[str] = None,
s3_region_name: Optional[str] = None,
s3_api_version: Optional[str] = None,
s3_use_ssl: bool = True,
s3_verify: Optional[bool] = None,
s3_endpoint_url: Optional[str] = None,
s3_aws_access_key_id: Optional[str] = None,
s3_aws_secret_access_key: Optional[str] = None,
s3_aws_session_token: Optional[str] = None,
s3_aws_session_name: Optional[str] = None,
s3_aws_profile_name: Optional[str] = None,
s3_aws_role_name: Optional[str] = None,
s3_aws_web_identity_token: Optional[str] = None,
s3_aws_sts_endpoint: Optional[str] = None,
s3_flush_interval: Optional[int] = DEFAULT_S3_FLUSH_INTERVAL_SECONDS,
s3_batch_size: Optional[int] = DEFAULT_S3_BATCH_SIZE,
s3_config=None,
s3_use_team_prefix: bool = False,
**kwargs,
):
try:
verbose_logger.debug(
f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}"
)
# IMPORTANT: We use a concurrent limit of 1 to upload to s3
# Files should get uploaded BUT they should not impact latency of LLM calling logic
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback,
)
self._init_s3_params(
s3_bucket_name=s3_bucket_name,
s3_region_name=s3_region_name,
s3_api_version=s3_api_version,
s3_use_ssl=s3_use_ssl,
s3_verify=s3_verify,
s3_endpoint_url=s3_endpoint_url,
s3_aws_access_key_id=s3_aws_access_key_id,
s3_aws_secret_access_key=s3_aws_secret_access_key,
s3_aws_session_token=s3_aws_session_token,
s3_aws_session_name=s3_aws_session_name,
s3_aws_profile_name=s3_aws_profile_name,
s3_aws_role_name=s3_aws_role_name,
s3_aws_web_identity_token=s3_aws_web_identity_token,
s3_aws_sts_endpoint=s3_aws_sts_endpoint,
s3_config=s3_config,
s3_path=s3_path,
s3_use_team_prefix=s3_use_team_prefix,
)
verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}")
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
verbose_logger.debug(
f"s3 flush interval: {s3_flush_interval}, s3 batch size: {s3_batch_size}"
)
# Call CustomLogger's __init__
CustomBatchLogger.__init__(
self,
flush_lock=self.flush_lock,
flush_interval=s3_flush_interval,
batch_size=s3_batch_size,
)
self.log_queue: List[s3BatchLoggingElement] = []
# Call BaseAWSLLM's __init__
BaseAWSLLM.__init__(self)
except Exception as e:
print_verbose(f"Got exception on init s3 client {str(e)}")
raise e
def _init_s3_params(
self,
s3_bucket_name: Optional[str] = None,
s3_region_name: Optional[str] = None,
s3_api_version: Optional[str] = None,
s3_use_ssl: bool = True,
s3_verify: Optional[bool] = None,
s3_endpoint_url: Optional[str] = None,
s3_aws_access_key_id: Optional[str] = None,
s3_aws_secret_access_key: Optional[str] = None,
s3_aws_session_token: Optional[str] = None,
s3_aws_session_name: Optional[str] = None,
s3_aws_profile_name: Optional[str] = None,
s3_aws_role_name: Optional[str] = None,
s3_aws_web_identity_token: Optional[str] = None,
s3_aws_sts_endpoint: Optional[str] = None,
s3_config=None,
s3_path: Optional[str] = None,
s3_use_team_prefix: bool = False,
):
"""
Initialize the s3 params for this logging callback
"""
litellm.s3_callback_params = litellm.s3_callback_params or {}
# read in .env variables - example os.environ/AWS_BUCKET_NAME
for key, value in litellm.s3_callback_params.items():
if isinstance(value, str) and value.startswith("os.environ/"):
litellm.s3_callback_params[key] = litellm.get_secret(value)
self.s3_bucket_name = (
litellm.s3_callback_params.get("s3_bucket_name") or s3_bucket_name
)
self.s3_region_name = (
litellm.s3_callback_params.get("s3_region_name") or s3_region_name
)
self.s3_api_version = (
litellm.s3_callback_params.get("s3_api_version") or s3_api_version
)
self.s3_use_ssl = (
litellm.s3_callback_params.get("s3_use_ssl", True) or s3_use_ssl
)
self.s3_verify = litellm.s3_callback_params.get("s3_verify") or s3_verify
self.s3_endpoint_url = (
litellm.s3_callback_params.get("s3_endpoint_url") or s3_endpoint_url
)
self.s3_aws_access_key_id = (
litellm.s3_callback_params.get("s3_aws_access_key_id")
or s3_aws_access_key_id
)
self.s3_aws_secret_access_key = (
litellm.s3_callback_params.get("s3_aws_secret_access_key")
or s3_aws_secret_access_key
)
self.s3_aws_session_token = (
litellm.s3_callback_params.get("s3_aws_session_token")
or s3_aws_session_token
)
self.s3_aws_session_name = (
litellm.s3_callback_params.get("s3_aws_session_name") or s3_aws_session_name
)
self.s3_aws_profile_name = (
litellm.s3_callback_params.get("s3_aws_profile_name") or s3_aws_profile_name
)
self.s3_aws_role_name = (
litellm.s3_callback_params.get("s3_aws_role_name") or s3_aws_role_name
)
self.s3_aws_web_identity_token = (
litellm.s3_callback_params.get("s3_aws_web_identity_token")
or s3_aws_web_identity_token
)
self.s3_aws_sts_endpoint = (
litellm.s3_callback_params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint
)
self.s3_config = litellm.s3_callback_params.get("s3_config") or s3_config
self.s3_path = litellm.s3_callback_params.get("s3_path") or s3_path
# done reading litellm.s3_callback_params
self.s3_use_team_prefix = (
bool(litellm.s3_callback_params.get("s3_use_team_prefix", False))
or s3_use_team_prefix
)
return
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
verbose_logger.debug(
f"s3 Logging - Enters logging function for model {kwargs}"
)
s3_batch_logging_element = self.create_s3_batch_logging_element(
start_time=start_time,
standard_logging_payload=kwargs.get("standard_logging_object", None),
)
if s3_batch_logging_element is None:
raise ValueError("s3_batch_logging_element is None")
verbose_logger.debug(
"\ns3 Logger - Logging payload = %s", s3_batch_logging_element
)
self.log_queue.append(s3_batch_logging_element)
verbose_logger.debug(
"s3 logging: queue length %s, batch size %s",
len(self.log_queue),
self.batch_size,
)
except Exception as e:
verbose_logger.exception(f"s3 Layer Error - {str(e)}")
pass
async def async_upload_data_to_s3(
self, batch_logging_element: s3BatchLoggingElement
):
try:
import hashlib
import requests
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
try:
from litellm.litellm_core_utils.asyncify import asyncify
asyncified_get_credentials = asyncify(self.get_credentials)
credentials = await asyncified_get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
aws_session_token=self.s3_aws_session_token,
aws_region_name=self.s3_region_name,
aws_session_name=self.s3_aws_session_name,
aws_profile_name=self.s3_aws_profile_name,
aws_role_name=self.s3_aws_role_name,
aws_web_identity_token=self.s3_aws_web_identity_token,
aws_sts_endpoint=self.s3_aws_sts_endpoint,
)
verbose_logger.debug(
f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}"
)
# Prepare the URL
url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}"
if self.s3_endpoint_url:
url = self.s3_endpoint_url + "/" + batch_logging_element.s3_object_key
# Convert JSON to string
json_string = json.dumps(batch_logging_element.payload)
# Calculate SHA256 hash of the content
content_hash = hashlib.sha256(json_string.encode("utf-8")).hexdigest()
# Prepare the request
headers = {
"Content-Type": "application/json",
"x-amz-content-sha256": content_hash,
"Content-Language": "en",
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
}
req = requests.Request("PUT", url, data=json_string, headers=headers)
prepped = req.prepare()
# Sign the request
aws_request = AWSRequest(
method=prepped.method,
url=prepped.url,
data=prepped.body,
headers=prepped.headers,
)
SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
# Prepare the signed headers
signed_headers = dict(aws_request.headers.items())
# Make the request
response = await self.async_httpx_client.put(
url, data=json_string, headers=signed_headers
)
response.raise_for_status()
except Exception as e:
verbose_logger.exception(f"Error uploading to s3: {str(e)}")
async def async_send_batch(self):
"""
Sends runs from self.log_queue
Returns: None
Raises: Does not raise an exception, will only verbose_logger.exception()
"""
verbose_logger.debug(f"s3_v2 logger - sending batch of {len(self.log_queue)}")
if not self.log_queue:
return
#########################################################
# Flush the log queue to s3
# the log queue can be bounded by DEFAULT_S3_BATCH_SIZE
# see custom_batch_logger.py which triggers the flush
#########################################################
for payload in self.log_queue:
asyncio.create_task(self.async_upload_data_to_s3(payload))
def create_s3_batch_logging_element(
self,
start_time: datetime,
standard_logging_payload: Optional[StandardLoggingPayload],
) -> Optional[s3BatchLoggingElement]:
"""
Helper function to create an s3BatchLoggingElement.
Args:
start_time (datetime): The start time of the logging event.
standard_logging_payload (Optional[StandardLoggingPayload]): The payload to be logged.
s3_path (Optional[str]): The S3 path prefix.
Returns:
Optional[s3BatchLoggingElement]: The created s3BatchLoggingElement, or None if payload is None.
"""
if standard_logging_payload is None:
return None
team_alias = standard_logging_payload["metadata"].get("user_api_key_team_alias")
team_alias_prefix = ""
if (
litellm.enable_preview_features
and self.s3_use_team_prefix
and team_alias is not None
):
team_alias_prefix = f"{team_alias}/"
s3_file_name = (
litellm.utils.get_logging_id(start_time, standard_logging_payload) or ""
)
s3_object_key = get_s3_object_key(
s3_path=cast(Optional[str], self.s3_path) or "",
team_alias_prefix=team_alias_prefix,
start_time=start_time,
s3_file_name=s3_file_name,
)
s3_object_download_filename = (
"time-"
+ start_time.strftime("%Y-%m-%dT%H-%M-%S-%f")
+ "_"
+ standard_logging_payload["id"]
+ ".json"
)
s3_object_download_filename = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json"
return s3BatchLoggingElement(
payload=dict(standard_logging_payload),
s3_object_key=s3_object_key,
s3_object_download_filename=s3_object_download_filename,
)
def upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement):
try:
import hashlib
import requests
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
try:
verbose_logger.debug(
f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}"
)
credentials: Credentials = self.get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
aws_session_token=self.s3_aws_session_token,
aws_region_name=self.s3_region_name,
)
# Prepare the URL
url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}"
if self.s3_endpoint_url:
url = self.s3_endpoint_url + "/" + batch_logging_element.s3_object_key
# Convert JSON to string
json_string = json.dumps(batch_logging_element.payload)
# Calculate SHA256 hash of the content
content_hash = hashlib.sha256(json_string.encode("utf-8")).hexdigest()
# Prepare the request
headers = {
"Content-Type": "application/json",
"x-amz-content-sha256": content_hash,
"Content-Language": "en",
"Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"',
"Cache-Control": "private, immutable, max-age=31536000, s-maxage=0",
}
req = requests.Request("PUT", url, data=json_string, headers=headers)
prepped = req.prepare()
# Sign the request
aws_request = AWSRequest(
method=prepped.method,
url=prepped.url,
data=prepped.body,
headers=prepped.headers,
)
SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request)
# Prepare the signed headers
signed_headers = dict(aws_request.headers.items())
httpx_client = _get_httpx_client()
# Make the request
response = httpx_client.put(url, data=json_string, headers=signed_headers)
response.raise_for_status()
except Exception as e:
verbose_logger.exception(f"Error uploading to s3: {str(e)}")

View file

@ -135,6 +135,7 @@ from ..integrations.opik.opik import OpikLogger
from ..integrations.prometheus import PrometheusLogger
from ..integrations.prompt_layer import PromptLayerLogger
from ..integrations.s3 import S3Logger
from ..integrations.s3_v2 import S3Logger as S3V2Logger
from ..integrations.supabase import Supabase
from ..integrations.traceloop import TraceloopLogger
from ..integrations.weights_biases import WeightsBiasesLogger
@ -2699,7 +2700,9 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915
sentry_sdk_instance.init(
dsn=os.environ.get("SENTRY_DSN"),
traces_sample_rate=float(sentry_trace_rate), # type: ignore
sample_rate=float(sentry_sample_rate),
sample_rate=float(
sentry_sample_rate if sentry_sample_rate else 1.0
),
)
capture_exception = sentry_sdk_instance.capture_exception
add_breadcrumb = sentry_sdk_instance.add_breadcrumb
@ -2867,6 +2870,14 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_gcs_bucket_logger = GCSBucketLogger()
_in_memory_loggers.append(_gcs_bucket_logger)
return _gcs_bucket_logger # type: ignore
elif logging_integration == "s3_v2":
for callback in _in_memory_loggers:
if isinstance(callback, S3V2Logger):
return callback # type: ignore
_s3_v2_logger = S3V2Logger()
_in_memory_loggers.append(_s3_v2_logger)
return _s3_v2_logger # type: ignore
elif logging_integration == "azure_storage":
for callback in _in_memory_loggers:
if isinstance(callback, AzureBlobStorageLogger):
@ -2962,7 +2973,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
galileo_logger = GalileoObserve()
_in_memory_loggers.append(galileo_logger)
return galileo_logger # type: ignore
elif logging_integration == "deepeval":
for callback in _in_memory_loggers:
if isinstance(callback, DeepEvalLogger):
@ -2970,7 +2981,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
deepeval_logger = DeepEvalLogger()
_in_memory_loggers.append(deepeval_logger)
return deepeval_logger # type: ignore
elif logging_integration == "logfire":
if "LOGFIRE_TOKEN" not in os.environ:
raise ValueError("LOGFIRE_TOKEN not found in environment variables")
@ -3172,6 +3183,10 @@ def get_custom_logger_compatible_class( # noqa: PLR0915
for callback in _in_memory_loggers:
if isinstance(callback, GCSBucketLogger):
return callback
elif logging_integration == "s3_v2":
for callback in _in_memory_loggers:
if isinstance(callback, S3V2Logger):
return callback
elif logging_integration == "azure_storage":
for callback in _in_memory_loggers:
if isinstance(callback, AzureBlobStorageLogger):

View file

@ -0,0 +1,13 @@
from typing import Dict
from pydantic import BaseModel
class s3BatchLoggingElement(BaseModel):
"""
Type of element stored in self.log_queue in S3Logger
"""
payload: Dict
s3_object_key: str
s3_object_download_filename: str

View file

@ -74,6 +74,59 @@ async def test_basic_s3_logging(sync_mode, streaming):
s3.delete_object(Bucket="load-testing-oct", Key=key)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"streaming", [(True)]
)
async def test_basic_s3_v2_logging(streaming):
from blockbuster import BlockBuster
from litellm.integrations.s3_v2 import S3Logger
s3_v2_logger = S3Logger(s3_flush_interval=1)
litellm.callbacks = [s3_v2_logger]
blockbuster = BlockBuster()
blockbuster.activate()
litellm._turn_on_debug()
litellm.callbacks = ["s3_v2"]
litellm.s3_callback_params = {
"s3_bucket_name": "load-testing-oct",
"s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY",
"s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID",
"s3_region_name": "us-west-2",
}
litellm.set_verbose = True
response_id = None
response = await litellm.acompletion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "This is a test"}],
stream=streaming,
)
if streaming:
async for chunk in response:
print(chunk)
response_id = chunk.id
else:
response_id = response.id
await asyncio.sleep(30)
print(f"response: {response}")
# stop blockbuster
blockbuster.deactivate()
total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct")
print(f"all_s3_keys: {all_s3_keys}")
#assert that atlest one key has response.id in it
assert any(response_id in key for key in all_s3_keys)
s3 = boto3.client("s3")
# delete all objects
for key in all_s3_keys:
s3.delete_object(Bucket="load-testing-oct", Key=key)
def list_all_s3_objects(bucket_name):
s3 = boto3.client("s3")

View file

@ -34,6 +34,7 @@ from litellm.integrations.opentelemetry import OpenTelemetry
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.s3_v2 import S3Logger
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVectorStore
from litellm.integrations.langfuse.langfuse_prompt_management import (
@ -88,6 +89,7 @@ callback_class_str_to_classType = {
"resend_email": ResendEmailLogger,
"smtp_email": SMTPEmailLogger,
"deepeval": DeepEvalLogger,
"s3_v2": S3Logger,
}
expected_env_vars = {
@ -113,6 +115,10 @@ expected_env_vars = {
"GCS_PUBSUB_PROJECT_ID": "gcs_pubsub_project_id",
"CONFIDENT_API_KEY": "confident_api_key",
"LITELM_ENVIRONMENT": "development",
"AWS_BUCKET_NAME": "aws_bucket_name",
"AWS_SECRET_ACCESS_KEY": "aws_secret_access_key",
"AWS_ACCESS_KEY_ID": "aws_access_key_id",
"AWS_REGION": "aws_region",
}