handler.py

"""RunPod Serverless handler for PP-OCRv6 (paddleocr, GPU).

Accepts a single image or a PDF and returns page-wise markdown plus per-line
layout. Models are loaded at import time, before RunPod starts dispatching, so
a warm worker answers in milliseconds.

Besides the OCR request itself the handler takes two tuning shapes:

* ``{"params": {...}}`` alongside a document forwards detection parameters to
  PaddleOCR's ``predict()`` for that one request.
* ``{"tune": {...}}`` rebuilds the model pool in place, so a sweep over pool
  size, batch size or weights costs one cold start instead of one per point.
  Gated by OCR_ALLOW_TUNE because it lets any caller change what the worker
  serves; leave it off in production.
"""

from __future__ import annotations

import time

# Before the heavy imports, so the boot breakdown below is measured from as
# close to process start as Python allows. A cold start is dominated by things
# this process cannot see -- image pull, container create -- and the only way to
# tell how large that invisible part is, is to subtract the part we can time:
# delayTime minus process_ready_s is everything that happened before Python ran.
_PROCESS_START = time.time()

import io  # noqa: E402
import os  # noqa: E402
import socket  # noqa: E402

import runpod  # noqa: E402
from PIL import Image  # noqa: E402

import ocr_engine  # noqa: E402
from image_input import is_pdf, resolve_payload  # noqa: E402
from pdf_input import render_pages  # noqa: E402

_IMPORT_SECONDS = time.time() - _PROCESS_START

CONCURRENCY = max(1, int(os.environ.get("CONCURRENCY", "2")))
PDF_MAX_EDGE = int(os.environ.get("PDF_MAX_EDGE", "1600"))
PDF_MAX_PAGES = int(os.environ.get("PDF_MAX_PAGES", "50"))
MAX_INPUT_BYTES = int(os.environ.get("MAX_INPUT_BYTES", str(10 * 1024 * 1024)))
ALLOW_TUNE = os.environ.get("OCR_ALLOW_TUNE", "0") == "1"
WORKER_ID = os.environ.get("RUNPOD_POD_ID", socket.gethostname())

print(f"imports done in {_IMPORT_SECONDS:.1f}s; loading PP-OCRv6 models...", flush=True)
_BOOT_LOAD_SECONDS = ocr_engine.warm_up()
_READY_SECONDS = time.time() - _PROCESS_START
print(
    f"models ready in {_BOOT_LOAD_SECONDS:.1f}s "
    f"(process ready {_READY_SECONDS:.1f}s after start)",
    flush=True,
)


def _mean(values: list) -> float | None:
    present = [value for value in values if value is not None]
    return round(sum(present) / len(present), 4) if present else None


def _meta(started: float, pages: int, truncated: bool, extra: dict) -> dict:
    cfg = ocr_engine.config()
    return {
        "model": cfg["rec_model"].replace("_rec", ""),
        "det_model": cfg["det_model"],
        "rec_model": cfg["rec_model"],
        "device": ocr_engine.DEVICE,
        "page_count": pages,
        "pages_truncated": truncated,
        "total_ms": round((time.perf_counter() - started) * 1000, 2),
        "model_load_seconds": round(ocr_engine.load_seconds(), 2),
        "boot_load_seconds": round(_BOOT_LOAD_SECONDS, 2),
        # Boot breakdown. process_ready_s is everything from the first line of
        # Python to the pool being built; subtracting it from the job's
        # delayTime leaves image pull plus container create, which is the part
        # only the platform controls.
        "import_seconds": round(_IMPORT_SECONDS, 2),
        "process_ready_seconds": round(_READY_SECONDS, 2),
        "uptime_seconds": round(time.time() - _PROCESS_START, 2),
        "worker_concurrency": CONCURRENCY,
        "recognition_batch_size": cfg["rec_batch_size"],
        "pool_size": cfg["pool_size"],
        "limit_side_len": cfg["limit_side_len"],
        "input_max_edge": cfg["input_max_edge"],
        "input_mode": cfg["input_mode"],
        "worker_id": WORKER_ID,
        **extra,
    }


