mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
[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:
parent
d408814978
commit
3db272b6d2
9 changed files with 536 additions and 3 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
438
litellm/integrations/s3_v2.py
Normal file
438
litellm/integrations/s3_v2.py
Normal 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)}")
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
13
litellm/types/integrations/s3_v2.py
Normal file
13
litellm/types/integrations/s3_v2.py
Normal 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
|
||||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue