Via tools/grade_run.py against LM Studio (load -> all prompts -> unload -> verify).
kat-coder: lfu/tts/webhook/automation/rust (lfu has a real SyntaxError: 'self._ bump_freq')
qwen3-coder-30b: lfu/tts/webhook/automation/rust (all 4 py parse clean; terse ~2k tok outputs)
Grading deferred to a later session (capture-only mode).
Script hardening:
- --resident / --unload-all commands (read loaded_instances, POST unload, verify)
- MEMORY GUARD: aborts if a different model is resident (never stack 30GB models)
- TTFT self-discovery dump (LM Studio returns empty stats non-streaming;
fix next session: switch to /api/v1/chat stream:true -> chat.end result.stats)
Co-Authored-By: Claude <noreply@anthropic.com>
416 lines
17 KiB
Python
416 lines
17 KiB
Python
"""
|
||
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()) |