feat(guardrails): add run_in_parallel opt-in for concurrent pre_call and post_call guardrails - #33770
Conversation
…guardrails Pre-call guardrails run sequentially because each may mutate the request payload and later guardrails depend on earlier mutations. Deployments with several slow block-only pre_call guardrails (external moderation, Bedrock, LLM-judge) therefore pay the sum of their latencies. during_call guardrails run concurrently but alongside the LLM call, so a violating payload has already been sent, which is unacceptable when the request must never reach the model. This adds a per-guardrail run_in_parallel flag (default off). Guardrails that opt in are pulled out of the sequential loop and run concurrently via asyncio.gather after every sequential (payload-mutating) guardrail has run, so they observe the mutated payload and still form a hard barrier before the LLM call; the first to raise blocks the request. Their returned data is discarded since they are declared block-only. The flag is wired from LitellmParams onto the guardrail instance at the same generic choke point in initialize_guardrail that already sets skip_system_message_in_guardrail, so no per-provider initializer needs to change.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Greptile SummaryThis PR adds an opt-in
Confidence Score: 5/5Safe to merge. The feature is strictly opt-in (default off), all existing sequential behaviour is preserved, and the race conditions the design might raise are correctly avoided by asyncio's cooperative scheduling and the unified hook's synchronous pop of guardrail_to_apply before its first I/O yield. The two previously-noted concerns (orphaned tasks and config-absent overwriting of constructor defaults) are both addressed in the code: return_exceptions=True awaits every sibling before processing results, and the registry writes run_in_parallel only when the value is not None. The ContextVar fix is correct — asyncio copies context into each gathered task. Tests cover the parallel/sequential ordering, block-wins-over-reroute priority, all-siblings-awaited guarantee, and the ContextVar regression for both hooks. No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/utils.py | Core parallel guardrail dispatch added to pre_call_hook and post_call_success_hook; blocking/flow-changing exception priority logic is correct; return_exceptions=True avoids orphaned tasks |
| litellm/integrations/custom_guardrail.py | ContextVar-based self-record guard correctly isolates each concurrent guardrail task; run_in_parallel attribute added with safe default=False; token reset in finally block prevents context leak |
| litellm/proxy/guardrails/guardrail_registry.py | run_in_parallel wired from config to instance only when explicitly provided (not None), preserving constructor-level defaults correctly |
| litellm/types/guardrails.py | run_in_parallel Optional[bool] field added to BaseLitellmParams with default=None and clear documentation about mutation-safety restrictions |
| tests/proxy_unit_tests/test_proxy_utils.py | Comprehensive new tests for both hooks covering concurrency, sequential vs parallel ordering, blocking behavior, block-wins-over-reroute priority, and all-siblings-awaited guarantees |
| tests/test_litellm/integrations/test_custom_guardrail.py | New regression test for the ContextVar fix verifying two concurrent guardrails each record their own lifecycle entry |
| tests/test_litellm/proxy/guardrails/test_guardrail_registry.py | Parametrized test covering None/True/False config values; correctly asserts constructor default is preserved when key is absent from config |
| tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py | Existing mock updated with run_in_parallel=False to keep testing the sequential path; not a weakening of coverage |
| tests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.py | Existing mock updated with run_in_parallel=False for the same reason as test_guardrail_pipeline; sequential assertions are preserved |
Reviews (3): Last reviewed commit: "test(guardrails): declare run_in_paralle..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
post_call_success_hook ran guardrails sequentially for the same reason pre_call did: response-modifying guardrails thread the response forward. But block-only output scanners (which read the response and reject on violation without changing it) serialize for no benefit and add latency. This reuses the existing run_in_parallel flag for the post_call hook. Opted-in post_call guardrails are pulled out of the sequential loop and run concurrently via asyncio.gather after the sequential (response-modifying) guardrails and before the non-guardrail CustomLogger callbacks, so they inspect the final response and still block it from reaching the client if any raises. Their returned response is discarded since they are block-only. The apply_guardrail path sets data["guardrail_to_apply"] immediately before awaiting, and unified_guardrail pops it before its first suspension point, so concurrent guardrails never race on that key under asyncio's cooperative scheduling.
noahnistler
changed the title
…over reroutes Addresses review feedback on the run_in_parallel opt-in. asyncio.gather propagated the first exception without cancelling or awaiting the siblings, so a block at t=0 left the other guardrails running as unobserved background tasks (wasted external calls plus event-loop warnings), and a fast SensitiveDataRouteException/ModifyResponseException could return a reroute or passthrough before a slower block finished, letting crafted input bypass the block. Both the pre_call and post_call parallel batches now gather with return_exceptions=True so every guardrail runs to completion, then raise any blocking exception ahead of a flow-changing one. The registry choke point wrote bool(None)==False onto every instance when the config omitted run_in_parallel, silently disabling a constructor-set default; it now only writes when the config provides an explicit value.
|
Thanks for the reviews. Pushed e33b9df addressing all three points. Orphaned tasks (Greptile) and the routing bypass (Veria) share a root cause: bare asyncio.gather propagates the first exception without awaiting the siblings, so a block could leave the others running unobserved, and a fast SensitiveDataRouteException/ModifyResponseException could return a reroute/passthrough before a slower block finished. Both the pre_call and post_call parallel batches now gather with return_exceptions=True, so every guardrail runs to completion (no orphans), then raise any blocking exception ahead of a flow-changing one so a block can never be bypassed by a faster reroute. I chose await-all over cancel-survivors precisely because cancelling on the first raise would reintroduce the bypass if that first raise were the reroute. Config-absent clobber (Greptile): the registry choke point now only writes run_in_parallel when the config provides an explicit value, so a constructor-set default is preserved. New regression tests: a slower block wins over a faster reroute, both pre_call and post_call await every sibling to completion when one blocks, and the registry preserves a True constructor default when the key is absent. Each fails on the pre-fix code. |
…rdrail The log_guardrail_information decorator skipped its auto-record when it saw that the count of standard_logging_guardrail_information entries in the shared request_data had grown during the wrapped call, taking that as proof the wrapped function had recorded its own richer entry. That heuristic breaks the moment guardrails run concurrently (parallel pre_call/post_call, during_call): a sibling guardrail's append inflates the shared count, so a guardrail that did not self-record wrongly concludes it already did and drops its own entry. The result is that enabling run_in_parallel silently loses per-guardrail lifecycle logs, so the Admin UI Request Lifecycle timeline and downstream loggers (Datadog, Langfuse, OTEL, spend logs) show only one of the concurrent guardrails. Replace the shared-count heuristic with a ContextVar flag set when a guardrail records its own entry. asyncio copies the context into each gathered task, so the flag is isolated per concurrent guardrail while still catching the self-record-then-skip-auto-record case within a single invocation.
The post_call partition reads run_in_parallel on every CustomGuardrail callback. A MagicMock(spec=CustomGuardrail) has no run_in_parallel (it is set in __init__, not on the class) so the attribute access raised, and even a class-level default would return a truthy child mock that wrongly routes the double into the parallel batch. Declare the flag False on the shared mock factories so these pre-existing hook tests exercise the sequential path they assert on.
yassin-berriai
left a comment
yassin-berriai
left a comment
There was a problem hiding this comment.
Reviewed the concurrency-critical paths against the base code rather than taking the description at its word, and this holds up well. The design is sound and the tests would catch regressions of each property. Findings below are all minor; nothing merge-blocking
Verified correct
The guardrail_to_apply race is genuinely safe. Each parallel _run_one sets data["guardrail_to_apply"] = callback then awaits, and the window from that write to the pop runs in a single event-loop step with no yield: _run_guardrail_with_metrics does no await before return await coro, and async_post_call_success_hook pops guardrail_to_apply (unified_guardrail.py:220) well before its first suspension at line 262. So concurrent _run_one coroutines under asyncio.gather cannot race on that key
Block-over-reroute plus await-all is right. return_exceptions=True then raising a non-flow-changing exception ahead of a reroute/passthrough preserves the barrier, and results follows gather's input order so raised[0] is deterministic and matches the sequential "first wins" semantics. The block-wins-over-reroute and awaits-all-when-one-blocks tests would both fail under the naive gather the first commit shipped
The ContextVar swap for the logging decorator is the correct mechanism, since asyncio copies the context into each gathered task so siblings cannot inflate each other's count. Good catch that this also fixes the latent during_call drop
A reroute raised by a parallel guardrail is still handled. The raw raise raised[0] propagates to the outer except SensitiveDataRouteException in pre_call_hook, so routing still applies rather than surfacing as an error
Minor findings
-
Redundant in-function import.
_run_parallel_post_call_guardrailsaddsfrom litellm.types.guardrails import GuardrailEventHooksinside the function body, but it is already imported at the top oflitellm/proxy/utils.py. Drop the local import and use the module-level name -
Unguarded attribute read. Both partitions read
callback.run_in_parallelon every guardrail callback, not only opted-in ones. RealCustomGuardrailinstances get it from the new__init__default, which is exactly why theMagicMock(spec=CustomGuardrail)doubles had to be patched in two test files. A third-partyCustomGuardrailsubclass that does not chainsuper().__init__()would now raiseAttributeErroron a path that previously worked.getattr(callback, "run_in_parallel", False)would keep the read total without changing behavior -
Timing upper-bound assertions may flake.
assert elapsed < 0.2withsleep=0.1across three guardrails leaves roughly 2x headroom, which can be tight on a loaded runner. Thestarts_before_first_end == 3overlap assertion is the timing-independent signal that actually proves concurrency, so the wall-clock upper bound is somewhat redundant flake surface -
Reroute precedence inversion (contract-consistent, worth a note). The sequential path defers a reroute and applies it before the parallel batch runs; if a parallel guardrail then also raises a reroute, the outer handler re-runs
_handle_sensitive_data_route_exceptionand overrides the sequential one'sdata["model"], so a parallel reroute beats a sequential one. This only matters ifrun_in_parallelis enabled on a rerouting guardrail, which the PR documents as unsupported (block-only), so it is not a bug. Since the path technically accepts a parallel reroute, a one-line docstring note, or an explicit guard rejecting a flow-changing exception from the parallel batch, would remove the surprise. Related: the concurrent_process_guardrail_callbackcalls all mutate the samedatadict in place, so the safety rests entirely on callers honoring the block-only contract with no runtime enforcement
CI
The two red checks look unrelated to this diff. proxy-infra fails on four test_streaming_* tests in tests/test_litellm/proxy/test_budget_reservation.py (object MagicMock can't be used in 'await' expression); that file and the streaming/reservation code it exercises are not in this PR, and both tests exist unchanged on the base branch, so this reads as base-branch drift on the merge ref. A rebase onto current staging will likely clear it. osv-scan is red repo-wide from default-branch advisories, not this diff
Net: approve-worthy after the trivial import cleanup, ideally with the rebase to shake off the unrelated proxy-infra red
…back Read run_in_parallel via getattr(..., False) in the pre_call and post_call partitions so a third-party CustomGuardrail subclass that overrides __init__ without chaining super().__init__() no longer raises AttributeError on a path that previously worked. Drop the redundant in-function GuardrailEventHooks import in _run_parallel_post_call_guardrails (already imported module-level). Remove the flaky wall-clock upper-bound assertions from the two concurrency tests; the all-start-before-any-end overlap assertion is the timing-independent signal that actually proves concurrency.
|
Thanks for the close read. Pushed db243b6 addressing the actionable findings. Redundant import (1): dropped the in-function Unguarded attribute read (2): good catch. Timing flake (3): removed the Reroute precedence (4): left as-is deliberately. Enabling On CI: agreed the two reds are unrelated. Re-validated end to end against a live proxy with a Postgres backend after these changes: a single request exercising pre and post, sequential and parallel, returns in ~18s (6s pre sequential + 3s pre parallel + 6s post sequential + 3s post parallel) versus ~24s if the parallel batches ran serially, and all eight guardrails persist their own lifecycle entry to spend logs with the two parallel guardrails in each hook sharing a start offset |
|
All four points are addressed cleanly:
One small remaining observation on Overall the implementation is correct: |
…ng' into litellm_parallel_pre_call_guardrails # Conflicts: # litellm/integrations/custom_guardrail.py
8177230
into
BerriAI:litellm_internal_staging
…and post_call guardrails (BerriAI#33770) * feat(guardrails): add run_in_parallel opt-in for concurrent pre_call guardrails Pre-call guardrails run sequentially because each may mutate the request payload and later guardrails depend on earlier mutations. Deployments with several slow block-only pre_call guardrails (external moderation, Bedrock, LLM-judge) therefore pay the sum of their latencies. during_call guardrails run concurrently but alongside the LLM call, so a violating payload has already been sent, which is unacceptable when the request must never reach the model. This adds a per-guardrail run_in_parallel flag (default off). Guardrails that opt in are pulled out of the sequential loop and run concurrently via asyncio.gather after every sequential (payload-mutating) guardrail has run, so they observe the mutated payload and still form a hard barrier before the LLM call; the first to raise blocks the request. Their returned data is discarded since they are declared block-only. The flag is wired from LitellmParams onto the guardrail instance at the same generic choke point in initialize_guardrail that already sets skip_system_message_in_guardrail, so no per-provider initializer needs to change. * feat(guardrails): extend run_in_parallel opt-in to post_call guardrails post_call_success_hook ran guardrails sequentially for the same reason pre_call did: response-modifying guardrails thread the response forward. But block-only output scanners (which read the response and reject on violation without changing it) serialize for no benefit and add latency. This reuses the existing run_in_parallel flag for the post_call hook. Opted-in post_call guardrails are pulled out of the sequential loop and run concurrently via asyncio.gather after the sequential (response-modifying) guardrails and before the non-guardrail CustomLogger callbacks, so they inspect the final response and still block it from reaching the client if any raises. Their returned response is discarded since they are block-only. The apply_guardrail path sets data["guardrail_to_apply"] immediately before awaiting, and unified_guardrail pops it before its first suspension point, so concurrent guardrails never race on that key under asyncio's cooperative scheduling. * fix(guardrails): await all parallel guardrails and prioritize blocks over reroutes Addresses review feedback on the run_in_parallel opt-in. asyncio.gather propagated the first exception without cancelling or awaiting the siblings, so a block at t=0 left the other guardrails running as unobserved background tasks (wasted external calls plus event-loop warnings), and a fast SensitiveDataRouteException/ModifyResponseException could return a reroute or passthrough before a slower block finished, letting crafted input bypass the block. Both the pre_call and post_call parallel batches now gather with return_exceptions=True so every guardrail runs to completion, then raise any blocking exception ahead of a flow-changing one. The registry choke point wrote bool(None)==False onto every instance when the config omitted run_in_parallel, silently disabling a constructor-set default; it now only writes when the config provides an explicit value. * fix(guardrails): record lifecycle logs for every concurrently-run guardrail The log_guardrail_information decorator skipped its auto-record when it saw that the count of standard_logging_guardrail_information entries in the shared request_data had grown during the wrapped call, taking that as proof the wrapped function had recorded its own richer entry. That heuristic breaks the moment guardrails run concurrently (parallel pre_call/post_call, during_call): a sibling guardrail's append inflates the shared count, so a guardrail that did not self-record wrongly concludes it already did and drops its own entry. The result is that enabling run_in_parallel silently loses per-guardrail lifecycle logs, so the Admin UI Request Lifecycle timeline and downstream loggers (Datadog, Langfuse, OTEL, spend logs) show only one of the concurrent guardrails. Replace the shared-count heuristic with a ContextVar flag set when a guardrail records its own entry. asyncio copies the context into each gathered task, so the flag is isolated per concurrent guardrail while still catching the self-record-then-skip-auto-record case within a single invocation. * test(guardrails): declare run_in_parallel on post_call guardrail mocks The post_call partition reads run_in_parallel on every CustomGuardrail callback. A MagicMock(spec=CustomGuardrail) has no run_in_parallel (it is set in __init__, not on the class) so the attribute access raised, and even a class-level default would return a truthy child mock that wrongly routes the double into the parallel batch. Declare the flag False on the shared mock factories so these pre-existing hook tests exercise the sequential path they assert on. * fix(guardrails): harden run_in_parallel reads and address review feedback Read run_in_parallel via getattr(..., False) in the pre_call and post_call partitions so a third-party CustomGuardrail subclass that overrides __init__ without chaining super().__init__() no longer raises AttributeError on a path that previously worked. Drop the redundant in-function GuardrailEventHooks import in _run_parallel_post_call_guardrails (already imported module-level). Remove the flaky wall-clock upper-bound assertions from the two concurrency tests; the all-start-before-any-end overlap assertion is the timing-independent signal that actually proves concurrency.
2 participants
Relevant issues
Linear ticket
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Type
🆕 New Feature
Changes
Guardrails that must block around the LLM call run serially and add latency. In
ProxyLogging.pre_call_hookandpost_call_success_hookthe guardrails run in a for-loop with await because a guardrail can modify the request or response and later guardrails depend on that (PII masking, prompt rewriting, sensitive-data routing, output redaction). A deployment with several slow block-only guardrails (external moderation, Bedrock, an LLM-judge) therefore pays the sum of their latencies. Theduring_callstage already runs guardrails concurrently, but it runs alongside the LLM call, so a violating payload has already been sent; that is unusable when the request must never reach the model, or when a violating response must never reach the clientThis adds an opt-in per-guardrail
run_in_parallelflag (default off, so existing setups are unchanged) covering both the pre_call and post_call hooks. Guardrails that opt in are pulled out of the sequential loop and run concurrently viaasyncio.gatherafter every sequential (mutating) guardrail has run, so they observe the final payload or response and still form a hard barrier. Their returned value is discarded since they are declared block-only, and per-guardrail latency metrics still come from the existing guardrail-execution pathThe parallel batch gathers with
return_exceptions=Trueso every opted-in guardrail is awaited to completion; a raise by one never leaves the others running as unobserved background tasks. Once the batch settles, a guardrail that blocks (any exception other than a reroute or passthrough) is raised ahead of one that only changes the request or response flow, so a fastSensitiveDataRouteExceptionorModifyResponseExceptioncan never let crafted input slip past a slower block. Awaiting all rather than cancelling on the first raise is deliberate: cancelling would reintroduce that bypass whenever the first raise happened to be the rerouteEnabling
run_in_parallelon a guardrail that mutates the request or response is unsafe (concurrent guardrails share one snapshot and their writes would race), which the config field description and the docstrings call out explicitly. For theapply_guardrailpath, each per-guardrail coroutine setsdata["guardrail_to_apply"]immediately before awaiting andunified_guardrailpops it before its first suspension point, so under asyncio's cooperative scheduling concurrent guardrails never race on that keyThe flag is read from
LitellmParamsand set on the guardrail instance at the same generic choke point ininitialize_guardrailthat already setsskip_system_message_in_guardrail, so no per-provider initializer needs to change. It is written only when the config provides an explicit value, so arun_in_parallel=Truedefault set in a guardrail's own constructor is preserved rather than silently resetModes deliberately left as-is:
during_callis already parallel; the streaming iterator hook is a sequential wrapping chain that cannot be parallelized; the per-chunk streaming hook and the failure hook have per-chunk mutation and first-exception-wins ordering that make parallelizing low-value or semantically riskyRunning guardrails concurrently also surfaced a pre-existing observability bug that this PR fixes. The
log_guardrail_informationdecorator decided whether to auto-record a guardrail's lifecycle entry by countingstandard_logging_guardrail_informationentries in the sharedrequest_databefore and after the wrapped call; a growth meant "the wrapped function recorded its own richer entry, so skip". Under concurrency a sibling guardrail's append inflates that shared count, so a guardrail that did not self-record wrongly concludes it already did and drops its own entry, leaving the Admin UI Request Lifecycle timeline and downstream loggers (Datadog, Langfuse, OTEL, spend logs) showing only one of the concurrent guardrails. The count heuristic is replaced with aContextVarflag set when a guardrail records its own entry; asyncio copies the context into each gathered task, so the flag is isolated per concurrent guardrail while still catching the self-record case within a single invocation. This also fixes the same latent drop forduring_call, which has always run guardrails concurrentlyFiles:
litellm/types/guardrails.py:run_in_parallelfield onBaseLitellmParamslitellm/proxy/guardrails/guardrail_registry.py: wire the config flag onto the instance, only when explicitly configuredlitellm/proxy/utils.py: partition the pre_call and post_call loops, add_run_parallel_pre_call_guardrailsand_run_parallel_post_call_guardrailswith await-all plus block-over-reroute prioritization. The partitions readrun_in_parallelviagetattr(callback, "run_in_parallel", False)so a third-partyCustomGuardrailsubclass that overrides__init__without chainingsuper().__init__()cannotAttributeErroron a path that previously workedlitellm/integrations/custom_guardrail.py: typedrun_in_parallelattribute onCustomGuardrail, and a concurrency-safeContextVarinlog_guardrail_informationso every concurrently-run guardrail records its own lifecycle entrytests/proxy_unit_tests/test_proxy_utils.py,tests/test_litellm/proxy/guardrails/test_init_guardrails.py,tests/test_litellm/proxy/guardrails/test_guardrail_registry.py, andtests/test_litellm/integrations/test_custom_guardrail.pytests/test_litellm/proxy/utils/proxy_logging/test_post_call_success_hook.pyandtests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py: existing hook tests buildMagicMock(spec=CustomGuardrail)doubles; the post_call partition now readsrun_in_parallelon every guardrail, so these mocks declare itFalseto keep exercising the sequential path they assert onTests cover, for both hooks, parallel guardrails overlapping and finishing in about max rather than sum, default guardrails staying strictly sequential, and a raising parallel guardrail blocking the request or response; pre_call additionally covers sequential mutations being visible to the parallel batch (ordering) and a should-not-run guardrail being skipped. Regression tests for the review fixes: a slower block wins over a faster reroute, both hooks await every sibling to completion when one blocks (no orphaned tasks), the config to instance wiring for True/False/unset, the registry preserving a True constructor default when the key is absent, and two guardrails run concurrently through the logging decorator each recording their own lifecycle entry (would drop one under the old shared-count guard)
Final Attestation