worker_common.py

from __future__ import annotations

import base64
import json
import logging
import os
import shutil
import tempfile
import time
import traceback
import uuid
from pathlib import Path
from typing import Any, Callable, Protocol, TypeVar

import gradio as gr


LOGGER = logging.getLogger("ocr-worker")
logging.basicConfig(
    level=os.getenv("LOG_LEVEL", "INFO"),
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)

F = TypeVar("F", bound=Callable[..., Any])


class Adapter(Protocol):
    model_id: str
    label: str

    def is_loaded(self) -> bool: ...

    def infer(self, image_path: Path, prompt: str, options: dict[str, Any]) -> dict[str, Any]: ...

    def runtime_metadata(self) -> dict[str, Any]: ...


def configure_cache() -> None:
    data = Path(os.getenv("DATA_DIR", "/data"))
    if data.exists():
        cache_root = data / ".cache"
        cache_root.mkdir(parents=True, exist_ok=True)
        os.environ.setdefault("HF_HOME", str(cache_root / "huggingface"))
        hub_cache = str(cache_root / "huggingface" / "hub")
        os.environ.setdefault("HF_HUB_CACHE", hub_cache)
        os.environ.setdefault("HUGGINGFACE_HUB_CACHE", hub_cache)
        os.environ.setdefault("TRANSFORMERS_CACHE", str(cache_root / "huggingface" / "transformers"))
        os.environ.setdefault("TORCH_HOME", str(cache_root / "torch"))
        os.environ.setdefault("PADDLEOCR_HOME", str(cache_root / "paddleocr"))
        os.environ.setdefault("PADDLE_PDX_CACHE_HOME", str(cache_root / "paddlex"))
        os.environ.setdefault("PADDLE_HOME", str(cache_root / "paddle"))


def gpu_task(function: F) -> F:
    if os.getenv("WORKER_RUNTIME", "gradio-zerogpu") != "gradio-zerogpu":
        return function
    try:
        import spaces
    except Exception:
        return function
    duration = int(os.getenv("ZERO_GPU_DURATION", "120"))
    size = os.getenv("ZERO_GPU_SIZE", "large").strip() or "large"
    return spaces.GPU(duration=duration, size=size)(function)  # type: ignore[return-value]


def _require_token(provided: str | None) -> None:
    expected = os.getenv("WORKER_API_TOKEN", "").strip()
    if expected and provided != expected:
        raise gr.Error("Invalid or missing worker token.")


def _copy_input(file_path: str | Path, destination: Path, max_mb: int) -> int:
    source = Path(file_path)
    if not source.exists() or not source.is_file():
        raise gr.Error("Uploaded file was not received by the worker.")
    size = source.stat().st_size
    if size <= 0:
        raise gr.Error("Uploaded file is empty.")
    if size > max_mb * 1024 * 1024:
        raise gr.Error(f"Upload exceeds worker limit of {max_mb} MB.")
    shutil.copyfile(source, destination)
    return size


def _safe_suffix(file_path: str | Path) -> str:
    suffix = Path(file_path).suffix.lower()
    if suffix not in {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}:
        return ".png"
    return suffix


def _image_payload(path_value: Any) -> tuple[str | None, str | None]:
    if not path_value:
        return None, None
    path = Path(str(path_value))
    if not path.exists() or not path.is_file():
        return None, None
    suffix = path.suffix.lower()
    mime = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
    }.get(suffix, "image/png")
    encoded = base64.b64encode(path.read_bytes()).decode("ascii")
    return encoded, mime


def health_payload(adapter: Adapter) -> dict[str, Any]:
    payload = {
        "status": "ok",
        "model": adapter.model_id,
        "label": adapter.label,
        "loaded": bool(adapter.is_loaded()),
        "runtime": os.getenv("WORKER_RUNTIME", "gradio-zerogpu"),
        "zero_gpu_duration": int(os.getenv("ZERO_GPU_DURATION", "120")),
        "zero_gpu_size": os.getenv("ZERO_GPU_SIZE", "large"),
    }
    metadata = getattr(adapter, "runtime_metadata", None)
    if callable(metadata):
        payload["runtime_metadata"] = metadata()
    return payload


