""" Async TTS Job Pipeline — stdlib-only, Python 3.11+ A bounded-concurrency job pipeline with retry, cancellation, backpressure, and lifecycle callbacks. Includes a self-testing ``main()``. """ from __future__ import annotations import asyncio import random import time from dataclasses import dataclass, field from enum import Enum from typing import Awaitable, Callable # ────────────────────────────────────────────────────────────────────── # Mock synthesizer # ────────────────────────────────────────────────────────────────────── async def mock_synthesize(text: str) -> bytes: """Simulate TTS synthesis. * Sleeps a random 50–300 ms. * Returns ``b"\\x00" * len(text)`` on success. * Raises ``RuntimeError`` ~10 % of the time to exercise retry logic. """ await asyncio.sleep(random.uniform(0.05, 0.30)) if random.random() < 0.10: raise RuntimeError(f"synth-error on {len(text)}-char payload") return b"\x00" * len(text) # ────────────────────────────────────────────────────────────────────── # Job model # ────────────────────────────────────────────────────────────────────── class JobState(str, Enum): QUEUED = "queued" STARTED = "started" COMPLETED = "completed" FAILED = "failed" CANCELLED = "cancelled" @dataclass class Job: job_id: str text: str voice: str state: JobState = JobState.QUEUED attempt: int = 0 result: bytes | None = None error: str | None = None _cancelled: bool = False # internal flag, protected by pipeline lock # ────────────────────────────────────────────────────────────────────── # Pipeline # ────────────────────────────────────────────────────────────────────── class TTSPipeline: """Async TTS job pipeline with bounded concurrency and retry.""" def __init__( self, max_concurrency: int = 4, max_queue_size: int = 100, ) -> None: self._max_concurrency = max_concurrency self._max_queue_size = max_queue_size # Bounded queue enforces backpressure on submit() self._queue: asyncio.Queue[Job] = asyncio.Queue(maxsize=max_queue_size) # Limits how many jobs are synthesising simultaneously self._semaphore = asyncio.Semaphore(max_concurrency) # All jobs ever submitted (for drain / lookup) self._jobs: dict[str, Job] = {} self._job_counter: int = 0 # Callbacks registered via on_event() self._callbacks: list[Callable[[str, str], Awaitable[None]]] = [] # Concurrency-tracking helpers (for tests) self._active_count: int = 0 self._active_lock = asyncio.Lock() # Internal lock serialises mutations to shared state self._lock = asyncio.Lock() # Background worker that pulls from the queue and dispatches tasks self._worker_task: asyncio.Task[None] | None = None # ── lifecycle ──────────────────────────────────────────────────── async def start(self) -> None: """Spawn the background worker. Call once before submitting.""" if self._worker_task is None: self._worker_task = asyncio.create_task(self._worker_loop()) async def stop(self) -> None: """Gracefully shut down the worker (for cleanup in tests).""" if self._worker_task is not None: self._worker_task.cancel() try: await self._worker_task except asyncio.CancelledError: pass self._worker_task = None # ── public API ─────────────────────────────────────────────────── async def submit(self, text: str, voice: str) -> str: """Enqueue a job and return its id immediately (non-blocking). Raises ``RuntimeError`` with a clear message when the queue is full (backpressure — does *not* block or grow unbounded). """ async with self._lock: self._job_counter += 1 job_id = f"job-{self._job_counter:04d}" job = Job(job_id=job_id, text=text, voice=voice) self._jobs[job_id] = job # Try to enqueue; on failure roll back the registration try: self._queue.put_nowait(job) except asyncio.QueueFull: async with self._lock: del self._jobs[job_id] raise RuntimeError( f"Queue full (max {self._max_queue_size}); job rejected: {job_id}" ) from None await self._notify(job, "queued") return job_id async def drain(self) -> None: """Block until every submitted job has reached a terminal state (completed / failed / cancelled) and the internal queue is empty. Jobs submitted *after* this call returns are **not** waited for. """ while True: async with self._lock: total = len(self._jobs) done = sum( 1 for j in self._jobs.values() if j.state in (JobState.COMPLETED, JobState.FAILED, JobState.CANCELLED) ) if total > 0 and total == done and self._queue.empty(): return # Yield control so workers can make progress await asyncio.sleep(0.005) async def cancel(self, job_id: str) -> None: """Cancel a job. * Queued jobs are removed from processing immediately. * In-flight jobs are marked for cancellation and checked after synthesis completes (or on the next retry boundary). """ async with self._lock: job = self._jobs.get(job_id) if job is None: raise KeyError(f"Unknown job: {job_id}") if job.state in (JobState.COMPLETED, JobState.FAILED): return # already terminal — no-op job._cancelled = True if job.state == JobState.QUEUED: # Already in the queue; the worker will see the flag and skip it. job.state = JobState.CANCELLED # If in-flight, the worker checks _cancelled after synthesis. await self._notify(job, "cancelled") def on_event(self, callback: Callable[[str, str], Awaitable[None]]) -> None: """Register a lifecycle-event callback. Events emitted: ``queued``, ``started``, ``completed``, ``failed``, ``cancelled``. """ self._callbacks.append(callback) # ── internal ───────────────────────────────────────────────────── async def _notify(self, job: Job, event: str) -> None: for cb in self._callbacks: try: await cb(job.job_id, event) except Exception: pass # one bad callback must not break the pipeline async def _worker_loop(self) -> None: """Single dispatcher: pulls jobs from the queue and spawns a task for each. The task itself acquires the concurrency semaphore.""" while True: job = await self._queue.get() # task_done is called by _process_job when it finishes (including retries) asyncio.create_task(self._process_job(job)) async def _process_job(self, job: Job) -> None: """Run one job with bounded-concurrency synthesis and retry logic.""" max_attempts = 3 backoffs = [0.1, 0.2, 0.4] # exponential: 0.1 → 0.2 → 0.4 s for attempt in range(max_attempts): job.attempt = attempt + 1 # Respect cancellation before even trying to acquire a slot if job._cancelled and job.state == JobState.CANCELLED: break # ── Acquire a concurrency slot for the actual synthesis ── async with self._semaphore: # Double-check cancellation after waiting for a slot if job._cancelled and job.state == JobState.CANCELLED: break async with self._active_lock: self._active_count += 1 job.state = JobState.STARTED await self._notify(job, "started") try: result = await mock_synthesize(job.text) if job._cancelled and job.state == JobState.CANCELLED: job.state = JobState.CANCELLED else: job.result = result job.state = JobState.COMPLETED await self._notify(job, "completed") except Exception as exc: if job._cancelled and job.state == JobState.CANCELLED: job.state = JobState.CANCELLED elif attempt < max_attempts - 1: # Will fall through to backoff below (outside semaphore) pass else: job.error = str(exc) job.state = JobState.FAILED await self._notify(job, "failed") finally: async with self._active_lock: self._active_count -= 1 # ── Retry backoff happens *outside* the semaphore so it # doesn't hold a concurrency slot while sleeping ────── if job.state == JobState.FAILED and attempt < max_attempts - 1: await asyncio.sleep(backoffs[attempt]) else: break # completed, cancelled, or exhausted retries self._queue.task_done() # ────────────────────────────────────────────────────────────────────── # Tests # ────────────────────────────────────────────────────────────────────── async def main() -> None: print("=" * 60) print("Async TTS Pipeline — Self-Test") print("=" * 60) # ── Helper: collect lifecycle events ─────────────────────────── events: list[tuple[str, str]] = [] async def record(event_id: str, event_type: str) -> None: events.append((event_id, event_type)) # ── Test A: bounded concurrency ──────────────────────────────── print("\n[Test A] Bounded concurrency (50 jobs, max 4 simultaneous)") print("-" * 50) pipeline_a = TTSPipeline(max_concurrency=4, max_queue_size=100) await pipeline_a.start() pipeline_a.on_event(record) peak_concurrency = 0 current_active = 0 active_lock = asyncio.Lock() async def track_events(job_id: str, event_type: str) -> None: nonlocal peak_concurrency, current_active if event_type == "started": async with active_lock: current_active += 1 peak_concurrency = max(peak_concurrency, current_active) elif event_type in ("completed", "failed", "cancelled"): async with active_lock: current_active -= 1 pipeline_a.on_event(track_events) ids_a = [await pipeline_a.submit(f"text-{i}", "mock-voice") for i in range(50)] await pipeline_a.drain() print(f" Submitted : {len(ids_a)}") print(f" Peak concurrency: {peak_concurrency}") assert peak_concurrency <= 4, f"Concurrency exceeded! Peak was {peak_concurrency}" print(" ✓ Concurrency never exceeded 4") # Verify no duplicates / drops states = {jid: next((j.state for j in pipeline_a._jobs.values() if j.job_id == jid)) for jid in ids_a} terminal = sum(1 for s in states.values() if s in (JobState.COMPLETED, JobState.FAILED)) assert terminal == 50, f"Expected 50 terminal jobs, got {terminal}" print(f" ✓ All 50 jobs reached a terminal state") await pipeline_a.stop() # ── Test B: backpressure (queue cap = 100) ───────────────────── print("\n[Test B] Backpressure — 101st job must be rejected") print("-" * 50) pipeline_b = TTSPipeline(max_concurrency=1, max_queue_size=100) await pipeline_b.start() # Fill the queue to capacity by submitting 100 fast jobs ids_b = [] for i in range(100): jid = await pipeline_b.submit(f"fill-{i}", "voice") ids_b.append(jid) # The 101st must raise immediately (queue is full) rejected = False try: await pipeline_b.submit("overflow", "voice") except RuntimeError as exc: rejected = True print(f" Rejection message: {exc}") assert rejected, "Expected RuntimeError for over-full queue" print(" ✓ 101st job was correctly rejected") await pipeline_b.drain() print(" ✓ Drain completed after backpressure test") await pipeline_b.stop() # ── Test C: cancellation mid-flight + pool health ────────────── print("\n[Test C] Cancel mid-flight, drain, verify pool health") print("-" * 50) pipeline_c = TTSPipeline(max_concurrency=4, max_queue_size=100) await pipeline_c.start() # Patch mock_synthesize to be slow so we can cancel mid-flight original_synthesize = mock_synthesize async def slow_synthesize(text: str) -> bytes: await asyncio.sleep(2.0) # deliberately slow return b"\x00" * len(text) import __main__ __main__.mock_synthesize = slow_synthesize # temporarily replace ids_c = [await pipeline_c.submit(f"cancel-me-{i}", "voice") for i in range(20)] await asyncio.sleep(0.05) # let a few start # Cancel one that should be in-flight await pipeline_c.cancel(ids_c[0]) print(f" Cancelled job {ids_c[0]}") # Also cancel one that is still queued await pipeline_c.cancel(ids_c[1]) print(f" Cancelled job {ids_c[1]} (queued)") # Restore original synthesizer so remaining jobs finish quickly __main__.mock_synthesize = original_synthesize await pipeline_c.drain() # Verify no duplicates: every submitted id appears exactly once seen_ids = [e[0] for e in events if e[1] == "queued"] assert len(seen_ids) == len(set(seen_ids)), "Duplicate job IDs detected!" print(" ✓ No duplicate jobs") # Verify no drops: every submitted id has a terminal state final_states: dict[str, JobState] = {} for jid in ids_c: job = pipeline_c._jobs[jid] final_states[jid] = job.state completed = sum(1 for s in final_states.values() if s == JobState.COMPLETED) cancelled = sum(1 for s in final_states.values() if s == JobState.CANCELLED) failed = sum(1 for s in final_states.values() if s == JobState.FAILED) print(f" Results: {completed} completed, {cancelled} cancelled, {failed} failed") assert completed + cancelled + failed == 20, "Jobs were dropped!" print(" ✓ No jobs dropped or duplicated") # Pool must still be usable after drain + cancel jid_new = await pipeline_c.submit("post-drain-job", "voice") assert jid_new.startswith("job-"), "Pipeline unusable after drain" await pipeline_c.drain() assert pipeline_c._jobs[jid_new].state == JobState.COMPLETED print(" ✓ Pipeline still usable after drain + cancellation") await pipeline_c.stop() # ── Summary ──────────────────────────────────────────────────── print("\n" + "=" * 60) print("All tests passed ✓") print("=" * 60) if __name__ == "__main__": asyncio.run(main())