From 8ec5c7ddf22f0490a7776de2496df5cc11da701c Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 14:12:32 -0700 Subject: [PATCH] perf(proxy): record applied guardrails in time linear in the names a request carries _process_guardrail_metadata checked each requested name against the growing applied_guardrails list, so a body naming 50k guardrails spent seconds of event loop time. Names are now filtered against sets and appended in one pass, with the same order, dedup and metadata bucket as before. --- litellm/proxy/utils.py | 52 ++++++++++---- .../proxy_logging/test_guardrail_pipeline.py | 70 ++++++++----------- 2 files changed, 67 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c7efbf9dacf..6e091391842 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -815,6 +815,29 @@ def _drop_response_steps_that_did_not_run( _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) +def _add_applied_guardrails_in_place( + data: dict[str, object], # mutable-ok: the names are recorded in the live request body's metadata + names: Sequence[str], +) -> None: + """Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names""" + from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header + + if not names: + return + metadata: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str-keyed bucket + if "applied_guardrails" not in metadata: + metadata["applied_guardrails"] = list(dict.fromkeys(names)) + return + applied: Final = metadata["applied_guardrails"] + if not isinstance(applied, list): + for name in names: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=name) + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + present: Final = frozenset(name for name in entries if isinstance(name, str)) + entries.extend(dict.fromkeys(name for name in names if name not in present)) + + def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: if not isinstance(bucket, dict): return @@ -2831,10 +2854,6 @@ class ProxyLogging: policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on another request, so every step is added. """ - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) - metadata_standard: Final = data.get("metadata") or {} metadata_litellm: Final = data.get("litellm_metadata") or {} @@ -2854,18 +2873,21 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] - not_run: Final = ( - frozenset[str]() - if skip_guardrails - else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload + request: Final = cast("dict[str, object]", data) # cast-ok: request payload + not_run: Final = frozenset[str]() if skip_guardrails else _request_steps_that_have_not_run(request) + already_applied: Final = frozenset( + name + for name in cast("Sequence[object]", applied_guardrails) # cast-ok: checked to be a list above + if isinstance(name, str) + ) + _add_applied_guardrails_in_place( + request, + [ + name + for name in guardrails_in_metadata + if isinstance(name, str) and name not in already_applied and name not in not_run + ], ) - for guardrail_name in guardrails_in_metadata: - if ( - isinstance(guardrail_name, str) - and guardrail_name not in applied_guardrails - and guardrail_name not in not_run - ): - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( self, diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5fd40696224..a46b56294bf 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import json +import time from collections.abc import AsyncGenerator, Iterator from contextlib import AbstractContextManager, nullcontext from copy import deepcopy @@ -275,54 +276,43 @@ async def test_process_guardrail_callback_enriches_and_reraises_http_exception( # --------------------------------------------------------------------------- -def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypatch): - calls: List[Dict[str, Any]] = [] - - def fake_add(request_data, guardrail_name): - calls.append({"data": request_data, "name": guardrail_name}) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) - data = {"metadata": {"guardrails": ["g1", "g2"]}} +def test_process_guardrail_metadata_adds_each_guardrail_once_in_order(proxy_logging): + data = {"metadata": {"guardrails": ["g1", "g2", "g1"]}} proxy_logging._process_guardrail_metadata(data) - snapshot = { - "call_count": len(calls), - "first_name": calls[0]["name"], - "second_name": calls[1]["name"], - "data_passed_is_input": all(c["data"] is data for c in calls), - } - assert snapshot == { - "call_count": 2, - "first_name": "g1", - "second_name": "g2", - "data_passed_is_input": True, - } + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeypatch): - calls: List[str] = [] - - def fake_add(request_data, guardrail_name): - calls.append(guardrail_name) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) +def test_process_guardrail_metadata_skips_already_applied(proxy_logging): data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} proxy_logging._process_guardrail_metadata(data) - assert calls == ["g2"] + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypatch): - from litellm.proxy.common_utils import callback_utils +def test_process_guardrail_metadata_records_in_the_litellm_metadata_bucket(proxy_logging): + data = {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2"]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2", "g1"]}} - monkeypatch.setattr( - callback_utils, - "add_guardrail_to_applied_guardrails_header", - MagicMock(side_effect=AssertionError("should not be called")), - ) - proxy_logging._process_guardrail_metadata({}) + +def test_process_guardrail_metadata_handles_many_guardrail_names_quickly(proxy_logging): + names = [f"guard-{i}" for i in range(50_000)] + data = {"metadata": {"guardrails": names, "applied_guardrails": []}} + started = time.perf_counter() + proxy_logging._process_guardrail_metadata(data) + assert time.perf_counter() - started < 2 + assert data["metadata"]["applied_guardrails"] == names + + +def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging): + data: Dict[str, Any] = {} + proxy_logging._process_guardrail_metadata(data) + assert data == {} + + +def test_process_guardrail_metadata_with_only_settings_entries_leaves_data_unchanged(proxy_logging): + data = {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} def test_process_guardrail_metadata_invalid_data_raises(proxy_logging):