src/pipeline/distributed_processor.py

"""Distributed processing module for scaling to 5M+ samples using Ray."""

import logging
import time
import hashlib
from typing import List, Dict, Any, Optional, Iterator
from pathlib import Path
import numpy as np
from PIL import Image
import io
import pickle
from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    import ray
    RAY_AVAILABLE = True
except ImportError:
    RAY_AVAILABLE = False
    ray = None

from .base import PipelineStage, CharacterAttributes, ProcessingResult
from .input_loader import DatasetItem

logger = logging.getLogger(__name__)

if RAY_AVAILABLE:
    @ray.remote
    class DistributedWorker:
        """Ray actor for distributed character attribute extraction."""
        
        def __init__(self, pipeline_config: Dict[str, Any]):
            """Initialize worker with pipeline configuration."""
            self.pipeline_config = pipeline_config
            self.pipeline = None
            self._initialize_pipeline()
        
        def _initialize_pipeline(self):
            """Initialize the character extraction pipeline on worker."""
            try:
                from character_pipeline import create_pipeline
                self.pipeline = create_pipeline(self.pipeline_config)
                logger.info(f"Worker {ray.get_runtime_context().get_worker_id()} initialized")
            except Exception as e:
                logger.error(f"Failed to initialize pipeline on worker: {e}")
                raise
        
        def process_batch(self, items: List[DatasetItem]) -> List[ProcessingResult]:
            """Process a batch of items on this worker."""
            results = []
            for item in items:
                try:
                    start_time = time.time()
                    result = self.pipeline.extract_from_image(item.image_path)
                    processing_time = time.time() - start_time
                    
                    results.append(ProcessingResult(
                        item_id=item.item_id,
                        attributes=result,
                        success=True,
                        processing_time=processing_time
                    ))
                except Exception as e:
                    logger.warning(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 get_worker_stats(self) -> Dict[str, Any]:
            """Get worker statistics and health status."""
            return {
                'worker_id': ray.get_runtime_context().get_worker_id(),
                'node_id': ray.get_runtime_context().get_node_id(),
                'memory_usage': ray.cluster_resources().get('memory', 0),
                'cpu_usage': ray.cluster_resources().get('CPU', 0)
            }
else:
    class DistributedWorker:
        """Fallback worker when Ray is not available."""
        def __init__(self, pipeline_config: Dict[str, Any]):
            raise ImportError("Ray is not available. Install with: pip install ray[default]")

class DistributedProcessor(PipelineStage):
    """Distributed processor for handling large-scale character extraction."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__("DistributedProcessor", config)
        
        # Configuration
        self.num_workers = config.get('num_workers', 4) if config else 4
        self.batch_size = config.get('batch_size', 32) if config else 32
        self.max_retries = config.get('max_retries', 3) if config else 3
        self.checkpoint_interval = config.get('checkpoint_interval', 1000) if config else 1000
        
        # Ray configuration
        self.ray_config = config.get('ray_config', {}) if config else {}
        
        # State
        self.workers = []
        self.is_initialized = False
        
    def initialize_cluster(self) -> bool:
        """Initialize Ray cluster and workers."""
        if not RAY_AVAILABLE:
            self.logger.error("Ray is not available. Install with: pip install ray[default]")
            return False
            
        try:
            if not ray.is_initialized():
                ray.init(**self.ray_config)
            
            # Create distributed workers
            pipeline_config = self.config.get('pipeline_config', {})
            self.workers = [
                DistributedWorker.remote(pipeline_config) 
                for _ in range(self.num_workers)
            ]
            
            # Test worker initialization
            health_checks = ray.get([worker.get_worker_stats.remote() for worker in self.workers])
            self.logger.info(f"Initialized {len(health_checks)} workers successfully")
            
            self.is_initialized = True
            return True
            
        except Exception as e:
            self.logger.error(f"Failed to initialize Ray cluster: {e}")
            return False
    
    def process_large_dataset(self, dataset_path: str, output_path: str) -> Dict[str, Any]:
        """Process large dataset with distributed workers."""
        if not self.is_initialized:
            if not self.initialize_cluster():
                raise RuntimeError("Failed to initialize distributed cluster")
        
        start_time = time.time()
        total_processed = 0
        total_successful = 0
        checkpoints = []
        
        try:
            # Load dataset in streaming fashion
            dataset_stream = self._create_dataset_stream(dataset_path)
            
            # Process in distributed batches
            batch_futures = []
            current_batch = []
            
            for item in dataset_stream:
                current_batch.append(item)
                
                if len(current_batch) >= self.batch_size:
                    # Distribute batch to available worker
                    worker = self.workers[len(batch_futures) % len(self.workers)]
                    future = worker.process_batch.remote(current_batch.copy())
                    batch_futures.append(future)
                    current_batch = []
                    
                    # Process completed batches
                    if len(batch_futures) >= self.num_workers * 2:
                        completed_results = self._collect_completed_batches(batch_futures)
                        total_processed += sum(len(results) for results in completed_results)
                        total_successful += sum(
                            sum(1 for r in results if r.success) 
                            for results in completed_results
                        )
                        
                        # Save checkpoint
                        if total_processed % self.checkpoint_interval == 0:
                            checkpoint = self._save_checkpoint(completed_results, output_path, total_processed)
                            checkpoints.append(checkpoint)
                            self.logger.info(f"Checkpoint saved: {total_processed} items processed")
            
            # Process remaining batch
            if current_batch:
                worker = self.workers[0]
                future = worker.process_batch.remote(current_batch)
                batch_futures.append(future)
            
            # Collect all remaining results
            if batch_futures:
                remaining_results = ray.get(batch_futures)
                total_processed += sum(len(results) for results in remaining_results)
                total_successful += sum(
                    sum(1 for r in results if r.success) 
                    for results in remaining_results
                )
                
                # Final checkpoint
                final_checkpoint = self._save_checkpoint(remaining_results, output_path, total_processed)
                checkpoints.append(final_checkpoint)
            
            processing_time = time.time() - start_time
            
            return {
                'total_processed': total_processed,
                'total_successful': total_successful,
                'success_rate': total_successful / total_processed if total_processed > 0 else 0,
                'processing_time': processing_time,
                'throughput': total_processed / processing_time if processing_time > 0 else 0,
                'checkpoints': checkpoints,
                'num_workers': len(self.workers)
            }
            
        except Exception as e:
            self.logger.error(f"Distributed processing failed: {e}")
            raise
    
    def _create_dataset_stream(self, dataset_path: str) -> Iterator[DatasetItem]:
        """Create streaming iterator for large dataset."""
        dataset_path = Path(dataset_path)
        
        if dataset_path.is_file():
            # Single file - assume it's a list of image paths
            with open(dataset_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
                        )
        else:
            # Directory - iterate through image files
            image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
            for idx, image_path in enumerate(dataset_path.rglob('*')):
                if image_path.suffix.lower() in image_extensions:
                    yield DatasetItem(
                        item_id=f"item_{idx}",
                        image_path=str(image_path)
                    )
    
    def _collect_completed_batches(self, batch_futures: List) -> List[List[ProcessingResult]]:
        """Collect completed batch results and remove from futures list."""
        completed_results = []
        ready_futures, batch_futures[:] = ray.wait(batch_futures, num_returns=len(batch_futures), timeout=0)
        
        if ready_futures:
            completed_results = ray.get(ready_futures)
        
        return completed_results
    
    def _save_checkpoint(self, results: List[List[ProcessingResult]], output_path: str, total_processed: int) -> str:
        """Save processing checkpoint."""
        checkpoint_path = f"{output_path}_checkpoint_{total_processed}.pkl"
        
        # Flatten results
        flat_results = []
        for batch_results in results:
            flat_results.extend(batch_results)
        
        with open(checkpoint_path, 'wb') as f:
            pickle.dump(flat_results, f)
        
        return checkpoint_path
    
    def estimate_scalability(self, sample_size: int = 1000) -> Dict[str, Any]:
        """Estimate processing capabilities for 5M scale."""
        if not self.is_initialized:
            if not self.initialize_cluster():
                raise RuntimeError("Failed to initialize distributed cluster")
        
        # Create sample dataset
        sample_items = [
            DatasetItem(item_id=f"sample_{i}", image_path="sample_image.jpg")
            for i in range(sample_size)
        ]
        
        # Measure processing time
        start_time = time.time()
        
        # Distribute sample processing
        batch_size = min(self.batch_size, sample_size // self.num_workers)
        futures = []
        
        for i in range(0, sample_size, batch_size):
            batch = sample_items[i:i + batch_size]
            worker = self.workers[i // batch_size % len(self.workers)]
            future = worker.process_batch.remote(batch)
            futures.append(future)
        
        # Collect results
        results = ray.get(futures)
        processing_time = time.time() - start_time
        
        # Calculate metrics
        total_items = sum(len(batch_results) for batch_results in results)
        throughput = total_items / processing_time if processing_time > 0 else 0
        
        # Estimate 5M scale
        estimated_5m_time = 5_000_000 / throughput if throughput > 0 else float('inf')
        estimated_5m_hours = estimated_5m_time / 3600
        
        # Memory estimation
        cluster_resources = ray.cluster_resources()
        total_memory_gb = cluster_resources.get('memory', 0) / (1024**3)
        total_cpus = cluster_resources.get('CPU', 0)
        
        return {
            'sample_size': sample_size,
            'processing_time': processing_time,
            'throughput_per_second': throughput,
            'estimated_5m_processing_time_hours': estimated_5m_hours,
            'estimated_5m_processing_time_days': estimated_5m_hours / 24,
            'cluster_resources': {
                'total_memory_gb': total_memory_gb,
                'total_cpus': total_cpus,
                'num_workers': len(self.workers)
            },
            'scalability_recommendations': self._generate_scalability_recommendations(throughput, total_memory_gb, total_cpus)
        }
    
    def _generate_scalability_recommendations(self, throughput: float, memory_gb: float, cpus: int) -> List[str]:
        """Generate recommendations for scaling to 5M samples."""
        recommendations = []
        
        if throughput < 100:  # Less than 100 items/second
            recommendations.append("Consider increasing worker count or optimizing pipeline")
        
        if memory_gb < 32:
            recommendations.append("Increase cluster memory for better caching and model loading")
        
        if cpus < 16:
            recommendations.append("Add more CPU cores for parallel processing")
        
        recommendations.extend([
            "Implement data sharding across multiple storage systems",
            "Use Redis cluster for distributed caching",
            "Consider GPU acceleration for CLIP model inference",
            "Implement progressive loading to reduce memory footprint",
            "Use Apache Beam or Spark for even larger scale processing"
        ])
        
        return recommendations
    
    def shutdown(self):
        """Shutdown distributed cluster."""
        if self.is_initialized and RAY_AVAILABLE:
            ray.shutdown()
            self.is_initialized = False
            self.logger.info("Distributed cluster shutdown complete")
    
    def process(self, input_data: Any) -> Any:
        """Process input using distributed workers."""
        if not RAY_AVAILABLE:
            raise RuntimeError("Ray is not available. Install with: pip install ray[default]")
            
        if isinstance(input_data, str):
            # Assume it's a dataset path
            return self.process_large_dataset(input_data, "distributed_output")
        else:
            raise ValueError("DistributedProcessor expects dataset path as input")
    
    def validate_input(self, input_data: Any) -> bool:
        """Validate input data."""
        return isinstance(input_data, str) and Path(input_data).exists()