def run_ocr_request(
    adapter: Adapter,
    file_path: str | Path,
    model: str,
    prompt: str,
    options_json: str,
    worker_token: str,
) -> dict[str, Any]:
    _require_token(worker_token)
    if model != adapter.model_id:
        raise gr.Error(f"This worker serves {adapter.model_id!r}, not {model!r}.")
    try:
        options = json.loads(options_json or "{}")
    except json.JSONDecodeError as exc:
        raise gr.Error(f"Invalid options_json: {exc}") from exc
    if not isinstance(options, dict):
        raise gr.Error("options_json must be a JSON object.")

    request_id = uuid.uuid4().hex
    request_dir = Path(tempfile.mkdtemp(prefix=f"ocr_worker_{request_id}_"))
    image_path = request_dir / f"input{_safe_suffix(file_path)}"
    input_bytes = _copy_input(file_path, image_path, int(os.getenv("MAX_UPLOAD_MB", "40")))

    started = time.perf_counter()
    try:
        result = adapter.infer(image_path, prompt, options)
        if not isinstance(result, dict):
            raise TypeError("Adapter returned a non-dict result.")

        annotated_b64, annotated_mime = _image_payload(result.pop("annotated_path", None))
        elapsed = time.perf_counter() - started
        metrics = dict(result.pop("metrics", {}) or {})
        metrics["elapsed_seconds"] = round(elapsed, 4)
        metrics["input_bytes"] = input_bytes

        payload: dict[str, Any] = {
            "schema_version": "1.0",
            "request_id": request_id,
            "model": adapter.model_id,
            "text": str(result.pop("text", "")),
            "markdown": str(result.pop("markdown", "")),
            "annotated_image_base64": annotated_b64,
            "annotated_image_mime": annotated_mime,
            "raw": result.pop("raw", {}),
            "warnings": list(result.pop("warnings", []) or []),
            "metrics": metrics,
        }
        if result:
            payload["adapter_extra"] = result
        return payload
    except gr.Error:
        raise
    except Exception as exc:
        LOGGER.error("Request %s failed: %s\n%s", request_id, exc, traceback.format_exc())
        detail: dict[str, Any] = {
            "error": type(exc).__name__,
            "message": str(exc),
            "request_id": request_id,
        }
        if os.getenv("DEBUG_ERRORS", "0") == "1":
            detail["traceback"] = traceback.format_exc()
        raise gr.Error(json.dumps(detail, ensure_ascii=False)) from exc
    finally:
        if os.getenv("KEEP_REQUEST_DIRS", "0") != "1":
            shutil.rmtree(request_dir, ignore_errors=True)


def build_demo(adapter: Adapter) -> gr.Blocks:
    @gpu_task
    def ocr(file_path: str, model: str, prompt: str, options_json: str, worker_token: str) -> dict[str, Any]:
        return run_ocr_request(adapter, file_path, model, prompt, options_json, worker_token)

    def health() -> dict[str, Any]:
        return health_payload(adapter)

    with gr.Blocks(title=f"{adapter.label} OCR Worker") as demo:
        gr.Markdown(f"# {adapter.label} Worker\n\nModel ID: `{adapter.model_id}`")
        with gr.Row():
            with gr.Column():
                file_input = gr.File(label="Image", type="filepath", file_count="single")
                model_input = gr.Textbox(label="Model", value=adapter.model_id)
                prompt_input = gr.Textbox(label="Prompt", lines=3)
                options_input = gr.Textbox(label="Options JSON", value="{}")
                token_input = gr.Textbox(label="Worker token", type="password")
                run_button = gr.Button("Run OCR", variant="primary")
            with gr.Column():
                result_output = gr.JSON(label="OCR response")
                health_output = gr.JSON(label="Health")
                health_button = gr.Button("Health")
        run_button.click(
            fn=ocr,
            inputs=[file_input, model_input, prompt_input, options_input, token_input],
            outputs=result_output,
            api_name="ocr",
        )
        health_button.click(fn=health, inputs=[], outputs=health_output, api_name="health")
        demo.load(fn=health, inputs=[], outputs=health_output)
    return demo