mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(logging)!: fire one CustomLogger failure hook per async call outside the async flag list (#47124)
* fix(logging): fire one CustomLogger hook per async call outside the async flag list aspeech, atranscription, aimage_generation, aimage_edit, aocr and asearch are missing from is_sync_litellm_request's flag list, so the sync success and failure handlers called a CustomLogger's sync hook on top of the async hook the async handler had already called for the same request. The sync handlers now skip a CustomLogger's log_success_event and log_failure_event when that run is paired with an async handler run (for_async_request) and the async handler gives the logger its own async hook. The async client wrapper, dispatch_failure_handlers, the proxy-only failure dispatch and the native bridge's sync failure step for an async call mark their sync failure run as paired. Resolves LIT-9486 * fix(logging): keep the sync CustomLogger hook when no async success runs The async client wrapper skips async success logging for calls carrying a fallbacks kwarg, and the native bridge skips it for internal and fallbacks calls, yet the paired sync success run still dropped a both-hooks logger's sync hook, so that logger got no success event at all. The sync success run now suppresses the sync hook only when the async success handler was actually scheduled for the call: the client wrapper passes async_handler_runs=not is_completion_with_fallbacks, and the native host records per call whether it enqueued or deferred the async success. Tests that pinned the doubled failure hook on native aocr, messages and responses calls and on proxy-rejected requests now expect the async hook alone. The integration test adds /v1/messages and /v1/responses cells and runs the key-auth rejection on every route. * fix(logging): type dispatch_success_handlers kwargs The new async_handler_runs keyword on success_handler made the untyped **kwargs that dispatch_success_handlers forwards an unknown argument for it, which grew the reportUnknownArgumentType count past its merge-base ceiling. Typing the forwarded kwargs as Any, as _success_handler_body already does, removes those diagnostics without changing what is forwarded. * test(integration): cover CustomLogger hook counts across modes, chaos and router strategies Moves the hook count rig into tests/integration/_support/hook_counts.py and extends the cells it drives: - streamed chat, Messages and Responses requests, success and failure - loggers the async handler never reaches (sync-only overrides, sync event lists only) keep one sync hook per request on every surface outside the async flag list, key-auth rejections included - a concurrent burst while the logger's destination stalls and then fails, and a burst while a worker is killed, never give a request two hooks - least-busy keeps counting a request in flight while async speech requests finish around it, and usage-based routing v2 counts an async image request's tokens once * test(integration): Datadog and LangSmith deliver one record per async request Points the built-in Datadog, LangSmith and OpenTelemetry loggers at a local sink and checks every async surface outside the flag list: one Datadog log entry and one LangSmith run per request (two on the merge base), and never two OpenTelemetry spans, which its own emit-once guard already ensured. * fix(logging): give a both-hooks logger its sync hook when the async run never reaches it The sync hook was skipped whenever the logger sat in the async callback list, so a logger the async run never reached got no hook at all: a sibling stalled past the logging worker timeout, a batch poll that skips default logging, or an async failure run cut off. The sync run now holds the hook and the async run hands it back unless it reached the logger * fix(logging): run the held sync success hooks when the proxy drops a blocked response's success A post-call guardrail block drops the parked async success, which left the sync success hook held for it unrun. Base gave a both-hooks logger that hook on these call types, so the drop now ends the pairing and the held hook runs as before * fix(logging): run held sync hooks only after every sync run of the request has closed A held hook released by the async run could run on another thread while the sync run was still walking its other callbacks. The pairing now waits until the async run has ended and every sync run has closed, and the last sync run to close runs the held hooks itself * fix(logging): end the async run when its coroutine is dropped before it starts A logging queue that drops the async success coroutine, or a deferred success that fails to schedule, left the sync hooks held for it unrun. The enqueued async run now ends when its coroutine is discarded, and the proxy ends it when firing the deferred success raises * fix(logging): leave a still-queued async run to the exit flush The discard finalizer ran at interpreter exit before the logging worker's exit flush, so a queued request could get its held sync hook and then its async hook * test: keep the router and file-stream tests in step with main Main gives each owned proxy a private Redis by default, so the router cells share the worker's Redis to read its counters, and the file-stream stub takes the hook pairing the stream now passes * fix(logging): run held sync hooks when a during_call guardrail raises or the exit flush gives up A during_call guardrail that raised after the LLM call parked its deferred success skipped the proxy's flush, and async runs the exit flush never started were never ended, so a both-hooks logger lost its sync hook in both cases * fix(logging): skip the sync failure hook when the async failure run logs to the logger An async call outside is_sync_litellm_request's flag list ran both failure handlers, so a CustomLogger with both failure hooks got both. The sync run of an async request now skips log_failure_event for a logger that overrides async_log_failure_event and is registered for async failures. The four sync runs of async failures (dispatch_failure_handlers, the @client async wrapper, the proxy-only failure thread and the Rust host) pass for_async_request=True. Drops the HookPairing fallback. Successes are left to #47192, which already gives such a logger only async_log_success_event. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
ffb7133d32
commit
dd511a2b22
20 changed files with 1532 additions and 20 deletions
|
|
@ -2466,7 +2466,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
await self.async_failure_handler(exception, traceback_exception)
|
||||
finally:
|
||||
if self._should_run_sync_failure_callbacks_for_async_calls():
|
||||
executor.submit(self.failure_handler, exception, traceback_exception)
|
||||
executor.submit(self.failure_handler, exception, traceback_exception, for_async_request=True)
|
||||
|
||||
def should_run_logging(
|
||||
self,
|
||||
|
|
@ -4090,15 +4090,21 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
*,
|
||||
for_async_request: bool = False,
|
||||
) -> None:
|
||||
"""Restores trace_id/session_id contextvars once this attempt's own failure
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
logging (including any nested calls its callbacks trigger) is fully done.
|
||||
|
||||
``for_async_request`` marks the sync run of a failure whose async run also logs it, so a
|
||||
CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook."""
|
||||
try:
|
||||
return self._failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
for_async_request=for_async_request,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
|
@ -4109,6 +4115,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
for_async_request: bool = False,
|
||||
) -> None:
|
||||
verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback)
|
||||
if not self.should_run_logging(event_type="sync_failure"): # prevent double logging
|
||||
|
|
@ -4201,6 +4208,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
isinstance(callback, CustomLogger)
|
||||
and is_sync_request
|
||||
and self.call_type != CallTypes.pass_through.value
|
||||
and not (for_async_request and self._async_handler_logs_to(callback))
|
||||
): # custom logger class
|
||||
callback.log_failure_event(
|
||||
start_time=start_time,
|
||||
|
|
@ -4563,6 +4571,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
cb_name: Final = self._get_callback_name(cb)
|
||||
return any(prefix in cb_name for prefix in INTERNAL_PREFIXES)
|
||||
|
||||
def _async_handler_logs_to(self, logger: CustomLogger) -> bool:
|
||||
"""True when the async failure run of this call hands ``logger`` its own ``async_log_failure_event``."""
|
||||
if type(logger).async_log_failure_event is CustomLogger.async_log_failure_event: # pyright: ignore[reportUnknownMemberType] # CustomLogger's hook signatures are untyped
|
||||
return False
|
||||
return logger in (self.dynamic_async_failure_callbacks or []) or logger in litellm._async_failure_callback # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # callback lists are untyped and have no public accessor
|
||||
|
||||
def _remove_internal_custom_logger_callbacks(self, callbacks: list) -> list:
|
||||
"""
|
||||
Removes internal custom logger callbacks from the list.
|
||||
|
|
|
|||
|
|
@ -3865,6 +3865,7 @@ class ProxyLogging:
|
|||
original_exception,
|
||||
traceback_str,
|
||||
),
|
||||
kwargs={"for_async_request": True},
|
||||
daemon=True,
|
||||
).start()
|
||||
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import contextvars
|
|||
import datetime
|
||||
import traceback
|
||||
import uuid
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Coroutine, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import (
|
||||
|
|
@ -61,9 +62,21 @@ def setup(
|
|||
}
|
||||
supplied: Final = arguments.get("litellm_logging_obj")
|
||||
if isinstance(supplied, Logging):
|
||||
return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous)
|
||||
return _begin_call(CallSetup(supplied, arguments), asynchronous)
|
||||
logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments)
|
||||
return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous)
|
||||
return _begin_call(CallSetup(logger, prepared), asynchronous)
|
||||
|
||||
|
||||
# Loggers whose current native call is async: their sync failure step is followed by the async one
|
||||
_ASYNC_CALLS: Final[weakref.WeakSet[Logging]] = weakref.WeakSet()
|
||||
|
||||
|
||||
def _begin_call(call_setup: CallSetup, asynchronous: bool) -> CallSetup:
|
||||
if asynchronous:
|
||||
_ASYNC_CALLS.add(call_setup.logger)
|
||||
else:
|
||||
_ASYNC_CALLS.discard(call_setup.logger)
|
||||
return _claim_budget_reservation(call_setup, asynchronous)
|
||||
|
||||
|
||||
def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup:
|
||||
|
|
@ -135,6 +148,8 @@ class LoggingSurface(Protocol):
|
|||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
*,
|
||||
for_async_request: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
def async_failure_handler(
|
||||
|
|
@ -261,7 +276,7 @@ def failure_handler(
|
|||
if not is_internal_call():
|
||||
unbind_budget_reservation_from_callbacks(logger.litellm_params)
|
||||
return logger.async_failure_handler(error, trace, start, end)
|
||||
logger.failure_handler(error, trace, start, end)
|
||||
logger.failure_handler(error, trace, start, end, for_async_request=logger in _ASYNC_CALLS)
|
||||
return None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2371,7 +2371,7 @@ def client(original_function):
|
|||
unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
|
||||
try:
|
||||
logging_obj.failure_handler(
|
||||
e, traceback_exception, start_time, end_time
|
||||
e, traceback_exception, start_time, end_time, for_async_request=True
|
||||
) # DO NOT MAKE THREADED - router retry fallback relies on this!
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
|
|
|||
91
tests/integration/_support/hook_count_probe.py
Normal file
91
tests/integration/_support/hook_count_probe.py
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
|
||||
def _record(hook: str, kwargs: Mapping[str, object]) -> bytes:
|
||||
return json.dumps(
|
||||
{"hook": hook, "litellm_call_id": kwargs.get("litellm_call_id"), "call_type": kwargs.get("call_type")}
|
||||
).encode()
|
||||
|
||||
|
||||
def _send(hook: str, kwargs: Mapping[str, object]) -> None:
|
||||
with httpx.Client(trust_env=False) as client:
|
||||
client.post(os.environ["HOOK_COUNT_PROBE_URL"], content=_record(hook, kwargs)).raise_for_status()
|
||||
|
||||
|
||||
async def _send_async(hook: str, kwargs: Mapping[str, object]) -> None:
|
||||
async with httpx.AsyncClient(trust_env=False) as client:
|
||||
response: Final = await client.post(os.environ["HOOK_COUNT_PROBE_URL"], content=_record(hook, kwargs))
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
class HookCountProbe(CustomLogger):
|
||||
"""Reports every sync and async success and failure hook it receives, one record per call."""
|
||||
|
||||
def log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send("sync_success", kwargs)
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
await _send_async("async_success", kwargs)
|
||||
|
||||
def log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send("sync_failure", kwargs)
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
await _send_async("async_failure", kwargs)
|
||||
|
||||
|
||||
class SyncOnlyHookCountProbe(CustomLogger):
|
||||
"""Overrides only the sync success and failure hooks."""
|
||||
|
||||
def log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send("sync_success", kwargs)
|
||||
|
||||
def log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send("sync_failure", kwargs)
|
||||
|
||||
|
||||
class BlockingPostCallGuardrail(CustomGuardrail):
|
||||
"""Rejects every response with a 400 after the provider answered."""
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self, data: Mapping[str, object], user_api_key_dict: object, response: object
|
||||
) -> None:
|
||||
raise HTTPException(status_code=400, detail="blocked by the hook-count guardrail")
|
||||
|
||||
|
||||
probe: Final = HookCountProbe()
|
||||
sync_only_probe: Final = SyncOnlyHookCountProbe()
|
||||
|
||||
|
||||
def sync_success_function(
|
||||
kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send("function_success", kwargs)
|
||||
|
||||
|
||||
def sync_failure_function(
|
||||
kwargs: Mapping[str, object], completion_response: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
_send("function_failure", kwargs)
|
||||
329
tests/integration/_support/hook_counts.py
Normal file
329
tests/integration/_support/hook_counts.py
Normal file
|
|
@ -0,0 +1,329 @@
|
|||
"""Scripted provider, proxy config and request helpers shared by the CustomLogger hook count tests."""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import yaml
|
||||
from integration._support.client import Gateway, held
|
||||
from integration._support.wire import Reply, Request, Wire
|
||||
|
||||
WAV_BYTES: Final = (
|
||||
b"RIFF\x30\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00"
|
||||
b"\x40\x1f\x00\x00\x40\x1f\x00\x00\x01\x00\x08\x00data\x08\x00\x00\x00"
|
||||
b"\x00\x01\x02\x03\x04\x05\x06\x07"
|
||||
)
|
||||
PNG_BYTES: Final = (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00"
|
||||
b"\x1f\x15\xc4\x89\x00\x00\x00\rIDAT\x08\xd7c\xf8\xcf\xc0\xf0\x1f\x00\x05\x00\x01\xff"
|
||||
b"\x89\x99=\x1d\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
_IMAGE_REPLY: Final = json.dumps({"created": 1700000000, "data": [{"b64_json": base64.b64encode(PNG_BYTES).decode()}]})
|
||||
_CHAT_REPLY: Final = json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-hook-count",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
||||
}
|
||||
)
|
||||
_EMBEDDING_REPLY: Final = json.dumps(
|
||||
{
|
||||
"object": "list",
|
||||
"model": "text-embedding-3-small",
|
||||
"data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}],
|
||||
"usage": {"prompt_tokens": 3, "total_tokens": 3},
|
||||
}
|
||||
)
|
||||
_MESSAGE_REPLY: Final = json.dumps(
|
||||
{
|
||||
"id": "msg_hook_count",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-haiku-4-5",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1},
|
||||
}
|
||||
)
|
||||
_RESPONSE_REPLY: Final = json.dumps(
|
||||
{
|
||||
"id": "resp_hook_count",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_hook_count",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 3, "output_tokens": 1, "total_tokens": 4},
|
||||
}
|
||||
)
|
||||
_OCR_REPLY: Final = json.dumps(
|
||||
{
|
||||
"pages": [
|
||||
{"index": 0, "markdown": "page", "images": [], "dimensions": {"dpi": 200, "height": 1100, "width": 850}}
|
||||
],
|
||||
"model": "mistral-ocr-latest",
|
||||
"usage_info": {"pages_processed": 1, "doc_size_bytes": 64},
|
||||
"document_annotation": None,
|
||||
}
|
||||
)
|
||||
_SEARCH_REPLY: Final = json.dumps(
|
||||
{"results": [{"title": "hook count", "url": "https://result.invalid/a", "snippet": "hit"}]}
|
||||
)
|
||||
_CHAT_FRAMES: Final = (
|
||||
'data: {"id":"chatcmpl-hook-count","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
|
||||
'"choices":[{"index":0,"delta":{"role":"assistant","content":"ok"},"finish_reason":null}]}',
|
||||
'data: {"id":"chatcmpl-hook-count","object":"chat.completion.chunk","created":1,"model":"gpt-4o-mini",'
|
||||
'"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],'
|
||||
'"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}',
|
||||
"data: [DONE]",
|
||||
)
|
||||
_MESSAGE_FRAMES: Final = (
|
||||
"event: message_start\n"
|
||||
'data: {"type":"message_start","message":{"id":"msg_hook_count","type":"message","role":"assistant",'
|
||||
'"model":"claude-haiku-4-5","content":[],"stop_reason":null,"stop_sequence":null,'
|
||||
'"usage":{"input_tokens":3,"output_tokens":0}}}',
|
||||
'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}',
|
||||
"event: content_block_delta\n"
|
||||
'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}',
|
||||
'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}',
|
||||
"event: message_delta\n"
|
||||
'data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":1}}',
|
||||
'event: message_stop\ndata: {"type":"message_stop"}',
|
||||
)
|
||||
_RESPONSE_FRAMES: Final = (
|
||||
"event: response.created\n"
|
||||
'data: {"type":"response.created","sequence_number":0,"response":{"id":"resp_hook_count","object":"response",'
|
||||
'"created_at":1,"status":"in_progress","model":"gpt-4o-mini","output":[],"usage":null}}',
|
||||
"event: response.output_text.delta\n"
|
||||
'data: {"type":"response.output_text.delta","sequence_number":1,"item_id":"msg_hook_count","output_index":0,'
|
||||
'"content_index":0,"delta":"ok"}',
|
||||
"event: response.completed\n"
|
||||
'data: {"type":"response.completed","sequence_number":2,"response":{"id":"resp_hook_count","object":"response",'
|
||||
'"created_at":1,"status":"completed","model":"gpt-4o-mini","output":[{"type":"message","id":"msg_hook_count",'
|
||||
'"status":"completed","role":"assistant","content":[{"type":"output_text","text":"ok","annotations":[]}]}],'
|
||||
'"usage":{"input_tokens":3,"output_tokens":1,"total_tokens":4}}}',
|
||||
)
|
||||
_STREAM_FRAMES: Final = {
|
||||
"/v1/chat/completions": _CHAT_FRAMES,
|
||||
"/v1/messages": _MESSAGE_FRAMES,
|
||||
"/v1/responses": _RESPONSE_FRAMES,
|
||||
}
|
||||
SURFACE_CALL_TYPES: Final = {
|
||||
"chat": "acompletion",
|
||||
"embedding": "aembedding",
|
||||
"messages": "anthropic_messages",
|
||||
"responses": "aresponses",
|
||||
"speech": "aspeech",
|
||||
"transcription": "atranscription",
|
||||
"image_generation": "aimage_generation",
|
||||
"image_edit": "aimage_edit",
|
||||
"ocr": "aocr",
|
||||
"search": "asearch",
|
||||
}
|
||||
IN_THE_ASYNC_FLAG_LIST: Final = frozenset({"chat", "embedding", "messages", "responses"})
|
||||
STREAMING_SURFACES: Final = ("chat", "messages", "responses")
|
||||
|
||||
|
||||
def _streams(request: Request) -> bool:
|
||||
if request.headers.get("content-type", "").startswith("application/json"):
|
||||
return json.loads(request.body or b"{}").get("stream") is True
|
||||
return False
|
||||
|
||||
|
||||
def provider(request: Request) -> Reply:
|
||||
"""Answers every surface the hook count tests drive; any target under /fail/ is a provider 400."""
|
||||
if request.method == "GET":
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
if request.target.startswith("/fail/"):
|
||||
return Reply(
|
||||
status=400,
|
||||
body=b'{"error":{"message":"scripted provider rejection","type":"invalid_request_error","code":null}}',
|
||||
)
|
||||
if request.target in _STREAM_FRAMES and _streams(request):
|
||||
return Reply(
|
||||
chunks=tuple(f"{frame}\n\n".encode() for frame in _STREAM_FRAMES[request.target]),
|
||||
content_type="text/event-stream",
|
||||
)
|
||||
replies: Final[dict[str, Callable[[], Reply]]] = {
|
||||
"/v1/chat/completions": lambda: Reply(body=_CHAT_REPLY.encode()),
|
||||
"/v1/embeddings": lambda: Reply(body=_EMBEDDING_REPLY.encode()),
|
||||
"/v1/messages": lambda: Reply(body=_MESSAGE_REPLY.encode()),
|
||||
"/v1/responses": lambda: Reply(body=_RESPONSE_REPLY.encode()),
|
||||
"/v1/audio/speech": lambda: Reply(body=WAV_BYTES, content_type="audio/wav"),
|
||||
"/v1/audio/transcriptions": lambda: Reply(body=b'{"text":"hook count transcript"}'),
|
||||
"/v1/images/generations": lambda: Reply(body=_IMAGE_REPLY.encode()),
|
||||
"/v1/images/edits": lambda: Reply(body=_IMAGE_REPLY.encode()),
|
||||
"/v1/ocr": lambda: Reply(body=_OCR_REPLY.encode()),
|
||||
"/perplexity/search": lambda: Reply(body=_SEARCH_REPLY.encode()),
|
||||
}
|
||||
return replies[request.target]()
|
||||
|
||||
|
||||
@dataclass
|
||||
class HookRecords:
|
||||
"""Every hook record the probe reported, kept across drains so cells can read their own call id."""
|
||||
|
||||
wire: Wire
|
||||
function_callbacks: bool
|
||||
seen: list[dict[str, str]] = field(default_factory=list)
|
||||
lock: threading.Lock = field(default_factory=threading.Lock)
|
||||
|
||||
def drained(self) -> tuple[dict[str, str], ...]:
|
||||
with self.lock:
|
||||
self.seen.extend(json.loads(request.body) for request in self.wire.drain())
|
||||
return tuple(self.seen)
|
||||
|
||||
def for_call(self, call_id: str) -> tuple[dict[str, str], ...]:
|
||||
return tuple(record for record in self.drained() if record["litellm_call_id"] == call_id)
|
||||
|
||||
|
||||
def _model(name: str, model: str, api_base: str, **extra: object) -> dict[str, object]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": model, "api_base": api_base, "api_key": "synthetic-provider-key"},
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def write_hook_count_config(
|
||||
path: Path,
|
||||
provider_url: str,
|
||||
litellm_settings: Mapping[str, object],
|
||||
guardrails: tuple[Mapping[str, object], ...] = (),
|
||||
) -> Path:
|
||||
"""One success and one failure deployment per surface, all answered by `provider`."""
|
||||
models: Final = []
|
||||
tools: Final = []
|
||||
for outcome, prefix in (("success", ""), ("failure", "/fail")):
|
||||
base: str = provider_url + prefix
|
||||
models.extend(
|
||||
(
|
||||
_model(f"hooks-chat-{outcome}", "openai/gpt-4o-mini", base + "/v1"),
|
||||
_model(f"hooks-embedding-{outcome}", "openai/text-embedding-3-small", base + "/v1"),
|
||||
_model(f"hooks-messages-{outcome}", "anthropic/claude-haiku-4-5", base),
|
||||
_model(f"hooks-responses-{outcome}", "openai/gpt-4o-mini", base + "/v1"),
|
||||
_model(f"hooks-speech-{outcome}", "openai/gpt-4o-mini-tts", base + "/v1"),
|
||||
_model(f"hooks-transcription-{outcome}", "openai/whisper-1", base + "/v1"),
|
||||
_model(f"hooks-image_generation-{outcome}", "openai/gpt-image-1", base + "/v1"),
|
||||
_model(f"hooks-image_edit-{outcome}", "openai/gpt-image-1", base + "/v1"),
|
||||
_model(f"hooks-ocr-{outcome}", "mistral/mistral-ocr-latest", base + "/v1", model_info={"mode": "ocr"}),
|
||||
)
|
||||
)
|
||||
tools.append(
|
||||
{
|
||||
"search_tool_name": f"hooks-search-{outcome}",
|
||||
"litellm_params": {
|
||||
"search_provider": "perplexity",
|
||||
"api_key": "synthetic-provider-key",
|
||||
"api_base": base + "/perplexity",
|
||||
},
|
||||
}
|
||||
)
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": models,
|
||||
"search_tools": tools,
|
||||
"litellm_settings": dict(litellm_settings),
|
||||
**({"guardrails": [dict(guardrail) for guardrail in guardrails]} if guardrails else {}),
|
||||
"general_settings": {
|
||||
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
||||
"database_url": "os.environ/DATABASE_URL",
|
||||
"disable_model_info_refresh": True,
|
||||
},
|
||||
"router_settings": {"num_retries": 0},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def send(
|
||||
proxy: Gateway, surface: str, outcome: str, call_id: str, *, key: str | None = None, stream: bool = False
|
||||
) -> httpx.Response:
|
||||
model: Final = f"hooks-{surface}-{outcome}"
|
||||
nonce: Final = f"hook-count-{call_id}"
|
||||
headers: Final = {"Authorization": f"Bearer {key or proxy.key}", "x-litellm-call-id": call_id}
|
||||
if surface == "transcription":
|
||||
return proxy.client.post(
|
||||
"/v1/audio/transcriptions",
|
||||
data={"model": model, "prompt": nonce},
|
||||
files={"file": ("a.wav", WAV_BYTES + call_id.encode(), "audio/wav")},
|
||||
headers=headers,
|
||||
)
|
||||
if surface == "image_edit":
|
||||
return proxy.client.post(
|
||||
"/v1/images/edits",
|
||||
data={"model": model, "prompt": nonce},
|
||||
files={"image": ("a.png", PNG_BYTES, "image/png")},
|
||||
headers=headers,
|
||||
)
|
||||
streaming: Final[dict[str, object]] = {"stream": True} if stream else {}
|
||||
routes: Final[dict[str, tuple[str, dict[str, object]]]] = {
|
||||
"chat": (
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": nonce}], **streaming},
|
||||
),
|
||||
"embedding": ("/v1/embeddings", {"model": model, "input": nonce}),
|
||||
"messages": (
|
||||
"/v1/messages",
|
||||
{"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": nonce}], **streaming},
|
||||
),
|
||||
"responses": ("/v1/responses", {"model": model, "input": nonce, **streaming}),
|
||||
"speech": ("/v1/audio/speech", {"model": model, "input": nonce, "voice": "alloy", "response_format": "wav"}),
|
||||
"image_generation": ("/v1/images/generations", {"model": model, "prompt": nonce, "n": 1}),
|
||||
"ocr": (
|
||||
"/v1/ocr",
|
||||
{
|
||||
"model": model,
|
||||
"document": {
|
||||
"type": "document_url",
|
||||
"document_url": "data:application/pdf;base64," + base64.b64encode(nonce.encode()).decode(),
|
||||
},
|
||||
},
|
||||
),
|
||||
"search": ("/v1/search", {"search_tool_name": f"hooks-search-{outcome}", "query": nonce}),
|
||||
}
|
||||
path, body = routes[surface]
|
||||
return proxy.client.post(path, json=body, headers=headers)
|
||||
|
||||
|
||||
def counts_function_hooks(records: HookRecords, surface: str) -> bool:
|
||||
return records.function_callbacks and surface not in IN_THE_ASYNC_FLAG_LIST
|
||||
|
||||
|
||||
def probe_hooks(records: HookRecords, surface: str, found: tuple[dict[str, str], ...]) -> list[str]:
|
||||
counted: Final = counts_function_hooks(records, surface)
|
||||
return sorted(record["hook"] for record in found if counted or not record["hook"].startswith("function_"))
|
||||
|
||||
|
||||
def one_async_hook(records: HookRecords, surface: str, outcome: str) -> list[str]:
|
||||
function_hook: Final = (f"function_{outcome}",) if counts_function_hooks(records, surface) else ()
|
||||
return sorted((f"async_{outcome}", *function_hook))
|
||||
|
||||
|
||||
def assert_hooks(records: HookRecords, surface: str, call_id: str, expected: list[str]) -> None:
|
||||
settled: Final = held(
|
||||
lambda: records.for_call(call_id),
|
||||
lambda found: probe_hooks(records, surface, found) == expected,
|
||||
holding=3,
|
||||
seconds=20,
|
||||
)
|
||||
assert {record["call_type"] for record in settled} == {SURFACE_CALL_TYPES[surface]}, settled
|
||||
|
|
@ -0,0 +1,96 @@
|
|||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.hook_counts import (
|
||||
STREAMING_SURFACES,
|
||||
SURFACE_CALL_TYPES,
|
||||
HookRecords,
|
||||
assert_hooks,
|
||||
one_async_hook,
|
||||
provider,
|
||||
send,
|
||||
write_hook_count_config,
|
||||
)
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, wire_server
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=["cache", "cache_and_sync_callback"])
|
||||
def hook_proxy(
|
||||
request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[tuple[Gateway, HookRecords]]:
|
||||
tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-counts")
|
||||
legacy_callbacks: Final = (
|
||||
{
|
||||
"success_callback": ["integration._support.hook_count_probe.sync_success_function"],
|
||||
"failure_callback": ["integration._support.hook_count_probe.sync_failure_function"],
|
||||
}
|
||||
if request.param == "cache_and_sync_callback"
|
||||
else {}
|
||||
)
|
||||
with (
|
||||
wire_server(provider) as upstream,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = write_hook_count_config(
|
||||
tmp_path / "hook-counts.yaml",
|
||||
upstream.url,
|
||||
{
|
||||
"callbacks": ["integration._support.hook_count_probe.probe"],
|
||||
**legacy_callbacks,
|
||||
"cache": True,
|
||||
"cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"},
|
||||
},
|
||||
)
|
||||
with owned_proxy(
|
||||
rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2
|
||||
) as proxy:
|
||||
yield proxy, HookRecords(sink, function_callbacks=request.param == "cache_and_sync_callback")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["success", "failure"])
|
||||
@pytest.mark.parametrize("surface", sorted(SURFACE_CALL_TYPES))
|
||||
def test_custom_logger_gets_one_async_hook_per_request_with_the_response_cache_on(
|
||||
hook_proxy: tuple[Gateway, HookRecords], surface: str, outcome: str
|
||||
) -> None:
|
||||
proxy, records = hook_proxy
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
response: Final = send(proxy, surface, outcome, call_id)
|
||||
assert response.status_code == (200 if outcome == "success" else 400), response.text
|
||||
assert_hooks(records, surface, call_id, one_async_hook(records, surface, outcome))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["success", "failure"])
|
||||
@pytest.mark.parametrize("surface", STREAMING_SURFACES)
|
||||
def test_custom_logger_gets_one_async_hook_per_streamed_request(
|
||||
hook_proxy: tuple[Gateway, HookRecords], surface: str, outcome: str
|
||||
) -> None:
|
||||
proxy, records = hook_proxy
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
response: Final = send(proxy, surface, outcome, call_id, stream=True)
|
||||
assert response.status_code == (200 if outcome == "success" else 400), response.text
|
||||
if outcome == "success":
|
||||
assert "ok" in response.text, response.text
|
||||
assert_hooks(records, surface, call_id, one_async_hook(records, surface, outcome))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", sorted(SURFACE_CALL_TYPES))
|
||||
def test_custom_logger_gets_one_async_hook_per_request_rejected_by_key_auth(
|
||||
hook_proxy: tuple[Gateway, HookRecords], surface: str
|
||||
) -> None:
|
||||
proxy, records = hook_proxy
|
||||
known_calls: Final = {record["litellm_call_id"] for record in records.drained()}
|
||||
response: Final = send(proxy, surface, "success", str(uuid.uuid4()), key="sk-not-a-virtual-key")
|
||||
assert response.status_code == 401, response.text
|
||||
|
||||
def new_calls(found: tuple[dict[str, str], ...]) -> set[str]:
|
||||
return {record["litellm_call_id"] for record in found if record["litellm_call_id"] not in known_calls}
|
||||
|
||||
(rejected_call,) = eventually(lambda: new_calls(records.drained()), bool, seconds=20)
|
||||
assert_hooks(records, surface, rejected_call, one_async_hook(records, surface, "failure"))
|
||||
|
|
@ -0,0 +1,181 @@
|
|||
"""One hook per request under a concurrent burst while the logger's destination stalls and fails, and while a
|
||||
worker is killed."""
|
||||
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, held
|
||||
from integration._support.hook_counts import (
|
||||
STREAMING_SURFACES,
|
||||
SURFACE_CALL_TYPES,
|
||||
HookRecords,
|
||||
assert_hooks,
|
||||
one_async_hook,
|
||||
probe_hooks,
|
||||
provider,
|
||||
send,
|
||||
write_hook_count_config,
|
||||
)
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 300)
|
||||
|
||||
RecordProperty = Callable[[str, object], None]
|
||||
|
||||
_SETTINGS: Final = {"callbacks": ["integration._support.hook_count_probe.probe"]}
|
||||
_HOLD_SECONDS: Final = 2.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Call:
|
||||
surface: str
|
||||
outcome: str
|
||||
stream: bool
|
||||
call_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Outcome:
|
||||
call: _Call
|
||||
status: int | None
|
||||
error: str | None
|
||||
|
||||
|
||||
def _burst(rounds: int) -> tuple[_Call, ...]:
|
||||
shapes: Final = (
|
||||
*((surface, outcome, False) for surface in sorted(SURFACE_CALL_TYPES) for outcome in ("success", "failure")),
|
||||
*((surface, "success", True) for surface in STREAMING_SURFACES),
|
||||
)
|
||||
return tuple(
|
||||
_Call(surface, outcome, stream, str(uuid.uuid4())) for _ in range(rounds) for surface, outcome, stream in shapes
|
||||
)
|
||||
|
||||
|
||||
def _attempt(proxy: Gateway, call: _Call) -> _Outcome:
|
||||
try:
|
||||
response: Final = send(proxy, call.surface, call.outcome, call.call_id, stream=call.stream)
|
||||
except httpx.HTTPError as error:
|
||||
return _Outcome(call, None, repr(error))
|
||||
return _Outcome(call, response.status_code, None)
|
||||
|
||||
|
||||
def _expected_status(call: _Call) -> int:
|
||||
return 200 if call.outcome == "success" else 400
|
||||
|
||||
|
||||
def _done(futures: tuple[Future[_Outcome], ...]) -> int:
|
||||
return sum(future.done() for future in futures)
|
||||
|
||||
|
||||
class _Destination:
|
||||
"""The probe's sink: answers at once, holds each report until released, or refuses it with a 503."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.mode = "ok"
|
||||
self.release = threading.Event()
|
||||
|
||||
def __call__(self, request: Request) -> Reply:
|
||||
if self.mode == "hold":
|
||||
self.release.wait(timeout=_HOLD_SECONDS + 1)
|
||||
if self.mode == "fail":
|
||||
return Reply(status=503, body=b'{"error":"destination unavailable"}')
|
||||
return Reply()
|
||||
|
||||
|
||||
def test_logger_destination_stalling_then_failing_mid_burst_gets_one_hook_per_request(
|
||||
tmp_path: Path, record_property: RecordProperty
|
||||
) -> None:
|
||||
destination: Final = _Destination()
|
||||
with (
|
||||
wire_server(provider) as upstream,
|
||||
wire_server(destination) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = write_hook_count_config(tmp_path / "hook-count-chaos.yaml", upstream.url, _SETTINGS)
|
||||
with owned_proxy(
|
||||
rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2
|
||||
) as proxy:
|
||||
records: Final = HookRecords(sink, function_callbacks=False)
|
||||
calls: Final = _burst(2)
|
||||
destination.mode = "hold"
|
||||
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
|
||||
futures: Final = tuple(pool.submit(_attempt, proxy, call) for call in calls)
|
||||
eventually(lambda: _done(futures), lambda done: done >= 1, seconds=60)
|
||||
held_liveliness: Final = proxy.client.get("/health/liveliness").status_code
|
||||
time.sleep(_HOLD_SECONDS)
|
||||
destination.mode = "fail"
|
||||
destination.release.set()
|
||||
time.sleep(_HOLD_SECONDS)
|
||||
destination.mode = "ok"
|
||||
outcomes: Final = tuple(future.result(timeout=180) for future in futures)
|
||||
record_property("burst_size", len(calls))
|
||||
assert held_liveliness == 200, held_liveliness
|
||||
assert [outcome for outcome in outcomes if outcome.status != _expected_status(outcome.call)] == []
|
||||
for call in calls:
|
||||
assert_hooks(records, call.surface, call.call_id, one_async_hook(records, call.surface, call.outcome))
|
||||
|
||||
|
||||
def test_worker_killed_mid_burst_leaves_at_most_one_hook_per_request(
|
||||
tmp_path: Path, record_property: RecordProperty
|
||||
) -> None:
|
||||
with (
|
||||
wire_server(provider) as upstream,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = write_hook_count_config(tmp_path / "hook-count-kill.yaml", upstream.url, _SETTINGS)
|
||||
with owned_proxy_process(
|
||||
rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2
|
||||
) as owned:
|
||||
proxy: Final = owned.gateway
|
||||
records: Final = HookRecords(sink, function_callbacks=False)
|
||||
victim: Final = min(eventually(owned.worker_pids, lambda found: len(found) == 2, seconds=60))
|
||||
calls: Final = _burst(2)
|
||||
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
|
||||
futures: Final = tuple(pool.submit(_attempt, proxy, call) for call in calls)
|
||||
eventually(lambda: _done(futures), lambda done: done >= 1, seconds=60)
|
||||
with owned.replacing_worker(victim) as kill:
|
||||
kill()
|
||||
done_at_kill: Final = _done(futures)
|
||||
survived: Final = eventually(
|
||||
lambda: _attempt(proxy, _Call("speech", "success", False, str(uuid.uuid4()))),
|
||||
lambda outcome: outcome.status == 200,
|
||||
seconds=30,
|
||||
)
|
||||
outcomes: Final = tuple(future.result(timeout=180) for future in futures)
|
||||
respawned: Final = _Call("speech", "success", False, str(uuid.uuid4()))
|
||||
after_respawn: Final = _attempt(proxy, respawned)
|
||||
completed: Final = tuple(outcome for outcome in outcomes if outcome.status is not None)
|
||||
interrupted: Final = tuple(outcome for outcome in outcomes if outcome.status is None)
|
||||
assert after_respawn.status == 200, after_respawn
|
||||
assert [outcome for outcome in completed if outcome.status != _expected_status(outcome.call)] == []
|
||||
for call in (survived.call, respawned):
|
||||
assert_hooks(records, call.surface, call.call_id, one_async_hook(records, call.surface, call.outcome))
|
||||
burst_hooks: Final = held(
|
||||
lambda: _hooks_of(records, calls),
|
||||
lambda found: all(
|
||||
hooks in ([], [f"async_{call.outcome}"]) for call, hooks in zip(calls, found, strict=True)
|
||||
),
|
||||
holding=3,
|
||||
seconds=20,
|
||||
)
|
||||
record_property("burst_size", len(calls))
|
||||
record_property("killed_worker_pid", victim)
|
||||
record_property("done_at_kill", done_at_kill)
|
||||
record_property("interrupted_calls", str(tuple(outcome.error for outcome in interrupted)))
|
||||
record_property("burst_calls_without_a_hook", sum(not hooks for hooks in burst_hooks))
|
||||
|
||||
|
||||
def _hooks_of(records: HookRecords, calls: tuple[_Call, ...]) -> tuple[list[str], ...]:
|
||||
"""A burst call the killed worker served may lose its logging with the worker, whether or not its response
|
||||
arrived first, so a burst call gets its one async hook or none, never two."""
|
||||
return tuple(probe_hooks(records, call.surface, records.for_call(call.call_id)) for call in calls)
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
"""A logger the async handler never reaches keeps its sync hook on async calls outside the async flag list.
|
||||
|
||||
A logger with both hooks registered only under `success_callback` and `failure_callback` is never reached by the async
|
||||
failure handler, so it keeps its sync failure hook. The async success handler does reach it, and gives it its async
|
||||
success hook instead of the sync one.
|
||||
|
||||
The sync success handler only runs for an async call when a legacy sync callback is configured, so every
|
||||
registration here carries the legacy function callbacks too."""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.hook_counts import (
|
||||
IN_THE_ASYNC_FLAG_LIST,
|
||||
SURFACE_CALL_TYPES,
|
||||
HookRecords,
|
||||
assert_hooks,
|
||||
provider,
|
||||
send,
|
||||
write_hook_count_config,
|
||||
)
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, wire_server
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_PROBES: Final = "integration._support.hook_count_probe"
|
||||
_REGISTRATIONS: Final[dict[str, dict[str, object]]] = {
|
||||
"sync_only_logger": {
|
||||
"callbacks": [f"{_PROBES}.sync_only_probe"],
|
||||
"success_callback": [f"{_PROBES}.sync_success_function"],
|
||||
"failure_callback": [f"{_PROBES}.sync_failure_function"],
|
||||
},
|
||||
"sync_event_lists": {
|
||||
"success_callback": [f"{_PROBES}.probe", f"{_PROBES}.sync_success_function"],
|
||||
"failure_callback": [f"{_PROBES}.probe", f"{_PROBES}.sync_failure_function"],
|
||||
},
|
||||
}
|
||||
_OUTSIDE_THE_ASYNC_FLAG_LIST: Final = sorted(SURFACE_CALL_TYPES.keys() - IN_THE_ASYNC_FLAG_LIST)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module", params=sorted(_REGISTRATIONS))
|
||||
def sync_hook_proxy(
|
||||
request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[tuple[Gateway, HookRecords, str]]:
|
||||
tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-count-modes")
|
||||
with (
|
||||
wire_server(provider) as upstream,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = write_hook_count_config(
|
||||
tmp_path / "hook-count-modes.yaml", upstream.url, _REGISTRATIONS[request.param]
|
||||
)
|
||||
with owned_proxy(
|
||||
rig, tmp_path, {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}, config=config, workers=2
|
||||
) as proxy:
|
||||
yield proxy, HookRecords(sink, function_callbacks=True), request.param
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["success", "failure"])
|
||||
@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST)
|
||||
def test_logger_the_async_handler_skips_gets_one_sync_hook_per_request(
|
||||
sync_hook_proxy: tuple[Gateway, HookRecords, str], surface: str, outcome: str
|
||||
) -> None:
|
||||
proxy, records, registration = sync_hook_proxy
|
||||
reached_by_the_async_success: Final = registration == "sync_event_lists" and outcome == "success"
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
response: Final = send(proxy, surface, outcome, call_id)
|
||||
assert response.status_code == (200 if outcome == "success" else 400), response.text
|
||||
assert_hooks(
|
||||
records,
|
||||
surface,
|
||||
call_id,
|
||||
sorted([f"function_{outcome}", f"async_{outcome}" if reached_by_the_async_success else f"sync_{outcome}"]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST)
|
||||
def test_logger_the_async_handler_skips_gets_one_sync_hook_per_request_rejected_by_key_auth(
|
||||
sync_hook_proxy: tuple[Gateway, HookRecords, str], surface: str
|
||||
) -> None:
|
||||
proxy, records, _ = sync_hook_proxy
|
||||
known_calls: Final = {record["litellm_call_id"] for record in records.drained()}
|
||||
response: Final = send(proxy, surface, "success", str(uuid.uuid4()), key="sk-not-a-virtual-key")
|
||||
assert response.status_code == 401, response.text
|
||||
|
||||
def new_calls(found: tuple[dict[str, str], ...]) -> set[str]:
|
||||
return {record["litellm_call_id"] for record in found if record["litellm_call_id"] not in known_calls}
|
||||
|
||||
(rejected_call,) = eventually(lambda: new_calls(records.drained()), bool, seconds=20)
|
||||
assert_hooks(records, surface, rejected_call, ["function_failure", "sync_failure"])
|
||||
|
|
@ -0,0 +1,75 @@
|
|||
"""A CustomLogger with both hooks gets one failure hook when a post-call guardrail blocks the response.
|
||||
|
||||
The proxy drops the async success logging it parked for the guardrails and logs the block as a failure through both
|
||||
the async and the sync failure handlers. The logger gets only its async failure hook from them. Speech and
|
||||
transcription responses never reach post-call guardrails, and image generation logs its success before they run, so
|
||||
neither parks.
|
||||
|
||||
A legacy sync callback is configured, because that is what makes the sync handler run for an async call."""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment, held
|
||||
from integration._support.hook_counts import (
|
||||
SURFACE_CALL_TYPES,
|
||||
HookRecords,
|
||||
probe_hooks,
|
||||
provider,
|
||||
send,
|
||||
write_hook_count_config,
|
||||
)
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, wire_server
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_PROBES: Final = "integration._support.hook_count_probe"
|
||||
_PARKED_FOR_POST_CALL_GUARDRAILS: Final = sorted(
|
||||
SURFACE_CALL_TYPES.keys() - {"speech", "transcription", "image_generation"}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def post_call_block(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, HookRecords]]:
|
||||
tmp_path: Final = tmp_path_factory.mktemp("hook-count-post-call-block")
|
||||
guardrail: Final = {
|
||||
"guardrail_name": "hook-count-post-call-block",
|
||||
"litellm_params": {
|
||||
"guardrail": f"{_PROBES}.BlockingPostCallGuardrail",
|
||||
"mode": "post_call",
|
||||
"default_on": True,
|
||||
},
|
||||
}
|
||||
with (
|
||||
wire_server(provider) as upstream,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = write_hook_count_config(
|
||||
tmp_path / "hook-count-post-call-block.yaml",
|
||||
upstream.url,
|
||||
{"callbacks": [f"{_PROBES}.probe"], "success_callback": [f"{_PROBES}.sync_success_function"]},
|
||||
(guardrail,),
|
||||
)
|
||||
overrides: Final = {"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"}
|
||||
with owned_proxy(rig, tmp_path, overrides, config=config, workers=2) as proxy:
|
||||
yield proxy, HookRecords(sink, function_callbacks=False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", _PARKED_FOR_POST_CALL_GUARDRAILS)
|
||||
def test_response_a_post_call_guardrail_blocks_gives_the_logger_one_failure_hook(
|
||||
post_call_block: tuple[Gateway, HookRecords], surface: str
|
||||
) -> None:
|
||||
proxy, records = post_call_block
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
response: Final = send(proxy, surface, "success", call_id)
|
||||
assert response.status_code == 400, response.text
|
||||
held(
|
||||
lambda: probe_hooks(records, surface, records.for_call(call_id)),
|
||||
lambda hooks: hooks == ["async_failure"],
|
||||
holding=3,
|
||||
seconds=20,
|
||||
)
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
"""Built-in loggers that override both hooks deliver one record per async request outside the async flag list:
|
||||
Datadog one log entry and LangSmith one run. OpenTelemetry keeps its own emit-once guard, so it never exports two
|
||||
spans for a request either way; its speech, OCR and search spans carry no `litellm.call_id`, so only "never two"
|
||||
is checked for it.
|
||||
|
||||
A legacy sync callback is configured, because that is what makes the sync handler run for an async call."""
|
||||
|
||||
import gzip
|
||||
import json
|
||||
import uuid
|
||||
import zlib
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, gateway_from_environment, held
|
||||
from integration._support.hook_counts import IN_THE_ASYNC_FLAG_LIST, SURFACE_CALL_TYPES, provider, send
|
||||
from integration._support.hook_counts import write_hook_count_config
|
||||
from integration._support.otlp_sink import decode_spans
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_PROBES: Final = "integration._support.hook_count_probe"
|
||||
_FLUSH_SECONDS: Final = 1
|
||||
_OUTSIDE_THE_ASYNC_FLAG_LIST: Final = sorted(SURFACE_CALL_TYPES.keys() - IN_THE_ASYNC_FLAG_LIST)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Vendors:
|
||||
proxy: Gateway
|
||||
sink: Wire
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def vendors(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Vendors]:
|
||||
tmp_path: Final = tmp_path_factory.mktemp("custom-logger-hook-count-vendors")
|
||||
with (
|
||||
wire_server(provider) as upstream,
|
||||
wire_server(lambda request: Reply(status=202)) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = write_hook_count_config(
|
||||
tmp_path / "hook-count-vendors.yaml",
|
||||
upstream.url,
|
||||
{
|
||||
"callbacks": ["datadog", "langsmith", "otel"],
|
||||
"success_callback": [f"{_PROBES}.sync_success_function"],
|
||||
"failure_callback": [f"{_PROBES}.sync_failure_function"],
|
||||
},
|
||||
)
|
||||
overrides: Final = {
|
||||
"HOOK_COUNT_PROBE_URL": sink.url + "/probe",
|
||||
"STORE_MODEL_IN_DB": "False",
|
||||
"DEFAULT_FLUSH_INTERVAL_SECONDS": str(_FLUSH_SECONDS),
|
||||
"DD_API_KEY": "synthetic-datadog-key",
|
||||
"DD_SITE": "us5.datadoghq.com",
|
||||
"DD_BASE_URL": sink.url + "/datadog",
|
||||
"LANGSMITH_API_KEY": "synthetic-langsmith-key",
|
||||
"LANGSMITH_PROJECT": "hook-counts",
|
||||
"LANGSMITH_BASE_URL": sink.url + "/langsmith",
|
||||
"OTEL_EXPORTER": "otlp_http",
|
||||
"OTEL_ENDPOINT": sink.url + "/otel/v1/traces",
|
||||
}
|
||||
with owned_proxy(rig, tmp_path, overrides, config=config, workers=2) as proxy:
|
||||
yield _Vendors(proxy, sink)
|
||||
|
||||
|
||||
def _unzipped(body: bytes) -> bytes:
|
||||
if body[:2] == b"\x1f\x8b":
|
||||
return gzip.decompress(body)
|
||||
try:
|
||||
return zlib.decompress(body, zlib.MAX_WBITS | 16)
|
||||
except zlib.error:
|
||||
return body
|
||||
|
||||
|
||||
def _entries(body: bytes) -> tuple[object, ...]:
|
||||
decoded: Final = json.loads(_unzipped(body))
|
||||
return tuple(decoded) if isinstance(decoded, list) else (decoded,)
|
||||
|
||||
|
||||
def _records(request: Request) -> tuple[tuple[str, str], ...]:
|
||||
"""Every record one delivery carries, as (destination, serialized record). A LangSmith batch carries its new
|
||||
runs under `post`, and an OpenTelemetry export carries every span of the request's trace."""
|
||||
if request.target.startswith("/datadog"):
|
||||
return tuple(("datadog", json.dumps(entry)) for entry in _entries(request.body))
|
||||
if request.target.startswith("/langsmith") and request.method == "POST":
|
||||
runs: Final = tuple(
|
||||
run
|
||||
for entry in _entries(request.body)
|
||||
for run in (entry.get("post", []) if isinstance(entry, dict) and "post" in entry else [entry])
|
||||
)
|
||||
return tuple(("langsmith", json.dumps(run)) for run in runs)
|
||||
if request.target.startswith("/otel"):
|
||||
spans: Final = decode_spans(request.body, request.headers.get("content-type", ""))
|
||||
return tuple(("otel", json.dumps(span)) for span in spans if "litellm.call_id" in span["attributes"])
|
||||
return ()
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Deliveries:
|
||||
sink: Wire
|
||||
seen: tuple[tuple[str, str], ...] = ()
|
||||
|
||||
def per_destination(self, call_id: str) -> dict[str, int]:
|
||||
self.seen = (*self.seen, *(record for request in self.sink.drain() for record in _records(request)))
|
||||
destinations: Final = ("datadog", "langsmith", "otel")
|
||||
return {
|
||||
destination: sum(1 for name, record in self.seen if name == destination and call_id in record)
|
||||
for destination in destinations
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("surface", _OUTSIDE_THE_ASYNC_FLAG_LIST)
|
||||
def test_built_in_loggers_deliver_one_record_per_async_request(vendors: _Vendors, surface: str) -> None:
|
||||
deliveries: Final = _Deliveries(vendors.sink)
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
response: Final = send(vendors.proxy, surface, "success", call_id)
|
||||
assert response.status_code == 200, response.text
|
||||
held(
|
||||
lambda: deliveries.per_destination(call_id),
|
||||
lambda counts: counts["datadog"] == 1 and counts["langsmith"] == 1 and counts["otel"] <= 1,
|
||||
holding=4 * _FLUSH_SECONDS,
|
||||
seconds=60,
|
||||
)
|
||||
|
|
@ -0,0 +1,214 @@
|
|||
"""Router strategies that keep both a sync and an async hook count a finished async request once: least-busy
|
||||
keeps counting a request still in flight while others finish around it, and usage-based routing v2 adds a
|
||||
request's tokens to its deployment once.
|
||||
|
||||
A legacy sync callback is configured, because that is what makes the sync handler run for an async call."""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, held
|
||||
from integration._support.hook_counts import PNG_BYTES, WAV_BYTES, HookRecords
|
||||
from integration._support.process import OwnedState, owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from redis import Redis
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_MODEL: Final = "router-strategy-model"
|
||||
_PROBES: Final = "integration._support.hook_count_probe"
|
||||
_REJECTION: Final = b'{"error":{"message":"scripted provider rejection","type":"invalid_request_error"}}'
|
||||
_IMAGE_TOKENS: Final = 12
|
||||
_IMAGE_REPLY: Final = json.dumps(
|
||||
{
|
||||
"created": 1700000000,
|
||||
"data": [{"b64_json": base64.b64encode(PNG_BYTES).decode()}],
|
||||
"usage": {
|
||||
"input_tokens": 5,
|
||||
"output_tokens": 7,
|
||||
"total_tokens": _IMAGE_TOKENS,
|
||||
"input_tokens_details": {"text_tokens": 5, "image_tokens": 0},
|
||||
},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Rig:
|
||||
proxy: Gateway
|
||||
records: HookRecords
|
||||
release: threading.Event
|
||||
deployment_id: str
|
||||
|
||||
|
||||
def _provider(release: threading.Event) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.target == "/v1/images/generations":
|
||||
return Reply(body=_IMAGE_REPLY)
|
||||
spoken: Final = str(json.loads(request.body or b"{}").get("input", ""))
|
||||
if spoken.startswith("fail"):
|
||||
return Reply(status=400, body=_REJECTION)
|
||||
if spoken.startswith("hold"):
|
||||
return Reply(chunks=(WAV_BYTES,), content_type="audio/wav", gate_before_first=release)
|
||||
return Reply(body=WAV_BYTES, content_type="audio/wav")
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _write_config(path: Path, provider_url: str, strategy: str, model: str, deployment_id: str) -> Path:
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": _MODEL,
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_base": provider_url + "/v1",
|
||||
"api_key": "synthetic-provider-key",
|
||||
},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
],
|
||||
"litellm_settings": {
|
||||
"success_callback": [f"{_PROBES}.sync_success_function"],
|
||||
"failure_callback": [f"{_PROBES}.sync_failure_function"],
|
||||
},
|
||||
"general_settings": {
|
||||
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
||||
"database_url": "os.environ/DATABASE_URL",
|
||||
"disable_model_info_refresh": True,
|
||||
},
|
||||
"router_settings": {
|
||||
"routing_strategy": strategy,
|
||||
"routing_strategy_args": {"ttl": 600},
|
||||
"num_retries": 0,
|
||||
"disable_cooldowns": True,
|
||||
"redis_host": "os.environ/REDIS_HOST",
|
||||
"redis_port": "os.environ/REDIS_PORT",
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _router_proxy(tmp_path: Path, strategy: str, model: str) -> Iterator[_Rig]:
|
||||
release: Final = threading.Event()
|
||||
deployment_id: Final = f"{strategy}-{uuid.uuid4().hex}"
|
||||
with (
|
||||
wire_server(_provider(release)) as upstream,
|
||||
wire_server(lambda request: Reply()) as sink,
|
||||
gateway_from_environment() as rig,
|
||||
):
|
||||
config: Final = _write_config(tmp_path / f"{strategy}.yaml", upstream.url, strategy, model, deployment_id)
|
||||
try:
|
||||
with owned_proxy(
|
||||
rig,
|
||||
tmp_path,
|
||||
{"HOOK_COUNT_PROBE_URL": sink.url, "STORE_MODEL_IN_DB": "False"},
|
||||
config=config,
|
||||
workers=2,
|
||||
state=OwnedState.SHARED_WITH_WORKER,
|
||||
) as proxy:
|
||||
yield _Rig(proxy, HookRecords(sink, function_callbacks=True), release, deployment_id)
|
||||
finally:
|
||||
release.set()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def least_busy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
with _router_proxy(tmp_path_factory.mktemp("least-busy"), "least-busy", "openai/gpt-4o-mini-tts") as rig:
|
||||
yield rig
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def usage_based(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
|
||||
with _router_proxy(tmp_path_factory.mktemp("usage-based"), "usage-based-routing-v2", "openai/gpt-image-1") as rig:
|
||||
yield rig
|
||||
|
||||
|
||||
def _redis() -> Redis:
|
||||
return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
|
||||
|
||||
|
||||
def _in_flight(rig: _Rig) -> int:
|
||||
with _redis() as cache:
|
||||
stored: Final = cache.get(f"{_MODEL}_request_count:{rig.deployment_id}")
|
||||
return 0 if stored is None else int(stored)
|
||||
|
||||
|
||||
def _tokens_counted(rig: _Rig) -> int:
|
||||
with _redis() as cache:
|
||||
keys: Final = tuple(cache.scan_iter(match=f"{rig.deployment_id}:*:tpm:*"))
|
||||
return sum(int(cache.get(key) or 0) for key in keys)
|
||||
|
||||
|
||||
def _post(rig: _Rig, path: str, body: dict[str, object], call_id: str) -> httpx.Response:
|
||||
return rig.proxy.client.post(
|
||||
path,
|
||||
json={"model": _MODEL, **body},
|
||||
headers={"Authorization": f"Bearer {rig.proxy.key}", "x-litellm-call-id": call_id},
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
|
||||
def _speak(rig: _Rig, text: str, call_id: str) -> httpx.Response:
|
||||
return _post(rig, "/v1/audio/speech", {"input": text, "voice": "alloy", "response_format": "wav"}, call_id)
|
||||
|
||||
|
||||
def _logged(rig: _Rig, call_id: str) -> None:
|
||||
"""The request's success or failure logging, the legacy sync callback included, has run."""
|
||||
eventually(lambda: rig.records.for_call(call_id), bool, seconds=20)
|
||||
|
||||
|
||||
def _finish_a_speech_request(rig: _Rig, outcome: str) -> None:
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
finished: Final = _speak(rig, f"{'ok' if outcome == 'success' else 'fail'} {call_id}", call_id)
|
||||
assert finished.status_code == (200 if outcome == "success" else 400), finished.text
|
||||
_logged(rig, call_id)
|
||||
held(lambda: _in_flight(rig), lambda count: count == 1, holding=3, seconds=20)
|
||||
|
||||
|
||||
def _generate_an_image(rig: _Rig) -> None:
|
||||
call_id: Final = str(uuid.uuid4())
|
||||
response: Final = _post(rig, "/v1/images/generations", {"prompt": f"tokens {call_id}", "n": 1}, call_id)
|
||||
assert response.status_code == 200, response.text
|
||||
_logged(rig, call_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("outcome", ["success", "failure"])
|
||||
def test_least_busy_keeps_counting_a_request_in_flight_while_async_speech_requests_finish(
|
||||
least_busy: _Rig, outcome: str
|
||||
) -> None:
|
||||
held(lambda: _in_flight(least_busy), lambda count: count == 0, holding=1, seconds=30)
|
||||
least_busy.release.clear()
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
in_flight: Final = pool.submit(_speak, least_busy, f"hold {uuid.uuid4().hex}", str(uuid.uuid4()))
|
||||
eventually(lambda: _in_flight(least_busy), lambda count: count == 1, seconds=30)
|
||||
try:
|
||||
for _ in range(3):
|
||||
_finish_a_speech_request(least_busy, outcome)
|
||||
finally:
|
||||
least_busy.release.set()
|
||||
assert in_flight.result(timeout=60).status_code == 200
|
||||
held(lambda: _in_flight(least_busy), lambda count: count == 0, holding=3, seconds=30)
|
||||
|
||||
|
||||
def test_usage_based_routing_counts_an_async_image_request_s_tokens_once(usage_based: _Rig) -> None:
|
||||
for _ in range(3):
|
||||
_generate_an_image(usage_based)
|
||||
held(lambda: _tokens_counted(usage_based), lambda tokens: tokens == 3 * _IMAGE_TOKENS, holding=3, seconds=30)
|
||||
|
|
@ -67,7 +67,7 @@ async def test_success_hook_fires_the_expected_number_of_times(
|
|||
pytest.param("async_log_failure_event", 1),
|
||||
pytest.param("async_post_call_failure_deployment_hook", 1),
|
||||
pytest.param("log_model_group_rate_limit_error", 1),
|
||||
pytest.param("log_failure_event", 0, marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
|
||||
pytest.param("log_failure_event", 0),
|
||||
),
|
||||
)
|
||||
async def test_failure_hook_fires_the_expected_number_of_times(
|
||||
|
|
|
|||
|
|
@ -237,7 +237,7 @@ async def test_native_aocr_failure_callbacks_receive_call_type_error_and_no_resp
|
|||
with pytest.raises(litellm.InternalServerError):
|
||||
await call_aocr(server, callbacks=[Observe()])
|
||||
|
||||
assert [observation[0] for observation in observations] == ["sync", "async"]
|
||||
assert [observation[0] for observation in observations] == ["async"]
|
||||
assert all(observation[1] == "aocr" for observation in observations)
|
||||
assert all(isinstance(observation[2], litellm.InternalServerError) for observation in observations)
|
||||
assert all(observation[3] is None for observation in observations)
|
||||
|
|
@ -278,7 +278,7 @@ async def test_native_aocr_failure_callbacks_receive_state_added_by_pre_call_cal
|
|||
with pytest.raises(litellm.InternalServerError):
|
||||
await call_aocr(server, callbacks=[TrackInFlightRequest()])
|
||||
|
||||
assert [event for event, _ in observed] == ["sync", "async"]
|
||||
assert [event for event, _ in observed] == ["async"]
|
||||
assert all(observed_token is token for _, observed_token in observed)
|
||||
|
||||
|
||||
|
|
@ -445,7 +445,7 @@ async def test_native_aocr_state_stashed_before_a_blocking_hook_raises_reaches_f
|
|||
await call_aocr(server)
|
||||
await drain_logging()
|
||||
|
||||
assert observed == [("sync", token, raised.value), ("async", token, raised.value)]
|
||||
assert observed == [("async", token, raised.value)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -144,7 +144,7 @@ async def test_mapped_failure_identity_and_deployment_snapshot(server: Recording
|
|||
with pytest.raises(litellm.InternalServerError) as caught:
|
||||
await call_aocr(server, callbacks=[recorder])
|
||||
failures: Final = tuple(event for event in recorder.events if "failure" in event.name)
|
||||
assert [event.name for event in failures] == ["log_failure_event", "async_log_failure_event"]
|
||||
assert [event.name for event in failures] == ["async_log_failure_event"]
|
||||
assert all(event.kwargs["exception"] is caught.value for event in failures)
|
||||
assert caught.value.status_code == 500
|
||||
assert snapshots[0] is not caught.value
|
||||
|
|
@ -324,24 +324,26 @@ async def test_failure_callbacks_continue_within_both_families(server: Recording
|
|||
server.enqueue(ResponseSpec(body={"message": "failed"}, status=500))
|
||||
observed: Final = []
|
||||
|
||||
class Broken(CustomLogger):
|
||||
class BrokenSync(CustomLogger):
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
observed.append(("broken-sync", kwargs["exception"]))
|
||||
raise RuntimeError("sync observer")
|
||||
|
||||
class FollowingSync(CustomLogger):
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
observed.append(("following-sync", kwargs["exception"]))
|
||||
|
||||
class BrokenAsync(CustomLogger):
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
observed.append(("broken-async", kwargs["exception"]))
|
||||
raise RuntimeError("async observer")
|
||||
|
||||
class Following(CustomLogger):
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
observed.append(("following-sync", kwargs["exception"]))
|
||||
|
||||
class FollowingAsync(CustomLogger):
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
observed.append(("following-async", kwargs["exception"]))
|
||||
|
||||
with pytest.raises(litellm.InternalServerError) as caught:
|
||||
await call_aocr(server, callbacks=[Broken(), Following()])
|
||||
await call_aocr(server, callbacks=[BrokenSync(), FollowingSync(), BrokenAsync(), FollowingAsync()])
|
||||
assert [name for name, _ in observed] == ["broken-sync", "following-sync", "broken-async", "following-async"]
|
||||
assert all(error is caught.value for _, error in observed)
|
||||
|
||||
|
|
|
|||
|
|
@ -95,7 +95,7 @@ async def test_success_hook_fires_the_expected_number_of_times(
|
|||
"async_post_call_failure_deployment_hook", True, 1, id="async-async_post_call_failure_deployment_hook"
|
||||
),
|
||||
pytest.param("log_model_group_rate_limit_error", True, 1, id="async-log_model_group_rate_limit_error"),
|
||||
pytest.param("log_failure_event", True, 0, id="async-log_failure_event", marks=SYNC_HANDLERS_ON_ASYNC_CALLS),
|
||||
pytest.param("log_failure_event", True, 0, id="async-log_failure_event"),
|
||||
pytest.param("log_failure_event", False, 1, id="sync-log_failure_event"),
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3088,7 +3088,7 @@ async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_c
|
|||
prefer_async_handlers=True,
|
||||
)
|
||||
|
||||
mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception)
|
||||
mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception, for_async_request=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8188,6 +8188,155 @@ def test_cache_callback_alone_does_not_start_sync_callbacks_for_async_requests(
|
|||
assert logging_obj._should_run_sync_callbacks_for_async_calls() is expected
|
||||
|
||||
|
||||
class _FailureHookRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.hooks: list[str] = []
|
||||
|
||||
def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
|
||||
self.hooks.append("sync_failure")
|
||||
|
||||
|
||||
class _BothFailureHooksRecorder(_FailureHookRecorder):
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: object, response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self.hooks.append("async_failure")
|
||||
|
||||
|
||||
def _failure_hook_count_logging_obj(
|
||||
call_type: str, recorder: _FailureHookRecorder, *, async_failure_listed: bool
|
||||
) -> LitellmLogging:
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini-tts",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type=call_type,
|
||||
start_time=datetime.datetime.now(),
|
||||
litellm_call_id=f"failure-hook-count-{call_type}",
|
||||
function_id=f"failure-hook-count-{call_type}",
|
||||
dynamic_failure_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder] if async_failure_listed else [],
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_global_callbacks(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
for name in (
|
||||
"callbacks",
|
||||
"success_callback",
|
||||
"failure_callback",
|
||||
"_async_success_callback",
|
||||
"_async_failure_callback",
|
||||
):
|
||||
monkeypatch.setattr(litellm, name, [])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("recorder_type", "async_listed", "expected"),
|
||||
[
|
||||
(_BothFailureHooksRecorder, True, ["async_failure"]),
|
||||
(_FailureHookRecorder, True, ["sync_failure"]),
|
||||
(_BothFailureHooksRecorder, False, ["sync_failure"]),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("call_type", ["aspeech", "atranscription", "aimage_edit", "asearch"])
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
async def test_async_call_outside_the_async_flag_list_gets_one_custom_logger_failure_hook(
|
||||
call_type: str, recorder_type: type[_FailureHookRecorder], async_listed: bool, expected: list[str]
|
||||
) -> None:
|
||||
recorder: Final = recorder_type()
|
||||
logging_obj: Final = _failure_hook_count_logging_obj(call_type, recorder, async_failure_listed=async_listed)
|
||||
|
||||
logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True)
|
||||
await logging_obj.async_failure_handler(ValueError("upstream rejected"), "traceback")
|
||||
|
||||
assert recorder.hooks == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
async def test_logger_in_the_global_failure_callback_lists_gets_one_failure_hook(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder: Final = _BothFailureHooksRecorder()
|
||||
monkeypatch.setattr(litellm, "failure_callback", [recorder])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [recorder])
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini-tts",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="aspeech",
|
||||
start_time=datetime.datetime.now(),
|
||||
litellm_call_id="global-failure",
|
||||
function_id="global-failure",
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
|
||||
|
||||
logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True)
|
||||
await logging_obj.async_failure_handler(ValueError("upstream rejected"), "traceback")
|
||||
|
||||
assert recorder.hooks == ["async_failure"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
def test_failure_handler_run_without_an_async_failure_run_keeps_the_sync_hook() -> None:
|
||||
recorder: Final = _BothFailureHooksRecorder()
|
||||
|
||||
_failure_hook_count_logging_obj("aspeech", recorder, async_failure_listed=True).failure_handler(
|
||||
ValueError("upstream rejected"), "traceback"
|
||||
)
|
||||
|
||||
assert recorder.hooks == ["sync_failure"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
def test_async_call_outside_the_async_flag_list_still_runs_sync_failure_function_callbacks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm_logging, "customLogger", CustomLogger())
|
||||
calls: Final[list[str]] = []
|
||||
|
||||
def failure_function(kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
|
||||
calls.append("function_failure")
|
||||
|
||||
recorder: Final = _BothFailureHooksRecorder()
|
||||
logging_obj: Final = _failure_hook_count_logging_obj("aspeech", recorder, async_failure_listed=True)
|
||||
logging_obj.dynamic_failure_callbacks = [recorder, failure_function]
|
||||
|
||||
logging_obj.failure_handler(ValueError("upstream rejected"), "traceback", for_async_request=True)
|
||||
|
||||
assert calls == ["function_failure"]
|
||||
assert recorder.hooks == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
async def test_client_wrapped_failure_gives_a_logger_in_litellm_callbacks_one_hook(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
recorder: Final = _BothFailureHooksRecorder()
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
|
||||
@litellm.utils.client
|
||||
async def aspeech(model: str, input: str, voice: str, **kwargs: object) -> None:
|
||||
raise litellm.BadRequestError(message="upstream rejected", model=model, llm_provider="openai")
|
||||
|
||||
@litellm.utils.client
|
||||
def speech(model: str, input: str, voice: str, **kwargs: object) -> None:
|
||||
raise litellm.BadRequestError(message="upstream rejected", model=model, llm_provider="openai")
|
||||
|
||||
with pytest.raises(litellm.BadRequestError):
|
||||
await aspeech(model="gpt-4o-mini-tts", input="hi", voice="alloy")
|
||||
with pytest.raises(litellm.BadRequestError):
|
||||
speech(model="gpt-4o-mini-tts", input="hi", voice="alloy")
|
||||
|
||||
assert recorder.hooks == ["async_failure", "sync_failure"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("turn_off_message_logging", [False, True])
|
||||
@pytest.mark.parametrize("make_response", [ModelResponse, _responses_api_response_with_text])
|
||||
|
|
|
|||
|
|
@ -209,4 +209,4 @@ async def test_chat_completion_request_with_redaction(route, body, router, monke
|
|||
|
||||
assert _test_logger.reaches_async_failure_event is True
|
||||
|
||||
assert _test_logger.reaches_sync_failure_event is True
|
||||
assert _test_logger.reaches_sync_failure_event is False
|
||||
|
|
|
|||
|
|
@ -2652,6 +2652,58 @@ async def test_post_call_failure_hook_auth_error_llm_api_route():
|
|||
mock_handle_logging.assert_called_once()
|
||||
|
||||
|
||||
class _InlineThread:
|
||||
def __init__(self, target, args=(), kwargs=None, daemon=None) -> None:
|
||||
self._run = lambda: target(*args, **(kwargs or {}))
|
||||
|
||||
def start(self) -> None:
|
||||
self._run()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_only_failure_gives_a_logger_with_both_failure_hooks_only_the_async_one(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import types
|
||||
|
||||
import litellm.proxy.utils as proxy_utils
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
hooks: Final[list[str]] = []
|
||||
|
||||
class Recorder(CustomLogger):
|
||||
def log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
hooks.append("sync_failure")
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
hooks.append("async_failure")
|
||||
|
||||
for name in ("failure_callback", "_async_failure_callback"):
|
||||
monkeypatch.setattr(litellm, name, [])
|
||||
monkeypatch.setattr(proxy_utils, "threading", types.SimpleNamespace(Thread=_InlineThread))
|
||||
recorder: Final = Recorder()
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-4o-mini-tts",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="aspeech",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="proxy-only-failure",
|
||||
function_id="proxy-only-failure",
|
||||
dynamic_failure_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
|
||||
|
||||
await proxy_utils.ProxyLogging._dispatch_proxy_only_failure_handlers(
|
||||
litellm_logging_obj=logging_obj,
|
||||
original_exception=HTTPException(status_code=401, detail="Authentication Error: invalid user key"),
|
||||
)
|
||||
|
||||
assert hooks == ["async_failure"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"request_data, route, expected_call_type",
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import pytest
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._internal_context import is_internal_call
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.rust_bridge.host import callbacks_legacy_python as legacy
|
||||
from litellm.rust_bridge.host.callbacks_legacy_python import failure_handler, setup
|
||||
|
|
@ -167,6 +168,75 @@ def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation
|
|||
pending.close()
|
||||
|
||||
|
||||
class _Hooks(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.hooks: list[str] = []
|
||||
|
||||
def log_success_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
|
||||
self.hooks.append("sync_success")
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: object, response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self.hooks.append("async_success")
|
||||
|
||||
def log_failure_event(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None:
|
||||
self.hooks.append("sync_failure")
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: object, response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self.hooks.append("async_failure")
|
||||
|
||||
|
||||
def _logger_with_hooks(hooks: _Hooks) -> Logging:
|
||||
logger: Final = Logging(
|
||||
model="mistral/mistral-ocr-latest",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="aocr",
|
||||
start_time=datetime.datetime.now(),
|
||||
litellm_call_id="hooks",
|
||||
function_id="hooks",
|
||||
dynamic_success_callbacks=[hooks],
|
||||
dynamic_async_success_callbacks=[hooks],
|
||||
dynamic_failure_callbacks=[hooks],
|
||||
dynamic_async_failure_callbacks=[hooks],
|
||||
)
|
||||
logger.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
|
||||
return logger
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_of_an_async_call_gives_a_logger_with_both_failure_hooks_only_the_async_one() -> None:
|
||||
hooks: Final = _Hooks()
|
||||
now: Final = datetime.datetime.now()
|
||||
kwargs: Final = {**_OCR_KWARGS, "litellm_logging_obj": _logger_with_hooks(hooks)}
|
||||
result: Final = setup("aocr", (), kwargs, now, asynchronous=True)
|
||||
|
||||
assert failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=False) is None
|
||||
pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True)
|
||||
assert pending is not None
|
||||
await pending
|
||||
|
||||
assert hooks.hooks == ["async_failure"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("earlier_async_call", [False, True])
|
||||
def test_failure_of_a_sync_call_keeps_the_sync_failure_hook(earlier_async_call: bool) -> None:
|
||||
hooks: Final = _Hooks()
|
||||
now: Final = datetime.datetime.now()
|
||||
kwargs: Final = {**_OCR_KWARGS, "litellm_logging_obj": _logger_with_hooks(hooks)}
|
||||
if earlier_async_call:
|
||||
setup("aocr", (), kwargs, now, asynchronous=True)
|
||||
result: Final = setup("ocr", (), kwargs, now, asynchronous=False)
|
||||
|
||||
assert failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=False) is None
|
||||
|
||||
assert hooks.hooks == ["sync_failure"]
|
||||
|
||||
|
||||
CONTRACT_PATH: Final = Path(__file__).parents[4] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json"
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue