"""Streaming data processor for memory-efficient handling of large datasets.""" import logging import time import json import csv from typing import Iterator, Dict, Any, Optional, List, Callable from pathlib import Path import threading import queue from concurrent.futures import ThreadPoolExecutor, as_completed import gc import psutil from dataclasses import asdict import ray from .base import PipelineStage, CharacterAttributes, ProcessingResult from .input_loader import DatasetItem from .retry_manager import RetryManager logger = logging.getLogger(__name__) class MemoryMonitor: """Monitor memory usage and trigger cleanup when needed.""" def __init__(self, max_memory_gb: float = 8.0): self.max_memory_bytes = max_memory_gb * 1024 * 1024 * 1024 self.process = psutil.Process() def get_memory_usage(self) -> Dict[str, float]: """Get current memory usage statistics.""" memory_info = self.process.memory_info() return { 'rss_gb': memory_info.rss / (1024**3), 'vms_gb': memory_info.vms / (1024**3), 'percent': self.process.memory_percent() } def should_cleanup(self) -> bool: """Check if memory cleanup is needed.""" return self.process.memory_info().rss > self.max_memory_bytes def force_cleanup(self): """Force garbage collection and memory cleanup.""" gc.collect() logger.info(f"Memory cleanup performed. Current usage: {self.get_memory_usage()}") class StreamingDataLoader: """Memory-efficient data loader for large datasets.""" def __init__(self, batch_size: int = 32, prefetch_size: int = 2): self.batch_size = batch_size self.prefetch_size = prefetch_size def load_from_directory(self, directory_path: str, file_extensions: List[str] = None) -> Iterator[DatasetItem]: """Stream dataset items from directory.""" if file_extensions is None: file_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff'] directory = Path(directory_path) if not directory.exists(): raise ValueError(f"Directory does not exist: {directory_path}") # Use generator to avoid loading all paths into memory for idx, file_path in enumerate(directory.rglob('*')): if file_path.suffix.lower() in file_extensions: # Check for corresponding text file text_path = file_path.with_suffix('.txt') text_file = str(text_path) if text_path.exists() else None yield DatasetItem( item_id=f"item_{idx}_{file_path.stem}", image_path=str(file_path), text_path=text_file ) def load_from_manifest(self, manifest_path: str) -> Iterator[DatasetItem]: """Stream dataset items from manifest file.""" manifest_file = Path(manifest_path) if manifest_file.suffix.lower() == '.json': yield from self._load_from_json_manifest(manifest_file) elif manifest_file.suffix.lower() == '.csv': yield from self._load_from_csv_manifest(manifest_file) elif manifest_file.suffix.lower() == '.txt': yield from self._load_from_text_manifest(manifest_file) else: raise ValueError(f"Unsupported manifest format: {manifest_file.suffix}") def _load_from_json_manifest(self, manifest_path: Path) -> Iterator[DatasetItem]: """Load from JSON Lines manifest.""" with open(manifest_path, 'r') as f: for line_num, line in enumerate(f): try: data = json.loads(line.strip()) yield DatasetItem( item_id=data.get('id', f"item_{line_num}"), image_path=data['image_path'], text_path=data.get('text_path') ) except (json.JSONDecodeError, KeyError) as e: logger.warning(f"Skipping invalid line {line_num} in manifest: {e}") def _load_from_csv_manifest(self, manifest_path: Path) -> Iterator[DatasetItem]: """Load from CSV manifest.""" with open(manifest_path, 'r') as f: reader = csv.DictReader(f) for row_num, row in enumerate(reader): try: yield DatasetItem( item_id=row.get('id', f"item_{row_num}"), image_path=row['image_path'], text_path=row.get('text_path') ) except KeyError as e: logger.warning(f"Skipping invalid row {row_num} in CSV: missing {e}") def _load_from_text_manifest(self, manifest_path: Path) -> Iterator[DatasetItem]: """Load from simple text file (one image path per line).""" with open(manifest_path, 'r') as f: for line_num, line in enumerate(f): image_path = line.strip() if image_path and Path(image_path).exists(): yield DatasetItem( item_id=f"item_{line_num}", image_path=image_path ) def create_batches(self, data_stream: Iterator[DatasetItem]) -> Iterator[List[DatasetItem]]: """Create batches from data stream.""" batch = [] for item in data_stream: batch.append(item) if len(batch) >= self.batch_size: yield batch batch = [] # Yield remaining items if batch: yield batch class StreamingProcessor(PipelineStage): """Streaming processor for memory-efficient large-scale processing.""" def __init__(self, config=None): """ Initializes the streaming processor. Args: config (dict): Configuration dictionary. """ if config: self.pipeline = config.get("pipeline") self.batch_size = config.get("batch_size", 16) self.parallelism = config.get("parallelism", 4) self.max_retries = config.get("max_retries", 3) self.initial_retry_delay = config.get("initial_retry_delay", 1) self.max_retry_delay = config.get("max_retry_delay", 60) self.use_ray = config.get("use_ray", False) else: self.pipeline = None self.batch_size = 16 self.parallelism = 4 self.max_retries = 3 self.initial_retry_delay = 1 self.max_retry_delay = 60 self.use_ray = False if self.pipeline is None: raise ValueError("A pipeline must be provided for processing.") if self.use_ray: if not ray.is_initialized(): ray.init(num_cpus=self.parallelism) self.pipeline = ray.remote(self.pipeline) self.retry_manager = RetryManager( max_retries=self.max_retries, initial_delay=self.initial_retry_delay, max_delay=self.max_retry_delay, ) def set_pipeline(self, pipeline): """Set the character extraction pipeline.""" self.pipeline = pipeline def process_stream(self, data_source: str, output_path: str, progress_callback: Optional[Callable] = None) -> Dict[str, Any]: """Process data stream with memory-efficient streaming.""" if self.pipeline is None: raise ValueError("Pipeline not set. Call set_pipeline() first.") start_time = time.time() output_file = Path(output_path) output_file.parent.mkdir(parents=True, exist_ok=True) # Determine data source type source_path = Path(data_source) if source_path.is_dir(): data_stream = self.data_loader.load_from_directory(data_source) elif source_path.is_file(): data_stream = self.data_loader.load_from_manifest(data_source) else: raise ValueError(f"Invalid data source: {data_source}") # Create batch stream batch_stream = self.data_loader.create_batches(data_stream) # Process batches with open(output_file, 'w') as output_f: self._write_output_header(output_f) for batch_num, batch in enumerate(batch_stream): try: # Process batch batch_results = self._process_batch(batch) # Write results for result in batch_results: self._write_result(output_f, result) if result.success: self.success_count += 1 else: self.error_count += 1 self.processed_count += 1 # Progress callback if progress_callback: progress_callback({ 'processed': self.processed_count, 'success': self.success_count, 'errors': self.error_count, 'batch_num': batch_num }) # Memory management if self.memory_monitor.should_cleanup(): self.memory_monitor.force_cleanup() # Checkpoint if self.processed_count % self.checkpoint_interval == 0: self._create_checkpoint(output_path, batch_num) logger.info(f"Checkpoint: {self.processed_count} items processed") except Exception as e: logger.error(f"Failed to process batch {batch_num}: {e}") self.error_count += len(batch) self.processed_count += len(batch) processing_time = time.time() - start_time return { 'total_processed': self.processed_count, 'successful': self.success_count, 'errors': self.error_count, 'success_rate': self.success_count / self.processed_count if self.processed_count > 0 else 0, 'processing_time': processing_time, 'throughput': self.processed_count / processing_time if processing_time > 0 else 0, 'memory_usage': self.memory_monitor.get_memory_usage(), 'output_file': str(output_file) } def _process_batch(self, batch: List[DatasetItem]) -> List[ProcessingResult]: """Process a batch of items using thread pool.""" results = [] with ThreadPoolExecutor(max_workers=self.num_workers) as executor: # Submit all items in batch future_to_item = { executor.submit(self._process_single_item, item): item for item in batch } # Collect results as they complete for future in as_completed(future_to_item): item = future_to_item[future] try: result = future.result() results.append(result) except Exception as e: logger.error(f"Failed to process {item.item_id}: {e}") results.append(ProcessingResult( item_id=item.item_id, attributes=CharacterAttributes(), success=False, error_message=str(e) )) return results def _process_single_item(self, item: DatasetItem) -> ProcessingResult: """Process a single dataset item.""" start_time = time.time() try: # Extract attributes using pipeline attributes = self.pipeline.extract_from_image(item.image_path) processing_time = time.time() - start_time return ProcessingResult( item_id=item.item_id, attributes=attributes, success=True, processing_time=processing_time ) except Exception as e: processing_time = time.time() - start_time return ProcessingResult( item_id=item.item_id, attributes=CharacterAttributes(), success=False, error_message=str(e), processing_time=processing_time ) def _write_output_header(self, file_handle): """Write output file header based on format.""" if self.output_format == 'csv': writer = csv.writer(file_handle) writer.writerow([ 'item_id', 'success', 'age', 'gender', 'ethnicity', 'hair_style', 'hair_color', 'hair_length', 'eye_color', 'body_type', 'dress', 'confidence_score', 'processing_time', 'error_message' ]) def _write_result(self, file_handle, result: ProcessingResult): """Write processing result to output file.""" if self.output_format == 'jsonl': # JSON Lines format result_dict = { 'item_id': result.item_id, 'success': result.success, 'attributes': asdict(result.attributes), 'processing_time': result.processing_time, 'error_message': result.error_message } file_handle.write(json.dumps(result_dict) + '\n') elif self.output_format == 'csv': # CSV format writer = csv.writer(file_handle) attrs = result.attributes writer.writerow([ result.item_id, result.success, attrs.age, attrs.gender, attrs.ethnicity, attrs.hair_style, attrs.hair_color, attrs.hair_length, attrs.eye_color, attrs.body_type, attrs.dress, attrs.confidence_score, result.processing_time, result.error_message ]) def _create_checkpoint(self, output_path: str, batch_num: int): """Create processing checkpoint.""" checkpoint_data = { 'processed_count': self.processed_count, 'success_count': self.success_count, 'error_count': self.error_count, 'batch_num': batch_num, 'timestamp': time.time(), 'memory_usage': self.memory_monitor.get_memory_usage() } checkpoint_path = f"{output_path}.checkpoint_{self.processed_count}.json" with open(checkpoint_path, 'w') as f: json.dump(checkpoint_data, f, indent=2) def estimate_processing_time(self, data_source: str, sample_size: int = 100) -> Dict[str, Any]: """Estimate processing time for full dataset based on sample.""" if self.pipeline is None: raise ValueError("Pipeline not set. Call set_pipeline() first.") # Load sample data source_path = Path(data_source) if source_path.is_dir(): data_stream = self.data_loader.load_from_directory(data_source) elif source_path.is_file(): data_stream = self.data_loader.load_from_manifest(data_source) else: raise ValueError(f"Invalid data source: {data_source}") # Process sample sample_items = [] for i, item in enumerate(data_stream): if i >= sample_size: break sample_items.append(item) if not sample_items: return {'error': 'No items found in data source'} # Time sample processing start_time = time.time() sample_results = self._process_batch(sample_items) sample_time = time.time() - start_time # Calculate metrics successful_results = [r for r in sample_results if r.success] avg_time_per_item = sample_time / len(sample_items) success_rate = len(successful_results) / len(sample_items) # Estimate full dataset size if source_path.is_dir(): # Count files in directory file_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff'] total_files = sum(1 for f in source_path.rglob('*') if f.suffix.lower() in file_extensions) else: # Estimate from manifest file size with open(source_path, 'r') as f: total_files = sum(1 for _ in f) # Projections estimated_total_time = total_files * avg_time_per_item estimated_total_hours = estimated_total_time / 3600 estimated_total_days = estimated_total_hours / 24 return { 'sample_size': len(sample_items), 'sample_processing_time': sample_time, 'avg_time_per_item': avg_time_per_item, 'success_rate': success_rate, 'estimated_total_files': total_files, 'estimated_total_time_seconds': estimated_total_time, 'estimated_total_time_hours': estimated_total_hours, 'estimated_total_time_days': estimated_total_days, 'throughput_items_per_second': 1 / avg_time_per_item if avg_time_per_item > 0 else 0, 'memory_usage': self.memory_monitor.get_memory_usage(), 'recommendations': self._generate_processing_recommendations(total_files, avg_time_per_item) } def _generate_processing_recommendations(self, total_files: int, avg_time_per_item: float) -> List[str]: """Generate recommendations for processing optimization.""" recommendations = [] if total_files > 1_000_000: recommendations.append("Consider distributed processing with Ray for datasets > 1M items") if avg_time_per_item > 1.0: recommendations.append("Processing time > 1s per item - consider GPU acceleration") if total_files > 100_000: recommendations.append("Enable Redis caching for large datasets") recommendations.append("Use database sharding for better performance") memory_usage = self.memory_monitor.get_memory_usage() if memory_usage['percent'] > 80: recommendations.append("High memory usage detected - reduce batch size") recommendations.extend([ "Monitor memory usage during processing", "Use checkpointing for long-running jobs", "Consider preprocessing to filter out edge cases" ]) return recommendations def process(self, input_data: Any) -> Any: """Process streaming data.""" if isinstance(input_data, dict): operation = input_data.get('operation') if operation == 'process_stream': return self.process_stream( input_data['data_source'], input_data['output_path'], input_data.get('progress_callback') ) elif operation == 'estimate': return self.estimate_processing_time( input_data['data_source'], input_data.get('sample_size', 100) ) raise ValueError("StreamingProcessor expects operation dict as input") def validate_input(self, input_data: Any) -> bool: """Validate input data.""" return isinstance(input_data, dict) and 'operation' in input_data