feat: add parallel execution handling in during_call_hook (#16279)

litellm/proxy/utils.py

@@ -1053,51 +1053,63 @@ class ProxyLogging:
         ],
     ):
         """
-        Runs the CustomGuardrail's async_moderation_hook()
+        Runs the CustomGuardrail's async_moderation_hook() in parallel
         """
+        # Step 1: Collect all guardrail tasks to run in parallel
+        guardrail_tasks = []
+
         for callback in litellm.callbacks:
-            try:
-                if isinstance(callback, CustomGuardrail):
-                    ################################################################
-                    # Check if guardrail should be run for GuardrailEventHooks.during_call hook
-                    ################################################################
+            if isinstance(callback, CustomGuardrail):
+                ################################################################
+                # Check if guardrail should be run for GuardrailEventHooks.during_call hook
+                ################################################################
 
-                    # V1 implementation - backwards compatibility
-                    if callback.event_hook is None and hasattr(
-                        callback, "moderation_check"
-                    ):
-                        if callback.moderation_check == "pre_call":  # type: ignore
-                            return
-                    else:
-                        # Main - V2 Guardrails implementation
-                        from litellm.types.guardrails import GuardrailEventHooks
+                # V1 implementation - backwards compatibility
+                if callback.event_hook is None and hasattr(
+                    callback, "moderation_check"
+                ):
+                    if callback.moderation_check == "pre_call":  # type: ignore
+                        return
+                else:
+                    # Main - V2 Guardrails implementation
+                    from litellm.types.guardrails import GuardrailEventHooks
 
-                        event_type = GuardrailEventHooks.during_call
-                        if call_type == "mcp_call":
-                            event_type = GuardrailEventHooks.during_mcp_call
-
-                        if (
-                            callback.should_run_guardrail(
-                                data=data, event_type=event_type
-                            )
-                            is not True
-                        ):
-                            continue
-                    # Convert user_api_key_dict to proper format for async_moderation_hook
+                    event_type = GuardrailEventHooks.during_call
                     if call_type == "mcp_call":
-                        user_api_key_auth_dict = (
-                            self._convert_user_api_key_auth_to_dict(user_api_key_dict)
-                        )
-                    else:
-                        user_api_key_auth_dict = user_api_key_dict
+                        event_type = GuardrailEventHooks.during_mcp_call
 
-                    await callback.async_moderation_hook(
+                    if (
+                        callback.should_run_guardrail(
+                            data=data, event_type=event_type
+                        )
+                        is not True
+                    ):
+                        continue
+                # Convert user_api_key_dict to proper format for async_moderation_hook
+                if call_type == "mcp_call":
+                    user_api_key_auth_dict = (
+                        self._convert_user_api_key_auth_to_dict(user_api_key_dict)
+                    )
+                else:
+                    user_api_key_auth_dict = user_api_key_dict
+
+                # Add task to list for parallel execution
+                guardrail_tasks.append(
+                    callback.async_moderation_hook(
                         data=data,
                         user_api_key_dict=user_api_key_auth_dict,  # type: ignore
                         call_type=call_type,
                     )
+                )
+
+        # Step 2: Run all guardrail tasks in parallel
+        if guardrail_tasks:
+            try:
+                await asyncio.gather(*guardrail_tasks)
             except Exception as e:
+                # If any guardrail raises an exception, it will propagate here
                 raise e
+
         return data
 
     async def failed_tracking_alert(

tests/proxy_unit_tests/test_proxy_utils.py

@@ -2154,3 +2154,99 @@ async def test_post_call_failure_hook_auth_error_llm_api_route():
         # Assert that _handle_logging_proxy_only_error WAS called
         # because /v1/chat/completions is an LLM API route
         mock_handle_logging.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_during_call_hook_parallel_execution():
+    """
+    Test that multiple guardrails in during_call_hook are executed in parallel.
+    Verifies parallel execution by checking timing and execution order.
+    """
+    from litellm.caching.caching import DualCache
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.types.guardrails import GuardrailEventHooks
+
+    cache = DualCache()
+    proxy_logging = ProxyLogging(user_api_key_cache=cache)
+    execution_order = []
+    
+    class TestGuardrail(CustomGuardrail):
+        def __init__(self, name):
+            super().__init__(
+                guardrail_name=name,
+                event_hook=GuardrailEventHooks.during_call,
+                default_on=True
+            )
+            self.name = name
+        
+        async def async_moderation_hook(self, data, user_api_key_dict, call_type):
+            execution_order.append(f"{self.name}_start")
+            await asyncio.sleep(0.1)
+            execution_order.append(f"{self.name}_end")
+            return data
+    
+    original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
+    
+    try:
+        litellm.callbacks = [TestGuardrail(f"g{i}") for i in range(3)]
+        
+        start_time = asyncio.get_event_loop().time()
+        result = await proxy_logging.during_call_hook(
+            data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]},
+            user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
+            call_type="completion",
+        )
+        execution_time = asyncio.get_event_loop().time() - start_time
+        
+        # Verify parallel execution: all start before any end
+        first_end_idx = next(i for i, item in enumerate(execution_order) if "end" in item)
+        starts_before_end = sum(1 for item in execution_order[:first_end_idx] if "start" in item)
+        assert starts_before_end == 3, f"Expected 3 starts before first end, got {starts_before_end}"
+        
+        # Verify timing: parallel ~0.1s vs sequential ~0.3s
+        assert execution_time < 0.2, f"Parallel execution took {execution_time}s, expected < 0.2s"
+        assert result["model"] == "gpt-4"
+    finally:
+        litellm.callbacks = original_callbacks
+
+
+@pytest.mark.asyncio
+async def test_during_call_hook_parallel_execution_with_error():
+    """
+    Test that exceptions from guardrails are properly raised in parallel execution.
+    """
+    from litellm.caching.caching import DualCache
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.proxy.utils import ProxyLogging
+    from litellm.types.guardrails import GuardrailEventHooks
+
+    cache = DualCache()
+    proxy_logging = ProxyLogging(user_api_key_cache=cache)
+    
+    class FailingGuardrail(CustomGuardrail):
+        def __init__(self):
+            super().__init__(
+                guardrail_name="failing_guardrail",
+                event_hook=GuardrailEventHooks.during_call,
+                default_on=True
+            )
+        
+        async def async_moderation_hook(self, data, user_api_key_dict, call_type):
+            raise ValueError("Guardrail violation detected!")
+    
+    original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
+    
+    try:
+        litellm.callbacks = [FailingGuardrail()]
+        
+        with pytest.raises(ValueError) as exc_info:
+            await proxy_logging.during_call_hook(
+                data={"model": "gpt-4", "messages": [{"role": "user", "content": "test"}]},
+                user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
+                call_type="completion",
+            )
+        
+        assert "Guardrail violation detected!" in str(exc_info.value)
+    finally:
+        litellm.callbacks = original_callbacks