From 63e855a03336a7e7c6e4032ff99a60105c7e3c31 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Wed, 28 Jan 2026 12:06:33 -0800 Subject: [PATCH] fix(presidio): fix 97s latency spike under load by implementing session pooling and PII caching - Implement loop-bound session pooling to reuse TCP connections in background threads (logging hooks) - Add local memory caching for PII results to skip redundant network calls for identical text - Add presidio_ad_hoc_recognizers_on_server flag to prevent redundant server-side registry reloads - Add comprehensive latency and optimization tests in tests/test_presidio_latency.py --- .../guardrails/guardrail_hooks/presidio.py | 51 +++++-- tests/test_presidio_latency.py | 139 ++++++++++++++++++ 2 files changed, 178 insertions(+), 12 deletions(-) create mode 100644 tests/test_presidio_latency.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index fbe6d753679..2a67bb93d57 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -84,6 +84,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): presidio_score_thresholds: Optional[ Dict[Union[PiiEntityType, str], float] ] = None, + presidio_ad_hoc_recognizers_on_server: Optional[bool] = None, **kwargs, ): if logging_only is True: @@ -104,6 +105,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): presidio_score_thresholds or {} ) self.presidio_language = presidio_language or "en" + self.presidio_ad_hoc_recognizers_on_server = ( + presidio_ad_hoc_recognizers_on_server or False + ) # Shared HTTP session to prevent memory leaks (issue #14540) self._http_session: Optional[aiohttp.ClientSession] = None # Lock to prevent race conditions when creating session under concurrent load @@ -114,6 +118,12 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self._main_thread_id = threading.get_ident() + # Loop-bound session cache for background threads + self._loop_sessions: Dict[asyncio.AbstractEventLoop, aiohttp.ClientSession] = {} + + # Result cache to avoid redundant network calls + self.pii_cache = DualCache() + if mock_testing is True: # for testing purposes only return @@ -189,9 +199,9 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): Logic: 1. If running in the main thread (where the object was initialized/destined to live normally), use the shared `self._http_session` (protected by a lock). - 2. If running in a background thread (e.g. logging hook), yield a NEW ephemeral session - and ensure it is closed after use. + 2. If running in a background thread (e.g. logging hook), use a cached session for that loop. """ + current_loop = asyncio.get_running_loop() # Check if we are in the stored main thread if threading.get_ident() == self._main_thread_id: @@ -201,22 +211,27 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): self._http_session = aiohttp.ClientSession() yield self._http_session else: - # Background thread -> create ephemeral session + # Background thread/loop -> use loop-bound session cache # This avoids "attached to a different loop" or "no running event loop" errors # when accessing the shared session created in the main loop - session = aiohttp.ClientSession() - try: - yield session - finally: - if not session.closed: - await session.close() + if ( + current_loop not in self._loop_sessions + or self._loop_sessions[current_loop].closed + ): + self._loop_sessions[current_loop] = aiohttp.ClientSession() + yield self._loop_sessions[current_loop] async def _close_http_session(self) -> None: - """Close the shared HTTP session if it exists.""" + """Close all cached HTTP sessions.""" if self._http_session is not None and not self._http_session.closed: await self._http_session.close() self._http_session = None + for session in self._loop_sessions.values(): + if not session.closed: + await session.close() + self._loop_sessions.clear() + def __del__(self): """Cleanup: we try to close, but doing async cleanup in __del__ is risky.""" pass @@ -239,7 +254,10 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ################################################################## ###### Check if user has configured any params for this guardrail ################################################################ - if self.ad_hoc_recognizers is not None: + if ( + self.ad_hoc_recognizers is not None + and self.presidio_ad_hoc_recognizers_on_server is False + ): analyze_payload["ad_hoc_recognizers"] = self.ad_hoc_recognizers if self.pii_entities_config: @@ -470,6 +488,13 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): """ Calls Presidio Analyze + Anonymize endpoints for PII Analysis + Masking """ + # Cache check + cache_key = f"presidio:pii:{text}:{output_parse_pii}:{presidio_config}:{self.pii_entities_config}" + cached_result = self.pii_cache.get_cache(cache_key) + if cached_result is not None: + verbose_proxy_logger.debug("PII Cache Hit for text: %s", text) + return cached_result + start_time = datetime.now() analyze_results: Optional[Union[List[PresidioAnalyzeResponseItem], Dict]] = None status: GuardrailStatus = "success" @@ -501,12 +526,14 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) # Then anonymize the text using the analysis results - return await self.anonymize_text( + anonymized_text = await self.anonymize_text( text=text, analyze_results=analyze_results, output_parse_pii=output_parse_pii, masked_entity_count=masked_entity_count, ) + self.pii_cache.set_cache(cache_key, anonymized_text) + return anonymized_text return redacted_text["text"] except Exception as e: status = "guardrail_failed_to_respond" diff --git a/tests/test_presidio_latency.py b/tests/test_presidio_latency.py new file mode 100644 index 00000000000..79c567e90ed --- /dev/null +++ b/tests/test_presidio_latency.py @@ -0,0 +1,139 @@ + +import asyncio +import aiohttp +import pytest +from unittest.mock import MagicMock, patch +from litellm.proxy.guardrails.guardrail_hooks.presidio import _OPTIONAL_PresidioPIIMasking + +@pytest.mark.asyncio +async def test_sanity_presidio_session_reuse_main_thread(): + """ + SANITY CHECK: + Verify that Presidio guardrail reuses sessions in the main thread. + This ensures we don't break existing session pooling functionality. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_analyzer_api_base="http://mock-analyzer", + presidio_anonymizer_api_base="http://mock-anonymizer" + ) + + session_creations = 0 + original_init = aiohttp.ClientSession.__init__ + + def mocked_init(self, *args, **kwargs): + nonlocal session_creations + session_creations += 1 + original_init(self, *args, **kwargs) + + with patch.object(aiohttp.ClientSession, "__init__", side_effect=mocked_init, autospec=True): + for _ in range(10): + async with presidio._get_session_iterator() as session: + pass + + # Expected: Only 1 session created for all 10 calls. + assert session_creations == 1 + + await presidio._close_http_session() + +@pytest.mark.asyncio +async def test_bug_presidio_session_explosion_background_thread_causes_latency(): + """ + BUG REPRODUCTION: + Verify that background threads (like logging hooks) REUSE sessions. + Previously, each call in a background loop created a NEW ephemeral session, + leading to socket exhaustion and the reported 97s latency spike. + """ + import threading + presidio = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_analyzer_api_base="http://mock-analyzer", + presidio_anonymizer_api_base="http://mock-anonymizer" + ) + + # Force the code to think it's in a background thread + presidio._main_thread_id = threading.get_ident() + 1 + + session_creations = 0 + original_init = aiohttp.ClientSession.__init__ + + def mocked_init(self, *args, **kwargs): + nonlocal session_creations + session_creations += 1 + original_init(self, *args, **kwargs) + + with patch.object(aiohttp.ClientSession, "__init__", side_effect=mocked_init, autospec=True): + for _ in range(10): + async with presidio._get_session_iterator() as session: + pass + + # FIX VERIFICATION: Should now be 1 session (reused) instead of 10. + assert session_creations == 1 + + await presidio._close_http_session() + +@pytest.mark.asyncio +async def test_optimization_presidio_avoid_recognizer_reloads_on_server(): + """ + OPTIMIZATION VERIFICATION: + Verify that ad_hoc_recognizers are OMITTED from the payload when + presidio_ad_hoc_recognizers_on_server is True. + Sending them on every request forces the Presidio server to reload its + entire registry, spiking CPU and latency. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_analyzer_api_base="http://mock-analyzer", + presidio_anonymizer_api_base="http://mock-anonymizer", + presidio_ad_hoc_recognizers_on_server=True + ) + presidio.ad_hoc_recognizers = [{"name": "CustomRecognizer", "supported_entity": "CUSTOM"}] + + payload = presidio._get_presidio_analyze_request_payload( + text="some text", + presidio_config=None, + request_data={} + ) + + # Optimization Check: payload should NOT contain the redundant recognizers + assert "ad_hoc_recognizers" not in payload + +@pytest.mark.asyncio +async def test_optimization_presidio_pii_caching_skips_network_calls(): + """ + OPTIMIZATION VERIFICATION: + Verify that PII results are cached in local memory. + Identical text should return instantly from cache without making a + network round-trip to the Presidio service. + """ + presidio = _OPTIONAL_PresidioPIIMasking( + mock_testing=True, + presidio_analyzer_api_base="http://mock-analyzer", + presidio_anonymizer_api_base="http://mock-anonymizer" + ) + + # Mock network calls + presidio.analyze_text = MagicMock(return_value=asyncio.Future()) + presidio.analyze_text.return_value.set_result([]) + presidio.anonymize_text = MagicMock(return_value=asyncio.Future()) + presidio.anonymize_text.return_value.set_result("redacted text") + + # First call - Must hit the mock (Network) + result1 = await presidio.check_pii( + text="repetitive call center text", + output_parse_pii=False, + presidio_config=None, + request_data={} + ) + assert result1 == "redacted text" + assert presidio.analyze_text.call_count == 1 + + # Second call - Must hit the CACHE (No Network) + result2 = await presidio.check_pii( + text="repetitive call center text", + output_parse_pii=False, + presidio_config=None, + request_data={} + ) + assert result2 == "redacted text" + assert presidio.analyze_text.call_count == 1 # Counter should NOT increase