def _handle_pdf(data: bytes, started: float, overrides: dict) -> dict:
    total_pages, page_images = render_pages(data, PDF_MAX_EDGE, PDF_MAX_PAGES)
    pages = []
    processing_ms = 0.0
    preprocess_ms = 0.0
    for index, image in page_images:
        result = ocr_engine.recognize(image, overrides)
        image.close()
        processing_ms += result["processing_ms"]
        preprocess_ms += result["preprocess_ms"]
        pages.append(
            {
                "page_index": index,
                "markdown": result["markdown"],
                "layout_json": result["layout_json"],
                "text_count": result["text_count"],
                "mean_score": result["mean_score"],
                "processing_ms": result["processing_ms"],
            }
        )

    markdown = "\n\n".join(page["markdown"] for page in pages)
    return {
        "markdown": markdown,
        "layout_json": [line for page in pages for line in page["layout_json"]],
        "pages": pages,
        "meta": _meta(
            started,
            len(pages),
            total_pages > len(pages),
            {
                "input_type": "pdf",
                "pdf_total_pages": total_pages,
                "processing_ms": round(processing_ms, 2),
                "preprocess_ms": round(preprocess_ms, 2),
                "text_count": sum(page["text_count"] for page in pages),
                "mean_score": _mean([page["mean_score"] for page in pages]),
                "pdf_render_max_edge": PDF_MAX_EDGE,
                "predict_overrides": overrides,
            },
        ),
    }


def _handle_image(data: bytes, started: float, overrides: dict) -> dict:
    with Image.open(io.BytesIO(data)) as image:
        result = ocr_engine.recognize(image.convert("RGB"), overrides)
    return {
        "markdown": result["markdown"],
        "layout_json": result["layout_json"],
        "pages": [
            {
                "page_index": 0,
                "markdown": result["markdown"],
                "layout_json": result["layout_json"],
                "text_count": result["text_count"],
                "mean_score": result["mean_score"],
                "processing_ms": result["processing_ms"],
            }
        ],
        "meta": _meta(
            started,
            1,
            False,
            {
                "input_type": "image",
                "processing_ms": result["processing_ms"],
                "preprocess_ms": result["preprocess_ms"],
                "original_size": result["original_size"],
                "processed_size": result["processed_size"],
                "text_count": result["text_count"],
                "mean_score": result["mean_score"],
                "predict_overrides": overrides,
            },
        ),
    }


def _handle_tune(request: dict, started: float) -> dict:
    if not ALLOW_TUNE:
        return {"error": "tuning is disabled (set OCR_ALLOW_TUNE=1)"}
    if not isinstance(request, dict):
        return {"error": "tune must be an object"}
    config = ocr_engine.reconfigure(request)
    return {
        "tuned": True,
        "config": config,
        "reload_seconds": round(ocr_engine.load_seconds(), 2),
        "meta": _meta(started, 0, False, {"input_type": "tune"}),
    }


def handler(event):
    started = time.perf_counter()
    job_input = event.get("input") or {}
    if not isinstance(job_input, dict):
        return {"error": "input must be an object"}

    try:
        if "tune" in job_input:
            return _handle_tune(job_input["tune"], started)

        overrides = job_input.get("params") or {}
        if not isinstance(overrides, dict):
            return {"error": "params must be an object"}

        data = resolve_payload(job_input, MAX_INPUT_BYTES)
        if is_pdf(data):
            return _handle_pdf(data, started, overrides)
        return _handle_image(data, started, overrides)
    except ValueError as exc:
        return {"error": str(exc)}
    except Exception as exc:  # surface the real cause instead of a bare timeout
        return {"error": f"{type(exc).__name__}: {exc}"}


def adjust_concurrency(_current: int) -> int:
    return CONCURRENCY


runpod.serverless.start({"handler": handler, "concurrency_modifier": adjust_concurrency})