TieMeasureFlow/src/vision_worker/main.py

"""The heavy container: the runner behind an internal API.

Separate from the FastAPI server on purpose. The API image stays light, a
VisionSuite upgrade does not restart production traffic, and an execution that
crashes does not take the other tablets' requests down with it.
"""
from __future__ import annotations

import asyncio
import io
import json
import os

import numpy as np
from fastapi import FastAPI, File, Form, HTTPException, UploadFile, status
from PIL import Image, UnidentifiedImageError

from src.vision.runner import GraphBuildError, RunOutcome, engine_version, run_graph

app = FastAPI(title="TieMeasureFlow Vision Worker", version="0.1.0")

# Below the backend's own 120s budget on purpose (#34): the worker must give
# up and answer 504 before the backend's httpx client times out on its own,
# so the operator sees a reason instead of a bare timeout.
_RUN_TIMEOUT_S = float(os.environ.get("VISION_RUN_TIMEOUT_S", "100"))
# Two slow graphs at once must not also stall a third tablet's request past
# any reasonable wait - reject it instead, with a status the backend can turn
# into a 503 the operator understands ("try again"), not a queued request
# that quietly outlasts anyone's patience.
_MAX_CONCURRENCY = int(os.environ.get("VISION_MAX_CONCURRENCY", "2"))
_run_slots = asyncio.Semaphore(_MAX_CONCURRENCY)


class _ImageDecodeError(Exception):
    """The upload is not a decodable image - raised inside the worker thread
    so it survives the trip through `asyncio.to_thread`."""


@app.get("/health")
async def health() -> dict:
    """Deliberately just this: reading `engine_version()` (cached after the
    first call) is the only work here, so this answers even while `/run` is
    busy decoding an image or running a graph in its own thread."""
    return {"status": "ok", "engine_version": engine_version()}


def _decode_and_run(content: bytes, graph: dict) -> RunOutcome:
    """The synchronous, CPU/IO-bound half of a request: PIL decode plus
    `run_graph`. Runs in a worker thread (see `run` below) so the event loop -
    and therefore `/health` - stays responsive while it executes. Once this
    thread is past its budget it is not interruptible: it keeps computing to
    completion in the background even after `run` has already answered 504
    for it (see the timeout handling there).
    """
    try:
        frame = np.array(Image.open(io.BytesIO(content)).convert("L"))
    except (UnidentifiedImageError, OSError) as exc:
        raise _ImageDecodeError(str(exc)) from exc
    return run_graph(frame, graph)


@app.post("/run")
async def run(
    image: UploadFile = File(...),
    graph: str = Form(...),
) -> dict:
    try:
        parsed_graph = json.loads(graph)
    except json.JSONDecodeError as exc:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail=f"graph is not valid JSON: {exc}",
        ) from exc

    if not isinstance(parsed_graph, dict):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail=f"graph must be a JSON object, got {type(parsed_graph).__name__}",
        )

    content = await image.read()

    if _run_slots.locked():
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=f"worker busy: already running {_MAX_CONCURRENCY} job(s)",
        )

    async with _run_slots:
        try:
            outcome = await asyncio.wait_for(
                asyncio.to_thread(_decode_and_run, content, parsed_graph),
                timeout=_RUN_TIMEOUT_S,
            )
        except asyncio.TimeoutError as exc:
            raise HTTPException(
                status_code=status.HTTP_504_GATEWAY_TIMEOUT,
                detail=f"vision run exceeded {_RUN_TIMEOUT_S}s",
            ) from exc
        except _ImageDecodeError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=f"image is not decodable: {exc}",
            ) from exc
        except GraphBuildError as exc:
            # The graph itself is bad (schema version, a tool missing
            # "kind", ...) - a bad request, not a server fault.
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=str(exc),
            ) from exc
        except Exception as exc:
            # The graph built fine; something failed while actually running
            # it - an engine fault, not the caller's. Still named, not a
            # naked 500.
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"vision run failed: {type(exc).__name__}: {exc}",
            ) from exc

    return {
        # `outcome.outputs` never carries NaN/inf/bool (see runner.run_graph):
        # a JSON encoder that accepts non-finite floats (Python's stdlib json
        # does, by default) must never get the chance to emit them here.
        "outputs": outcome.outputs,
        "failures": [
            {"tool_id": f.tool_id, "tool_name": f.tool_name, "error": f.error}
            for f in outcome.failures
        ],
        "engine_version": outcome.engine_version,
        "duration_ms": outcome.duration_ms,
        "invalid_outputs": outcome.invalid_outputs,
    }