device_manager.py

"""
Device Management Module
Handles hardware detection, optimization, and device switching
"""

import torch
import logging
import platform
import subprocess
from typing import Optional, Dict, Any, List
from exceptions import DeviceError

logger = logging.getLogger(__name__)

class DeviceManager:
    """
    Manages device detection, validation, and optimization for video processing
    """
    
    def __init__(self):
        self._optimal_device = None
        self._device_info = {}
        self._cuda_tested = False
        self._mps_tested = False
        self._initialize_device_info()
    
    def _initialize_device_info(self):
        """Initialize comprehensive device information"""
        self._device_info = {
            'platform': platform.system(),
            'python_version': platform.python_version(),
            'pytorch_version': torch.__version__,
            'cuda_available': torch.cuda.is_available(),
            'cuda_version': torch.version.cuda if torch.cuda.is_available() else None,
            'mps_available': self._check_mps_availability(),
            'cpu_count': torch.get_num_threads(),
        }
        
        if self._device_info['cuda_available']:
            self._device_info.update(self._get_cuda_info())
        
        if self._device_info['mps_available']:
            self._device_info.update(self._get_mps_info())
        
        logger.debug(f"Device info initialized: {self._device_info}")
    
    def _check_mps_availability(self) -> bool:
        """Check if Metal Performance Shaders (MPS) is available on macOS"""
        try:
            if platform.system() == 'Darwin':  # macOS
                return hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
        except Exception:
            pass
        return False
    
    def _get_cuda_info(self) -> Dict[str, Any]:
        """Get detailed CUDA information"""
        cuda_info = {}
        try:
            if torch.cuda.is_available():
                cuda_info.update({
                    'cuda_device_count': torch.cuda.device_count(),
                    'cuda_current_device': torch.cuda.current_device(),
                    'cuda_devices': []
                })
                
                for i in range(torch.cuda.device_count()):
                    device_props = torch.cuda.get_device_properties(i)
                    device_info = {
                        'index': i,
                        'name': device_props.name,
                        'memory_total_gb': device_props.total_memory / (1024**3),
                        'memory_total_mb': device_props.total_memory / (1024**2),
                        'multiprocessor_count': device_props.multiprocessor_count,
                        'compute_capability': f"{device_props.major}.{device_props.minor}"
                    }
                    
                    # Get current memory usage
                    try:
                        memory_allocated = torch.cuda.memory_allocated(i) / (1024**3)
                        memory_reserved = torch.cuda.memory_reserved(i) / (1024**3)
                        device_info.update({
                            'memory_allocated_gb': memory_allocated,
                            'memory_reserved_gb': memory_reserved,
                            'memory_free_gb': device_info['memory_total_gb'] - memory_reserved
                        })
                    except Exception as e:
                        logger.warning(f"Could not get memory info for CUDA device {i}: {e}")
                    
                    cuda_info['cuda_devices'].append(device_info)
        
        except Exception as e:
            logger.error(f"Error getting CUDA info: {e}")
        
        return cuda_info
    
    def _get_mps_info(self) -> Dict[str, Any]:
        """Get Metal Performance Shaders information"""
        mps_info = {}
        try:
            if self._device_info['mps_available']:
                # Get system memory as MPS uses unified memory
                try:
                    result = subprocess.run(['sysctl', 'hw.memsize'], 
                                          capture_output=True, text=True, timeout=5)
                    if result.returncode == 0:
                        memory_bytes = int(result.stdout.split(':')[1].strip())
                        mps_info['mps_system_memory_gb'] = memory_bytes / (1024**3)
                except Exception as e:
                    logger.warning(f"Could not get system memory info: {e}")
                
                mps_info['mps_device'] = 'Apple Silicon GPU'
        
        except Exception as e:
            logger.error(f"Error getting MPS info: {e}")
        
        return mps_info
    
    def get_optimal_device(self) -> torch.device:
        """
        Get the optimal device for video processing with comprehensive testing
        """
        if self._optimal_device is not None:
            return self._optimal_device
        
        logger.info("Determining optimal device for video processing...")
        
        # Try CUDA first (most common for AI workloads)
        if self._device_info['cuda_available'] and not self._cuda_tested:
            cuda_device = self._test_cuda_device()
            if cuda_device is not None:
                self._optimal_device = cuda_device
                logger.info(f"Selected CUDA device: {self._get_device_name(cuda_device)}")
                return self._optimal_device
        
        # Try MPS on Apple Silicon
        if self._device_info['mps_available'] and not self._mps_tested:
            mps_device = self._test_mps_device()
            if mps_device is not None:
                self._optimal_device = mps_device
                logger.info(f"Selected MPS device: {self._get_device_name(mps_device)}")
                return self._optimal_device
        
        # Fallback to CPU
        self._optimal_device = torch.device("cpu")
        logger.info("Using CPU device (no suitable GPU found or GPU tests failed)")
        return self._optimal_device
    
    def _test_cuda_device(self) -> Optional[torch.device]:
        """Test CUDA device functionality"""
        self._cuda_tested = True
        
        try:
            # Find best CUDA device (highest memory)
            best_device_idx = 0
            best_memory = 0
            
            for device_info in self._device_info.get('cuda_devices', []):
                if device_info['memory_free_gb'] > best_memory:
                    best_memory = device_info['memory_free_gb']
                    best_device_idx = device_info['index']
            
            device = torch.device(f"cuda:{best_device_idx}")
            
            # Test basic functionality
            test_tensor = torch.tensor([1.0], device=device)
            result = test_tensor * 2
            
            # Test memory operations
            large_tensor = torch.randn(1000, 1000, device=device)
            del large_tensor, test_tensor, result
            torch.cuda.empty_cache()
            torch.cuda.synchronize()
            
            logger.info(f"CUDA device {best_device_idx} passed functionality tests")
            return device
            
        except Exception as e:
            logger.warning(f"CUDA device test failed: {e}")
            return None
    
    def _test_mps_device(self) -> Optional[torch.device]:
        """Test MPS device functionality"""
        self._mps_tested = True
        
        try:
            device = torch.device("mps")
            
            # Test basic functionality
            test_tensor = torch.tensor([1.0], device=device)
            result = test_tensor * 2
            
            # Test memory operations
            large_tensor = torch.randn(1000, 1000, device=device)
            del large_tensor, test_tensor, result
            
            # MPS doesn't have explicit cache clearing like CUDA
            logger.info("MPS device passed functionality tests")
            return device
            
        except Exception as e:
            logger.warning(f"MPS device test failed: {e}")
            return None
    
    def _get_device_name(self, device: torch.device) -> str:
        """Get human-readable device name"""
        if device.type == 'cuda':
            if self._device_info.get('cuda_devices'):
                device_idx = device.index or 0
                for cuda_device in self._device_info['cuda_devices']:
                    if cuda_device['index'] == device_idx:
                        return cuda_device['name']
            return f"CUDA Device {device.index or 0}"
        elif device.type == 'mps':
            return "Apple Silicon GPU (MPS)"
        else:
            return "CPU"
    
    def get_device_capabilities(self, device: Optional[torch.device] = None) -> Dict[str, Any]:
        """Get capabilities of the specified device"""
        if device is None:
            device = self.get_optimal_device()
        
        capabilities = {
            'device_type': device.type,
            'device_name': self._get_device_name(device),
            'supports_mixed_precision': False,
            'recommended_batch_size': 1,
            'memory_efficiency': 'medium'
        }
        
        if device.type == 'cuda':
            device_idx = device.index or 0
            for cuda_device in self._device_info.get('cuda_devices', []):
                if cuda_device['index'] == device_idx:
                    # Check compute capability for mixed precision
                    compute_version = float(cuda_device.get('compute_capability', '0.0'))
                    capabilities['supports_mixed_precision'] = compute_version >= 7.0
                    
                    # Estimate batch size based on memory
                    memory_gb = cuda_device.get('memory_free_gb', 0)
                    if memory_gb >= 24:
                        capabilities['recommended_batch_size'] = 4
                        capabilities['memory_efficiency'] = 'high'
                    elif memory_gb >= 12:
                        capabilities['recommended_batch_size'] = 2
                        capabilities['memory_efficiency'] = 'high'
                    elif memory_gb >= 6:
                        capabilities['recommended_batch_size'] = 1
                        capabilities['memory_efficiency'] = 'medium'
                    else:
                        capabilities['memory_efficiency'] = 'low'
                    
                    capabilities['memory_available_gb'] = memory_gb
                    break
        
        elif device.type == 'mps':
            capabilities['supports_mixed_precision'] = True  # MPS supports fp16
            capabilities['memory_efficiency'] = 'high'  # Unified memory
            system_memory = self._device_info.get('mps_system_memory_gb', 8)
            if system_memory >= 16:
                capabilities['recommended_batch_size'] = 2
            capabilities['memory_available_gb'] = system_memory * 0.7  # Rough estimate
        
        else:  # CPU
            capabilities['memory_efficiency'] = 'low'
            capabilities['supports_mixed_precision'] = False
            
        return capabilities
    
    def switch_device(self, device_type: str) -> torch.device:
        """
        Switch to a specific device type
        
        Args:
            device_type: 'cuda', 'mps', or 'cpu'
        """
        try:
            if device_type.lower() == 'cuda':
                if not self._device_info['cuda_available']:
                    raise DeviceError('cuda', 'CUDA not available on this system')
                
                device = self._test_cuda_device()
                if device is None:
                    raise DeviceError('cuda', 'CUDA device failed functionality tests')
                
            elif device_type.lower() == 'mps':
                if not self._device_info['mps_available']:
                    raise DeviceError('mps', 'MPS not available on this system')
                
                device = self._test_mps_device()
                if device is None:
                    raise DeviceError('mps', 'MPS device failed functionality tests')
                
            elif device_type.lower() == 'cpu':
                device = torch.device('cpu')
                
            else:
                raise DeviceError('unknown', f'Unknown device type: {device_type}')
            
            self._optimal_device = device
            logger.info(f"Switched to device: {self._get_device_name(device)}")
            return device
            
        except DeviceError:
            raise
        except Exception as e:
            raise DeviceError(device_type, f"Failed to switch to {device_type}: {str(e)}")
    
    def get_available_devices(self) -> List[str]:
        """Get list of available device types"""
        devices = ['cpu']  # CPU always available
        
        if self._device_info['cuda_available']:
            devices.append('cuda')
        
        if self._device_info['mps_available']:
            devices.append('mps')
        
        return devices
    
    def get_device_status(self) -> Dict[str, Any]:
        """Get comprehensive device status"""
        current_device = self.get_optimal_device()
        
        status = {
            'current_device': str(current_device),
            'current_device_name': self._get_device_name(current_device),
            'available_devices': self.get_available_devices(),
            'device_info': self._device_info.copy(),
            'capabilities': self.get_device_capabilities(current_device)
        }
        
        # Add current memory usage if on GPU
        if current_device.type == 'cuda':
            try:
                device_idx = current_device.index or 0
                status['current_memory_usage'] = {
                    'allocated_gb': torch.cuda.memory_allocated(device_idx) / (1024**3),
                    'reserved_gb': torch.cuda.memory_reserved(device_idx) / (1024**3),
                    'max_allocated_gb': torch.cuda.max_memory_allocated(device_idx) / (1024**3),
                    'max_reserved_gb': torch.cuda.max_memory_reserved(device_idx) / (1024**3)
                }
            except Exception as e:
                logger.warning(f"Could not get current memory usage: {e}")
        
        return status
    
    def optimize_for_processing(self) -> Dict[str, Any]:
        """Optimize device settings for video processing"""
        device = self.get_optimal_device()
        optimizations = {
            'device': str(device),
            'optimizations_applied': []
        }
        
        try:
            if device.type == 'cuda':
                # Enable cuDNN benchmarking for consistent input sizes
                torch.backends.cudnn.benchmark = True
                optimizations['optimizations_applied'].append('cudnn_benchmark')
                
                # Enable cuDNN deterministic mode if needed for reproducibility
                # torch.backends.cudnn.deterministic = True
                
                # Set memory allocation strategy
                # os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
                optimizations['optimizations_applied'].append('cuda_memory_strategy')
            
            elif device.type == 'mps':
                # MPS-specific optimizations would go here
                optimizations['optimizations_applied'].append('mps_optimized')
            
            else:  # CPU
                # Set optimal number of threads for CPU processing
                torch.set_num_threads(min(torch.get_num_threads(), 8))
                optimizations['optimizations_applied'].append('cpu_thread_optimization')
            
            logger.info(f"Applied optimizations for {device}: {optimizations['optimizations_applied']}")
            
        except Exception as e:
            logger.warning(f"Some optimizations failed: {e}")
            optimizations['optimization_errors'] = str(e)
        
        return optimizations
    
    def cleanup_device_memory(self):
        """Clean up device memory"""
        device = self.get_optimal_device()
        
        if device.type == 'cuda':
            try:
                torch.cuda.empty_cache()
                torch.cuda.synchronize()
                logger.debug("CUDA memory cache cleared")
            except Exception as e:
                logger.warning(f"CUDA memory cleanup failed: {e}")
        
        elif device.type == 'mps':
            try:
                # MPS uses unified memory, less explicit cleanup needed
                # But we can still run garbage collection
                import gc
                gc.collect()
                logger.debug("MPS memory cleanup completed")
            except Exception as e:
                logger.warning(f"MPS memory cleanup failed: {e}")
        
        else:  # CPU
            try:
                import gc
                gc.collect()
                logger.debug("CPU memory cleanup completed")
            except Exception as e:
                logger.warning(f"CPU memory cleanup failed: {e}")