Skip to content

Commit 58b9a5a

Browse files
codelioncodelion
committed
Cache expensive resources in privacy plugin
Refactored privacy_plugin to use singleton instances for AnalyzerEngine and AnonymizerEngine, preventing repeated initialization and improving performance. Added tests to verify resource caching, singleton reuse, and performance across multiple invocations.
1 parent 6d09721 commit 58b9a5a

File tree

2 files changed

+250
-10
lines changed

2 files changed

+250
-10
lines changed

‎optillm/plugins/privacy_plugin.py‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,15 @@
33
from presidio_anonymizer import AnonymizerEngine, DeanonymizeEngine, OperatorConfig
44
from presidio_anonymizer.operators import Operator, OperatorType
55

6-
from typing import Dict, Tuple
6+
from typing import Dict, Tuple, Optional
77

88
SLUG = "privacy"
99

10+
# Singleton instances for expensive resources
11+
_analyzer_engine: Optional[AnalyzerEngine] = None
12+
_anonymizer_engine: Optional[AnonymizerEngine] = None
13+
_model_downloaded: bool = False
14+
1015
class InstanceCounterAnonymizer(Operator):
1116
"""
1217
Anonymizer which replaces the entity value
@@ -67,11 +72,14 @@ def operator_type(self) -> OperatorType:
6772
return OperatorType.Anonymize
6873

6974
def download_model(model_name):
70-
if not spacy.util.is_package(model_name):
71-
print(f"Downloading {model_name} model...")
72-
spacy.cli.download(model_name)
73-
else:
74-
print(f"{model_name} model already downloaded.")
75+
global _model_downloaded
76+
if not _model_downloaded:
77+
if not spacy.util.is_package(model_name):
78+
print(f"Downloading {model_name} model...")
79+
spacy.cli.download(model_name)
80+
else:
81+
print(f"{model_name} model already downloaded.")
82+
_model_downloaded = True
7583

7684
def replace_entities(entity_map, text):
7785
# Create a reverse mapping of placeholders to entity names
@@ -92,17 +100,32 @@ def replace_placeholder(match):
92100

93101
return replaced_text
94102

103+
def get_analyzer_engine() -> AnalyzerEngine:
104+
"""Get or create singleton AnalyzerEngine instance."""
105+
global _analyzer_engine
106+
if _analyzer_engine is None:
107+
_analyzer_engine = AnalyzerEngine()
108+
return _analyzer_engine
109+
110+
def get_anonymizer_engine() -> AnonymizerEngine:
111+
"""Get or create singleton AnonymizerEngine instance."""
112+
global _anonymizer_engine
113+
if _anonymizer_engine is None:
114+
_anonymizer_engine = AnonymizerEngine()
115+
_anonymizer_engine.add_anonymizer(InstanceCounterAnonymizer)
116+
return _anonymizer_engine
117+
95118
def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
96119
# Use the function
97120
model_name = "en_core_web_lg"
98121
download_model(model_name)
99122

100-
analyzer = AnalyzerEngine()
123+
# Use singleton instances
124+
analyzer = get_analyzer_engine()
101125
analyzer_results = analyzer.analyze(text=initial_query, language="en")
102126

103-
# Create Anonymizer engine and add the custom anonymizer
104-
anonymizer_engine = AnonymizerEngine()
105-
anonymizer_engine.add_anonymizer(InstanceCounterAnonymizer)
127+
# Use singleton anonymizer engine
128+
anonymizer_engine = get_anonymizer_engine()
106129

107130
# Create a mapping between entity types and counters
108131
entity_mapping = dict()

‎tests/test_privacy_plugin_performance.py‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Test to ensure privacy plugin resources are properly cached and not reloaded on each request.
4+
This test will fail if resources are being recreated on every call, preventing performance regressions.
5+
"""
6+
7+
import time
8+
import sys
9+
import os
10+
from unittest.mock import Mock, patch, MagicMock
11+
import importlib
12+
13+
# Add parent directory to path
14+
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
15+
16+
def test_privacy_plugin_resource_caching():
17+
"""
18+
Test that expensive resources (AnalyzerEngine, AnonymizerEngine) are created only once
19+
and reused across multiple plugin invocations.
20+
"""
21+
print("Testing privacy plugin resource caching...")
22+
23+
# Need to reset the module state before testing
24+
if 'optillm.plugins.privacy_plugin' in sys.modules:
25+
del sys.modules['optillm.plugins.privacy_plugin']
26+
27+
# Mock the expensive AnalyzerEngine and AnonymizerEngine at the module level before import
28+
with patch('presidio_analyzer.AnalyzerEngine') as MockAnalyzerEngine, \
29+
patch('presidio_anonymizer.AnonymizerEngine') as MockAnonymizerEngine, \
30+
patch('spacy.util.is_package', return_value=True):
31+
32+
# Set up mock instances
33+
mock_analyzer_instance = MagicMock()
34+
mock_analyzer_instance.analyze.return_value = []
35+
MockAnalyzerEngine.return_value = mock_analyzer_instance
36+
37+
mock_anonymizer_instance = MagicMock()
38+
mock_anonymizer_instance.anonymize.return_value = MagicMock(text="anonymized text")
39+
mock_anonymizer_instance.add_anonymizer = MagicMock()
40+
MockAnonymizerEngine.return_value = mock_anonymizer_instance
41+
42+
# Import the module with mocks in place
43+
import optillm.plugins.privacy_plugin as privacy_plugin
44+
45+
# Mock client for the run function
46+
mock_client = Mock()
47+
mock_response = Mock()
48+
mock_response.choices = [Mock(message=Mock(content="response"))]
49+
mock_response.usage.completion_tokens = 10
50+
mock_client.chat.completions.create.return_value = mock_response
51+
52+
# First invocation
53+
print("First invocation...")
54+
result1, tokens1 = privacy_plugin.run("system", "query 1", mock_client, "model")
55+
56+
# Check that resources were created once
57+
assert MockAnalyzerEngine.call_count == 1, f"AnalyzerEngine created {MockAnalyzerEngine.call_count} times, expected 1"
58+
assert MockAnonymizerEngine.call_count == 1, f"AnonymizerEngine created {MockAnonymizerEngine.call_count} times, expected 1"
59+
60+
# Second invocation
61+
print("Second invocation...")
62+
result2, tokens2 = privacy_plugin.run("system", "query 2", mock_client, "model")
63+
64+
# Check that resources were NOT created again
65+
assert MockAnalyzerEngine.call_count == 1, f"AnalyzerEngine created {MockAnalyzerEngine.call_count} times after 2nd call, expected 1"
66+
assert MockAnonymizerEngine.call_count == 1, f"AnonymizerEngine created {MockAnonymizerEngine.call_count} times after 2nd call, expected 1"
67+
68+
# Third invocation to be extra sure
69+
print("Third invocation...")
70+
result3, tokens3 = privacy_plugin.run("system", "query 3", mock_client, "model")
71+
72+
# Still should be 1
73+
assert MockAnalyzerEngine.call_count == 1, f"AnalyzerEngine created {MockAnalyzerEngine.call_count} times after 3rd call, expected 1"
74+
assert MockAnonymizerEngine.call_count == 1, f"AnonymizerEngine created {MockAnonymizerEngine.call_count} times after 3rd call, expected 1"
75+
76+
print("✅ Privacy plugin resource caching test PASSED - Resources are properly cached!")
77+
return True
78+
79+
def test_privacy_plugin_performance():
80+
"""
81+
Test that multiple invocations of the privacy plugin don't have degraded performance.
82+
This catches the actual performance issue even without mocking.
83+
"""
84+
print("\nTesting privacy plugin performance (real execution)...")
85+
86+
try:
87+
# Try to import the actual plugin
88+
import optillm.plugins.privacy_plugin as privacy_plugin
89+
90+
# Check if required dependencies are available
91+
try:
92+
import spacy
93+
from presidio_analyzer import AnalyzerEngine
94+
from presidio_anonymizer import AnonymizerEngine
95+
except ImportError as e:
96+
print(f"⚠️ Skipping performance test - dependencies not installed: {e}")
97+
return True
98+
99+
# Mock client
100+
mock_client = Mock()
101+
mock_response = Mock()
102+
mock_response.choices = [Mock(message=Mock(content="response"))]
103+
mock_response.usage.completion_tokens = 10
104+
mock_client.chat.completions.create.return_value = mock_response
105+
106+
# Warm-up call (might include model download)
107+
print("Warm-up call...")
108+
start = time.time()
109+
privacy_plugin.run("system", "warm up query", mock_client, "model")
110+
warmup_time = time.time() - start
111+
print(f"Warm-up time: {warmup_time:.2f}s")
112+
113+
# First real measurement
114+
print("First measurement call...")
115+
start = time.time()
116+
privacy_plugin.run("system", "test query 1", mock_client, "model")
117+
first_time = time.time() - start
118+
print(f"First call time: {first_time:.2f}s")
119+
120+
# Second measurement - should be fast if caching works
121+
print("Second measurement call...")
122+
start = time.time()
123+
privacy_plugin.run("system", "test query 2", mock_client, "model")
124+
second_time = time.time() - start
125+
print(f"Second call time: {second_time:.2f}s")
126+
127+
# Third measurement
128+
print("Third measurement call...")
129+
start = time.time()
130+
privacy_plugin.run("system", "test query 3", mock_client, "model")
131+
third_time = time.time() - start
132+
print(f"Third call time: {third_time:.2f}s")
133+
134+
# Performance assertions
135+
# Second and third calls should be much faster than first (at least 10x faster)
136+
# Allow some tolerance for the first call as it might still be initializing
137+
max_acceptable_time = 2.0 # 2 seconds max for subsequent calls
138+
139+
if second_time > max_acceptable_time:
140+
raise AssertionError(f"Second call took {second_time:.2f}s, expected < {max_acceptable_time}s. Resources might not be cached!")
141+
142+
if third_time > max_acceptable_time:
143+
raise AssertionError(f"Third call took {third_time:.2f}s, expected < {max_acceptable_time}s. Resources might not be cached!")
144+
145+
print(f"✅ Privacy plugin performance test PASSED - Subsequent calls are fast ({second_time:.2f}s, {third_time:.2f}s)!")
146+
return True
147+
148+
except Exception as e:
149+
print(f"❌ Performance test failed: {e}")
150+
raise
151+
152+
def test_singleton_instances_are_reused():
153+
"""
154+
Direct test that singleton instances are the same object across calls.
155+
"""
156+
print("\nTesting singleton instance reuse...")
157+
158+
try:
159+
import optillm.plugins.privacy_plugin as privacy_plugin
160+
importlib.reload(privacy_plugin)
161+
162+
# Get first instances
163+
analyzer1 = privacy_plugin.get_analyzer_engine()
164+
anonymizer1 = privacy_plugin.get_anonymizer_engine()
165+
166+
# Get second instances
167+
analyzer2 = privacy_plugin.get_analyzer_engine()
168+
anonymizer2 = privacy_plugin.get_anonymizer_engine()
169+
170+
# They should be the exact same object
171+
assert analyzer1 is analyzer2, "AnalyzerEngine instances are not the same object!"
172+
assert anonymizer1 is anonymizer2, "AnonymizerEngine instances are not the same object!"
173+
174+
print("✅ Singleton instance test PASSED - Same objects are reused!")
175+
return True
176+
177+
except ImportError as e:
178+
print(f"⚠️ Skipping singleton test - dependencies not installed: {e}")
179+
return True
180+
except Exception as e:
181+
print(f"❌ Singleton test failed: {e}")
182+
raise
183+
184+
if __name__ == "__main__":
185+
print("=" * 60)
186+
print("Privacy Plugin Performance & Caching Tests")
187+
print("=" * 60)
188+
189+
all_passed = True
190+
191+
try:
192+
test_privacy_plugin_resource_caching()
193+
except Exception as e:
194+
all_passed = False
195+
print(f"❌ Resource caching test failed: {e}")
196+
197+
try:
198+
test_singleton_instances_are_reused()
199+
except Exception as e:
200+
all_passed = False
201+
print(f"❌ Singleton instance test failed: {e}")
202+
203+
try:
204+
test_privacy_plugin_performance()
205+
except Exception as e:
206+
all_passed = False
207+
print(f"❌ Performance test failed: {e}")
208+
209+
print("\n" + "=" * 60)
210+
if all_passed:
211+
print("✅ ALL TESTS PASSED!")
212+
print("Privacy plugin resources are properly cached.")
213+
sys.exit(0)
214+
else:
215+
print("❌ SOME TESTS FAILED!")
216+
print("Privacy plugin may have performance issues.")
217+
sys.exit(1)

0 commit comments

Comments
 (0)