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
118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
"""Cancellation on disconnect, the time limit, and GPU queue overflow."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from app import media_io
|
|
from app.errors import Busy, Cancelled
|
|
from app.runtime import Gate, run_job
|
|
|
|
from .conftest import AUTH
|
|
from .fakes import BlockingEmbedder, fake_backend
|
|
|
|
|
|
class FakeRequest:
|
|
def __init__(self, disconnect_after: int | None) -> None:
|
|
self.calls = 0
|
|
self.disconnect_after = disconnect_after
|
|
|
|
async def is_disconnected(self) -> bool:
|
|
self.calls += 1
|
|
return self.disconnect_after is not None and self.calls > self.disconnect_after
|
|
|
|
|
|
def _spin(seen: dict[str, bool]): # type: ignore[no-untyped-def]
|
|
def work(cancel: threading.Event) -> str:
|
|
for _ in range(400):
|
|
if cancel.wait(0.01):
|
|
seen["cancelled"] = True
|
|
raise Cancelled()
|
|
return "done"
|
|
|
|
return work
|
|
|
|
|
|
def test_disconnect_cancels_work() -> None:
|
|
seen: dict[str, bool] = {}
|
|
gate = Gate(1, 0)
|
|
with pytest.raises(Cancelled) as error:
|
|
asyncio.run(run_job(FakeRequest(disconnect_after=1), gate, _spin(seen), max_seconds=60))
|
|
assert error.value.reason == "client disconnected"
|
|
assert seen["cancelled"]
|
|
assert gate.inflight == 0
|
|
|
|
|
|
def test_time_limit_cancels_work() -> None:
|
|
seen: dict[str, bool] = {}
|
|
with pytest.raises(Cancelled) as error:
|
|
asyncio.run(run_job(FakeRequest(None), Gate(1, 0), _spin(seen), max_seconds=0.3))
|
|
assert error.value.reason == "exceeded MEDIA_MAX_SECONDS"
|
|
|
|
|
|
def test_completed_work_returns_result() -> None:
|
|
assert asyncio.run(run_job(FakeRequest(None), Gate(1, 0), lambda _c: "ok", max_seconds=5)) == "ok"
|
|
|
|
|
|
def test_gate_overflow_is_busy() -> None:
|
|
gate = Gate(1, 1)
|
|
gate.reserve()
|
|
gate.reserve()
|
|
with pytest.raises(Busy):
|
|
gate.reserve()
|
|
gate.release_reservation()
|
|
gate.reserve()
|
|
|
|
|
|
def test_cancelled_transform_leaves_no_output(media_root: Path) -> None:
|
|
cancel = threading.Event()
|
|
path = str(media_root / "testsrc.mp4")
|
|
info = media_io.probe(path, threading.Event())
|
|
out = media_root / "cancelled" / "out.mkv"
|
|
out.parent.mkdir(exist_ok=True)
|
|
|
|
def frame_fn(index: int, _t: int, frame: np.ndarray) -> np.ndarray:
|
|
if index == 5:
|
|
cancel.set()
|
|
return frame
|
|
|
|
with pytest.raises(Cancelled):
|
|
media_io.transform_video(path, str(out), info, cancel, frame_fn)
|
|
assert list(out.parent.iterdir()) == []
|
|
|
|
|
|
def test_http_busy_when_gpu_queue_is_full(make_client, media_root: Path) -> None:
|
|
blocking = BlockingEmbedder()
|
|
client = make_client(
|
|
backends={"embedFrames": fake_backend(blocking, "fake-embed")}, gpu_concurrency=1, queue_size=0
|
|
)
|
|
body = {"path": f"{media_root}/testsrc.mp4", "timesMs": [0]}
|
|
results: list[int] = []
|
|
first = threading.Thread(
|
|
target=lambda: results.append(client.post("/v1/embedFrames", headers=AUTH, json=body).status_code)
|
|
)
|
|
first.start()
|
|
assert blocking.entered.wait(20)
|
|
second = client.post("/v1/embedFrames", headers=AUTH, json=body)
|
|
assert second.status_code == 503
|
|
assert second.json()["error"] == {
|
|
"code": "busy",
|
|
"message": "inference queue is full",
|
|
"operation": "embedFrames",
|
|
}
|
|
# CPU-bound operations use their own gate and still run.
|
|
assert (
|
|
client.post(
|
|
"/v1/saliency", headers=AUTH, json={"path": f"{media_root}/spot.mp4", "timesMs": [0]}
|
|
).status_code
|
|
== 200
|
|
)
|
|
blocking.release.set()
|
|
first.join(20)
|
|
assert results == [200]
|