scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix(presidio): fix 97s latency spike under load by implementing session pooling and PII caching

Browse source 
- 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
This commit is contained in:
Alexsander Hamir 2026-01-28 12:06:33 -08:00
parent 4c1b24eed9
commit 63e855a033
2 changed files with 178 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -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"

139
tests/test_presidio_latency.py Normal file
View file

@ -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