diff --git a/backend/app.py b/backend/app.py index ea8411c..e77df60 100644 --- a/backend/app.py +++ b/backend/app.py @@ -103,6 +103,32 @@ if not os.environ.get("MIOPEN_LOG_LEVEL"): os.environ["MIOPEN_LOG_LEVEL"] = "4" import torch + + +def _bounded_torch_thread_env(name: str, fallback: int, maximum: int) -> int: + """Read one bounded PyTorch pool size from the process environment.""" + try: + configured = int(os.environ.get(name, str(fallback)) or fallback) + except (TypeError, ValueError): + configured = fallback + return min(maximum, max(1, configured)) + + +_voicebox_cpu_threads = _bounded_torch_thread_env("VOICEBOX_CPU_THREADS", 4, 32) +_voicebox_interop_threads = _bounded_torch_thread_env("VOICEBOX_INTEROP_THREADS", 1, 8) +torch.set_num_threads(_voicebox_cpu_threads) +try: + torch.set_num_interop_threads(_voicebox_interop_threads) +except RuntimeError as error: + logger.warning( + "PyTorch inter-op thread pool was already initialized; keeping current setting: %s", + error, + ) +logger.info( + "PyTorch thread pools bounded: intra-op=%d inter-op=%d", + _voicebox_cpu_threads, + _voicebox_interop_threads, +) from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from urllib.parse import quote @@ -112,7 +138,7 @@ from .services import tts, transcribe, llm from .database import get_db from .utils.platform_detect import get_backend_type from .utils.progress import get_progress_manager -from .services.task_queue import create_background_task, init_queue +from .services.task_queue import create_background_task, init_queue, shutdown_queue from .routes import register_routers @@ -359,6 +385,10 @@ async def _run_startup(application: FastAPI) -> None: async def _run_shutdown() -> None: """Unload models on lifespan exit.""" logger.info("Voicebox server shutting down...") + try: + await shutdown_queue() + except Exception: + logger.exception("Failed to stop generation workers") try: tts.unload_tts_model() except Exception: diff --git a/backend/backends/chatterbox_backend.py b/backend/backends/chatterbox_backend.py index e7a025b..674919d 100644 --- a/backend/backends/chatterbox_backend.py +++ b/backend/backends/chatterbox_backend.py @@ -48,6 +48,10 @@ class ChatterboxTTSBackend: self.model_size = "default" self._device = None self._model_load_lock = asyncio.Lock() + # Generation installs attention hooks on the shared model. The lock + # lives in the worker thread so task cancellation cannot release it + # while asyncio.to_thread() is still executing. + self._generation_lock = threading.Lock() def _get_device(self) -> str: return get_torch_device(force_cpu_on_mac=True, allow_xpu=True) @@ -198,20 +202,21 @@ class ChatterboxTTSBackend: def _generate_sync(): import torch - if seed is not None: - manual_seed(seed, self._device) + with self._generation_lock: + if seed is not None: + manual_seed(seed, self._device) - logger.info(f"[Chatterbox] Generating: lang={language}") + logger.info(f"[Chatterbox] Generating: lang={language}") - wav = self.model.generate( - text, - language_id=language, - audio_prompt_path=ref_audio, - exaggeration=lang_defaults["exaggeration"], - cfg_weight=lang_defaults["cfg_weight"], - temperature=lang_defaults["temperature"], - repetition_penalty=lang_defaults["repetition_penalty"], - ) + wav = self.model.generate( + text, + language_id=language, + audio_prompt_path=ref_audio, + exaggeration=lang_defaults["exaggeration"], + cfg_weight=lang_defaults["cfg_weight"], + temperature=lang_defaults["temperature"], + repetition_penalty=lang_defaults["repetition_penalty"], + ) # Convert tensor -> numpy if isinstance(wav, torch.Tensor): @@ -223,4 +228,14 @@ class ChatterboxTTSBackend: return audio, sample_rate - return await asyncio.to_thread(_generate_sync) + physical_generation = asyncio.create_task(asyncio.to_thread(_generate_sync)) + try: + return await asyncio.shield(physical_generation) + except asyncio.CancelledError: + # Cancelling to_thread() cannot stop its worker. Drain it before + # propagating cancellation so the queue cannot overlap models. + try: + await asyncio.shield(physical_generation) + except Exception: + logger.exception("Chatterbox generation failed while draining cancellation") + raise diff --git a/backend/backends/chatterbox_turbo_backend.py b/backend/backends/chatterbox_turbo_backend.py index 6f7d6b9..8e10011 100644 --- a/backend/backends/chatterbox_turbo_backend.py +++ b/backend/backends/chatterbox_turbo_backend.py @@ -8,7 +8,9 @@ Forces CPU on macOS due to known MPS tensor issues. import asyncio import logging +import os import threading +from contextlib import nullcontext from pathlib import Path from typing import ClassVar, List, Optional, Tuple @@ -37,6 +39,20 @@ _TURBO_WEIGHT_FILES = [ ] +def _bounded_replica_count() -> int: + """Return the bounded number of isolated Turbo model replicas. + + Chatterbox mutates model-local attention state while generating, so one + model object must never serve overlapping calls. Parallelism is provided + by independent replicas instead of weakening that safety boundary. + """ + try: + configured = int(os.environ.get("VOICEBOX_CHATTERBOX_TURBO_REPLICAS", "2") or 2) + except (TypeError, ValueError): + configured = 2 + return min(4, max(1, configured)) + + class ChatterboxTurboTTSBackend: """Chatterbox Turbo TTS backend — fast, English-only, with paralinguistic tags.""" @@ -48,12 +64,21 @@ class ChatterboxTurboTTSBackend: self.model_size = "default" self._device = None self._model_load_lock = asyncio.Lock() + self._replica_count = _bounded_replica_count() + self._replicas = [] + self._available_replicas = None + self._seeded_generation_lock = threading.Lock() + self._active_replica_ids = set() + self._unload_requested = False def _get_device(self) -> str: return get_torch_device(force_cpu_on_mac=True, allow_xpu=True) def is_loaded(self) -> bool: - return self.model is not None + return ( + len(self._replicas) == self._replica_count + and self._available_replicas is not None + ) def _get_model_path(self, model_size: str = "default") -> str: return CHATTERBOX_TURBO_HF_REPO @@ -62,23 +87,69 @@ class ChatterboxTurboTTSBackend: return is_model_cached(CHATTERBOX_TURBO_HF_REPO, required_files=_TURBO_WEIGHT_FILES) async def load_model(self, model_size: str = "default") -> None: - """Load the Chatterbox Turbo model.""" - if self.model is not None: + """Load every configured Chatterbox Turbo model replica.""" + if self.is_loaded(): return async with self._model_load_lock: - if self.model is not None: + if self.is_loaded(): return - await asyncio.to_thread(self._load_model_sync) - - def _load_model_sync(self): - """Synchronous model loading.""" + while len(self._replicas) < self._replica_count: + replica_index = len(self._replicas) + physical_load = asyncio.create_task( + asyncio.to_thread( + self._load_model_sync, + replica_index, + ) + ) + cancelled = False + while True: + try: + model, device, cuda_stream = await asyncio.shield(physical_load) + break + except asyncio.CancelledError: + # asyncio cannot stop a native model load. Keep the + # load lock leased, retain its resulting model/VRAM, + # and only then propagate cancellation. + cancelled = True + if physical_load.cancelled(): + raise + self._replicas.append({ + "id": replica_index, + "model": model, + "device": device, + "cuda_stream": cuda_stream, + }) + if cancelled: + if len(self._replicas) == self._replica_count: + self._publish_replica_pool() + raise asyncio.CancelledError + + self._publish_replica_pool() + + def _publish_replica_pool(self) -> None: + """Publish a pool only when every physical replica is retained.""" + if len(self._replicas) != self._replica_count: + return + available = asyncio.Queue(maxsize=self._replica_count) + for replica in self._replicas: + available.put_nowait(replica) + self._available_replicas = available + self.model = self._replicas[0]["model"] + self._device = self._replicas[0]["device"] + + def _load_model_sync(self, replica_index: int = 0): + """Synchronously construct one physically independent model replica.""" model_name = "chatterbox-turbo" is_cached = self._is_model_cached() with model_load_progress(model_name, is_cached): device = self._get_device() - self._device = device - logger.info(f"Loading Chatterbox Turbo TTS on {device}...") + logger.info( + "Loading Chatterbox Turbo TTS replica %d/%d on %s...", + replica_index + 1, + self._replica_count, + device, + ) import torch from huggingface_hub import snapshot_download @@ -107,19 +178,54 @@ class ChatterboxTurboTTSBackend: model = ChatterboxTurboTTS.from_local(local_path, device) patch_chatterbox_f32(model) - self.model = model - - logger.info("Chatterbox Turbo TTS loaded successfully") + cuda_stream = None + if str(device).startswith("cuda") and torch.cuda.is_available(): + # PyTorch's default stream can serialize otherwise independent + # host threads. Give each model replica an explicit stream so + # kernels from separate live calls may overlap safely. + cuda_stream = torch.cuda.Stream(device=device) + + logger.info( + "Chatterbox Turbo TTS replica %d/%d loaded successfully", + replica_index + 1, + self._replica_count, + ) + return model, device, cuda_stream def unload_model(self) -> None: - """Unload model to free memory.""" - if self.model is not None: - device = self._device - del self.model - self.model = None - self._device = None + """Unload replicas only after every physical generation releases one.""" + if self._active_replica_ids: + self._unload_requested = True + logger.info( + "Deferring Chatterbox Turbo unload until %d active replica(s) finish", + len(self._active_replica_ids), + ) + return + self._unload_replicas() + + def _unload_replicas(self) -> None: + devices = {replica["device"] for replica in self._replicas} + self._available_replicas = None + self._replicas.clear() + self.model = None + self._device = None + self._unload_requested = False + for device in devices: empty_device_cache(device) - logger.info("Chatterbox Turbo unloaded") + logger.info("Chatterbox Turbo replicas unloaded") + + def generation_capacity(self) -> dict: + """Expose bounded physical capacity for readiness/health adapters.""" + return { + "configured": self._replica_count, + "loaded": len(self._replicas), + "active": len(self._active_replica_ids), + "available": ( + self._available_replicas.qsize() + if self._available_replicas is not None + else 0 + ), + } async def create_voice_prompt( self, @@ -171,36 +277,85 @@ class ChatterboxTurboTTSBackend: """ await self.load_model() - ref_audio = voice_prompt.get("ref_audio") - if ref_audio and not Path(ref_audio).exists(): - logger.warning(f"Reference audio not found: {ref_audio}") - ref_audio = None - - def _generate_sync(): - import torch - - if seed is not None: - manual_seed(seed, self._device) - - logger.info("[Chatterbox Turbo] Generating (English)") - - wav = self.model.generate( - text, - audio_prompt_path=ref_audio, - temperature=0.8, - top_k=1000, - top_p=0.95, - repetition_penalty=1.2, - ) - - # Convert tensor -> numpy - if isinstance(wav, torch.Tensor): - audio = wav.squeeze().cpu().numpy().astype(np.float32) - else: - audio = np.asarray(wav, dtype=np.float32) - - sample_rate = getattr(self.model, "sr", None) or getattr(self.model, "sample_rate", 24000) - - return audio, sample_rate - - return await asyncio.to_thread(_generate_sync) + available = self._available_replicas + if available is None: + raise RuntimeError("Chatterbox Turbo replica pool is not ready") + replica = await available.get() + replica_id = replica["id"] + self._active_replica_ids.add(replica_id) + try: + ref_audio = voice_prompt.get("ref_audio") + if ref_audio and not Path(ref_audio).exists(): + logger.warning(f"Reference audio not found: {ref_audio}") + ref_audio = None + + def _generate_sync(): + import torch + + # A supplied seed controls process-global RNG state in the + # upstream model. Preserve deterministic seeded calls by + # serializing only those calls; ordinary live calls use + # independent replicas concurrently. + seed_guard = ( + self._seeded_generation_lock + if seed is not None + else nullcontext() + ) + with seed_guard: + if seed is not None: + manual_seed(seed, replica["device"]) + + logger.info( + "[Chatterbox Turbo] Generating on replica %d/%d (English)", + replica_id + 1, + self._replica_count, + ) + + stream = replica["cuda_stream"] + stream_guard = ( + torch.cuda.stream(stream) + if stream is not None + else nullcontext() + ) + with stream_guard: + wav = replica["model"].generate( + text, + audio_prompt_path=ref_audio, + temperature=0.8, + top_k=1000, + top_p=0.95, + repetition_penalty=1.2, + ) + if stream is not None: + stream.synchronize() + + # Convert tensor -> numpy + if isinstance(wav, torch.Tensor): + audio = wav.squeeze().cpu().numpy().astype(np.float32) + else: + audio = np.asarray(wav, dtype=np.float32) + + model = replica["model"] + sample_rate = getattr(model, "sr", None) or getattr(model, "sample_rate", 24000) + + return audio, sample_rate + + physical_generation = asyncio.create_task(asyncio.to_thread(_generate_sync)) + try: + return await asyncio.shield(physical_generation) + except asyncio.CancelledError: + # asyncio cannot stop an in-flight worker thread. Keep this + # replica leased until its physical call has drained. + try: + await asyncio.shield(physical_generation) + except Exception: + logger.exception( + "Chatterbox Turbo generation failed while draining cancellation" + ) + raise + finally: + self._active_replica_ids.discard(replica_id) + if self._unload_requested and not self._active_replica_ids: + self._unload_replicas() + elif self._available_replicas is available: + available.put_nowait(replica) diff --git a/backend/models.py b/backend/models.py index 7970ce4..2efe60e 100644 --- a/backend/models.py +++ b/backend/models.py @@ -445,6 +445,7 @@ class HealthResponse(BaseModel): backend_variant: Optional[str] = None # Binary variant (cpu, cuda, or rocm) supports_rocm: bool = False # AMD GPU on Windows — the ROCm backend is applicable gpu_compatibility_warning: Optional[str] = None # Warning if GPU arch unsupported + generation_capacity: Optional[dict] = None class DirectoryCheck(BaseModel): diff --git a/backend/routes/generations.py b/backend/routes/generations.py index 215c96c..f527065 100644 --- a/backend/routes/generations.py +++ b/backend/routes/generations.py @@ -13,7 +13,13 @@ from .. import config, models from ..services import history, personality, profiles, tts from ..database import Generation as DBGeneration, VoiceProfile as DBVoiceProfile, get_db from ..services.generation import run_generation -from ..services.task_queue import cancel_generation as cancel_generation_job, enqueue_generation +from ..services.task_queue import ( + GenerationQueueFull, + GenerationWorkersUnavailable, + assert_generation_workers_available, + cancel_generation as cancel_generation_job, + enqueue_generation, +) from ..utils.audio import load_audio from ..utils.tasks import get_task_manager @@ -53,6 +59,36 @@ def _resolve_generation_engine(data: models.GenerationRequest, profile) -> str: return data.engine or getattr(profile, "default_engine", None) or getattr(profile, "preset_engine", None) or "qwen" +async def _enqueue_generation_or_reject( + *, + db, + task_manager, + generation_id, + generation_coro, + engine, +): + try: + enqueue_generation(generation_id, generation_coro, engine=engine) + except GenerationQueueFull as error: + task_manager.complete_generation(generation_id) + await history.update_generation_status( + generation_id=generation_id, + status="failed", + db=db, + error=str(error), + ) + raise HTTPException(status_code=429, detail=str(error)) from error + except GenerationWorkersUnavailable as error: + task_manager.complete_generation(generation_id) + await history.update_generation_status( + generation_id=generation_id, + status="failed", + db=db, + error=str(error), + ) + raise HTTPException(status_code=503, detail=str(error)) from error + + @router.post("/generate", response_model=models.GenerationResponse) async def generate_speech( data: models.GenerationRequest, @@ -76,6 +112,11 @@ async def generate_speech( model_size = (data.model_size or "1.7B") if engine_has_model_sizes(engine) else None + try: + assert_generation_workers_available() + except GenerationWorkersUnavailable as error: + raise HTTPException(status_code=503, detail=str(error)) from error + text = data.text source = "manual" if data.personality and getattr(profile, "personality", None): @@ -123,9 +164,11 @@ async def generate_speech( except Exception: pass - enqueue_generation( - generation_id, - run_generation( + await _enqueue_generation_or_reject( + db=db, + task_manager=task_manager, + generation_id=generation_id, + generation_coro=run_generation( generation_id=generation_id, profile_id=data.profile_id, text=text, @@ -139,7 +182,8 @@ async def generate_speech( mode="generate", max_chunk_chars=data.max_chunk_chars, crossfade_ms=data.crossfade_ms, - ) + ), + engine=engine, ) return generation @@ -169,9 +213,11 @@ async def retry_generation(generation_id: str, db: Session = Depends(get_db)): text=gen.text, ) - enqueue_generation( - generation_id, - run_generation( + await _enqueue_generation_or_reject( + db=db, + task_manager=task_manager, + generation_id=generation_id, + generation_coro=run_generation( generation_id=generation_id, profile_id=gen.profile_id, text=gen.text, @@ -181,7 +227,8 @@ async def retry_generation(generation_id: str, db: Session = Depends(get_db)): seed=gen.seed, instruct=gen.instruct, mode="retry", - ) + ), + engine=gen.engine or "qwen", ) return models.GenerationResponse.model_validate(gen) @@ -213,9 +260,11 @@ async def regenerate_generation(generation_id: str, db: Session = Depends(get_db version_id = str(uuid.uuid4()) - enqueue_generation( - generation_id, - run_generation( + await _enqueue_generation_or_reject( + db=db, + task_manager=task_manager, + generation_id=generation_id, + generation_coro=run_generation( generation_id=generation_id, profile_id=gen.profile_id, text=gen.text, @@ -226,7 +275,8 @@ async def regenerate_generation(generation_id: str, db: Session = Depends(get_db instruct=gen.instruct, mode="regenerate", version_id=version_id, - ) + ), + engine=gen.engine or "qwen", ) return models.GenerationResponse.model_validate(gen) diff --git a/backend/routes/health.py b/backend/routes/health.py index 1568455..7eee0ae 100644 --- a/backend/routes/health.py +++ b/backend/routes/health.py @@ -12,6 +12,7 @@ from sqlalchemy.orm import Session from .. import config, models from ..services import tts +from ..services import task_queue from ..database import get_db from ..utils.platform_detect import get_backend_type, is_amd_gpu_windows @@ -176,6 +177,16 @@ async def health(): elif has_xpu: default_variant = "xpu" + turbo_backend = None + try: + from ..backends import get_tts_backend_for_engine + + turbo_backend = get_tts_backend_for_engine("chatterbox_turbo") + except Exception: + # Health remains available during optional backend import failures, + # but reports zero loaded replicas so clients fail closed to one lane. + turbo_backend = None + return models.HealthResponse( status="healthy", model_loaded=model_loaded, @@ -188,6 +199,7 @@ async def health(): backend_variant=os.environ.get("VOICEBOX_BACKEND_VARIANT", default_variant), supports_rocm=is_amd_gpu_windows(), gpu_compatibility_warning=gpu_compat_warning, + generation_capacity=task_queue.generation_capacity_snapshot(turbo_backend), ) diff --git a/backend/services/task_queue.py b/backend/services/task_queue.py index 3ec4237..9a4202c 100644 --- a/backend/services/task_queue.py +++ b/backend/services/task_queue.py @@ -1,9 +1,8 @@ -""" -Serial generation queue — ensures only one TTS inference runs at a time -to avoid GPU contention. -""" +"""Bounded, engine-aware generation worker pool.""" import asyncio +import logging +import os import traceback from dataclasses import dataclass from typing import Coroutine, Literal @@ -18,14 +17,160 @@ class GenerationJob: generation_id: str coro: Coroutine + engine: str -# Generation queue — serializes TTS inference to avoid GPU contention +# Generation queue — bounded worker pool. Backend-specific capacity controls +# whether jobs for one engine may execute physically in parallel. _generation_queue: asyncio.Queue = None # type: ignore # initialized at startup -_generation_worker_task: asyncio.Task | None = None +_generation_worker_tasks: set[asyncio.Task] = set() _queued_generation_ids: set[str] = set() _running_generation_tasks: dict[str, asyncio.Task] = {} +_running_generation_engines: dict[str, str] = {} _cancelled_generation_ids: set[str] = set() +_backend_entered_generation_ids: set[str] = set() +_engine_generation_semaphores: dict[str, asyncio.Semaphore] = {} +_generation_worker_epoch = 0 +_generation_worker_failures = 0 +_generation_worker_restarts = 0 +_generation_worker_pool_stopping = True + +logger = logging.getLogger(__name__) + + +class GenerationQueueFull(RuntimeError): + """Raised before admission when the bounded generation queue is full.""" + + +class GenerationWorkersUnavailable(RuntimeError): + """Raised before admission when no live worker can own a generation.""" + + +def generation_queue_capacity() -> int: + return _bounded_int_env("VOICEBOX_MAX_PENDING_GENERATIONS", 64, 1, 512) + + +def _bounded_int_env(name: str, fallback: int, minimum: int, maximum: int) -> int: + try: + configured = int(os.environ.get(name, str(fallback)) or fallback) + except (TypeError, ValueError): + configured = fallback + return min(maximum, max(minimum, configured)) + + +def generation_worker_count() -> int: + return _bounded_int_env("VOICEBOX_GENERATION_WORKERS", 2, 1, 8) + + +def _live_generation_workers() -> list[asyncio.Task]: + return [task for task in _generation_worker_tasks if not task.done()] + + +def generation_workers_ready() -> bool: + return ( + not _generation_worker_pool_stopping + and _generation_queue is not None + and len(_live_generation_workers()) > 0 + ) + + +def assert_generation_workers_available() -> int: + """Restore supervised capacity, then fail before persistence if none exists.""" + _ensure_generation_workers(restarted=True) + live = len(_live_generation_workers()) + if not generation_workers_ready(): + raise GenerationWorkersUnavailable( + f"Voicebox generation workers are unavailable " + f"({live}/{generation_worker_count()} live)" + ) + return live + + +def engine_generation_capacity(engine: str) -> int: + normalized = str(engine or "unknown").strip().lower() + if normalized == "chatterbox_turbo": + return _bounded_int_env( + "VOICEBOX_CHATTERBOX_TURBO_REPLICAS", + 2, + 1, + 4, + ) + # Other upstream backends still expose singleton model objects. They may + # run alongside another engine, but same-engine overlap remains closed + # until that backend implements independent replicas of its own. + return 1 + + +def engine_supports_physical_task_cancellation(engine: str) -> bool: + """Whether cancellation is proven to retain all native load/generate leases.""" + return str(engine or "").strip().lower() == "chatterbox_turbo" + + +def generation_capacity_snapshot(turbo_backend=None) -> dict: + turbo_capacity = { + "configured": engine_generation_capacity("chatterbox_turbo"), + "loaded": 0, + "active": 0, + "available": 0, + } + capacity_reader = getattr(turbo_backend, "generation_capacity", None) + if callable(capacity_reader): + reported = capacity_reader() + if isinstance(reported, dict): + for key in turbo_capacity: + value = reported.get(key) + if isinstance(value, int) and value >= 0: + turbo_capacity[key] = value + live_workers = len(_live_generation_workers()) + configured_workers = generation_worker_count() + return { + "workers": { + "configured": configured_workers, + "live": live_workers, + # Retain the old field for clients deployed before the liveness + # contract was expanded. It means live worker tasks, not busy jobs. + "active": live_workers, + "failed": _generation_worker_failures, + "restarts": _generation_worker_restarts, + "restarting": ( + max(0, configured_workers - live_workers) + if not _generation_worker_pool_stopping + else 0 + ), + "ready": generation_workers_ready(), + "stopping": _generation_worker_pool_stopping, + }, + "queue": { + "queued": len(_queued_generation_ids), + "running": len(_running_generation_tasks), + "max_pending": generation_queue_capacity(), + }, + "chatterbox_turbo": turbo_capacity, + } + + +def _engine_semaphore(engine: str) -> asyncio.Semaphore: + normalized = str(engine or "unknown").strip().lower() or "unknown" + semaphore = _engine_generation_semaphores.get(normalized) + if semaphore is None: + semaphore = asyncio.Semaphore(engine_generation_capacity(normalized)) + _engine_generation_semaphores[normalized] = semaphore + return semaphore + + +async def _run_generation_job(job: GenerationJob) -> None: + entered_backend = False + try: + async with _engine_semaphore(job.engine): + entered_backend = True + _backend_entered_generation_ids.add(job.generation_id) + await job.coro + finally: + _backend_entered_generation_ids.discard(job.generation_id) + # A coroutine object that was cancelled while waiting for engine + # capacity was never awaited and must be explicitly closed. + if not entered_backend: + job.coro.close() def create_background_task(coro) -> asyncio.Task: @@ -46,12 +191,22 @@ async def _generation_worker(): job.coro.close() continue - task = asyncio.create_task(job.coro) + task = asyncio.create_task(_run_generation_job(job)) _running_generation_tasks[job.generation_id] = task + _running_generation_engines[job.generation_id] = str( + job.engine or "unknown" + ).strip().lower() _queued_generation_ids.discard(job.generation_id) try: await task except asyncio.CancelledError: + # Backends that dispatch worker threads drain those workers + # before allowing this task to reach its cancelled state. + # Preserve cancellation of the worker itself during pool + # shutdown; suppress only cancellation of the child job. + worker = asyncio.current_task() + if worker is not None and worker.cancelling(): + raise if not task.cancelled(): raise except Exception: @@ -62,10 +217,63 @@ async def _generation_worker(): ) finally: _running_generation_tasks.pop(job.generation_id, None) + _running_generation_engines.pop(job.generation_id, None) _queued_generation_ids.discard(job.generation_id) + _cancelled_generation_ids.discard(job.generation_id) _generation_queue.task_done() +def _generation_worker_finished(task: asyncio.Task, epoch: int) -> None: + """Replace a worker that exits outside an intentional pool generation.""" + global _generation_worker_failures + + _generation_worker_tasks.discard(task) + if epoch != _generation_worker_epoch or _generation_worker_pool_stopping: + return + _generation_worker_failures += 1 + if task.cancelled(): + logger.error("Voicebox generation worker was cancelled unexpectedly") + else: + error = task.exception() + if error is None: + logger.error("Voicebox generation worker exited unexpectedly") + else: + logger.exception( + "Voicebox generation worker crashed", + exc_info=(type(error), error, error.__traceback__), + ) + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + if loop.is_running(): + loop.call_soon(_ensure_generation_workers, True) + + +def _start_generation_worker(*, restarted: bool) -> None: + global _generation_worker_restarts + + epoch = _generation_worker_epoch + worker = create_background_task(_generation_worker()) + _generation_worker_tasks.add(worker) + worker.add_done_callback( + lambda completed, worker_epoch=epoch: _generation_worker_finished( + completed, + worker_epoch, + ) + ) + if restarted: + _generation_worker_restarts += 1 + + +def _ensure_generation_workers(restarted: bool = False) -> None: + if _generation_worker_pool_stopping or _generation_queue is None: + return + missing = generation_worker_count() - len(_live_generation_workers()) + for _ in range(max(0, missing)): + _start_generation_worker(restarted=restarted) + + async def _force_fail_if_active(generation_id: str, error: str) -> None: """Best-effort recovery — flip an active row to failed if the worker bailed before writing a terminal status. Catches the case where the gen @@ -93,20 +301,47 @@ async def _force_fail_if_active(generation_id: str, error: str) -> None: traceback.print_exc() -def enqueue_generation(generation_id: str, coro): - """Add a generation coroutine to the serial queue.""" - if _generation_queue is None: - raise RuntimeError("Generation queue has not been initialized") +def enqueue_generation(generation_id: str, coro, engine: str = "unknown"): + """Add generation work to the bounded, engine-aware worker pool.""" + try: + assert_generation_workers_available() + except GenerationWorkersUnavailable: + coro.close() + raise + job = GenerationJob( + generation_id=generation_id, + coro=coro, + engine=str(engine or "unknown"), + ) + try: + _generation_queue.put_nowait(job) + except asyncio.QueueFull as error: + # The coroutine was constructed by the route before admission. Close + # it explicitly so overload cannot leak an un-awaited coroutine. + job.coro.close() + raise GenerationQueueFull( + f"Voicebox generation queue is full ({generation_queue_capacity()} waiting)" + ) from error _queued_generation_ids.add(generation_id) - _generation_queue.put_nowait(GenerationJob(generation_id=generation_id, coro=coro)) def cancel_generation(generation_id: str) -> Literal["queued", "running"] | None: """Cancel a queued or running generation if it is still active.""" running_task = _running_generation_tasks.get(generation_id) if running_task is not None: - running_task.cancel() + entered_backend = generation_id in _backend_entered_generation_ids + engine = _running_generation_engines.get(generation_id, "unknown") + if not entered_backend: + _cancelled_generation_ids.add(generation_id) + running_task.cancel() + return "queued" + if engine_supports_physical_task_cancellation(engine): + _cancelled_generation_ids.add(generation_id) + running_task.cancel() + # Unsafe singleton backends intentionally keep running to terminal + # completion. Their semaphore remains physically leased, so a cancel + # request can never overlap an uninterruptible asyncio.to_thread call. return "running" if generation_id in _queued_generation_ids: @@ -122,18 +357,48 @@ def init_queue(force: bool = False): Must be called once during application startup (inside a running event loop). """ - global _generation_queue, _generation_worker_task - global _queued_generation_ids, _running_generation_tasks, _cancelled_generation_ids + global _generation_queue, _generation_worker_tasks + global _queued_generation_ids, _running_generation_tasks, _running_generation_engines + global _cancelled_generation_ids + global _backend_entered_generation_ids, _engine_generation_semaphores + global _generation_worker_epoch, _generation_worker_pool_stopping - if _generation_worker_task is not None and not _generation_worker_task.done(): + live_workers = _live_generation_workers() + if live_workers: if not force: + _generation_worker_pool_stopping = False + _ensure_generation_workers() return - _generation_worker_task.cancel() - for task in list(_running_generation_tasks.values()): - task.cancel() + if any(not task.done() for task in _running_generation_tasks.values()): + raise RuntimeError( + "Cannot reinitialize generation workers while physical jobs are active" + ) + _generation_worker_pool_stopping = True + _generation_worker_epoch += 1 + for worker in live_workers: + worker.cancel() - _generation_queue = asyncio.Queue() + _generation_queue = asyncio.Queue(maxsize=generation_queue_capacity()) _queued_generation_ids = set() _running_generation_tasks = {} + _running_generation_engines = {} _cancelled_generation_ids = set() - _generation_worker_task = create_background_task(_generation_worker()) + _backend_entered_generation_ids = set() + _engine_generation_semaphores = {} + _generation_worker_tasks = set() + _generation_worker_pool_stopping = False + _ensure_generation_workers() + + +async def shutdown_queue() -> None: + """Stop workers intentionally without allowing the supervisor to replace them.""" + global _generation_worker_epoch, _generation_worker_pool_stopping + + _generation_worker_pool_stopping = True + _generation_worker_epoch += 1 + workers = list(_generation_worker_tasks) + for worker in workers: + worker.cancel() + if workers: + await asyncio.gather(*workers, return_exceptions=True) + _generation_worker_tasks.clear()