""" GPU Utilization Optimizer Smooths out GPU usage spikes and reduces dwell time through optimized processing """ import torch import logging import time import threading from queue import Queue, Empty from typing import List, Dict, Any, Optional, Tuple from concurrent.futures import ThreadPoolExecutor, as_completed import asyncio from collections import deque class AsyncTTSProcessor: """Asynchronous TTS processor with GPU utilization smoothing""" def __init__(self, model, max_queue_size=10, prefetch_count=3, batch_size=4): self.model = model self.max_queue_size = max_queue_size self.prefetch_count = prefetch_count self.batch_size = batch_size self.logger = logging.getLogger(__name__) # Processing queues self.input_queue = Queue(maxsize=max_queue_size) self.processing_queue = Queue(maxsize=batch_size * 2) self.result_queue = Queue() # Worker threads self.preprocessor_thread = None self.gpu_worker_thread = None self.running = False # Performance tracking self.stats = { 'total_processed': 0, 'avg_gpu_utilization': 0.0, 'queue_efficiency': 0.0, 'processing_time_saved': 0.0 } # Prefetch buffer for smooth processing self.prefetch_buffer = deque(maxlen=prefetch_count) def start(self): """Start async processing threads""" if self.running: return self.running = True # Start preprocessing thread self.preprocessor_thread = threading.Thread(target=self._preprocessor_worker) self.preprocessor_thread.daemon = True self.preprocessor_thread.start() # Start GPU processing thread self.gpu_worker_thread = threading.Thread(target=self._gpu_worker) self.gpu_worker_thread.daemon = True self.gpu_worker_thread.start() self.logger.info("🚀 Async TTS processor started") def stop(self): """Stop async processing""" self.running = False if self.preprocessor_thread: self.preprocessor_thread.join(timeout=2) if self.gpu_worker_thread: self.gpu_worker_thread.join(timeout=2) self.logger.info("⏚ī¸ Async TTS processor stopped") def _preprocessor_worker(self): """Background text preprocessing worker""" while self.running: try: # Get text from input queue item = self.input_queue.get(timeout=0.1) if item is None: # Shutdown signal break text, task_id, params = item # Preprocessing operations (tokenization, text normalization, etc.) preprocessed = self._preprocess_text(text, params) # Add to processing queue self.processing_queue.put((preprocessed, task_id, params)) except Empty: continue except Exception as e: self.logger.error(f"Preprocessing error: {e}") def _gpu_worker(self): """Background GPU processing worker with batch optimization""" batch_buffer = [] last_batch_time = time.time() batch_timeout = 0.05 # 50ms timeout for batching while self.running: try: # Try to get items for batch processing try: item = self.processing_queue.get(timeout=0.01) batch_buffer.append(item) except Empty: pass # Process batch if conditions are met current_time = time.time() should_process = ( len(batch_buffer) >= self.batch_size or # Batch is full (len(batch_buffer) > 0 and current_time - last_batch_time > batch_timeout) or # Timeout not self.running # Shutdown ) if should_process and batch_buffer: self._process_gpu_batch(batch_buffer) batch_buffer.clear() last_batch_time = current_time if not batch_buffer: time.sleep(0.001) # Brief sleep when no work except Exception as e: self.logger.error(f"GPU worker error: {e}") batch_buffer.clear() def _preprocess_text(self, text: str, params: Dict) -> str: """Text preprocessing operations""" # Basic text normalization text = text.strip() # Could add more preprocessing here: # - Text tokenization # - Phoneme conversion # - Prosody analysis return text def _process_gpu_batch(self, batch_items: List[Tuple]): """Process a batch of items on GPU""" if not batch_items: return try: # Group by similar parameters for better batching param_groups = self._group_by_parameters(batch_items) for param_signature, items in param_groups.items(): # Extract texts and metadata texts = [item[0] for item in items] # preprocessed text task_ids = [item[1] for item in items] params = items[0][2] # Use first item's params (they should be similar) # Single GPU call for the group with torch.inference_mode(): if len(texts) == 1: # Single generation audio = self.model.generate(texts[0], **params) audios = [audio] else: # Try batch generation if available try: if hasattr(self.model, 'generate_batch'): audios = self.model.generate_batch(texts, **params) else: # Fallback to sequential with GPU persistence audios = [] for text in texts: audio = self.model.generate(text, **params) audios.append(audio) except Exception: # Final fallback audios = [] for text in texts: try: audio = self.model.generate(text, **params) audios.append(audio) except Exception as e: self.logger.error(f"Individual generation failed: {e}") audios.append(None) # Store results for task_id, audio in zip(task_ids, audios): self.result_queue.put((task_id, audio)) self.stats['total_processed'] += len(items) except Exception as e: self.logger.error(f"Batch processing failed: {e}") # Add None results for failed items for item in batch_items: task_id = item[1] self.result_queue.put((task_id, None)) def _group_by_parameters(self, items: List[Tuple]) -> Dict[str, List[Tuple]]: """Group items by similar TTS parameters""" groups = {} for item in items: _, task_id, params = item # Create parameter signature for grouping sig_items = sorted(params.items()) if params else [] param_signature = str(sig_items) if param_signature not in groups: groups[param_signature] = [] groups[param_signature].append(item) return groups def generate_async(self, text: str, task_id: Optional[str] = None, **params) -> str: """Add text to async processing queue""" if task_id is None: task_id = f"task_{int(time.time() * 1000000)}" try: self.input_queue.put((text, task_id, params), timeout=1.0) return task_id except: raise Exception("Processing queue full") def get_result(self, timeout: float = 5.0) -> Optional[Tuple[str, torch.Tensor]]: """Get next completed result""" try: return self.result_queue.get(timeout=timeout) except Empty: return None def get_stats(self) -> Dict[str, Any]: """Get processing statistics""" return self.stats.copy() class PipelinedTTSProcessor: """Pipelined TTS processor for smoother GPU utilization""" def __init__(self, model, pipeline_depth=3): self.model = model self.pipeline_depth = pipeline_depth self.logger = logging.getLogger(__name__) # Pipeline stages self.stages = ['preprocess', 'generate', 'postprocess'] self.stage_queues = {stage: Queue(maxsize=pipeline_depth) for stage in self.stages} self.result_queue = Queue() # Worker threads self.workers = {} self.running = False # Performance tracking self.pipeline_stats = { 'throughput': 0.0, 'gpu_efficiency': 0.0, 'stage_utilization': {stage: 0.0 for stage in self.stages} } def start(self): """Start pipelined processing""" if self.running: return self.running = True # Start workers for each stage self.workers['preprocess'] = threading.Thread(target=self._preprocess_worker) self.workers['generate'] = threading.Thread(target=self._generate_worker) self.workers['postprocess'] = threading.Thread(target=self._postprocess_worker) for worker in self.workers.values(): worker.daemon = True worker.start() self.logger.info("🔄 Pipelined TTS processor started") def stop(self): """Stop pipelined processing""" self.running = False for worker in self.workers.values(): if worker.is_alive(): worker.join(timeout=1) self.logger.info("⏸ī¸ Pipelined TTS processor stopped") def _preprocess_worker(self): """Preprocessing pipeline stage""" while self.running: try: item = self.stage_queues['preprocess'].get(timeout=0.1) if item is None: break text, task_id, params = item # Text preprocessing processed_text = text.strip() # Basic processing # Pass to next stage self.stage_queues['generate'].put((processed_text, task_id, params)) except Empty: continue except Exception as e: self.logger.error(f"Preprocess stage error: {e}") def _generate_worker(self): """GPU generation pipeline stage (critical path)""" while self.running: try: item = self.stage_queues['generate'].get(timeout=0.1) if item is None: break text, task_id, params = item # GPU TTS generation start_time = time.time() with torch.inference_mode(): audio = self.model.generate(text, **params) generation_time = time.time() - start_time # Pass to next stage self.stage_queues['postprocess'].put((audio, task_id, generation_time)) except Empty: continue except Exception as e: self.logger.error(f"Generate stage error: {e}") # Pass error result self.stage_queues['postprocess'].put((None, item[1] if 'item' in locals() else 'unknown', 0)) def _postprocess_worker(self): """Post-processing pipeline stage""" while self.running: try: item = self.stage_queues['postprocess'].get(timeout=0.1) if item is None: break audio, task_id, generation_time = item # Post-processing (normalization, format conversion, etc.) # For now, just pass through processed_audio = audio # Store final result self.result_queue.put((task_id, processed_audio, generation_time)) except Empty: continue except Exception as e: self.logger.error(f"Postprocess stage error: {e}") def generate_pipelined(self, text: str, task_id: Optional[str] = None, **params) -> str: """Add text to pipelined processing""" if task_id is None: task_id = f"pipeline_{int(time.time() * 1000000)}" try: self.stage_queues['preprocess'].put((text, task_id, params), timeout=1.0) return task_id except: raise Exception("Pipeline queue full") def get_result(self, timeout: float = 10.0) -> Optional[Tuple[str, torch.Tensor, float]]: """Get next completed result from pipeline""" try: return self.result_queue.get(timeout=timeout) except Empty: return None class SmoothGPUOptimizer: """Main GPU utilization optimizer combining multiple strategies""" def __init__(self, model, strategy='async'): self.model = model self.strategy = strategy self.logger = logging.getLogger(__name__) # Initialize chosen strategy if strategy == 'async': self.processor = AsyncTTSProcessor(model, max_queue_size=15, batch_size=3) elif strategy == 'pipeline': self.processor = PipelinedTTSProcessor(model, pipeline_depth=4) else: raise ValueError(f"Unknown strategy: {strategy}") # Start processing self.processor.start() # Performance tracking self.optimization_stats = { 'strategy': strategy, 'utilization_improvement': 0.0, 'throughput_improvement': 0.0, 'dwell_time_reduction': 0.0 } def generate_optimized(self, texts: List[str], **params) -> List[torch.Tensor]: """Generate audio with optimized GPU utilization""" if not texts: return [] results = [] task_ids = [] # Submit all texts for processing for i, text in enumerate(texts): if self.strategy == 'async': task_id = self.processor.generate_async(text, f"batch_{i}", **params) else: # pipeline task_id = self.processor.generate_pipelined(text, f"batch_{i}", **params) task_ids.append(task_id) # Collect results completed = {} while len(completed) < len(task_ids): result = self.processor.get_result(timeout=15.0) if result is None: self.logger.warning("Timeout waiting for results") break if self.strategy == 'async': task_id, audio = result else: # pipeline task_id, audio, gen_time = result completed[task_id] = audio # Order results to match input order for task_id in task_ids: results.append(completed.get(task_id)) return results def cleanup(self): """Clean up processor resources""" if hasattr(self.processor, 'stop'): self.processor.stop() def get_optimization_stats(self) -> Dict[str, Any]: """Get optimization statistics""" stats = self.optimization_stats.copy() if hasattr(self.processor, 'get_stats'): processor_stats = self.processor.get_stats() stats.update({f'processor_{k}': v for k, v in processor_stats.items()}) return stats def create_gpu_optimizer(model, strategy='async'): """Create GPU utilization optimizer""" return SmoothGPUOptimizer(model, strategy=strategy)