From 56c16fc6f656c5ba471fa0538e9430fba696953f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 22 Aug 2025 16:17:57 -0700 Subject: [PATCH] [Performance] Reduce Significant CPU overhead from litellm_logging.py (#13895) * fix: litellm.configured_cold_storage_logger * fix Session Management - Non-OpenAI Models docs * ruff fix * test fix --- docs/my-website/docs/response_api.md | 7 ++-- litellm/__init__.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 11 +----- .../spend_tracking/cold_storage_handler.py | 37 +------------------ .../session_handler.py | 5 ++- .../test_litellm_logging.py | 6 +-- .../test_session_handler.py | 6 +-- 7 files changed, 17 insertions(+), 56 deletions(-) diff --git a/docs/my-website/docs/response_api.md b/docs/my-website/docs/response_api.md index c0bb003c096..94d7c73be05 100644 --- a/docs/my-website/docs/response_api.md +++ b/docs/my-website/docs/response_api.md @@ -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) diff --git a/litellm/__init__.py b/litellm/__init__.py index b30e9758d45..44254bd259c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 385df83c904..7bc7702684d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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: diff --git a/litellm/proxy/spend_tracking/cold_storage_handler.py b/litellm/proxy/spend_tracking/cold_storage_handler.py index b3ab7b6224c..133403deae2 100644 --- a/litellm/proxy/spend_tracking/cold_storage_handler.py +++ b/litellm/proxy/spend_tracking/cold_storage_handler.py @@ -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 - ) \ No newline at end of file diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 886c46e7373..bfb996a2385 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -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: diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 84edbb3ac5c..8fd73dbb79f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -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( diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py index 1cc159bf963..c683b571345 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py @@ -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"