From 2182e8590ce9857f3caecdfb4697b62a251b6ce0 Mon Sep 17 00:00:00 2001 From: adc77 Date: Tue, 20 Oct 2026 01:02:16 +0530 Subject: [PATCH] fix(logging): preserve spend accounting when callbacks time out --- litellm/integrations/custom_logger.py | 1 + litellm/litellm_core_utils/litellm_logging.py | 124 ++++++++++++------ .../proxy/hooks/proxy_track_cost_callback.py | 4 +- .../test_litellm_logging.py | 95 +++++++++++++- 4 files changed, 184 insertions(+), 40 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..ca7877134db 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -66,6 +66,7 @@ _BASE64_INLINE_PATTERN: Final = re.compile( class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset() + is_internal_accounting_callback: ClassVar[bool] = False enforces_request_content: bool = False """ diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0ac216c2029..3292da43247 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -19,6 +19,7 @@ from functools import lru_cache from types import MappingProxyType, TracebackType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast +from anyio import fail_after from httpx import Response, Timeout from pydantic import BaseModel, JsonValue, TypeAdapter @@ -111,9 +112,11 @@ from litellm.litellm_core_utils.logging_utils import ( truncate_base64_in_messages, truncate_base64_in_messages_async, ) +from litellm.litellm_core_utils.logging_worker import optional_callback_budget from litellm.litellm_core_utils.model_param_helper import ModelParamHelper from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request from litellm.litellm_core_utils.redact_messages import ( + perform_redaction, redact_message_input_output_from_custom_logger, redact_message_input_output_from_logging, redact_streaming_responses_for_custom_logger, @@ -3526,6 +3529,68 @@ class Logging(LiteLLMLoggingBaseClass): e, ) + async def _run_async_success_hooks( + self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object + ) -> object: + for callback in callbacks: + 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, + ) + is not True + ): + continue + + self.shadow_eval_request_snapshot = None + 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 = cast( # cast-ok: legacy redaction returns heterogeneous response objects + object, + 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) + + return result + + async def _prepare_accounting_success_hooks( + self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object + ) -> object: + from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE + + try: + with fail_after(optional_callback_budget(LOGGING_WORKER_MAX_TIME_PER_COROUTINE, fraction=0.5)): + return await self._run_async_success_hooks(callbacks, result) + except TimeoutError: + self.shadow_eval_request_snapshot = None + self.model_call_details = redact_streaming_responses_for_custom_logger( + self.model_call_details, CustomLogger(message_logging=False) + ) + verbose_logger.warning("Logging hooks timed out; recording spend with request and response bodies redacted") + return cast( # cast-ok: legacy redaction returns heterogeneous response objects + object, perform_redaction(self.model_call_details, result) + ) + async def async_success_handler( self, result: object = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) @@ -3735,7 +3800,7 @@ class Logging(LiteLLMLoggingBaseClass): self._get_completed_success_callbacks(asynchronous=True) if not self.stream else callbacks ) - result = redact_message_input_output_from_logging( + redacted_result: Final = redact_message_input_output_from_logging( model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), result=result, ) @@ -3743,45 +3808,30 @@ class Logging(LiteLLMLoggingBaseClass): ## LOGGING HOOK ## hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks))) - for callback in hook_callbacks: - 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, - ) - is not True - ): - continue - - self.shadow_eval_request_snapshot = None - 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) + accounting_callbacks: Final = tuple( + { + type(callback): callback + for callback in delivery_callbacks + if isinstance(callback, CustomLogger) and callback.is_internal_accounting_callback + }.values() + ) + result = ( + await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result) + if accounting_callbacks + else await self._run_async_success_hooks(hook_callbacks, redacted_result) + ) + ordered_callbacks: Final = ( + *accounting_callbacks, + *( + callback + for callback in delivery_callbacks + if not (isinstance(callback, CustomLogger) and callback.is_internal_accounting_callback) + ), + ) self.mark_logging_complete(event_type="async_success") - for callback in delivery_callbacks: + for callback in ordered_callbacks: # check if callback can run for this request litellm_params = self.model_call_details.get("litellm_params", {}) should_run = self.should_run_callback( diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3195394b9aa..967ce1ea328 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -3,7 +3,7 @@ import traceback from collections.abc import Callable, Mapping, Sequence from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, cast +from typing import TYPE_CHECKING, ClassVar, Final, Protocol, cast from pydantic import TypeAdapter, ValidationError @@ -154,6 +154,8 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: class _ProxyDBLogger(CustomLogger): + is_internal_accounting_callback: ClassVar[bool] = True + def __init__( self, spend_event_producer: SpendEventProducer | None = None, diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d1258e43154..4e057b7ea61 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import contextvars import copy import datetime import importlib.abc @@ -48,7 +49,7 @@ from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, set_callbacks, ) -from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck @@ -2125,7 +2126,7 @@ async def test_aresponses_websocket_turn_reaches_a_custom_logger_through_the_asy async def test_aresponses_websocket_hands_back_the_provider_failure_without_a_success_log( monkeypatch: pytest.MonkeyPatch, ): - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker from litellm.responses.main import base_llm_http_handler success_events = [] @@ -12958,3 +12959,93 @@ def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_a assert s3_writer.log_event.call_count == expected_s3_writes if expected_s3_writes: assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("accounting_first", [False, True]) +@pytest.mark.parametrize("duplicate_accounting", [False, True]) +@pytest.mark.parametrize("blocked_phase", ["hook", "delivery", "neither", "redacted"]) +async def test_slow_optional_logging_cannot_cancel_accounting( + monkeypatch: pytest.MonkeyPatch, + stream: bool, + accounting_first: bool, + duplicate_accounting: bool, + blocked_phase: str, +) -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + from litellm.proxy.proxy_server import general_settings + from litellm.proxy.spend_tracking.spend_event import SpendEvent, decode_spend_event + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer, UnixAddress + + recorded: Final[asyncio.Queue[tuple[bytes, str]]] = asyncio.Queue() + hook_context: Final = contextvars.ContextVar("accounting_hook_context", default="outside") + delivered: Final = asyncio.Event() + + async def store_event(line: bytes) -> None: + recorded.put_nowait((line, hook_context.get())) + + producer: Final = SpendEventProducer( + address=UnixAddress("/unused"), + on_unavailable="fallback", + buffer_size=1, + connect_timeout=1, + fallback=store_event, + ) + await producer.close(drain_timeout=1) + monkeypatch.setitem(general_settings, "store_prompts_in_spend_logs", True) + cancelled: Final = asyncio.Event() + never: Final = asyncio.Event() + + class OptionalLogger(CustomLogger): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + hook_context.set("inside") + if blocked_phase == "hook": + try: + await never.wait() + finally: + cancelled.set() + return kwargs, result + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + delivered.set() + if blocked_phase == "delivery": + try: + await never.wait() + finally: + cancelled.set() + + optional: Final = OptionalLogger(message_logging=blocked_phase != "redacted") + accounting: Final = _ProxyDBLogger(producer) + accounting_callbacks: Final = [accounting, _ProxyDBLogger(producer)] if duplicate_accounting else [accounting] + callbacks: Final = [*accounting_callbacks, optional] if accounting_first else [optional, *accounting_callbacks] + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj: Final = _async_request_logging_obj("accounting-timeout", callbacks, stream=stream) + response: Final = ModelResponse( + choices=[{"message": {"role": "assistant", "content": "private-marker"}}], + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + worker: Final = LoggingWorker(timeout=0.1) + worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=response)) + try: + await worker.flush() + line, context = recorded.get_nowait() + assert context == "inside" + assert delivered.is_set() + event: Final = decode_spend_event(line) + assert isinstance(event, SpendEvent) + assert event.response is not None + assert Usage.model_validate(event.response["usage"]).total_tokens == 15 + assert recorded.empty() + assert cancelled.is_set() is (blocked_phase in ("hook", "delivery")) + assert ("private-marker" in line.decode()) is (blocked_phase not in ("hook", "redacted")) + assert (REDACTED_BY_LITELLM in line.decode()) is (blocked_phase in ("hook", "redacted")) + assert response.choices[0].message.content == "private-marker" + finally: + await worker.stop()