Spaces:
Sleeping
Sleeping
| """Serial async job queue for GPU-bound work. | |
| Why: TTS inference, diarization, demucs, Whisper, ffmpeg muxing all contend | |
| for the same GPU/CPU/fd budget. Without a gate, a second dub request can | |
| OOM VRAM or trigger posix_spawn EAGAIN on macOS. This module provides a | |
| single serialized worker with cancel semantics and queue-position | |
| introspection so the UI can tell users "2 jobs ahead of you". | |
| Complements `core.tasks.task_manager`: | |
| * `task_manager` — SSE event streaming (per-task history + listeners). | |
| * `job_queue` — resource gating, cancellation, position tracking. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import enum | |
| import logging | |
| import time | |
| import uuid | |
| from dataclasses import dataclass, field | |
| from typing import Any, Awaitable, Callable, Dict, List, Optional | |
| logger = logging.getLogger("omnivoice.jobs") | |
| class JobState(str, enum.Enum): | |
| QUEUED = "queued" | |
| RUNNING = "running" | |
| DONE = "done" | |
| FAILED = "failed" | |
| CANCELLED = "cancelled" | |
| class Job: | |
| id: str | |
| label: str | |
| fn: Callable[["Job"], Awaitable[Any]] | |
| state: JobState = JobState.QUEUED | |
| created_at: float = field(default_factory=time.time) | |
| started_at: Optional[float] = None | |
| finished_at: Optional[float] = None | |
| error: Optional[str] = None | |
| result: Any = None | |
| cancel_event: asyncio.Event = field(default_factory=asyncio.Event) | |
| done_event: asyncio.Event = field(default_factory=asyncio.Event) | |
| def is_cancelled(self) -> bool: | |
| return self.cancel_event.is_set() | |
| def to_dict(self) -> dict: | |
| return { | |
| "id": self.id, | |
| "label": self.label, | |
| "state": self.state.value, | |
| "created_at": self.created_at, | |
| "started_at": self.started_at, | |
| "finished_at": self.finished_at, | |
| "error": self.error, | |
| } | |
| class JobQueue: | |
| """Single-lane serial async worker with cancellation and introspection.""" | |
| def __init__(self, name: str = "jobs", concurrency: int = 1) -> None: | |
| self.name = name | |
| self._concurrency = max(1, concurrency) | |
| self._queue: asyncio.Queue[str] = asyncio.Queue() | |
| self._jobs: Dict[str, Job] = {} | |
| self._order: List[str] = [] | |
| self._workers: List[asyncio.Task] = [] | |
| self._lock = asyncio.Lock() | |
| async def start(self) -> None: | |
| if self._workers: | |
| return | |
| self._workers = [ | |
| asyncio.create_task(self._run_worker(i), name=f"{self.name}-worker-{i}") | |
| for i in range(self._concurrency) | |
| ] | |
| async def stop(self) -> None: | |
| for w in self._workers: | |
| w.cancel() | |
| if self._workers: | |
| await asyncio.gather(*self._workers, return_exceptions=True) | |
| self._workers = [] | |
| async def submit( | |
| self, | |
| label: str, | |
| fn: Callable[[Job], Awaitable[Any]], | |
| job_id: Optional[str] = None, | |
| ) -> Job: | |
| """Enqueue a coroutine-returning callable. Idempotent on job_id.""" | |
| jid = job_id or uuid.uuid4().hex[:12] | |
| async with self._lock: | |
| existing = self._jobs.get(jid) | |
| if existing is not None: | |
| return existing | |
| job = Job(id=jid, label=label, fn=fn) | |
| self._jobs[jid] = job | |
| self._order.append(jid) | |
| await self._queue.put(jid) | |
| return job | |
| def get(self, job_id: str) -> Optional[Job]: | |
| return self._jobs.get(job_id) | |
| def position(self, job_id: str) -> int: | |
| """Return 0-indexed position among queued jobs, or -1 if not queued.""" | |
| job = self._jobs.get(job_id) | |
| if not job or job.state != JobState.QUEUED: | |
| return -1 | |
| idx = 0 | |
| for jid in self._order: | |
| other = self._jobs.get(jid) | |
| if not other or other.state != JobState.QUEUED: | |
| continue | |
| if jid == job_id: | |
| return idx | |
| idx += 1 | |
| return -1 | |
| def list_jobs(self, include_finished: bool = False) -> List[Job]: | |
| if include_finished: | |
| return [self._jobs[j] for j in self._order if j in self._jobs] | |
| live = {JobState.QUEUED, JobState.RUNNING} | |
| return [self._jobs[j] for j in self._order if j in self._jobs and self._jobs[j].state in live] | |
| async def cancel(self, job_id: str) -> bool: | |
| job = self._jobs.get(job_id) | |
| if not job: | |
| return False | |
| if job.state == JobState.QUEUED: | |
| job.state = JobState.CANCELLED | |
| job.finished_at = time.time() | |
| job.cancel_event.set() | |
| job.done_event.set() | |
| return True | |
| if job.state == JobState.RUNNING: | |
| # Cooperative: worker polls `job.is_cancelled` between steps. | |
| job.cancel_event.set() | |
| return True | |
| return False | |
| async def wait(self, job_id: str, timeout: Optional[float] = None) -> Job: | |
| job = self._jobs.get(job_id) | |
| if not job: | |
| raise KeyError(job_id) | |
| if timeout is None: | |
| await job.done_event.wait() | |
| else: | |
| await asyncio.wait_for(job.done_event.wait(), timeout=timeout) | |
| return job | |
| def purge_finished(self, older_than_seconds: float = 3600.0) -> int: | |
| """Drop rows for jobs finished more than `older_than_seconds` ago.""" | |
| cutoff = time.time() - older_than_seconds | |
| drop = [ | |
| jid for jid in self._order | |
| if (j := self._jobs.get(jid)) | |
| and j.state in (JobState.DONE, JobState.FAILED, JobState.CANCELLED) | |
| and (j.finished_at or 0.0) < cutoff | |
| ] | |
| for jid in drop: | |
| self._jobs.pop(jid, None) | |
| self._order = [j for j in self._order if j in self._jobs] | |
| return len(drop) | |
| async def _run_worker(self, idx: int) -> None: | |
| while True: | |
| jid = await self._queue.get() | |
| try: | |
| job = self._jobs.get(jid) | |
| if job is None: | |
| continue | |
| if job.state == JobState.CANCELLED: | |
| # Cancelled while queued — skip execution. | |
| continue | |
| job.state = JobState.RUNNING | |
| job.started_at = time.time() | |
| try: | |
| result = await job.fn(job) | |
| if job.is_cancelled: | |
| job.state = JobState.CANCELLED | |
| else: | |
| job.state = JobState.DONE | |
| job.result = result | |
| except asyncio.CancelledError: | |
| job.state = JobState.CANCELLED | |
| raise | |
| except Exception as e: | |
| logger.exception("Job %s (%s) failed", job.id, job.label) | |
| job.state = JobState.FAILED | |
| job.error = str(e) | |
| finally: | |
| job.finished_at = time.time() | |
| job.done_event.set() | |
| finally: | |
| self._queue.task_done() | |
| # Shared GPU-bound queue — single lane, protects TTS/ASR/demucs from VRAM contention. | |
| gpu_queue = JobQueue(name="gpu", concurrency=1) | |