scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Merge pull request #41368 from BerriAI/litellm_shadow_streaming_multi_target

Browse source 
feat(router): stream shadow traffic and fan out silent_model to multiple targets
This commit is contained in:
Yassin Kortam 2026-09-16 09:08:24 -07:00 • committed by GitHub
commit 29a712b186
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 372 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
litellm/router.py
View file

@ -425,12 +425,34 @@ def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream])
_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({})
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
def _as_retry_skipped_deployment_ids(value: object) -> tuple[str, ...]:
return tuple(item for item in value if isinstance(item, str)) if isinstance(value, tuple) else ()
def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]:
if silent_model is None:
return ()
try:
targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model)
except ValidationError:
verbose_router_logger.warning(
"silent_model must be a model name or a list of model names, got %r; skipping shadow traffic",
silent_model,
)
return ()
return (targets,) if isinstance(targets, str) else tuple(targets)
def _silent_experiment_kwargs_snapshot(kwargs: Mapping[str, object]) -> Mapping[str, object]:
metadata: Final = kwargs.get("metadata")
if not isinstance(metadata, Mapping):
return MappingProxyType({**kwargs})
return MappingProxyType({**kwargs, "metadata": dict(metadata)})
def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]:
"""
Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still
@ -2455,18 +2477,17 @@ class Router:
)
silent_model: Final = litellm_params.pop("silent_model", None)
if silent_model is not None:
for silent_target in _silent_experiment_targets(silent_model):
# Mirroring traffic to a secondary model
# Use threading.Thread (not ThreadPoolExecutor) - executor.submit()
# requires pickling args, which fails when kwargs contain unpicklable
# objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment.
thread: Final = threading.Thread(
threading.Thread(
target=self._silent_experiment_completion,
args=(silent_model, messages),
kwargs=kwargs,
args=(silent_target, messages),
kwargs=_silent_experiment_kwargs_snapshot(kwargs),
daemon=True,
)
thread.start()
).start()
kwargs.setdefault("messages", messages)
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
@ -2567,9 +2588,6 @@ class Router:
silent_kwargs["metadata"]["is_silent_experiment"] = True
# Force stream=False so the response is fully consumed and callbacks fire
silent_kwargs["stream"] = False
# Pop logging objects and call IDs to ensure a fresh logging context
# This prevents collisions in the Proxy's database (spend_logs)
silent_kwargs.pop("litellm_call_id", None)
@ -2579,6 +2597,23 @@ class Router:
return silent_kwargs
async def _run_silent_experiment(
self, silent_model: str, messages: Sequence[Mapping[str, str]], silent_kwargs: Mapping[str, object]
) -> None:
remaining_kwargs: Final = MappingProxyType(
{key: value for key, value in silent_kwargs.items() if key != "stream"}
)
response: Final = await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
stream=bool(silent_kwargs.get("stream", False)),
**remaining_kwargs,
)
if not isinstance(response, CustomStreamWrapper):
return
async for _ in response:
pass
def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs):
"""
Run a silent experiment in the background (thread).
@ -2604,11 +2639,7 @@ class Router:
try:
async def _run_silent_completion():
await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
**silent_kwargs,
)
await self._run_silent_experiment(silent_model, messages, silent_kwargs)
# Drain any fire-and-forget tasks (e.g. alerting hooks)
# scheduled via asyncio.create_task during acompletion.
pending: Final = asyncio.all_tasks()
@ -3500,11 +3531,7 @@ class Router:
silent_kwargs["metadata"]["model_group"] = silent_model
# Trigger the silent request
await self.acompletion(
model=silent_model,
messages=cast(list[AllMessageValues], messages),
**silent_kwargs,
)
await self._run_silent_experiment(silent_model, messages, silent_kwargs)
except Exception as e:
verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e)
@ -3563,14 +3590,14 @@ class Router:
)
silent_model: Final = litellm_params.pop("silent_model", None)
if silent_model is not None:
for silent_target in _silent_experiment_targets(silent_model):
# Mirroring traffic to a secondary model
# This is a silent experiment, so we don't want to block the primary request
asyncio.create_task(
self._silent_experiment_acompletion(
silent_model=silent_model,
silent_model=silent_target,
messages=messages, # Use messages instead of *args
**kwargs,
**_silent_experiment_kwargs_snapshot(kwargs),
)
)

325
tests/test_litellm/test_router_silent_experiment.py
View file

