evOo-local/infra/media-worker/app/main.py

"""HTTP API v1 of the Media Studio sidecar.

GET /healthz is open; everything else needs the bearer token. POST
/v1/<operation> takes the TypeScript request object (minus `signal`) and
answers with the TypeScript result object.
"""

from __future__ import annotations

import logging
import threading
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import anyio
from fastapi import Depends, FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel, ValidationError
from starlette.exceptions import HTTPException as StarletteHTTPException

from . import VERSION, encoders, schemas
from .auth import require_token
from .config import Settings
from .errors import Cancelled, MediaError, error_response
from .models import OPERATIONS, Registry
from .ops import audio, embeddings, faces, inpaint, redact, saliency, scenes, segment, synthesize, transcribe
from .ops.common import Ctx, jsonable
from .ops.devices import effective_accelerator
from .runtime import Gate, run_job

log = logging.getLogger("media-worker")

MAX_BODY_BYTES = 64 * 1024 * 1024
# Status for a request whose client went away; nobody reads it.
CLIENT_CLOSED = 499


@dataclass(frozen=True)
class Route:
    schema: type[BaseModel]
    run: Callable[[Ctx, Any, threading.Event], dict[str, Any]]
    gpu: bool


ROUTES: dict[str, Route] = {
    "detectScenes": Route(schemas.PathRequest, scenes.run, False),
    "transcribe": Route(schemas.TranscribeRequest, transcribe.run, True),
    "detectFaces": Route(schemas.DetectFacesRequest, faces.detect_run, True),
    "trackFaces": Route(schemas.TrackFacesRequest, faces.track_run, True),
    "faceLandmarks": Route(schemas.TimesRequest, faces.landmarks_run, True),
    "embedFrames": Route(schemas.TimesRequest, embeddings.embed_run, True),
    "classifyContent": Route(schemas.ClassifyContentRequest, embeddings.classify_run, True),
    "segmentPrompted": Route(schemas.SegmentPromptedRequest, segment.run, True),
    "inpaintVideo": Route(schemas.InpaintVideoRequest, inpaint.run, True),
    "redact": Route(schemas.RedactRequest, redact.run, False),
    "synthesizeIdentity": Route(schemas.SynthesizeIdentityRequest, synthesize.run, True),
    "audioFeatures": Route(schemas.PathRequest, audio.run, False),
    "saliency": Route(schemas.TimesRequest, saliency.run, False),
}
assert set(ROUTES) == set(OPERATIONS)


def _operation_of(request: Request) -> str | None:
    name = request.url.path.rsplit("/", 1)[-1]
    return name if name in ROUTES else None


def _validation_message(error: ValidationError) -> str:
    first = error.errors()[0] if error.errors() else {}
    where = ".".join(str(part) for part in first.get("loc", ()))
    return (
        f"{where}: {first.get('msg', 'invalid request')}"
        if where
        else str(first.get("msg", "invalid request"))
    )


def create_app(settings: Settings | None = None, registry: Registry | None = None) -> FastAPI:
    settings = settings or Settings.from_env()
    registry = registry or Registry(settings)
    app = FastAPI(title="media-worker", version=VERSION, docs_url=None, redoc_url=None, openapi_url=None)
    app.state.settings = settings
    app.state.ctx = Ctx(settings, registry)
    app.state.gpu_gate = Gate(settings.gpu_concurrency, settings.queue_size)
    app.state.cpu_gate = Gate(settings.cpu_concurrency, settings.queue_size)

    @app.exception_handler(MediaError)
    async def media_error(request: Request, error: MediaError) -> JSONResponse:
        return error_response(error.code, error.message, error.operation or _operation_of(request))

    @app.exception_handler(RequestValidationError)
    async def validation_error(request: Request, error: RequestValidationError) -> JSONResponse:
        return error_response("bad_request", "invalid request", _operation_of(request))

    @app.exception_handler(StarletteHTTPException)
    async def http_error(request: Request, error: StarletteHTTPException) -> JSONResponse:
        code = "bad_request" if error.status_code < 500 else "internal"
        return error_response(code, str(error.detail), _operation_of(request), status=error.status_code)

    @app.exception_handler(Exception)
    async def unexpected(request: Request, error: Exception) -> JSONResponse:
        log.exception("unhandled error")
        return error_response("internal", "internal error", _operation_of(request))

    @app.get("/healthz")
    async def healthz() -> dict[str, Any]:
        return {"ok": True, "version": VERSION}

    @app.get("/v1/capabilities", dependencies=[Depends(require_token)])
    async def capabilities() -> dict[str, Any]:
        hardware = await anyio.to_thread.run_sync(encoders.hardware_encoders)
        return {
            "configured": True,
            "operations": registry.operation_capabilities(),
            "accelerator": effective_accelerator(settings),
            "hardwareEncoders": list(hardware),
        }

    @app.post("/v1/{operation}", dependencies=[Depends(require_token)])
    async def operate(operation: str, request: Request) -> Response:
        route = ROUTES.get(operation)
        if route is None:
            return error_response("bad_request", "unknown operation", None, status=404)
        raw = bytearray()
        async for chunk in request.stream():
            raw.extend(chunk)
            if len(raw) > MAX_BODY_BYTES:
                raise MediaError("too_large", "request body too large", operation)
        try:
            body = route.schema.model_validate_json(bytes(raw))
        except ValidationError as error:
            raise MediaError("bad_request", _validation_message(error), operation) from error
        ctx: Ctx = request.app.state.ctx
        gate: Gate = request.app.state.gpu_gate if route.gpu else request.app.state.cpu_gate
        try:
            result = await run_job(
                request, gate, lambda cancel: route.run(ctx, body, cancel), settings.max_seconds
            )
        except MediaError as error:
            error.operation = operation
            raise
        except Cancelled as error:
            if error.reason == "client disconnected":
                return Response(status_code=CLIENT_CLOSED)
            raise MediaError("internal", error.reason, operation) from error
        except Exception as error:
            log.exception("%s failed", operation)
            raise MediaError("internal", "internal error", operation) from error
        return JSONResponse(jsonable(result))

    return app