mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(router): leave callback-started calls unmarked and stop client flags from turning mirroring off
Logging handlers now run outside the mirrored-call context, so an LLM call a success or failure callback starts while logging a mirrored call is not flagged as mirrored. The guard against mirroring a mirror reads the same context variable instead of request metadata, so a client-sent is_silent_experiment no longer skips the mirrors. The unused module-level get_standard_logging_metadata stops copying resolved keys from raw metadata.
This commit is contained in:
parent
054f220e53
commit
e4bf3f356b
6 changed files with 160 additions and 47 deletions
|
|
@ -139,6 +139,16 @@ def in_mirrored_call() -> bool:
|
|||
return _mirrored_call.get()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def outside_mirrored_call() -> Generator[None]:
|
||||
"""Logging a mirrored call: the calls its callbacks start are not mirrors themselves."""
|
||||
token: Final = _mirrored_call.set(False)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_restore(_mirrored_call, token)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def pinned_billing_time(moment: datetime) -> Generator[None]:
|
||||
"""Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read."""
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from httpx import Response, Timeout
|
|||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import in_mirrored_call, post_response_phase
|
||||
from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase
|
||||
from litellm._logging import (
|
||||
is_debugging_on,
|
||||
redact_string,
|
||||
|
|
@ -3246,7 +3246,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Restores trace_id/session_id contextvars once this attempt's own success
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
with post_response_phase():
|
||||
with post_response_phase(), outside_mirrored_call():
|
||||
return self._success_handler_body(
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
|
|
@ -3746,7 +3746,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Restores trace_id/session_id contextvars once this attempt's own success
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
with post_response_phase():
|
||||
with post_response_phase(), outside_mirrored_call():
|
||||
return await self._async_success_handler_body(
|
||||
result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs
|
||||
)
|
||||
|
|
@ -4263,13 +4263,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
``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,
|
||||
)
|
||||
with outside_mirrored_call():
|
||||
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()
|
||||
|
||||
|
|
@ -4462,12 +4463,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Restores trace_id/session_id contextvars once this attempt's own failure
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
return await self._async_failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
with outside_mirrored_call():
|
||||
return await self._async_failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
|
|
@ -7714,9 +7716,8 @@ def get_standard_logging_metadata(
|
|||
)
|
||||
if isinstance(metadata, dict):
|
||||
# Update the clean_metadata with values from input metadata that match StandardLoggingMetadata fields
|
||||
for key in StandardLoggingMetadata.__annotations__:
|
||||
if key in metadata:
|
||||
clean_metadata[key] = metadata[key]
|
||||
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
|
||||
clean_metadata[key] = metadata[key]
|
||||
|
||||
if metadata.get("user_api_key") is not None:
|
||||
if is_valid_sha256_hash(str(metadata.get("user_api_key"))):
|
||||
|
|
|
|||
|
|
@ -52,7 +52,7 @@ from typing_extensions import assert_never, overload
|
|||
import litellm
|
||||
import litellm.litellm_core_utils.exception_mapping_utils
|
||||
from litellm import get_secret_str
|
||||
from litellm._internal_context import mirrored_call, service_target, with_service_target
|
||||
from litellm._internal_context import in_mirrored_call, mirrored_call, service_target, with_service_target
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import (
|
||||
|
|
@ -702,7 +702,7 @@ def _every_deployment_refused(
|
|||
|
||||
|
||||
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
|
||||
if silent_model is None:
|
||||
if silent_model is None or in_mirrored_call():
|
||||
return ()
|
||||
try:
|
||||
targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model)
|
||||
|
|
@ -3325,10 +3325,6 @@ class Router:
|
|||
Run a silent experiment in the background (thread).
|
||||
"""
|
||||
try:
|
||||
# Prevent infinite recursion if silent model also has a silent model
|
||||
if kwargs.get("metadata", {}).get("is_silent_experiment", False):
|
||||
return
|
||||
|
||||
messages = copy.deepcopy(messages)
|
||||
|
||||
verbose_router_logger.info("Starting silent experiment for model %s", silent_model)
|
||||
|
|
@ -4513,10 +4509,6 @@ class Router:
|
|||
Run a silent experiment in the background.
|
||||
"""
|
||||
try:
|
||||
# Prevent infinite recursion if silent model also has a silent model
|
||||
if kwargs.get("metadata", {}).get("is_silent_experiment", False):
|
||||
return
|
||||
|
||||
messages = copy.deepcopy(messages)
|
||||
|
||||
verbose_router_logger.info("Starting silent experiment for model %s", silent_model)
|
||||
|
|
|
|||
|
|
@ -82,9 +82,9 @@ class _ShadowRecorder(CustomLogger):
|
|||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
metadata: Final = _OBJECT.validate_python(_OBJECT.validate_python(kwargs["litellm_params"])["metadata"] or {})
|
||||
if metadata.get("is_silent_experiment") is True:
|
||||
self.shadow_groups.put(str(metadata.get("model_group")))
|
||||
payload: Final = _OBJECT.validate_python(kwargs["standard_logging_object"])
|
||||
if _OBJECT.validate_python(payload["metadata"]).get("is_silent_experiment") is True:
|
||||
self.shadow_groups.put(str(payload["model_group"]))
|
||||
|
||||
def drained(self) -> tuple[str, ...]:
|
||||
return tuple(self.shadow_groups.get_nowait() for _ in range(self.shadow_groups.qsize()))
|
||||
|
|
|
|||
|
|
@ -11975,6 +11975,16 @@ def test_get_standard_logging_metadata(metadata_key, metadata_value):
|
|||
assert standard_logging_metadata[metadata_key] == metadata_value
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
|
||||
def test_module_level_get_standard_logging_metadata_ignores_a_caller_sent_silent_experiment_flag():
|
||||
from litellm.litellm_core_utils.litellm_logging import get_standard_logging_metadata
|
||||
|
||||
result = get_standard_logging_metadata({"is_silent_experiment": True, "user_api_key_alias": "alias"})
|
||||
|
||||
assert "is_silent_experiment" not in result
|
||||
assert result["user_api_key_alias"] == "alias"
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown")
|
||||
def test_get_standard_logging_metadata_user_api_key_hash():
|
||||
valid_hash = "a" * 64 # 64 character string
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import asyncio
|
|||
import datetime
|
||||
import json
|
||||
import time
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Callable, Generator, Mapping
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -42,15 +42,43 @@ async def _settle_shared_logging_worker() -> None:
|
|||
await GLOBAL_LOGGING_WORKER.stop()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def recording_logger():
|
||||
class _JudgingLogger(_RecordingLogger):
|
||||
async def _judge(self, kwargs: Mapping[str, object]) -> None:
|
||||
metadata: Final = kwargs["litellm_params"]["metadata"]
|
||||
if "judge_of" in metadata:
|
||||
return
|
||||
await litellm.acompletion(
|
||||
model="openai/gpt-5.4-nano",
|
||||
api_key="fake-key",
|
||||
messages=[{"role": "user", "content": "score it"}],
|
||||
mock_response="ok",
|
||||
metadata={"judge_of": metadata["model_group"]},
|
||||
)
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await super().async_log_success_event(kwargs, response_obj, start_time, end_time)
|
||||
await self._judge(kwargs)
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await self._judge(kwargs)
|
||||
|
||||
def judge_flags(self) -> dict[str, object]:
|
||||
return {
|
||||
call["litellm_params"]["metadata"]["judge_of"]: call["standard_logging_object"]["metadata"].get(
|
||||
"is_silent_experiment"
|
||||
)
|
||||
for call in self.success_kwargs
|
||||
if "judge_of" in call["litellm_params"]["metadata"]
|
||||
}
|
||||
|
||||
|
||||
def _installed(logger: _RecordingLogger) -> Generator[_RecordingLogger]:
|
||||
settle_loop: Final = asyncio.new_event_loop()
|
||||
try:
|
||||
settle_loop.run_until_complete(_settle_shared_logging_worker())
|
||||
finally:
|
||||
settle_loop.close()
|
||||
original_callbacks: Final = litellm.callbacks
|
||||
logger: Final = _RecordingLogger()
|
||||
litellm.callbacks = [logger]
|
||||
try:
|
||||
yield logger
|
||||
|
|
@ -58,6 +86,16 @@ def recording_logger():
|
|||
litellm.callbacks = original_callbacks
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def recording_logger():
|
||||
yield from _installed(_RecordingLogger())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def judging_logger():
|
||||
yield from _installed(_JudgingLogger())
|
||||
|
||||
|
||||
async def _wait_for_shadow_successes(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
|
||||
|
|
@ -70,6 +108,21 @@ async def _wait_for_all_successes(logger: _RecordingLogger, expected: int, timeo
|
|||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
async def _wait_for_judges(logger: _JudgingLogger, judged: set[str], timeout: float = 5.0) -> None:
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
while not judged <= logger.judge_flags().keys() and time.monotonic() < deadline:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
def _silent_flag_by_group(logger: _RecordingLogger) -> dict[str, object]:
|
||||
return {
|
||||
call["standard_logging_object"]["model_group"]: call["standard_logging_object"]["metadata"].get(
|
||||
"is_silent_experiment"
|
||||
)
|
||||
for call in logger.success_kwargs
|
||||
}
|
||||
|
||||
|
||||
def _wait_for_shadow_successes_sync(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
|
||||
|
|
@ -806,15 +859,62 @@ async def test_aembedding_with_silent_model_sends_provider_body_without_it(
|
|||
request_body: Final = json.loads(route.calls.last.request.read())
|
||||
assert request_body == {"model": "embed-model", "input": ["black dresses"], "input_type": "query"}
|
||||
assert response.data[0]["embedding"] == [0.1, 0.2]
|
||||
@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS)
|
||||
def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment):
|
||||
router = Router(model_list=_streaming_model_list(["shadow-a"]))
|
||||
with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion:
|
||||
run_silent_experiment(
|
||||
router,
|
||||
silent_model="shadow-a",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
metadata={"is_silent_experiment": True},
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_caller_sent_silent_experiment_flag_still_mirrors_the_request_async(
|
||||
recording_logger: _RecordingLogger,
|
||||
) -> None:
|
||||
router: Final = Router(model_list=_streaming_model_list(["shadow-a"]))
|
||||
await router.acompletion(
|
||||
model="primary-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="pong",
|
||||
metadata={"is_silent_experiment": True},
|
||||
)
|
||||
await _wait_for_all_successes(recording_logger, expected=2)
|
||||
|
||||
acompletion.assert_not_awaited()
|
||||
assert _silent_flag_by_group(recording_logger) == {"primary-model": None, "shadow-a": True}
|
||||
|
||||
|
||||
def test_a_caller_sent_silent_experiment_flag_still_mirrors_the_request_sync(
|
||||
recording_logger: _RecordingLogger,
|
||||
) -> None:
|
||||
router: Final = Router(model_list=_streaming_model_list(["shadow-a"]))
|
||||
router.completion(
|
||||
model="primary-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="pong",
|
||||
metadata={"is_silent_experiment": True},
|
||||
)
|
||||
deadline: Final = time.monotonic() + 5.0
|
||||
while "shadow-a" not in _silent_flag_by_group(recording_logger) and time.monotonic() < deadline:
|
||||
time.sleep(0.05)
|
||||
|
||||
assert _silent_flag_by_group(recording_logger) == {"shadow-a": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("shadow_outcome", ["pong", "litellm.InternalServerError"])
|
||||
async def test_calls_a_callback_starts_while_logging_a_mirrored_call_are_not_marked_as_mirrored(
|
||||
judging_logger: _JudgingLogger, shadow_outcome: str
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "primary-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-5.4-mini",
|
||||
"api_key": "fake-key",
|
||||
"silent_model": "shadow-a",
|
||||
"mock_response": "pong",
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "shadow-a",
|
||||
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-key", "mock_response": shadow_outcome},
|
||||
},
|
||||
]
|
||||
)
|
||||
await router.acompletion(model="primary-model", messages=[{"role": "user", "content": "hi"}])
|
||||
await _wait_for_judges(judging_logger, {"primary-model", "shadow-a"})
|
||||
|
||||
assert judging_logger.judge_flags() == {"primary-model": None, "shadow-a": None}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue