Spaces:
Build error
Build error
| """ | |
| Background analysis jobs: a fair queue in front of the CPU, real per-step | |
| progress, cancellation, and a small result cache. | |
| POST /api/analyze/jobs starts a job and returns its id; the client polls | |
| GET /api/analyze/jobs/{id}. The Space has two vCPUs and one analysis | |
| already keeps both busy, so jobs take the CPU one at a time (CPU_SLOTS) in | |
| arrival order and everyone else sees their place in the queue. Network | |
| waits (YouTube download, the external FST model) happen outside that slot. | |
| A job belongs to whoever polls it. One that nobody has polled for a while | |
| when its turn comes (the tab was closed) is dropped instead of analysed. | |
| A reloaded page resumes polling the same id, so it isn't dropped. | |
| State lives in process memory: the Space runs one replica, and a job only | |
| has to outlive the request that started it, not a restart. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import os | |
| import time | |
| import uuid | |
| from collections import OrderedDict | |
| from contextlib import asynccontextmanager | |
| from dataclasses import dataclass, field | |
| from typing import Any, AsyncIterator, Awaitable, Callable, Dict, List, Optional | |
| JOB_TTL_SEC = 15 * 60 | |
| MAX_JOBS = 200 | |
| # Jobs waiting or running at once; beyond this, new ones get 503. | |
| MAX_PENDING = int(os.getenv("AURIS_MAX_PENDING", "24")) | |
| CPU_SLOTS = int(os.getenv("AURIS_CPU_SLOTS", "1")) | |
| # Unpolled this long when its turn comes: the visitor has gone. | |
| ABANDON_SEC = 180.0 | |
| CACHE_SIZE = 64 | |
| CACHE_TTL_SEC = 6 * 3600 | |
| STEP_PENDING = "pending" | |
| STEP_RUNNING = "running" | |
| STEP_DONE = "done" | |
| STEP_FAILED = "failed" | |
| STEP_SKIPPED = "skipped" | |
| FILE_STEPS = ["features", "vocals", "wav2vec2", "clap", "timeline", "fst", "xai", "meta"] | |
| URL_STEPS = ["download", *FILE_STEPS] | |
| class JobCancelled(Exception): | |
| """Raised inside a pipeline when its job was cancelled or abandoned.""" | |
| def __init__(self, code: str = "cancelled") -> None: | |
| super().__init__(code) | |
| self.code = code | |
| class Progress: | |
| """Ordered step states for one analysis, plus its cancel flag.""" | |
| def __init__(self, steps: List[str]) -> None: | |
| self._steps: Dict[str, str] = {s: STEP_PENDING for s in steps} | |
| self._started: Dict[str, float] = {} | |
| self._took: Dict[str, float] = {} | |
| # Steps made of many parts (the timeline scan) report how far along they are. | |
| self._fraction: Dict[str, float] = {} | |
| self.cancelled = False | |
| # queued: waiting for the CPU; running: downloading or analysing. | |
| self.phase = "queued" | |
| # Returns True when nobody is waiting for the result any more. | |
| self.gone: Callable[[], Awaitable[bool]] = _never | |
| def set(self, step: str, state: str) -> None: | |
| if step not in self._steps: | |
| return | |
| now = time.monotonic() | |
| if state == STEP_RUNNING: | |
| self._started.setdefault(step, now) | |
| elif step in self._started and step not in self._took: | |
| self._took[step] = round(now - self._started[step], 2) | |
| self._steps[step] = state | |
| def advance(self, step: str, done: int, total: int) -> None: | |
| if step in self._steps and total > 0: | |
| self._fraction[step] = round(min(1.0, done / total), 3) | |
| def running(self, *steps: str) -> None: | |
| for s in steps: | |
| self.set(s, STEP_RUNNING) | |
| def skip_pending(self) -> None: | |
| for step, state in self._steps.items(): | |
| if state in (STEP_PENDING, STEP_RUNNING): | |
| self._steps[step] = STEP_SKIPPED | |
| def check(self) -> None: | |
| if self.cancelled: | |
| raise JobCancelled("cancelled") | |
| def snapshot(self) -> List[Dict[str, Any]]: | |
| out: List[Dict[str, Any]] = [] | |
| for step, state in self._steps.items(): | |
| item: Dict[str, Any] = {"id": step, "state": state} | |
| if step in self._took: | |
| item["seconds"] = self._took[step] | |
| if step in self._fraction and state == STEP_RUNNING: | |
| item["progress"] = self._fraction[step] | |
| out.append(item) | |
| return out | |
| async def _never() -> bool: | |
| return False | |
| class CpuGate: | |
| """First come, first served access to the CPU-heavy stage.""" | |
| def __init__(self, slots: int) -> None: | |
| self._slots = max(1, slots) | |
| self._busy = 0 | |
| self._waiting: List[str] = [] | |
| self._events: Dict[str, asyncio.Event] = {} | |
| # Mean seconds a job holds the CPU; None until one has finished. | |
| self._mean: Optional[float] = None | |
| def position(self, ticket: str) -> Optional[int]: | |
| """1-based place in the queue, or None when not waiting.""" | |
| try: | |
| return self._waiting.index(ticket) + 1 | |
| except ValueError: | |
| return None | |
| def estimated_wait(self, position: int) -> Optional[float]: | |
| if self._mean is None: | |
| return None | |
| # Everyone ahead, plus whatever is on the CPU now (half done on average). | |
| return round(((position - 1) / self._slots + (0.5 if self._busy else 0)) * self._mean, 1) | |
| def waiting(self) -> int: | |
| return len(self._waiting) | |
| def busy(self) -> int: | |
| return self._busy | |
| def _wake(self) -> None: | |
| while self._waiting and self._busy < self._slots: | |
| ticket = self._waiting.pop(0) | |
| self._busy += 1 | |
| self._events.pop(ticket).set() | |
| async def slot(self, ticket: str, progress: Optional[Progress] = None) -> AsyncIterator[None]: | |
| if progress is not None: | |
| progress.phase = "queued" | |
| event = asyncio.Event() | |
| self._events[ticket] = event | |
| self._waiting.append(ticket) | |
| self._wake() | |
| try: | |
| await event.wait() | |
| except BaseException: | |
| if ticket in self._waiting: | |
| self._waiting.remove(ticket) | |
| self._events.pop(ticket, None) | |
| else: | |
| self._release() | |
| raise | |
| started = time.monotonic() | |
| try: | |
| if progress is not None: | |
| progress.phase = "running" | |
| progress.check() | |
| if await progress.gone(): | |
| raise JobCancelled("job_abandoned") | |
| yield | |
| finally: | |
| held = time.monotonic() - started | |
| if held > 1.0: | |
| self._mean = held if self._mean is None else self._mean * 0.7 + held * 0.3 | |
| self._release() | |
| def _release(self) -> None: | |
| self._busy = max(0, self._busy - 1) | |
| self._wake() | |
| cpu_gate = CpuGate(CPU_SLOTS) | |
| class ResultCache: | |
| """Recent successful responses by content key (same file, same link).""" | |
| def __init__(self, size: int, ttl: float) -> None: | |
| self._size = size | |
| self._ttl = ttl | |
| self._items: "OrderedDict[str, tuple[float, Dict[str, Any]]]" = OrderedDict() | |
| def get(self, key: Optional[str]) -> Optional[Dict[str, Any]]: | |
| if not key or key not in self._items: | |
| return None | |
| stored, value = self._items[key] | |
| if time.time() - stored > self._ttl: | |
| del self._items[key] | |
| return None | |
| self._items.move_to_end(key) | |
| return value | |
| def put(self, key: Optional[str], value: Dict[str, Any]) -> None: | |
| if not key or value.get("errors") or not value.get("result"): | |
| return | |
| self._items[key] = (time.time(), value) | |
| self._items.move_to_end(key) | |
| while len(self._items) > self._size: | |
| self._items.popitem(last=False) | |
| result_cache = ResultCache(CACHE_SIZE, CACHE_TTL_SEC) | |
| class Job: | |
| id: str | |
| source_type: str | |
| progress: Progress | |
| key: Optional[str] = None | |
| created: float = field(default_factory=time.time) | |
| updated: float = field(default_factory=time.time) | |
| last_seen: float = field(default_factory=time.monotonic) | |
| status: str = "active" # active | done | error | |
| cached: bool = False | |
| response: Optional[Dict[str, Any]] = None | |
| task: Optional[asyncio.Task] = None | |
| def finished(self) -> bool: | |
| return self.status in ("done", "error") | |
| def to_dict(self) -> Dict[str, Any]: | |
| data: Dict[str, Any] = { | |
| "jobId": self.id, | |
| "sourceType": self.source_type, | |
| # "running" until settled, whether waiting or working, so clients | |
| # that predate the queue keep polling; `phase` tells them apart. | |
| "status": self.status if self.finished else "running", | |
| "phase": self.status if self.finished else self.progress.phase, | |
| "steps": self.progress.snapshot(), | |
| "elapsedSec": round((self.updated if self.finished else time.time()) - self.created, 2), | |
| "response": self.response, | |
| } | |
| position = cpu_gate.position(self.id) | |
| if not self.finished and position is not None: | |
| data["queue"] = { | |
| "position": position, | |
| # Everyone in line before you, plus whoever is on the CPU now. | |
| "ahead": position - 1 + cpu_gate.busy, | |
| "estimatedWaitSec": cpu_gate.estimated_wait(position), | |
| } | |
| if self.cached: | |
| data["cached"] = True | |
| return data | |
| class QueueFull(Exception): | |
| pass | |
| Work = Callable[[Job], Awaitable[Dict[str, Any]]] | |
| class JobStore: | |
| def __init__(self) -> None: | |
| self._jobs: Dict[str, Job] = {} | |
| self._inflight: Dict[str, str] = {} | |
| def _prune(self) -> None: | |
| cutoff = time.time() - JOB_TTL_SEC | |
| for job_id in [j.id for j in self._jobs.values() if j.finished and j.updated < cutoff]: | |
| del self._jobs[job_id] | |
| if len(self._jobs) > MAX_JOBS: | |
| finished = sorted((j for j in self._jobs.values() if j.finished), key=lambda j: j.updated) | |
| for job in finished[: len(self._jobs) - MAX_JOBS]: | |
| del self._jobs[job.id] | |
| def pending_count(self) -> int: | |
| return sum(1 for j in self._jobs.values() if not j.finished) | |
| def running_count(self) -> int: | |
| return sum(1 for j in self._jobs.values() if not j.finished and j.progress.phase == "running") | |
| def start(self, source_type: str, steps: List[str], work: Work, key: Optional[str] = None) -> Job: | |
| """Queue a job; the same file or link already in flight is shared, a recent one is served from cache.""" | |
| self._prune() | |
| if key and key in self._inflight: | |
| shared = self._jobs.get(self._inflight[key]) | |
| if shared and not shared.finished: | |
| shared.last_seen = time.monotonic() | |
| return shared | |
| job = Job(id=uuid.uuid4().hex, source_type=source_type, progress=Progress(steps), key=key) | |
| hit = result_cache.get(key) | |
| if hit is not None: | |
| job.status, job.cached, job.response = "done", True, hit | |
| job.progress.skip_pending() | |
| for step in steps: | |
| job.progress.set(step, STEP_DONE) | |
| self._jobs[job.id] = job | |
| return job | |
| if self.pending_count() >= MAX_PENDING: | |
| raise QueueFull() | |
| async def gone() -> bool: | |
| return time.monotonic() - job.last_seen > ABANDON_SEC | |
| job.progress.gone = gone | |
| self._jobs[job.id] = job | |
| if key: | |
| self._inflight[key] = job.id | |
| async def runner() -> None: | |
| try: | |
| job.response = await work(job) | |
| failed = bool(job.response.get("errors")) and not job.response.get("result") | |
| job.status = "error" if failed else "done" | |
| result_cache.put(key, job.response) | |
| except JobCancelled as stop: | |
| job.progress.skip_pending() | |
| job.response = {"result": None, "warnings": [], "errors": [stop.code]} | |
| job.status = "error" | |
| except asyncio.CancelledError: | |
| # Cancelled while still waiting for the CPU. | |
| job.progress.skip_pending() | |
| job.response = {"result": None, "warnings": [], "errors": ["cancelled"]} | |
| job.status = "error" | |
| except Exception: # noqa: BLE001 — the job must always settle | |
| job.progress.skip_pending() | |
| job.response = {"result": None, "warnings": [], "errors": ["internal_error"]} | |
| job.status = "error" | |
| finally: | |
| job.updated = time.time() | |
| if key and self._inflight.get(key) == job.id: | |
| del self._inflight[key] | |
| job.task = None | |
| # Keep a reference: asyncio only holds weak references to tasks. | |
| job.task = asyncio.create_task(runner()) | |
| return job | |
| def get(self, job_id: str, touch: bool = True) -> Optional[Job]: | |
| self._prune() | |
| job = self._jobs.get(job_id) | |
| if job is not None and touch: | |
| job.last_seen = time.monotonic() | |
| return job | |
| def cancel(self, job_id: str) -> Optional[Job]: | |
| job = self._jobs.get(job_id) | |
| if job is None: | |
| return None | |
| if not job.finished: | |
| job.progress.cancelled = True | |
| # A job still waiting for the CPU never gets there. | |
| if cpu_gate.position(job.id) is not None and job.task is not None: | |
| job.task.cancel() | |
| return job | |
| job_store = JobStore() | |