73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Memory Leak Detection Script
|
|
Run this to verify that NO models are loaded at startup.
|
|
"""
|
|
import psutil
|
|
import os
|
|
import sys
|
|
|
|
def get_memory_mb():
|
|
"""Get current process memory in MB"""
|
|
process = psutil.Process(os.getpid())
|
|
return process.memory_info().rss / 1024 / 1024
|
|
|
|
def test_import_leak():
|
|
"""Test if importing modules causes memory spike"""
|
|
print("🧪 Testing for Import-Time Memory Leaks\n")
|
|
|
|
baseline = get_memory_mb()
|
|
print(f"1️⃣ Baseline Memory: {baseline:.1f} MB")
|
|
|
|
import importlib
|
|
|
|
# Test each module import
|
|
modules_to_test = [
|
|
"app.celery",
|
|
"app.modules.documents.tasks.celery_tasks",
|
|
"app.core.ssrf_protection",
|
|
"app.core.sanitization",
|
|
]
|
|
|
|
for module_path in modules_to_test:
|
|
before = get_memory_mb()
|
|
try:
|
|
# Security: Use import_module instead of exec() to prevent ACE
|
|
importlib.import_module(module_path)
|
|
after = get_memory_mb()
|
|
delta = after - before
|
|
|
|
status = "✅ PASS" if delta < 100 else "❌ LEAK DETECTED"
|
|
print(f"{status} | {module_path}")
|
|
print(f" Memory: {before:.1f} MB -> {after:.1f} MB (Δ {delta:+.1f} MB)")
|
|
|
|
if delta >= 100:
|
|
print(f" ⚠️ WARNING: {delta:.1f} MB spike suggests models loaded at import!")
|
|
return False
|
|
except Exception as e:
|
|
print(f"❌ ERROR | {module_path}")
|
|
print(f" {e}")
|
|
return False
|
|
|
|
final = get_memory_mb()
|
|
total_delta = final - baseline
|
|
print(f"\n📊 Total Memory Delta: {total_delta:+.1f} MB")
|
|
|
|
if total_delta < 200:
|
|
print("✅ PASSED: No significant memory leaks detected")
|
|
print(" Safe to deploy - models load on-demand only")
|
|
return True
|
|
else:
|
|
print(f"❌ FAILED: {total_delta:.1f} MB leaked during imports")
|
|
print(" DO NOT DEPLOY - fix memory leaks first")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
# Change to backend directory
|
|
backend_dir = os.path.join(os.path.dirname(__file__), '..')
|
|
os.chdir(backend_dir)
|
|
sys.path.insert(0, backend_dir)
|
|
|
|
success = test_import_leak()
|
|
sys.exit(0 if success else 1)
|