scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

[Performance] Reduce Significant CPU overhead from litellm_logging.py (#13895)

Browse source 
* fix: litellm.configured_cold_storage_logger

* fix Session Management - Non-OpenAI Models docs

* ruff fix

* test fix
This commit is contained in:
Ishaan Jaff 2025-08-22 16:17:57 -07:00
parent 3c1dcb64cc
commit 5210f97da6
7 changed files with 17 additions and 56 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
docs/my-website/docs/response_api.md
View file

@ -805,16 +805,17 @@ LiteLLM Proxy supports session management for non-OpenAI models. This allows you
Set `store_prompts_in_cold_storage: true` in your proxy config.yaml. When this is enabled, LiteLLM will store the request and response content in the s3 bucket you specify.
```yaml
```yaml showLineNumbers title="config.yaml with Session Continuity"
litellm_settings:
callbacks: ["s3_v2"]
cold_storage_custom_logger: s3_v2
s3_callback_params: # learn more https://docs.litellm.ai/docs/proxy/logging#s3-buckets
s3_bucket_name: litellm-logs # AWS Bucket Name for S3
s3_region_name: us-west-2
s3_region_name: us-west-2
general_settings:
cold_storage_custom_logger: s3_v2
store_prompts_in_cold_storage: true
store_prompts_in_spend_logs: true
```
2. Make request 1 with no `previous_response_id` (new session)

1
litellm/__init__.py
View file

@ -146,6 +146,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"vector_store_pre_call_hook",
"dotprompt",
]
configured_cold_storage_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None
_known_custom_logger_compatible_callbacks: List = list(
get_args(_custom_logger_compatible_callbacks_literal)

11
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4106,18 +4106,9 @@ class StandardLoggingPayloadSetup:
"""
# Generate object key in same format as S3Logger
from litellm.integrations.s3 import get_s3_object_key
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
# Only generate object key if cold storage is configured
try:
configured_cold_storage_logger = (
ColdStorageHandler._get_configured_cold_storage_custom_logger()
)
except Exception as e:
verbose_logger.debug(f"Cold storage custom logger unavailable: {e}")
return None
if configured_cold_storage_logger is None:
if litellm.configured_cold_storage_logger is None:
return None
try:

37
litellm/proxy/spend_tracking/cold_storage_handler.py
View file

@ -3,11 +3,10 @@ This module is responsible for handling Getting/Setting the proxy server request
It allows fetching a dict of the proxy server request from s3 or GCS bucket.
"""
from typing import Optional, cast
from typing import Optional
import litellm
from litellm import _custom_logger_compatible_callbacks_literal
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
@ -57,38 +56,6 @@ class ColdStorageHandler:
def _select_custom_logger_for_cold_storage(
self,
) -> Optional[_custom_logger_compatible_callbacks_literal]:
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = ColdStorageHandler._get_configured_cold_storage_custom_logger()
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = litellm.configured_cold_storage_logger
return cold_storage_custom_logger
@staticmethod
def _get_configured_cold_storage_custom_logger() -> Optional[_custom_logger_compatible_callbacks_literal]:
"""Return the configured cold storage custom logger.
During interpreter shutdown importing ``proxy_server`` can raise a
``RuntimeError`` (e.g. "can't register atexit after shutdown").
In these scenarios we gracefully return ``None`` instead of bubbling
the exception up the call stack.
"""
try:
from litellm.proxy.proxy_server import general_settings
except Exception as e:
verbose_proxy_logger.debug(
f"Unable to import proxy_server for cold storage logging: {e}"
)
return None
cold_storage_custom_logger: Optional[str] = general_settings.get(
"cold_storage_custom_logger"
)
if not cold_storage_custom_logger:
verbose_proxy_logger.debug(
"No cold storage custom logger found in general settings"
)
return None
return cast(
_custom_logger_compatible_callbacks_literal, cold_storage_custom_logger
)

5
litellm/responses/litellm_completion_transformation/session_handler.py
View file

@ -1,6 +1,7 @@
import json
from typing import TYPE_CHECKING, Any, List, Optional, Union, cast
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
@ -235,10 +236,10 @@ class ResponsesSessionHandler:
"""
Only check cold storage when both are true
1. `LITELLM_TRUNCATED_PAYLOAD_FIELD` is in the proxy server request dict
2. `ColdStorageHandler._get_configured_cold_storage_custom_logger()` is not None
2. `litellm.configured_cold_storage_logger` is not None
"""
from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD
configured_cold_storage_custom_logger = ColdStorageHandler._get_configured_cold_storage_custom_logger()
configured_cold_storage_custom_logger = litellm.configured_cold_storage_logger
if configured_cold_storage_custom_logger is None:
return False
if proxy_server_request_dict is None:

6
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -411,7 +411,7 @@ async def test_e2e_generate_cold_storage_object_key_successful():
response_id = "chatcmpl-test-12345"
team_alias = "test-team"
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"), \
with patch("litellm.configured_cold_storage_logger", return_value="s3"), \
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
# Mock the S3 object key generation to return a predictable result
@ -453,7 +453,7 @@ async def test_e2e_generate_cold_storage_object_key_not_configured():
response_id = "chatcmpl-test-67890"
team_alias = "another-team"
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value=None):
with patch("litellm.configured_cold_storage_logger", return_value=None):
# Call the function
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
@ -479,7 +479,7 @@ async def test_e2e_generate_cold_storage_object_key_runtime_error_handled():
team_alias = "team"
with patch(
"litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger",
"litellm.configured_cold_storage_logger",
side_effect=RuntimeError("can't register atexit after shutdown"),
):
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(

6
tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py
View file

@ -223,7 +223,7 @@ async def test_e2e_cold_storage_successful_retrieval():
new_callable=AsyncMock,
) as mock_get_spend_logs, \
patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, \
patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"):
patch("litellm.configured_cold_storage_logger", return_value="s3"):
# Setup mocks
mock_get_spend_logs.return_value = mock_spend_logs
@ -343,7 +343,7 @@ async def test_should_check_cold_storage_for_full_payload():
# Test case 4: None request (should return True)
proxy_request_none = None
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value="s3"):
with patch("litellm.configured_cold_storage_logger", return_value="s3"):
# Test case 1: Should return True for truncated content
result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf)
assert result1 == True, "Should return True for proxy request with truncated PDF content"
@ -361,6 +361,6 @@ async def test_should_check_cold_storage_for_full_payload():
assert result4 == True, "Should return True for None proxy request"
# Test case 5: Should return False when cold storage is not configured
with patch("litellm.proxy.spend_tracking.cold_storage_handler.ColdStorageHandler._get_configured_cold_storage_custom_logger", return_value=None):
with patch("litellm.configured_cold_storage_logger", return_value=None):
result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf)
assert result5 == False, "Should return False when cold storage is not configured, even with truncated content"