mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(logging): guarantee max_parallel_requests slot release when streaming logging fails (#39093)
* fix(logging): guarantee max_parallel_requests slot release when stream logging fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): cover guardrail branch of streaming logging hook failure isolation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
f9c6eda909
commit
5988d93fed
2 changed files with 136 additions and 26 deletions
|
|
@ -2959,13 +2959,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 later callbacks (slot release)
|
||||
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 later callbacks (slot release)
|
||||
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:
|
||||
|
|
@ -3005,32 +3017,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")
|
||||
|
||||
|
|
|
|||
|
|
@ -5479,6 +5479,97 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj):
|
|||
assert "key=*****" in raw_api_base
|
||||
|
||||
|
||||
def _streaming_logging_obj_with_callbacks(callbacks: list[CustomLogger]):
|
||||
import datetime
|
||||
|
||||
obj = LitellmLogging(
|
||||
model="anthropic/claude-opus-5",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=True,
|
||||
call_type="completion",
|
||||
start_time=datetime.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())
|
||||
|
||||
assert logging_obj.model_call_details.get("standard_logging_object") is None
|
||||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failure():
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
||||
skipping = CustomGuardrail(guardrail_name="skipping-guardrail")
|
||||
skipping.should_run_guardrail = MagicMock(return_value=False)
|
||||
skipping.async_logging_hook = AsyncMock()
|
||||
raising = CustomGuardrail(guardrail_name="raising-guardrail")
|
||||
raising.should_run_guardrail = MagicMock(return_value=True)
|
||||
raising.async_logging_hook = AsyncMock(side_effect=RuntimeError("guardrail hook failed"))
|
||||
releasing = CustomLogger()
|
||||
releasing.async_log_success_event = AsyncMock()
|
||||
|
||||
patcher, logging_obj = _streaming_logging_obj_with_callbacks([skipping, raising, releasing])
|
||||
with patcher:
|
||||
await logging_obj.async_success_handler(result=_assembled_stream_result())
|
||||
|
||||
skipping.async_logging_hook.assert_not_awaited()
|
||||
raising.async_logging_hook.assert_awaited_once()
|
||||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
def _resolve(custom_llm_provider, litellm_params, optional_params, model):
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_resolve_vertex_location_for_cost,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue