feat: implement hybrid image extraction and memory management

Backend:
- Add hybrid image extraction for Direct track (inline image blocks)
- Add render_inline_image_regions() fallback when OCR doesn't find images
- Add check_document_for_missing_images() for detecting missing images
- Add memory management system (MemoryGuard, ModelManager, ServicePool)
- Update pdf_generator_service to handle HYBRID processing track
- Add ElementType.LOGO for logo extraction

Frontend:
- Fix PDF viewer re-rendering issues with memoization
- Add TaskNotFound component and useTaskValidation hook
- Disable StrictMode due to react-pdf incompatibility
- Fix task detail and results page loading states

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
egg
2025-11-26 10:56:22 +08:00
parent ba8ddf2b68
commit 1afdb822c3
26 changed files with 8273 additions and 366 deletions

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,380 @@
"""
Tests for OCR Service Memory Integration
Tests the integration of MemoryGuard with OCRService patterns,
including pre-operation memory checks and CPU fallback logic.
"""
import pytest
from unittest.mock import Mock, patch, MagicMock
import sys
# Mock paddle before importing memory_manager
paddle_mock = MagicMock()
paddle_mock.is_compiled_with_cuda.return_value = False
paddle_mock.device.cuda.device_count.return_value = 0
paddle_mock.device.cuda.memory_allocated.return_value = 0
paddle_mock.device.cuda.memory_reserved.return_value = 0
paddle_mock.device.cuda.empty_cache = MagicMock()
sys.modules['paddle'] = paddle_mock
from app.services.memory_manager import (
MemoryGuard,
MemoryConfig,
MemoryStats,
)
class TestMemoryGuardIntegration:
"""Tests for MemoryGuard integration patterns used in OCRService"""
def setup_method(self):
"""Setup for each test"""
self.config = MemoryConfig(
warning_threshold=0.80,
critical_threshold=0.95,
emergency_threshold=0.98,
enable_cpu_fallback=True,
)
def teardown_method(self):
"""Cleanup after each test"""
pass
def test_memory_check_below_threshold_allows_processing(self):
"""Test that memory check returns True when below thresholds"""
guard = MemoryGuard(self.config)
# Mock stats below warning threshold
with patch.object(guard, 'get_memory_stats') as mock_stats:
mock_stats.return_value = MemoryStats(
gpu_used_ratio=0.50,
gpu_free_mb=4000,
gpu_total_mb=8000,
)
is_available, stats = guard.check_memory(required_mb=2000)
assert is_available is True
assert stats.gpu_free_mb >= 2000
guard.shutdown()
def test_memory_check_above_critical_blocks_processing(self):
"""Test that memory check returns False when above critical threshold"""
guard = MemoryGuard(self.config)
# Mock stats above critical threshold
with patch.object(guard, 'get_memory_stats') as mock_stats:
mock_stats.return_value = MemoryStats(
gpu_used_ratio=0.96,
gpu_free_mb=320,
gpu_total_mb=8000,
)
is_available, stats = guard.check_memory(required_mb=1000)
assert is_available is False
guard.shutdown()
def test_memory_check_insufficient_free_memory(self):
"""Test that memory check returns False when free memory < required"""
guard = MemoryGuard(self.config)
# Mock stats with insufficient free memory but below critical ratio
with patch.object(guard, 'get_memory_stats') as mock_stats:
mock_stats.return_value = MemoryStats(
gpu_used_ratio=0.70,
gpu_free_mb=500,
gpu_total_mb=8000,
)
is_available, stats = guard.check_memory(required_mb=1000)
# Should return False (not enough free memory)
assert is_available is False
guard.shutdown()
class TestCPUFallbackPattern:
"""Tests for CPU fallback pattern as used in OCRService"""
def test_cpu_fallback_activation_pattern(self):
"""Test the CPU fallback activation pattern"""
# Simulate the pattern used in OCRService._activate_cpu_fallback
class MockOCRService:
def __init__(self):
self._cpu_fallback_active = False
self.use_gpu = True
self.gpu_available = True
self.gpu_info = {'device_id': 0}
self._memory_guard = Mock()
def _activate_cpu_fallback(self):
if self._cpu_fallback_active:
return
self._cpu_fallback_active = True
self.use_gpu = False
self.gpu_info['cpu_fallback'] = True
self.gpu_info['fallback_reason'] = 'GPU memory insufficient'
if self._memory_guard:
self._memory_guard.clear_gpu_cache()
service = MockOCRService()
# Verify initial state
assert service._cpu_fallback_active is False
assert service.use_gpu is True
# Activate fallback
service._activate_cpu_fallback()
# Verify fallback state
assert service._cpu_fallback_active is True
assert service.use_gpu is False
assert service.gpu_info.get('cpu_fallback') is True
service._memory_guard.clear_gpu_cache.assert_called_once()
def test_cpu_fallback_idempotent(self):
"""Test that CPU fallback activation is idempotent"""
class MockOCRService:
def __init__(self):
self._cpu_fallback_active = False
self.use_gpu = True
self._memory_guard = Mock()
self.gpu_info = {}
def _activate_cpu_fallback(self):
if self._cpu_fallback_active:
return
self._cpu_fallback_active = True
self.use_gpu = False
if self._memory_guard:
self._memory_guard.clear_gpu_cache()
service = MockOCRService()
# Activate twice
service._activate_cpu_fallback()
service._activate_cpu_fallback()
# clear_gpu_cache should only be called once
assert service._memory_guard.clear_gpu_cache.call_count == 1
def test_gpu_mode_restoration_pattern(self):
"""Test the GPU mode restoration pattern"""
# Simulate the pattern used in OCRService._restore_gpu_mode
class MockOCRService:
def __init__(self):
self._cpu_fallback_active = True
self.use_gpu = False
self.gpu_available = True
self.gpu_info = {
'device_id': 0,
'cpu_fallback': True,
'fallback_reason': 'test'
}
self._memory_guard = Mock()
def _restore_gpu_mode(self):
if not self._cpu_fallback_active:
return
if not self.gpu_available:
return
# Check if GPU memory is now available
if self._memory_guard:
is_available, stats = self._memory_guard.check_memory(required_mb=2000)
if is_available:
self._cpu_fallback_active = False
self.use_gpu = True
self.gpu_info.pop('cpu_fallback', None)
self.gpu_info.pop('fallback_reason', None)
service = MockOCRService()
# Mock memory guard to indicate sufficient memory
mock_stats = Mock()
mock_stats.gpu_free_mb = 5000
service._memory_guard.check_memory.return_value = (True, mock_stats)
# Restore GPU mode
service._restore_gpu_mode()
# Verify GPU mode restored
assert service._cpu_fallback_active is False
assert service.use_gpu is True
assert 'cpu_fallback' not in service.gpu_info
def test_gpu_mode_not_restored_when_memory_still_low(self):
"""Test that GPU mode is not restored when memory is still low"""
class MockOCRService:
def __init__(self):
self._cpu_fallback_active = True
self.use_gpu = False
self.gpu_available = True
self.gpu_info = {'cpu_fallback': True}
self._memory_guard = Mock()
def _restore_gpu_mode(self):
if not self._cpu_fallback_active:
return
if not self.gpu_available:
return
if self._memory_guard:
is_available, stats = self._memory_guard.check_memory(required_mb=2000)
if is_available:
self._cpu_fallback_active = False
self.use_gpu = True
service = MockOCRService()
# Mock memory guard to indicate insufficient memory
mock_stats = Mock()
mock_stats.gpu_free_mb = 500
service._memory_guard.check_memory.return_value = (False, mock_stats)
# Try to restore GPU mode
service._restore_gpu_mode()
# Verify still in fallback mode
assert service._cpu_fallback_active is True
assert service.use_gpu is False
class TestPreOperationMemoryCheckPattern:
"""Tests for pre-operation memory check pattern as used in OCRService"""
def test_pre_operation_check_with_fallback(self):
"""Test the pre-operation memory check pattern with fallback"""
guard = MemoryGuard(MemoryConfig(
warning_threshold=0.80,
critical_threshold=0.95,
enable_cpu_fallback=True,
))
# Simulate the pattern:
# 1. Check if in CPU fallback mode
# 2. Try to restore GPU mode if memory available
# 3. Perform memory check for operation
class MockService:
def __init__(self):
self._cpu_fallback_active = False
self.use_gpu = True
self.gpu_available = True
self._memory_guard = guard
def _restore_gpu_mode(self):
pass # Simplified
def pre_operation_check(self, required_mb: int) -> bool:
# Try restore first
if self._cpu_fallback_active:
self._restore_gpu_mode()
# Perform memory check
if not self.use_gpu:
return True # CPU mode, no GPU check needed
is_available, stats = self._memory_guard.check_memory(required_mb=required_mb)
return is_available
service = MockService()
# Mock sufficient memory
with patch.object(guard, 'get_memory_stats') as mock_stats:
mock_stats.return_value = MemoryStats(
gpu_used_ratio=0.50,
gpu_free_mb=4000,
gpu_total_mb=8000,
)
result = service.pre_operation_check(required_mb=2000)
assert result is True
guard.shutdown()
def test_pre_operation_check_returns_true_in_cpu_mode(self):
"""Test that pre-operation check returns True when in CPU mode"""
class MockService:
def __init__(self):
self._cpu_fallback_active = True
self.use_gpu = False
self._memory_guard = Mock()
def pre_operation_check(self, required_mb: int) -> bool:
if not self.use_gpu:
return True # CPU mode, no GPU check needed
return False
service = MockService()
result = service.pre_operation_check(required_mb=5000)
# Should return True because we're in CPU mode
assert result is True
# Memory guard should not be called
service._memory_guard.check_memory.assert_not_called()
class TestMemoryCheckWithCleanup:
"""Tests for memory check with cleanup pattern"""
def test_memory_check_triggers_cleanup_on_failure(self):
"""Test that memory check triggers cleanup when insufficient"""
guard = MemoryGuard(MemoryConfig(
warning_threshold=0.80,
critical_threshold=0.95,
))
# Track cleanup calls
cleanup_called = False
def mock_cleanup():
nonlocal cleanup_called
cleanup_called = True
class MockService:
def __init__(self):
self._memory_guard = guard
self.cleanup_func = mock_cleanup
def check_gpu_memory(self, required_mb: int) -> bool:
# First check
with patch.object(self._memory_guard, 'get_memory_stats') as mock_stats:
# First call - low memory
mock_stats.return_value = MemoryStats(
gpu_used_ratio=0.96,
gpu_free_mb=300,
gpu_total_mb=8000,
)
is_available, stats = self._memory_guard.check_memory(required_mb=required_mb)
if not is_available:
# Trigger cleanup
self.cleanup_func()
self._memory_guard.clear_gpu_cache()
return False
return True
service = MockService()
result = service.check_gpu_memory(required_mb=1000)
# Cleanup should have been triggered
assert cleanup_called is True
assert result is False
guard.shutdown()
if __name__ == "__main__":
pytest.main([__file__, "-v"])

View File

@@ -0,0 +1,387 @@
"""
Tests for OCR Service Pool
Tests OCRServicePool functionality including acquire, release, and concurrency.
"""
import pytest
import threading
import time
from unittest.mock import Mock, patch, MagicMock
import sys
# Mock paddle before importing service_pool to avoid import errors
# when paddle is not installed in the test environment
paddle_mock = MagicMock()
paddle_mock.is_compiled_with_cuda.return_value = False
paddle_mock.device.cuda.device_count.return_value = 0
paddle_mock.device.cuda.memory_allocated.return_value = 0
paddle_mock.device.cuda.memory_reserved.return_value = 0
paddle_mock.device.cuda.empty_cache = MagicMock()
sys.modules['paddle'] = paddle_mock
from app.services.service_pool import (
OCRServicePool,
PooledService,
PoolConfig,
ServiceState,
get_service_pool,
shutdown_service_pool,
)
class TestPoolConfig:
"""Tests for PoolConfig class"""
def test_default_values(self):
"""Test default configuration values"""
config = PoolConfig()
assert config.max_services_per_device == 1
assert config.max_total_services == 2
assert config.acquire_timeout_seconds == 300.0
assert config.max_queue_size == 50
assert config.max_consecutive_errors == 3
def test_custom_values(self):
"""Test custom configuration values"""
config = PoolConfig(
max_services_per_device=2,
max_total_services=4,
acquire_timeout_seconds=60.0,
)
assert config.max_services_per_device == 2
assert config.max_total_services == 4
assert config.acquire_timeout_seconds == 60.0
class TestPooledService:
"""Tests for PooledService class"""
def test_creation(self):
"""Test PooledService creation"""
mock_service = Mock()
pooled = PooledService(
service=mock_service,
device="GPU:0",
)
assert pooled.service is mock_service
assert pooled.device == "GPU:0"
assert pooled.state == ServiceState.AVAILABLE
assert pooled.use_count == 0
assert pooled.error_count == 0
class TestOCRServicePool:
"""Tests for OCRServicePool class"""
def setup_method(self):
"""Reset singleton before each test"""
shutdown_service_pool()
OCRServicePool._instance = None
OCRServicePool._lock = threading.Lock()
def teardown_method(self):
"""Cleanup after each test"""
shutdown_service_pool()
OCRServicePool._instance = None
def test_singleton_pattern(self):
"""Test that OCRServicePool is a singleton"""
pool1 = OCRServicePool()
pool2 = OCRServicePool()
assert pool1 is pool2
pool1.shutdown()
def test_initialize_device(self):
"""Test device initialization"""
config = PoolConfig()
pool = OCRServicePool(config)
# Default device should be initialized
assert "GPU:0" in pool.services
assert "GPU:0" in pool.semaphores
# Test adding new device
pool._initialize_device("GPU:1")
assert "GPU:1" in pool.services
assert "GPU:1" in pool.semaphores
pool.shutdown()
def test_acquire_creates_service(self):
"""Test that acquire creates a new service if none available"""
config = PoolConfig(max_services_per_device=1)
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
mock_service.process = Mock()
mock_service.get_gpu_status = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
pooled = pool.acquire(device="GPU:0", timeout=5.0)
assert pooled is not None
assert pooled.state == ServiceState.IN_USE
assert pooled.use_count == 1
pool.shutdown()
def test_acquire_reuses_available_service(self):
"""Test that acquire reuses available services"""
config = PoolConfig(max_services_per_device=1)
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
# First acquire
pooled1 = pool.acquire(device="GPU:0")
service_id = id(pooled1.service)
pool.release(pooled1)
# Second acquire should get the same service
pooled2 = pool.acquire(device="GPU:0")
assert id(pooled2.service) == service_id
assert pooled2.use_count == 2
pool.shutdown()
def test_release_makes_service_available(self):
"""Test that release makes service available again"""
config = PoolConfig()
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
pooled = pool.acquire(device="GPU:0")
assert pooled.state == ServiceState.IN_USE
pool.release(pooled)
assert pooled.state == ServiceState.AVAILABLE
pool.shutdown()
def test_release_with_error(self):
"""Test that release with error increments error count"""
config = PoolConfig(max_consecutive_errors=3)
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
pooled = pool.acquire(device="GPU:0")
pool.release(pooled, error=Exception("Test error"))
assert pooled.error_count == 1
assert pooled.state == ServiceState.AVAILABLE
pool.shutdown()
def test_release_marks_unhealthy_after_errors(self):
"""Test that service is marked unhealthy after too many errors"""
config = PoolConfig(max_consecutive_errors=2)
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
pooled = pool.acquire(device="GPU:0")
pool.release(pooled, error=Exception("Error 1"))
pooled = pool.acquire(device="GPU:0")
pool.release(pooled, error=Exception("Error 2"))
assert pooled.state == ServiceState.UNHEALTHY
assert pooled.error_count == 2
pool.shutdown()
def test_acquire_context_manager(self):
"""Test context manager for acquire/release"""
config = PoolConfig()
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
with pool.acquire_context(device="GPU:0") as pooled:
assert pooled is not None
assert pooled.state == ServiceState.IN_USE
# After context, service should be available
assert pooled.state == ServiceState.AVAILABLE
pool.shutdown()
def test_acquire_context_manager_with_error(self):
"""Test context manager releases on error"""
config = PoolConfig()
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
with pytest.raises(ValueError):
with pool.acquire_context(device="GPU:0") as pooled:
raise ValueError("Test error")
# Service should still be available after error
assert pooled.error_count == 1
pool.shutdown()
def test_acquire_timeout(self):
"""Test that acquire times out when no service available"""
config = PoolConfig(
max_services_per_device=1,
max_total_services=1,
)
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
# Acquire the only service
pooled1 = pool.acquire(device="GPU:0")
assert pooled1 is not None
# Try to acquire another - should timeout
pooled2 = pool.acquire(device="GPU:0", timeout=0.5)
assert pooled2 is None
pool.shutdown()
def test_get_pool_stats(self):
"""Test pool statistics"""
config = PoolConfig()
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
# Acquire a service
pooled = pool.acquire(device="GPU:0")
stats = pool.get_pool_stats()
assert stats["total_services"] == 1
assert stats["in_use_services"] == 1
assert stats["available_services"] == 0
assert stats["metrics"]["total_acquisitions"] == 1
pool.release(pooled)
stats = pool.get_pool_stats()
assert stats["available_services"] == 1
assert stats["metrics"]["total_releases"] == 1
pool.shutdown()
def test_health_check(self):
"""Test health check functionality"""
config = PoolConfig()
pool = OCRServicePool(config)
# Pre-populate with a mock service
mock_service = Mock()
mock_service.process = Mock()
mock_service.get_gpu_status = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
# Acquire and release to update use_count
pooled = pool.acquire(device="GPU:0")
pool.release(pooled)
health = pool.health_check()
assert health["healthy"] is True
assert len(health["services"]) == 1
assert health["services"][0]["responsive"] is True
pool.shutdown()
def test_concurrent_acquire(self):
"""Test concurrent service acquisition"""
config = PoolConfig(
max_services_per_device=2,
max_total_services=2,
)
pool = OCRServicePool(config)
# Pre-populate with 2 mock services
for i in range(2):
mock_service = Mock()
pooled_service = PooledService(service=mock_service, device="GPU:0")
pool.services["GPU:0"].append(pooled_service)
results = []
def worker(worker_id):
pooled = pool.acquire(device="GPU:0", timeout=5.0, task_id=f"task_{worker_id}")
if pooled:
results.append((worker_id, pooled))
time.sleep(0.1) # Simulate work
pool.release(pooled)
threads = [threading.Thread(target=worker, args=(i,)) for i in range(4)]
for t in threads:
t.start()
for t in threads:
t.join()
# All workers should have acquired a service
assert len(results) == 4
pool.shutdown()
class TestGetServicePool:
"""Tests for get_service_pool helper function"""
def setup_method(self):
"""Reset singleton before each test"""
shutdown_service_pool()
OCRServicePool._instance = None
def teardown_method(self):
"""Cleanup after each test"""
shutdown_service_pool()
OCRServicePool._instance = None
def test_get_service_pool_creates_singleton(self):
"""Test that get_service_pool creates a singleton"""
pool1 = get_service_pool()
pool2 = get_service_pool()
assert pool1 is pool2
shutdown_service_pool()
def test_shutdown_service_pool(self):
"""Test shutdown_service_pool cleans up"""
pool = get_service_pool()
shutdown_service_pool()
# Should be able to create new pool
new_pool = get_service_pool()
assert new_pool._initialized is True
shutdown_service_pool()
if __name__ == "__main__":
pytest.main([__file__, "-v"])