@ -1,11 +1,73 @@
import asyncio
import time
from collections.abc import Callable, Mapping
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.router import Router
from litellm.router import _silent_experiment_kwargs_snapshot
from litellm.router import _silent_experiment_targets
class _RecordingLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.success_kwargs: list[dict[str, object]] = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.success_kwargs.append(kwargs)
def shadow_successes(self) -> list[dict[str, object]]:
return [
call
for call in self.success_kwargs
if call.get("litellm_params", {}).get("metadata", {}).get("is_silent_experiment") is True
]
@pytest.fixture
def recording_logger():
original_callbacks: Final = litellm.callbacks
logger: Final = _RecordingLogger()
litellm.callbacks = [logger]
try:
yield logger
finally:
litellm.callbacks = original_callbacks
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:
await asyncio.sleep(0.05)
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:
time.sleep(0.05)
def _streaming_model_list(silent_model: object) -> list[dict[str, object]]:
return [
{
"model_name": "primary-model",
"litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key", "silent_model": silent_model},
},
{
"model_name": "shadow-a",
"litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "fake-key", "silent_model": "shadow-b"},
},
{
"model_name": "shadow-b",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "fake-key"},
},
]
class _NonCopyableSpan:
@ -65,8 +127,7 @@ def test_get_silent_experiment_kwargs():
assert result["metadata"]["is_silent_experiment"] is True
assert result["metadata"]["foo"] == "bar"
assert "litellm_call_id" not in result
# stream must be forced to False so callbacks fire in background
assert result["stream"] is False
assert result["stream"] is True
# proxy_server_request must be preserved for spend log metadata
assert "proxy_server_request" in result
# CRITICAL: metadata must be a DIFFERENT dict object than the original,
@ -86,6 +147,247 @@ def test_get_silent_experiment_kwargs():
assert result["metadata"]["user_api_key_auth"] is mock_auth
def test_get_silent_experiment_kwargs_without_stream_stays_non_streaming():
router = Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k"}}])
result = router._get_silent_experiment_kwargs(metadata={"foo": "bar"}, stream=False)
assert result["stream"] is False
assert "stream" not in router._get_silent_experiment_kwargs(metadata={"foo": "bar"})
@pytest.mark.parametrize(
"silent_model, expected",
[
("shadow-a", ("shadow-a",)),
(["shadow-a", "shadow-b"], ("shadow-a", "shadow-b")),
([], ()),
(None, ()),
(42, ()),
(["shadow-a", 42], ()),
],
)
def test_silent_experiment_targets(silent_model, expected):
assert _silent_experiment_targets(silent_model) == expected
@pytest.mark.asyncio
async def test_streaming_shadow_is_streamed_and_drained_async(recording_logger):
router = Router(model_list=_streaming_model_list("shadow-a"))
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
stream_options={"include_usage": True},
mock_response="pong",
metadata={"foo": "bar"},
)
chunks = [chunk async for chunk in response]
assert chunks
await _wait_for_shadow_successes(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
shadow = shadow_successes[0]
assert shadow["stream"] is True
assert shadow["stream_options"] == {"include_usage": True}
assert shadow["litellm_params"]["metadata"]["model_group"] == "shadow-a"
assert shadow["async_complete_streaming_response"] is not None
def test_streaming_shadow_is_streamed_and_drained_sync(recording_logger):
router = Router(model_list=_streaming_model_list("shadow-a"))
response = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
mock_response="pong",
metadata={"foo": "bar"},
)
chunks = list(response)
assert chunks
_wait_for_shadow_successes_sync(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
assert shadow_successes[0]["stream"] is True
assert shadow_successes[0]["litellm_params"]["metadata"]["model_group"] == "shadow-a"
assert shadow_successes[0]["async_complete_streaming_response"] is not None
@pytest.mark.asyncio
async def test_multiple_shadow_targets_fan_out_async(recording_logger):
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
metadata = {"foo": "bar"}
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
stream=True,
mock_response="pong",
metadata=metadata,
)
assert [chunk async for chunk in response]
await _wait_for_shadow_successes(recording_logger, expected=2)
shadow_successes = recording_logger.shadow_successes()
model_groups = sorted(call["litellm_params"]["metadata"]["model_group"] for call in shadow_successes)
assert model_groups == ["shadow-a", "shadow-b"]
shadow_metadatas = [call["litellm_params"]["metadata"] for call in shadow_successes]
assert shadow_metadatas[0] is not shadow_metadatas[1]
assert all(call["stream"] is True for call in shadow_successes)
assert "is_silent_experiment" not in metadata
assert metadata.get("model_group") != "shadow-a"
primary_successes = [call for call in recording_logger.success_kwargs if call not in shadow_successes]
assert len(primary_successes) == 1
assert primary_successes[0]["litellm_params"]["metadata"]["model_group"] == "primary-model"
def test_multiple_shadow_targets_fan_out_sync(recording_logger):
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
response = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert response.choices[0].message.content == "pong"
_wait_for_shadow_successes_sync(recording_logger, expected=2)
shadow_successes = recording_logger.shadow_successes()
model_groups = sorted(call["litellm_params"]["metadata"]["model_group"] for call in shadow_successes)
assert model_groups == ["shadow-a", "shadow-b"]
assert all(call["stream"] is False for call in shadow_successes)
def _tagged_primary_model_list() -> list[dict[str, object]]:
return [
{
"model_name": "primary-model",
"litellm_params": {
"model": "openai/gpt-5.4-mini",
"api_key": "fake-key",
"silent_model": "shadow-b",
"tags": ["primary-only"],
},
},
{
"model_name": "shadow-b",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "fake-key"},
},
]
def test_silent_experiment_kwargs_snapshot_is_isolated_from_later_primary_mutations():
metadata = {"foo": "bar"}
kwargs: dict[str, object] = {"metadata": metadata, "stream": True}
snapshot = _silent_experiment_kwargs_snapshot(kwargs)
kwargs["messages"] = [{"role": "user", "content": "added by the primary"}]
metadata["tags"] = ["primary-only"]
assert dict(snapshot) == {"metadata": {"foo": "bar"}, "stream": True}
assert dict(_silent_experiment_kwargs_snapshot({"stream": False, "metadata": None})) == {
"stream": False,
"metadata": None,
}
def test_sync_shadow_gets_kwargs_snapshot_taken_before_primary_mutates_them(recording_logger):
deferred: list[Callable[[], None]] = []
class _DeferredThread:
def __init__(self, target, args, kwargs, daemon) -> None:
deferred.append(lambda: target(*args, **kwargs))
def start(self) -> None:
return None
router = Router(model_list=_tagged_primary_model_list())
with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race
"litellm.router.threading", SimpleNamespace(Thread=_DeferredThread)
):
response = router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert response.choices[0].message.content == "pong"
assert len(deferred) == 1
deferred[0]()
_wait_for_shadow_successes_sync(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
shadow_metadata = shadow_successes[0]["litellm_params"]["metadata"]
assert shadow_metadata["model_group"] == "shadow-b"
assert "primary-only" not in shadow_metadata.get("tags", [])
def test_sync_shadow_workers_do_not_share_metadata_with_each_other(recording_logger):
workers: list[tuple[Mapping[str, object], Callable[[], None]]] = []
class _DeferredThread:
def __init__(self, target, args, kwargs, daemon) -> None:
workers.append((kwargs, lambda: target(*args, **kwargs)))
def start(self) -> None:
return None
router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"]))
with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race
"litellm.router.threading", SimpleNamespace(Thread=_DeferredThread)
):
router.completion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert len(workers) == 2
(first_kwargs, run_first), (_, run_second) = workers
first_kwargs["metadata"].pop("foo")
run_second()
run_first()
_wait_for_shadow_successes_sync(recording_logger, expected=2)
metadata_by_group = {
call["litellm_params"]["metadata"]["model_group"]: call["litellm_params"]["metadata"]
for call in recording_logger.shadow_successes()
}
assert metadata_by_group["shadow-b"]["foo"] == "bar"
assert "foo" not in metadata_by_group["shadow-a"]
@pytest.mark.asyncio
async def test_async_shadow_does_not_inherit_primary_deployment_tags(recording_logger):
router = Router(model_list=_tagged_primary_model_list())
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
metadata={"foo": "bar"},
)
assert response.choices[0].message.content == "pong"
await _wait_for_shadow_successes(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
assert "primary-only" not in shadow_successes[0]["litellm_params"]["metadata"].get("tags", [])
@pytest.mark.asyncio
async def test_shadow_of_a_shadow_is_not_launched(recording_logger):
router = Router(model_list=_streaming_model_list(["shadow-a"]))
response = await router.acompletion(
model="primary-model",
messages=[{"role": "user", "content": "hi"}],
mock_response="pong",
)
assert response.choices[0].message.content == "pong"
await _wait_for_shadow_successes(recording_logger, expected=2, timeout=1.0)
model_groups = [call["litellm_params"]["metadata"]["model_group"] for call in recording_logger.shadow_successes()]
assert model_groups == ["shadow-a"]
def test_silent_experiment_completion_direct():
"""
Test _silent_experiment_completion directly (for router code coverage).
@ -127,6 +429,25 @@ async def test_silent_experiment_acompletion_direct():
)
@pytest.mark.asyncio
async def test_run_silent_experiment_drains_stream_so_callbacks_fire(recording_logger):
router = Router(model_list=_streaming_model_list(None))
silent_kwargs: Final = {
"stream": True,
"stream_options": {"include_usage": True},
"mock_response": "pong",
"metadata": {"is_silent_experiment": True, "model_group": "shadow-b"},
}
await router._run_silent_experiment("shadow-b", [{"role": "user", "content": "hi"}], silent_kwargs)
await _wait_for_shadow_successes(recording_logger, expected=1)
shadow_successes = recording_logger.shadow_successes()
assert len(shadow_successes) == 1
assert shadow_successes[0]["stream"] is True
assert shadow_successes[0]["async_complete_streaming_response"] is not None
assert silent_kwargs["stream"] is True
@pytest.mark.asyncio
async def test_router_silent_experiment_acompletion():
"""