Ran via tools/grade_run.py against LM Studio (no clipboard). Results: TTS: 80 Minor Flaws PASSES (real N-worker concurrency) <- Qwen 49, didn't parse Rust: 72 Minor Flaws COMPILES CLEAN (0 errs w/ deps) <- Qwen 50, 7 real errors Webhook: 55 Critical uses forbidden aiohttp (won't run) <- Qwen 75, passed Automation: 48 Critical SyntaxError (global-after-assign) <- first run for both DECISIVE head-to-head: Gemma generalizes where Qwen fails (TTS, Rust), but Qwen beats it on stdlib-discipline prompts (webhook). The two are COMPLEMENTARY local offloads, not redundant. Fixed: grade_run.py extractor (markdown/prose wrapping, multi-fence lang selection), TTFT-null handling in generator. TTFT capture from LM Studio API still needs the right stats key (left null + noted). Co-Authored-By: Claude <noreply@anthropic.com>
273 lines
9.2 KiB
Python
273 lines
9.2 KiB
Python
import asyncio
|
|
import random
|
|
import uuid
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, auto
|
|
from typing import Callable, Dict, Optional, Set
|
|
|
|
|
|
class JobStatus(Enum):
|
|
QUEUED = auto()
|
|
STARTED = auto()
|
|
COMPLETED = auto()
|
|
FAILED = auto()
|
|
CANCELLED = auto()
|
|
|
|
|
|
@dataclass
|
|
class Job:
|
|
id: str
|
|
text: str
|
|
voice: str
|
|
status: JobStatus = JobStatus.QUEUED
|
|
retries_attempted: int = 0
|
|
cancelled: bool = False
|
|
|
|
|
|
async def mock_synthesize(text: str) -> bytes:
|
|
"""
|
|
Simulates a TTS engine.
|
|
Sleeps 50-300ms, fails ~10% of the time.
|
|
"""
|
|
await asyncio.sleep(random.uniform(0.05, 0.3))
|
|
if random.random() < 0.1:
|
|
raise Exception("Synthetic Synthesis Error")
|
|
return b"\x00" * len(text)
|
|
|
|
|
|
class TTSPipeline:
|
|
def __init__(self, max_concurrency: int = 4, max_queue_size: int = 100):
|
|
self.max_concurrency = max_concurrency
|
|
self.max_queue_size = max_queue_size
|
|
|
|
self._queue: asyncio.Queue[str] = asyncio.Queue()
|
|
self._jobs: Dict[str, Job] = {}
|
|
self._workers: Set[asyncio.Task] = set()
|
|
self._callback: Optional[Callable[[str, JobStatus], None]] = None
|
|
|
|
self._semaphore = asyncio.Semaphore(max_concurrency)
|
|
self._running_count = 0 # Tracks jobs in queue + jobs in flight
|
|
self._lock = asyncio.Lock() # Protects shared state updates
|
|
|
|
# Metrics for testing
|
|
self._current_active_count = 0
|
|
self._peak_concurrency = 0
|
|
|
|
def set_callback(self, cb: Callable[[str, JobStatus], None]):
|
|
self._callback = cb
|
|
|
|
def _emit(self, job_id: str, status: JobStatus):
|
|
if self._callback:
|
|
self._callback(job_id, status)
|
|
|
|
async def start(self):
|
|
"""Initialize the worker pool."""
|
|
for _ in range(self.max_concurrency):
|
|
worker = asyncio.create_task(self._worker())
|
|
self._workers.add(worker)
|
|
|
|
async def stop(self):
|
|
"""Gracefully shut down workers."""
|
|
for worker in self._workers:
|
|
worker.cancel()
|
|
await asyncio.gather(*self._workers, return_exceptions=True)
|
|
self._workers.clear()
|
|
|
|
async def submit(self, text: str, voice: str) -> str:
|
|
"""Enqueue a job. Returns job_id immediately."""
|
|
if self._queue.qsize() >= self.max_queue_size:
|
|
raise BufferError("Pipeline queue is full (backpressure)")
|
|
|
|
job_id = str(uuid.uuid4())
|
|
job = Job(id=job_id, text=text, voice=voice)
|
|
|
|
async with self._lock:
|
|
self._jobs[job_id] = job
|
|
self._running_count += 1
|
|
|
|
await self._queue.put(job_id)
|
|
self._emit(job_id, JobStatus.QUEUED)
|
|
return job_id
|
|
|
|
async def cancel(self, job_id: str):
|
|
"""Cancel a queued or in-flight job."""
|
|
async with self._lock:
|
|
if job_id in self._jobs:
|
|
job = self._jobs[job_id]
|
|
if job.status not in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
|
|
job.cancelled = True
|
|
|
|
async def drain(self):
|
|
"""Wait until all queued and in-flight jobs are finished."""
|
|
# 1. Wait for the queue to be empty (all items processed by workers)
|
|
await self._queue.join()
|
|
# 2. Wait for the running count to hit zero (all workers finished their logic)
|
|
while True:
|
|
async with self._lock:
|
|
if self._running_count == 0:
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
async def _worker(self):
|
|
while True:
|
|
job_id = await self._queue.get()
|
|
try:
|
|
async with self._lock:
|
|
job = self._jobs.get(job_id)
|
|
|
|
if not job:
|
|
continue
|
|
|
|
# Check if cancelled while sitting in queue
|
|
if job.cancelled:
|
|
await self._finalize_job(job, JobStatus.CANCELLED)
|
|
continue
|
|
|
|
# Process job with bounded concurrency and retries
|
|
async with self._semaphore:
|
|
await self._process_with_retries(job)
|
|
|
|
except asyncio.CancelledError:
|
|
break
|
|
finally:
|
|
self._queue.task_done()
|
|
|
|
async def _process_with_retries(self, job: Job):
|
|
max_retries = 3
|
|
backoff = 0.1
|
|
|
|
for attempt in range(max_retries + 1):
|
|
# Check cancellation before starting/retrying
|
|
if job.cancelled:
|
|
await self._finalize_job(job, JobStatus.CANCELLED)
|
|
return
|
|
|
|
try:
|
|
# Track concurrency for testing purposes
|
|
async with self._lock:
|
|
self._current_active_count += 1
|
|
self._peak_concurrency = max(self._peak_concurrency, self._current_active_count)
|
|
|
|
self._emit(job.id, JobStatus.STARTED)
|
|
|
|
# The actual work
|
|
await mock_synthesize(job.text)
|
|
|
|
# Check cancellation after the await (in case it was cancelled during synthesis)
|
|
if job.cancelled:
|
|
await self._finalize_job(job, JobStatus.CANCELLED)
|
|
else:
|
|
await self._finalize_job(job, JobStatus.COMPLETED)
|
|
return
|
|
|
|
except Exception as e:
|
|
if job.cancelled:
|
|
await self._finalize_job(job, JobStatus.CANCELLED)
|
|
return
|
|
|
|
if attempt < max_retries:
|
|
# Exponential backoff
|
|
await asyncio.sleep(backoff * (2 ** attempt))
|
|
else:
|
|
# Exhausted retries
|
|
await self._finalize_job(job, JobStatus.FAILED)
|
|
finally:
|
|
async with self._lock:
|
|
self._current_active_count -= 1
|
|
|
|
async def _finalize_job(self, job: Job, status: JobStatus):
|
|
"""Updates job state and decrements the global running counter."""
|
|
async with self._lock:
|
|
if job.status not in (JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED):
|
|
job.status = status
|
|
self._running_count -= 1
|
|
self._emit(job.id, status)
|
|
|
|
|
|
# ==========================================
|
|
# TEST SUITE
|
|
# ==========================================
|
|
|
|
async def main():
|
|
print("--- Starting TTS Pipeline Tests ---\n")
|
|
pipeline = TTSPipeline(max_concurrency=4, max_queue_size=100)
|
|
await pipeline.start()
|
|
|
|
# Event tracking for assertions
|
|
events = []
|
|
def callback(job_id, status):
|
|
events.append((job_id, status))
|
|
|
|
pipeline.set_callback(callback)
|
|
|
|
# --- Test A: Bounded Concurrency & Drain ---
|
|
print("Test A: Submitting 50 jobs and checking concurrency...")
|
|
job_ids = []
|
|
for i in range(50):
|
|
jid = await pipeline.submit(f"Text {i}", "en-US")
|
|
job_ids.append(jid)
|
|
|
|
await pipeline.drain()
|
|
|
|
assert pipeline._peak_concurrency <= 4, f"Concurrency exceeded! Peak: {pipeline._peak_concurrency}"
|
|
print(f" [PASS] Max concurrency was {pipeline._peak_concurrency}/4")
|
|
|
|
# --- Test B: Backpressure ---
|
|
print("\nTest B: Testing backpressure (100 job cap)...")
|
|
# Fill the queue to 100. Note: drain() cleared previous jobs, so current count is 0.
|
|
# We submit 100 to fill it.
|
|
for i in range(100):
|
|
await pipeline.submit("Backpressure test", "en-US")
|
|
|
|
try:
|
|
await pipeline.submit("The breaking job", "en-US")
|
|
print(" [FAIL] Pipeline accepted 101st job without error.")
|
|
except BufferError:
|
|
print(" [PASS] Pipeline correctly rejected 101st job with BufferError.")
|
|
|
|
# Clear the queue for next test
|
|
await pipeline.drain()
|
|
|
|
# --- Test C: Cancellation & Recovery ---
|
|
print("\nTest C: Mid-flight cancellation and recovery...")
|
|
# Submit a batch
|
|
batch_ids = []
|
|
for i in range(10):
|
|
batch_ids.append(await pipeline.submit("Cancel me", "en-US"))
|
|
|
|
# Wait a tiny bit to ensure jobs are picked up by workers (in-flight)
|
|
await asyncio.sleep(0.1)
|
|
|
|
# Cancel the first job in the batch
|
|
target_id = batch_ids[0]
|
|
await pipeline.cancel(target_id)
|
|
print(f" Cancelled job: {target_id}")
|
|
|
|
await pipeline.drain()
|
|
|
|
# Verify the cancelled job is marked correctly in events
|
|
cancelled_events = [e for e in events if e[0] == target_id and e[1] == JobStatus.CANCELLED]
|
|
assert len(cancelled_events) > 0, "Job was not recorded as CANCELLED"
|
|
print(" [PASS] Job cancellation detected.")
|
|
|
|
# Verify pool is still usable
|
|
print(" Verifying pool usability...")
|
|
new_job = await pipeline.submit("Post-cancel job", "en-US")
|
|
await pipeline.drain()
|
|
|
|
# Check if the new job completed
|
|
new_job_events = [e for e in events if e[0] == new_job and e[1] == JobStatus.COMPLETED]
|
|
# Note: It might be FAILED if the 10% error hit, but it shouldn't be stuck.
|
|
assert any(e[1] in (JobStatus.COMPLETED, JobStatus.FAILED) for e in events if e[0] == new_job), "Pool stuck!"
|
|
print(" [PASS] Pool recovered and processed new jobs.")
|
|
|
|
await pipeline.stop()
|
|
print("\n--- All Tests Passed Successfully ---")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
asyncio.run(main())
|
|
except KeyboardInterrupt:
|
|
pass |