5e2094f2ad
FastAPI service in infra/media-worker implementing every MediaInferenceProvider operation behind HTTP API v1, plus capabilities. Bearer auth, MEDIA_ROOT realpath enforcement, strict request schemas, cancellation on disconnect, GPU/CPU gates with 503 overflow, input size/duration limits, lazy SHA256-pinned model downloads, fail-closed redaction with atomic output, offline pytest suite with fake backends, CPU/CUDA/ROCm Dockerfile and a license review of every model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0138YqDj3ttfayRwuMKKrqVj
94 lines
2.8 KiB
Python
94 lines
2.8 KiB
Python
"""Request execution: concurrency gates, cancellation on disconnect, time limit."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
from collections.abc import Callable
|
|
from typing import Any, Protocol
|
|
|
|
from .errors import Busy, Cancelled
|
|
|
|
POLL_SECONDS = 0.25
|
|
|
|
|
|
class Gate:
|
|
"""`concurrency` jobs run; up to `queue_size` more wait; anything beyond is busy.
|
|
|
|
Thread-based so it works no matter which event loop serves the request.
|
|
"""
|
|
|
|
def __init__(self, concurrency: int, queue_size: int) -> None:
|
|
self.concurrency = concurrency
|
|
self.capacity = concurrency + queue_size
|
|
self._slots = threading.Semaphore(concurrency)
|
|
self._lock = threading.Lock()
|
|
self.inflight = 0
|
|
|
|
def reserve(self) -> None:
|
|
with self._lock:
|
|
if self.inflight >= self.capacity:
|
|
raise Busy()
|
|
self.inflight += 1
|
|
|
|
def release_reservation(self) -> None:
|
|
with self._lock:
|
|
self.inflight -= 1
|
|
|
|
def run(self, fn: Callable[[threading.Event], Any], cancel: threading.Event) -> Any:
|
|
"""Wait for a slot (reservation already taken), run fn, free both."""
|
|
try:
|
|
while not self._slots.acquire(timeout=POLL_SECONDS):
|
|
if cancel.is_set():
|
|
raise Cancelled()
|
|
try:
|
|
if cancel.is_set():
|
|
raise Cancelled()
|
|
return fn(cancel)
|
|
finally:
|
|
self._slots.release()
|
|
finally:
|
|
self.release_reservation()
|
|
|
|
|
|
class DisconnectProbe(Protocol):
|
|
async def is_disconnected(self) -> bool: ...
|
|
|
|
|
|
async def run_job(
|
|
request: DisconnectProbe,
|
|
gate: Gate,
|
|
fn: Callable[[threading.Event], Any],
|
|
max_seconds: float,
|
|
) -> Any:
|
|
"""Run fn in a worker thread. Sets the cancel event when the client
|
|
disconnects or the deadline passes; fn checks it between frames/chunks."""
|
|
gate.reserve()
|
|
cancel = threading.Event()
|
|
task = asyncio.ensure_future(asyncio.to_thread(gate.run, fn, cancel))
|
|
deadline = time.monotonic() + max_seconds
|
|
reason: str | None = None
|
|
try:
|
|
while True:
|
|
done, _ = await asyncio.wait({task}, timeout=POLL_SECONDS)
|
|
if done:
|
|
break
|
|
if reason is None:
|
|
if time.monotonic() > deadline:
|
|
reason = "exceeded MEDIA_MAX_SECONDS"
|
|
cancel.set()
|
|
elif await request.is_disconnected():
|
|
reason = "client disconnected"
|
|
cancel.set()
|
|
except asyncio.CancelledError:
|
|
cancel.set()
|
|
raise
|
|
try:
|
|
result = task.result()
|
|
except Cancelled as error:
|
|
raise Cancelled(reason or error.reason) from error
|
|
if reason is not None:
|
|
raise Cancelled(reason)
|
|
return result
|