mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(logging): log spend once for large non-streaming requests on the chat to Responses bridge (#44508)
* fix(logging): claim the async success dedup flag before the first await Nested @client wrappers on the chat -> Responses bridge schedule two async success tasks on one logging object. The dedup flag was checked at the top of _async_success_handler_body but only set after several awaits; the worker-thread base64 offload for messages >= 256 KiB yields there, so the second task passed the same check and the request was logged twice (two SpendLogs rows, key/team/daily spend counted twice). Fixes #44500 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * fix(logging): serialise async success handlers per logging object instead of claiming early Claiming the dedup flag before the first await (previous commit) let a first handler that is cancelled before its callbacks run (e.g. by the logging worker's per-coroutine timeout) suppress the second one, losing the request's success accounting. Serialise non-streaming handlers on the same Logging object with an asyncio.Lock kept in a module-level WeakKeyDictionary: the second handler waits, then skips if the first logged, or logs itself if the first raised or was cancelled first. Streaming chunk logging is not locked. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> * test(integration): pin one spend row and one key charge for large bridged chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): wait for the large request's spend row before asserting Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): use fixed timestamps in async success dedup tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(logging): trim the async success dedup lock comment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): reset background interaction dedupe under the success dedup lock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): use a fixed start time in the interactions logging helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: claim the async success log at schedule time instead of locking the handler Nested @client wrappers (chat over the Responses bridge, Anthropic Messages over the chat adapter) exit with one shared Logging object and each queued their own async success handler. The handler checked has_logged_async_success before its awaited work and set it after, so both could pass the check and the request was logged and billed twice. _schedule_async_success_logging now claims the log on the Logging object synchronously through claim_async_success_log, and a later wrapper returns without enqueuing. There is no await between the check and the decision, so no lock is needed. The per-object asyncio.Lock, the WeakKeyDictionary that held it and the locked block in async_log_background_interaction_completion are gone. Tests: the regression test moves to tests/unit/test_utils.py and drives two _dispatch_success_logging exits with one Logging object through the logging worker, asserting one log with the inner result. It fails without the claim. The cancelled-handler test is dropped since a second wrapper no longer retries. The background interaction completion test stays. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): pin the background completion event order and name the innermost wrapper in the claim docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): send one bridged request and assert the inner Responses result is the logged one Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: kerry <kerry@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Hektor Jacynycz Garcia
GitHub
parent
7f5499a751
commit
157d96dcd0
5 changed files with 251 additions and 5 deletions
|
|
@ -723,6 +723,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
# enqueue closure here instead of firing it immediately.
|
||||
self._defer_async_logging: bool = False
|
||||
self._enqueue_deferred_logging: Callable[[], None] | None = None
|
||||
self._async_success_scheduled: bool = False
|
||||
self._on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
|
||||
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
|
||||
|
||||
|
|
@ -2261,6 +2262,15 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
except Exception:
|
||||
return True
|
||||
|
||||
def claim_async_success_log(self) -> bool:
|
||||
"""One async success log per request. The innermost @client wrapper always exits first,
|
||||
since the outer one is awaiting it, so it claims the log here and the outer wrapper gets
|
||||
False and schedules nothing."""
|
||||
if self._async_success_scheduled:
|
||||
return False
|
||||
self._async_success_scheduled = True
|
||||
return True
|
||||
|
||||
def mark_logging_complete(
|
||||
self,
|
||||
event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"],
|
||||
|
|
|
|||
|
|
@ -1373,11 +1373,14 @@ def _schedule_async_success_logging(
|
|||
|
||||
Nested @client wrappers (Anthropic Messages over the chat adapter, chat over the Responses
|
||||
bridge) each exit through here with the same logging object and their own shape of the same
|
||||
response. The immediate path already logs one request once, since the first task marks
|
||||
``has_logged_async_success`` and the later ones skip. The deferred slot keeps the same
|
||||
first-wins rule: the innermost wrapper's provider-shaped result is the one the spend log
|
||||
reads usage from, and a later wrapper never swaps in its client-shaped translation.
|
||||
response. The innermost wrapper always exits first, since the outer one is awaiting it, so it
|
||||
claims the async success log here, synchronously, and the outer wrapper returns without
|
||||
enqueuing: one request queues one handler, whatever the handler tasks do later. The deferred
|
||||
slot keeps the same innermost-wins rule: the provider-shaped result is the one the spend log
|
||||
reads usage from, and the outer wrapper never swaps in its client-shaped translation.
|
||||
"""
|
||||
if not logging_obj.claim_async_success_log():
|
||||
return
|
||||
|
||||
def _enqueue_async_logging() -> None:
|
||||
asyncio.create_task(
|
||||
|
|
|
|||
111
tests/integration/spend/test_responses_bridge_spend_once.py
Normal file
111
tests/integration/spend/test_responses_bridge_spend_once.py
Normal file
|
|
@ -0,0 +1,111 @@
|
|||
import json
|
||||
from base64 import b64decode
|
||||
from collections.abc import Callable
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
PROMPT_TOKENS: Final = 30
|
||||
COMPLETION_TOKENS: Final = 5
|
||||
INPUT_RATE: Final = 0.001
|
||||
OUTPUT_RATE: Final = 0.002
|
||||
EXPECTED_SPEND: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE
|
||||
|
||||
|
||||
def _responses_reply(prompt: str, response_id: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method == "GET" and request.target == "/v1/models":
|
||||
return Reply(
|
||||
body=json.dumps({"object": "list", "data": [{"id": "gpt-6.1-sol", "object": "model"}]}).encode()
|
||||
)
|
||||
assert request.method == "POST", request.method
|
||||
assert request.target == "/v1/responses", request.target
|
||||
body: Final = object_value(json.loads(request.body))
|
||||
encoded_input: Final = json.dumps(body["input"])
|
||||
assert prompt in encoded_input, encoded_input
|
||||
assert body.get("stream") is not True, body
|
||||
response: Final = {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"status": "completed",
|
||||
"model": body["model"],
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": f"msg_{uuid4().hex}",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"input_tokens": PROMPT_TOKENS,
|
||||
"output_tokens": COMPLETION_TOKENS,
|
||||
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
|
||||
},
|
||||
}
|
||||
return Reply(body=json.dumps(response).encode())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _rows_for_key(key: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT litellm_call_id, request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
def test_bridged_chat_completion_above_log_offload_threshold_logs_and_charges_once(gateway: Gateway) -> None:
|
||||
prompt_prefix: Final = f"spend once {uuid4().hex[:8]} "
|
||||
long_prompt: Final = (prompt_prefix + "lorem ipsum " * 30_000)[:300_000]
|
||||
upstream_id: Final = f"resp_{uuid4().hex}"
|
||||
with (
|
||||
wire_server(_responses_reply(long_prompt, upstream_id)) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
model: Final = scenario.model(
|
||||
model="openai/responses/gpt-6.1-sol",
|
||||
api_base=f"{wire.url}/v1",
|
||||
input_cost_per_token=INPUT_RATE,
|
||||
output_cost_per_token=OUTPUT_RATE,
|
||||
)
|
||||
key: Final = scenario.key(models=[model])
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": long_prompt}], "stream": False},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert "ok" in response.text, response.text
|
||||
assert response.json()["id"] == upstream_id, response.text
|
||||
call_id: Final = response.headers["x-litellm-call-id"]
|
||||
|
||||
rows: Final = eventually(
|
||||
lambda: _rows_for_key(key),
|
||||
lambda values: any(row["litellm_call_id"] == call_id for row in values),
|
||||
seconds=70,
|
||||
)
|
||||
call_rows: Final = [row for row in rows if row["litellm_call_id"] == call_id]
|
||||
assert len(call_rows) == 1, rows
|
||||
assert float(str(call_rows[0]["spend"])) == pytest.approx(EXPECTED_SPEND), rows
|
||||
logged_id: Final = b64decode(str(call_rows[0]["request_id"]).removeprefix("resp_")).decode()
|
||||
assert logged_id.startswith("litellm:") and f";response_id:{upstream_id}" in logged_id, rows
|
||||
|
||||
key_spend: Final = eventually(
|
||||
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
|
||||
lambda values: len(values) == 1 and float(str(values[0]["spend"])) >= EXPECTED_SPEND,
|
||||
seconds=70,
|
||||
)
|
||||
assert float(str(key_spend[0]["spend"])) == pytest.approx(EXPECTED_SPEND), key_spend
|
||||
|
|
@ -5975,7 +5975,7 @@ def _interactions_logging_obj(stream: bool, call_type: str = "acreate"):
|
|||
messages=[],
|
||||
stream=stream,
|
||||
call_type=call_type,
|
||||
start_time=time.time(),
|
||||
start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc).timestamp(),
|
||||
litellm_call_id="interactions-call-id",
|
||||
function_id="interactions-fn-id",
|
||||
)
|
||||
|
|
@ -9536,3 +9536,61 @@ def test_get_custom_logger_compatible_class_does_not_match_generic_api_logger(
|
|||
assert logging_module.get_custom_logger_compatible_class(integration) is None
|
||||
finally:
|
||||
logging_module._in_memory_loggers.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_background_interaction_completion_logs_while_in_progress_handler_is_parked(monkeypatch):
|
||||
"""
|
||||
The create's in_progress success handler can still be parked on awaited work when
|
||||
the settlement arrives. The completed result is the only event carrying usage and
|
||||
cost, so it has to reach the success callbacks whatever the one-success-log-per-request
|
||||
dedupe does, or the settlement is marked billed with no spend row behind it.
|
||||
"""
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.types.interactions import InteractionsAPIResponse
|
||||
|
||||
class CountingLogger(CustomLogger):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.logged_results: list[object] = []
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self.logged_results.append(response_obj)
|
||||
|
||||
counting_logger: Final = CountingLogger()
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger])
|
||||
|
||||
in_progress_parked: Final = asyncio.Event()
|
||||
release_in_progress: Final = asyncio.Event()
|
||||
real_truncate: Final = litellm_logging.truncate_base64_in_messages_async
|
||||
calls: list[int] = []
|
||||
|
||||
async def parked_then_real_truncate(messages):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
in_progress_parked.set()
|
||||
await release_in_progress.wait()
|
||||
return await real_truncate(messages)
|
||||
|
||||
monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", parked_then_real_truncate)
|
||||
|
||||
logging_obj: Final = _interactions_logging_obj(stream=False)
|
||||
in_progress: Final = InteractionsAPIResponse(id="interactions/abc", model="gemini-2.5-flash", status="in_progress")
|
||||
completed: Final = InteractionsAPIResponse(
|
||||
id="interactions/abc",
|
||||
model="gemini-2.5-flash",
|
||||
status="completed",
|
||||
steps=[],
|
||||
usage=dict(INTERACTIONS_USAGE_BLOCK),
|
||||
)
|
||||
|
||||
first: Final = asyncio.create_task(logging_obj.async_success_handler(result=in_progress))
|
||||
await in_progress_parked.wait()
|
||||
completion: Final = asyncio.create_task(logging_obj.async_log_background_interaction_completion(result=completed))
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
release_in_progress.set()
|
||||
await first
|
||||
await completion
|
||||
|
||||
assert counting_logger.logged_results == [completed, in_progress], counting_logger.logged_results
|
||||
|
|
|
|||
|
|
@ -6711,3 +6711,67 @@ def test_function_setup_never_logs_the_ocr_data_uri_payload() -> None:
|
|||
|
||||
assert logged == [{"role": "user", "content": f"data:application/pdf;base64 ({len(payload)} chars)"}]
|
||||
assert payload not in str(logged)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_wrapper_exits_schedule_one_async_success_log(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Chat over the Responses bridge exits two @client wrappers with one logging object. Issue
|
||||
#44500: both exits enqueued a success handler, and on a large prompt the first one yielded to
|
||||
the worker-thread base64 offload before it marked ``has_logged_async_success``, so the second
|
||||
passed the check too and the request was logged and billed twice. The schedule step claims the
|
||||
log for the object synchronously, so only the inner provider-shaped result is ever logged."""
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.utils import _dispatch_success_logging
|
||||
|
||||
class CountingLogger(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.logged_results: list[object] = []
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
if kwargs["litellm_call_id"] == "bridge-call-id":
|
||||
self.logged_results.append(response_obj)
|
||||
|
||||
counting_logger: Final = CountingLogger()
|
||||
monkeypatch.setattr(litellm, "_async_success_callback", [counting_logger])
|
||||
real_truncate: Final = litellm_logging.truncate_base64_in_messages_async
|
||||
|
||||
async def yielding_truncate(messages):
|
||||
await asyncio.sleep(0)
|
||||
return await real_truncate(messages)
|
||||
|
||||
monkeypatch.setattr(litellm_logging, "truncate_base64_in_messages_async", yielding_truncate)
|
||||
|
||||
messages: Final = [{"role": "user", "content": "hello"}]
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-5.6-luna",
|
||||
messages=messages,
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime(2026, 1, 1, tzinfo=timezone.utc).timestamp(),
|
||||
litellm_call_id="bridge-call-id",
|
||||
function_id="bridge-fn-id",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={}, optional_params={}, model="gpt-5.6-luna", custom_llm_provider="openai", input=messages
|
||||
)
|
||||
inner_result: Final = ModelResponse(id="inner")
|
||||
outer_result: Final = ModelResponse(id="outer")
|
||||
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
|
||||
for result in (inner_result, outer_result):
|
||||
_dispatch_success_logging(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
is_completion_with_fallbacks=False,
|
||||
is_litellm_internal_call=False,
|
||||
)
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
|
||||
|
||||
assert len(counting_logger.logged_results) == 1, counting_logger.logged_results
|
||||
assert counting_logger.logged_results[0] is inner_result
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue