From e66f38e447ef0291a684ade244ba9014762657b3 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 16:40:54 +0000 Subject: [PATCH] fix(guardrails): select one load-balanced callback per hook Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/guardrails/init_guardrails.py | 9 +- litellm/proxy/utils.py | 84 +++++++++++++------ .../test_guardrail_load_balancing.py | 6 +- .../test_guardrail_load_balancing.py | 61 ++++++++++---- 4 files changed, 112 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index ecdf99e49d0..a5e522f3779 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -47,10 +47,10 @@ def init_guardrails_v2( # verbose_proxy_logger.debug(f"\nGuardrail List:{guardrail_list}\n") # Populate router's guardrail_list for load balancing support - _populate_router_guardrail_list(guardrail_list=guardrail_list) + _populate_router_guardrail_list(guardrail_list=guardrail_list, router=llm_router) -def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None: +def _populate_router_guardrail_list(guardrail_list: list[Guardrail], router: Router | None = None) -> None: """ Populate the router's guardrail_list from initialized guardrails. @@ -61,7 +61,8 @@ def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None: from litellm.proxy.proxy_server import llm_router from litellm.types.router import GuardrailTypedDict - if llm_router is None: + target_router: Final = router if router is not None else llm_router + if target_router is None: verbose_proxy_logger.debug("Router not initialized yet, skipping guardrail_list population") return @@ -94,7 +95,7 @@ def _populate_router_guardrail_list(guardrail_list: list[Guardrail]) -> None: router_guardrail_list.append(router_guardrail) - llm_router.apply(ChangeSet(assign={"guardrail_list": router_guardrail_list})) + target_router.apply(ChangeSet(assign={"guardrail_list": router_guardrail_list})) verbose_proxy_logger.debug("Populated router guardrail_list with %s guardrails", len(router_guardrail_list)) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a2928bba287..9117f7417ac 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -30,7 +30,7 @@ from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from email.utils import formatdate from functools import partial, reduce -from itertools import takewhile +from itertools import chain, takewhile from operator import attrgetter from types import MappingProxyType from typing import ( @@ -889,6 +889,43 @@ def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> metadata.pop("applied_guardrails") +@overload +def _load_balanced_callbacks(callbacks: Sequence[CustomLogger]) -> tuple[CustomLogger, ...]: ... + + +@overload +def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: ... + + +def _load_balanced_callbacks(callbacks: Sequence[object]) -> tuple[object, ...]: + from litellm.proxy.proxy_server import llm_router + + if llm_router is None or not llm_router.settings().is_set("guardrail_list"): + return tuple(callbacks) + settings: Final = llm_router.settings() + balanced_names: Final = frozenset( + entry["guardrail_name"] + for entry in settings.guardrail_list + if sum(other["guardrail_name"] == entry["guardrail_name"] for other in settings.guardrail_list) > 1 + ) + + def selected(index: int, callback: object) -> tuple[object, ...]: + if not isinstance(callback, CustomGuardrail) or callback.guardrail_name not in balanced_names: + return (callback,) + if any( + isinstance(previous, CustomGuardrail) and previous.guardrail_name == callback.guardrail_name + for previous in callbacks[:index] + ): + return () + deployment: Final = settings.guardrail(callback.guardrail_name) + target: Final = deployment.get("callback") + if target is None: + raise ValueError(f"No callback found for guardrail: {callback.guardrail_name}") + return (target,) + + return tuple(chain.from_iterable(selected(index, callback) for index, callback in enumerate(callbacks))) + + def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: resolved: Final = tuple( litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -900,7 +937,7 @@ def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple else callback for callback in litellm.callbacks ) - present: Final = tuple(callback for callback in resolved if callback is not None) + present: Final = _load_balanced_callbacks(tuple(callback for callback in resolved if callback is not None)) guardrails: Final = tuple(callback for callback in present if isinstance(callback, CustomGuardrail)) others: Final = cast( # cast-ok: mirrors the legacy loop, which treated every non-guardrail entry as a CustomLogger "tuple[CustomLogger, ...]", @@ -2705,24 +2742,13 @@ class ProxyLogging: error_type = None try: - # Check if load balancing should be used - if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name): - response = await self._execute_guardrail_with_load_balancing( - guardrail_name=guardrail_name, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) - else: - # Single guardrail - execute directly - response = await self._execute_guardrail_hook( - callback=callback, - hook_type="pre_call", - data=data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - ) + response = await self._execute_guardrail_hook( + callback=callback, + hook_type="pre_call", + data=data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + ) # Process the response if one was returned if response is not None: @@ -3265,12 +3291,15 @@ class ProxyLogging: self._process_guardrail_metadata(data, skip_guardrails=skip_guardrails) return data + selected_callbacks: Final = ( + caps.resolved_callbacks if skip_guardrails else _load_balanced_callbacks(caps.resolved_callbacks) + ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( () if skip_guardrails else tuple( cb - for cb in caps.resolved_callbacks + for cb in selected_callbacks if isinstance(cb, CustomGuardrail) and getattr(cb, "run_in_parallel", False) and not (cb.guardrail_name and cb.guardrail_name in pipeline_managed) @@ -3278,7 +3307,7 @@ class ProxyLogging: ) deferred_route_exc: SensitiveDataRouteException | None = None - for _callback in caps.resolved_callbacks: + for _callback in selected_callbacks: start_time = time.time() try: if isinstance(_callback, CustomGuardrail) and data is not None: @@ -3750,7 +3779,7 @@ class ProxyLogging: # Step 1: Collect all guardrail tasks to run in parallel guardrail_tasks: Final = [] - for callback in litellm.callbacks: + for callback in _load_balanced_callbacks(tuple(litellm.callbacks)): if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) @@ -4724,7 +4753,7 @@ class ProxyLogging: stream_gated_guardrail_names(data, user_api_key_dict) if caps.has_guardrail else frozenset() ) - for callback in litellm.callbacks: + for callback in _load_balanced_callbacks(tuple(litellm.callbacks)): try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): @@ -4814,9 +4843,14 @@ class ProxyLogging: stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) + selected_iterator_callbacks: Final = _load_balanced_callbacks(tuple(cb for cb, _ in caps.iterator_overrides)) + selected_iterator_overrides: Final = tuple( + (cb, "override" if "async_post_call_streaming_iterator_hook" in type(cb).__dict__ else "apply_guardrail") + for cb in selected_iterator_callbacks + ) running_overrides: Final = tuple( (callback, kind) - for callback, kind in caps.iterator_overrides + for callback, kind in selected_iterator_overrides if not isinstance(callback, CustomGuardrail) or ( callback.guardrail_name not in pipeline_gated_names diff --git a/tests/integration/observability/test_guardrail_load_balancing.py b/tests/integration/observability/test_guardrail_load_balancing.py index 585cc7c3cc9..796e11f135c 100644 --- a/tests/integration/observability/test_guardrail_load_balancing.py +++ b/tests/integration/observability/test_guardrail_load_balancing.py @@ -42,10 +42,6 @@ def _deployment(name: str, api_base: str) -> dict[str, object]: def test_requests_naming_a_shared_guardrail_are_spread_over_its_deployments(gateway: Gateway, tmp_path: Path) -> None: - pytest.skip( - "BUG: init_guardrails_v2 populates llm_router.guardrail_list from the proxy_server global, which is still " - "None during load_config, so every deployment sharing a guardrail_name checks every request" - ) name: Final = f"lb-guardrail-{uuid4().hex}" prompts: Final = tuple(f"Hello request {index}" for index in range(_REQUESTS)) with ( @@ -75,5 +71,7 @@ def test_requests_naming_a_shared_guardrail_are_spread_over_its_deployments(gate (200, name) ] * _REQUESTS seen: Final = tuple(frozenset(json.loads(call.body)["texts"][0] for call in calls) for calls in checked) + assert sum(len(calls) for calls in checked) == _REQUESTS + assert seen[0].isdisjoint(seen[1]) assert all(0 < len(prompts_seen) < _REQUESTS for prompts_seen in seen), [len(prompts_seen) for prompts_seen in seen] assert seen[0] | seen[1] == frozenset(prompts) diff --git a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py index fd38e9e35db..88f7ef3b1bb 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py +++ b/tests/unit/proxy/guardrails/test_guardrail_load_balancing.py @@ -3,6 +3,7 @@ Test guardrail load balancing through the Router and ProxyLogging. """ from unittest.mock import MagicMock, patch, AsyncMock +from typing import Final import litellm @@ -31,9 +32,26 @@ class MockGuardrail(CustomGuardrail): MockGuardrail.call_count += 1 return None + async def async_moderation_hook(self, data, user_api_key_dict, call_type): + self.calls += 1 + + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + self.calls += 1 + return response + + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + self.calls += 1 + + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + self.calls += 1 + async for chunk in response: + yield chunk + @pytest.mark.asyncio -async def test_proxy_logging_pre_call_hook_load_balancing(): +@pytest.mark.parametrize("hook", ("pre_call", "during_call", "post_call", "streaming", "iterator")) +@pytest.mark.parametrize("deployment_count", (0, 1, 2)) +async def test_proxy_logging_hooks_load_balance_once_per_request(hook: str, deployment_count: int): """Test that async_pre_call_hook load balances across multiple guardrails.""" MockGuardrail.call_count = 0 @@ -62,33 +80,46 @@ async def test_proxy_logging_pre_call_hook_load_balancing(): "litellm_params": {"model": "gpt-5.5", "api_key": "fake-key"}, } ], - guardrail_list=guardrail_list, + guardrail_list=guardrail_list[:deployment_count], ) proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) original_callbacks = litellm.callbacks.copy() - litellm.callbacks = [guardrail_1] + litellm.callbacks = [guardrail_1, guardrail_2] try: - with patch("litellm.proxy.proxy_server.llm_router", router): + with patch("litellm.proxy.proxy_server.llm_router", router if deployment_count else None): for _ in range(50): - await proxy_logging.pre_call_hook( - user_api_key_dict=MagicMock(), - data={"messages": [{"role": "user", "content": "test"}]}, - call_type="completion", - ) + previous_calls: Final = guardrail_1.calls + guardrail_2.calls + data: Final = {"messages": [{"role": "user", "content": "test"}]} + auth: Final = MagicMock() + response: Final = litellm.ModelResponse(choices=[{"message": {"content": "test", "role": "assistant"}}]) + if hook == "pre_call": + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="completion") + elif hook == "during_call": + await proxy_logging.during_call_hook(user_api_key_dict=auth, data=data, call_type="completion") + elif hook == "post_call": + await proxy_logging.post_call_success_hook(user_api_key_dict=auth, data=data, response=response) + elif hook == "streaming": + await proxy_logging.async_post_call_streaming_hook(user_api_key_dict=auth, data=data, response=response) + else: + async def chunks(chunk=response): + yield chunk + + async for _chunk in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, request_data=data, response=chunks() + ): + pass + expected_calls: Final = 1 if deployment_count == 2 else 2 + assert guardrail_1.calls + guardrail_2.calls == previous_calls + expected_calls assert guardrail_1.calls > 0, "Guardrail 1 should have been called" assert guardrail_2.calls > 0, "Guardrail 2 should have been called" total = guardrail_1.calls + guardrail_2.calls - assert total == 50, f"Expected 50 total calls, got {total}" - - min_calls = min(guardrail_1.calls, guardrail_2.calls) - assert ( - min_calls >= 10 - ), f"Expected at least 10 calls to each guardrail, got min={min_calls}" + expected_total: Final = 50 if deployment_count == 2 else 100 + assert total == expected_total, f"Expected {expected_total} total calls, got {total}" finally: litellm.callbacks = original_callbacks