From 34fbbe30625a67d687de8432024da717b0ab28be Mon Sep 17 00:00:00 2001 From: milan Date: Wed, 19 Aug 2026 22:04:32 +0000 Subject: [PATCH] fix(proxy): guarantee parallel slot release when stream logging fails and make slot TTL configurable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 71 ++++++++++++------- .../hooks/parallel_request_limiter_v3.py | 2 +- .../test_litellm_logging.py | 68 ++++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 23 ++++++ 4 files changed, 137 insertions(+), 27 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 946110abf9e..edc2624ca05 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2740,13 +2740,25 @@ class Logging(LiteLLMLoggingBaseClass): "Model=%s not found in completion cost map. Setting 'response_cost' to None", self.model ) self.model_call_details["response_cost"] = None + except Exception: # noqa: BLE001 # cost calculation must never block slot-release callbacks + verbose_logger.exception( + "Error calculating streaming response cost for model=%s. Setting 'response_cost' to None", + self.model, + ) + self.model_call_details["response_cost"] = None self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + try: + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) + except Exception: # noqa: BLE001 # payload build must never block slot-release callbacks + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception building the standard logging payload " + "for a streaming response; callbacks still run without it" + ) # print standard logging payload if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: @@ -2786,32 +2798,39 @@ class Logging(LiteLLMLoggingBaseClass): ## LOGGING HOOK ## for callback in callbacks: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks + try: + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + continue + + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, ) - is not True - ): - continue - - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - result = redact_message_input_output_from_custom_logger( - result=result, litellm_logging_obj=self, custom_logger=callback - ) - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, + elif isinstance(callback, CustomLogger): + result = redact_message_input_output_from_custom_logger( + result=result, litellm_logging_obj=self, custom_logger=callback + ) + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", + traceback.format_exc(), ) + self._handle_callback_failure(callback=callback) self.has_run_logging(event_type="async_success") diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 2d799ded752..071eb07e265 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -369,7 +369,7 @@ PROJECT_OTPM_DESCRIPTOR_KEY: Final = "model_per_project_otpm" # considered leaked (worker crashed without any release callback firing) and # pruned. Also the longest request duration the gauge can track: a request # running longer than this stops occupying its slot. -PARALLEL_REQUEST_SLOT_TTL_SECONDS: Final = 3600 +PARALLEL_REQUEST_SLOT_TTL_SECONDS: Final = int(os.getenv("LITELLM_PARALLEL_REQUEST_SLOT_TTL_SECONDS", "3600")) CacheCounterValue: TypeAlias = int | float | str | bytes 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 54016470f8b..e459b500355 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -4958,3 +4958,71 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj): raw_api_base = logging_obj.model_call_details["raw_request_typed_dict"]["raw_request_api_base"] assert _GEMINI_KEY not in raw_api_base assert "key=*****" in raw_api_base + + +def _streaming_logging_obj_with_callbacks(callbacks): + from datetime import datetime + + obj = LitellmLogging( + model="anthropic/claude-opus-4-6", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="slot-leak-test", + function_id="slot-leak-test", + ) + obj.model_call_details["litellm_params"] = {"metadata": {}} + return patch.object(obj, "get_combined_callback_list", return_value=callbacks), obj + + +def _assembled_stream_result(): + response = ModelResponse() + response.choices[0].message.content = "hello" + return response + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_logging_hook_failure(): + """Regression for leaked max_parallel_requests slots: a raising + async_logging_hook must not abort the success-callback loop that + releases the rate-limiter slot.""" + broken = CustomLogger() + broken.async_logging_hook = AsyncMock(side_effect=RuntimeError("broken stream payload")) + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([broken, releasing]) + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_cost_calculation_failure(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + with patcher, patch.object( + logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block") + ): + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + assert logging_obj.model_call_details["response_cost"] is None + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_standard_logging_payload_failure(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + with patcher, patch.object( + logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream") + ): + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + releasing.async_log_success_event.assert_awaited_once() diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 54366226dfb..338103cf1e8 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5858,3 +5858,26 @@ async def test_conflicting_token_limits_cannot_bypass_tpm_reservation(): ) assert exc_info.value.status_code == 429 + + +def test_parallel_request_slot_ttl_env_override(): + """LITELLM_PARALLEL_REQUEST_SLOT_TTL_SECONDS shortens how long a leaked + slot can block a key before being pruned as stale. The constant is bound at + import time, so the override is checked in a fresh interpreter.""" + import subprocess + + assert PARALLEL_REQUEST_SLOT_TTL_SECONDS == 3600 + + output = subprocess.run( + [ + sys.executable, + "-c", + "from litellm.proxy.hooks.parallel_request_limiter_v3 import " + "PARALLEL_REQUEST_SLOT_TTL_SECONDS; print(PARALLEL_REQUEST_SLOT_TTL_SECONDS)", + ], + env={**os.environ, "LITELLM_PARALLEL_REQUEST_SLOT_TTL_SECONDS": "300"}, + capture_output=True, + text=True, + check=True, + ) + assert output.stdout.strip() == "300"