import asyncio import random import time from collections import defaultdict, deque from enum import Enum from typing import Optional, Callable, Dict, Any, Set from dataclasses import dataclass from concurrent.futures import ThreadPoolExecutor # Job lifecycle events class JobEvent(Enum): QUEUED = "queued" STARTED = "started" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" @dataclass class Job: id: str text: str voice: str status: JobEvent = JobEvent.QUEUED retries: int = 0 cancelled: bool = False class TTSJobPipeline: def __init__(self, max_concurrent: int = 4, queue_limit: int = 100): self.max_concurrent = max_concurrent self.queue_limit = queue_limit self.jobs: Dict[str, Job] = {} self.job_queue: deque = deque() self.semaphore = asyncio.Semaphore(max_concurrent) self.executor = ThreadPoolExecutor(max_workers=max_concurrent) self._callback: Optional[Callable[[str, JobEvent, Optional[str]], None]] = None self._task_lock = asyncio.Lock() self._active_tasks: Set[asyncio.Task] = set() self._max_concurrency = 0 self._shutdown_event = asyncio.Event() def register_callback(self, callback: Callable[[str, JobEvent, Optional[str]], None]): """Register a callback to receive job lifecycle events.""" self._callback = callback async def submit(self, text: str, voice: str) -> str: """Submit a new job. Raises an exception if queue is full.""" job_id = f"job_{int(time.time() * 1000000)}_{random.randint(1000, 9999)}" job = Job(id=job_id, text=text, voice=voice) async with self._task_lock: if len(self.job_queue) >= self.queue_limit: raise Exception("Queue is full") self.jobs[job_id] = job self.job_queue.append(job_id) self._callback_event(job_id, JobEvent.QUEUED) # Start processing if we're not already at max concurrency if len(self._active_tasks) < self.max_concurrent: self._start_worker() return job_id async def drain(self): """Wait for all queued and in-flight jobs to complete.""" while self.job_queue or self._active_tasks: await asyncio.sleep(0.01) # Yield control to avoid busy waiting self._shutdown_event.set() async def cancel(self, job_id: str): """Cancel a queued job or mark an in-flight job for cancellation.""" async with self._task_lock: if job_id not in self.jobs: return job = self.jobs[job_id] if job.status == JobEvent.QUEUED: # Remove from queue try: self.job_queue.remove(job_id) except ValueError: pass # Already removed job.status = JobEvent.CANCELLED self._callback_event(job_id, JobEvent.CANCELLED) elif job.status == JobEvent.STARTED: # Mark for cancellation job.cancelled = True def _callback_event(self, job_id: str, event: JobEvent, error: Optional[str] = None): """Notify callback of job lifecycle events.""" if self._callback: try: self._callback(job_id, event, error) except Exception: pass # Ignore callback errors def _start_worker(self): """Start a new worker task if needed.""" if not self._shutdown_event.is_set() and len(self._active_tasks) < self.max_concurrent: task = asyncio.create_task(self._worker()) self._active_tasks.add(task) task.add_done_callback(lambda t: self._active_tasks.discard(t)) async def _worker(self): """Worker that processes jobs from the queue.""" while not self._shutdown_event.is_set(): async with self.semaphore: try: # Get next job if not self.job_queue: break # No more jobs in queue job_id = self.job_queue.popleft() job = self.jobs[job_id] # Check if cancelled before starting if job.cancelled: self._callback_event(job_id, JobEvent.CANCELLED) continue # Update job status job.status = JobEvent.STARTED self._callback_event(job_id, JobEvent.STARTED) # Process job with retries success = await self._process_with_retry(job) if success: self._callback_event(job_id, JobEvent.COMPLETED) else: self._callback_event(job_id, JobEvent.FAILED) except Exception as e: # Handle unexpected errors pass async def _process_with_retry(self, job: Job) -> bool: """Process a job with retry logic.""" max_retries = 3 backoff_base = 0.1 for attempt in range(max_retries + 1): try: # Check if cancelled if job.cancelled: self._callback_event(job.id, JobEvent.CANCELLED) return False # Simulate synthesis result = await self._mock_synthesize(job.text) # If successful, update job status job.status = JobEvent.COMPLETED return True except Exception as e: if attempt < max_retries and not job.cancelled: # Exponential backoff delay = backoff_base * (2 ** attempt) await asyncio.sleep(delay) else: # Exhausted retries or cancelled job.status = JobEvent.FAILED return False async def _mock_synthesize(self, text: str) -> bytes: """Mock synthesizer that simulates processing with random delays and occasional failures.""" # Simulate delay delay = random.uniform(0.05, 0.3) await asyncio.sleep(delay) # Simulate failure ~10% of the time if random.random() < 0.1: raise Exception("Synthesis failed") return b"\x00" * len(text) # Test function async def main(): pipeline = TTSJobPipeline(max_concurrent=4, queue_limit=100) # Track concurrency max_concurrency = 0 concurrency_tracker = [] def callback(job_id: str, event: JobEvent, error: Optional[str] = None): nonlocal max_concurrency if event == JobEvent.STARTED: concurrency_tracker.append(1) max_concurrency = max(max_concurrency, len(concurrency_tracker)) elif event == JobEvent.COMPLETED or event == JobEvent.FAILED or event == JobEvent.CANCELLED: if concurrency_tracker: concurrency_tracker.pop() pipeline.register_callback(callback) print("Test a: Submitting 50 jobs and checking concurrency...") job_ids = [] for i in range(50): job_id = await pipeline.submit(f"Text {i}", "voice1") job_ids.append(job_id) await pipeline.drain() print(f"Max concurrency achieved: {max_concurrency}") assert max_concurrency <= 4, f"Expected at most 4 concurrent jobs, got {max_concurrency}" print("Test b: Testing backpressure...") try: for i in range(105): # Exceeds queue limit await pipeline.submit(f"Text {i}", "voice1") assert False, "Expected queue limit to be exceeded" except Exception as e: print(f"Correctly rejected due to queue limit: {e}") print("Test c: Submitting batch, cancelling one mid-flight...") job_ids = [] for i in range(10): job_id = await pipeline.submit(f"Batch {i}", "voice2") job_ids.append(job_id) # Cancel one job mid-flight await pipeline.cancel(job_ids[5]) await pipeline.drain() print("All tests passed!") if __name__ == "__main__": asyncio.run(main())