handler.py

"""RunPod Serverless handler for PaddleOCR-VL-0.9B served by vLLM.

Mirrors the PP-OCRv6 endpoint's request and response shape so the same
benchmark harness drives both, and mirrors the GLM-OCR endpoint's process
model (vLLM as a child process, handler forwards to it) so the two
vision-language endpoints are measured under the same conditions.
"""

from __future__ import annotations

import asyncio
import io
import logging
import os
import socket
import subprocess
import threading
import time

import requests
import runpod
from PIL import Image

from image_input import is_pdf, resolve_payload
from pdf_input import render_pages
from vl_client import recognise
from vllm_command import build_vllm_command

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("handler")

VLLM_PORT = int(os.environ.get("VLLM_PORT", "8080"))
VLLM_URL = f"http://localhost:{VLLM_PORT}"

MODEL_NAME = os.environ.get("MODEL_NAME", "PaddlePaddle/PaddleOCR-VL")
MODEL_PATH = os.environ.get("MODEL_PATH", MODEL_NAME)
MAX_MODEL_LEN = os.environ.get("MAX_MODEL_LEN", "16384")
GPU_MEMORY_UTILIZATION = os.environ.get("GPU_MEMORY_UTILIZATION", "0.90")
MAX_NUM_BATCHED_TOKENS = os.environ.get("MAX_NUM_BATCHED_TOKENS", "16384").strip()
ENFORCE_EAGER = os.environ.get("ENFORCE_EAGER", "0").lower() in {"1", "true", "yes"}

CONCURRENCY = max(1, int(os.environ.get("CONCURRENCY", "1")))
# vLLM must be allowed at least as many running sequences as the handler will
# hand it, otherwise the extra requests queue inside vLLM and the concurrency
# being measured is not real.
MAX_NUM_SEQS = os.environ.get("MAX_NUM_SEQS", "").strip() or str(max(CONCURRENCY, 4))

MAX_IMAGE_SIDE = int(os.environ.get("MAX_IMAGE_SIDE", "1900"))
MAX_OUTPUT_TOKENS = int(os.environ.get("MAX_OUTPUT_TOKENS", "4096"))
OCR_PROMPT = os.environ.get("OCR_PROMPT", "OCR:")
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)))
VLLM_START_TIMEOUT = int(os.environ.get("VLLM_START_TIMEOUT", "900"))
WORKER_ID = os.environ.get("RUNPOD_POD_ID", socket.gethostname())

_VLLM_READY_SECONDS = 0.0


def _stream_logs(pipe) -> None:
    try:
        for line in pipe:
            line = line.rstrip()
            if line:
                log.info("[vllm] %s", line)
    finally:
        pipe.close()


def start_vllm() -> subprocess.Popen:
    command = build_vllm_command(
        model_path=MODEL_PATH,
        model_name=MODEL_NAME,
        port=VLLM_PORT,
        max_model_len=MAX_MODEL_LEN,
        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
        max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS,
        max_num_seqs=MAX_NUM_SEQS,
        enforce_eager=ENFORCE_EAGER,
    )
    log.info("starting vLLM: %s", " ".join(command))
    process = subprocess.Popen(
        command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
    )
    threading.Thread(target=_stream_logs, args=(process.stdout,), daemon=True).start()
    return process


def wait_for_vllm(process: subprocess.Popen, timeout: int) -> float:
    """Block until /health answers, and fail loudly if vLLM dies first."""
    started = time.perf_counter()
    while time.perf_counter() - started < timeout:
        if process.poll() is not None:
            raise RuntimeError(f"vLLM exited with code {process.returncode} before becoming ready")
        try:
            if requests.get(f"{VLLM_URL}/health", timeout=2).status_code == 200:
                elapsed = time.perf_counter() - started
                log.info("vLLM ready in %.1fs", elapsed)
                return elapsed
        except requests.RequestException:
            pass
        time.sleep(2)
    raise TimeoutError(f"vLLM did not become ready within {timeout}s")


def _meta(started: float, pages: int, truncated: bool, extra: dict) -> dict:
    return {
        "model": MODEL_NAME,
        "backend": "vllm",
        "page_count": pages,
        "pages_truncated": truncated,
        "total_ms": round((time.perf_counter() - started) * 1000, 2),
        "worker_concurrency": CONCURRENCY,
        "max_num_seqs": int(MAX_NUM_SEQS),
        "max_image_side": MAX_IMAGE_SIDE,
        "prompt": OCR_PROMPT,
        "vllm_ready_seconds": round(_VLLM_READY_SECONDS, 2),
        "worker_id": WORKER_ID,
        **extra,
    }


def _recognise_page(image: Image.Image, prompt: str) -> dict:
    started = time.perf_counter()
    result = recognise(
        base_url=VLLM_URL,
        model_name=MODEL_NAME,
        image=image,
        prompt=prompt,
        max_side=MAX_IMAGE_SIDE,
        max_output_tokens=MAX_OUTPUT_TOKENS,
    )
    result["processing_ms"] = round((time.perf_counter() - started) * 1000, 2)
    return result


def _totals(pages: list[dict]) -> dict:
    def total(key: str) -> int:
        return sum(page.get(key) or 0 for page in pages)

    return {
        "processing_ms": round(sum(page["processing_ms"] for page in pages), 2),
        "prompt_tokens": total("prompt_tokens"),
        "completion_tokens": total("completion_tokens"),
        "total_tokens": total("total_tokens"),
    }


def _handle_pdf(data: bytes, prompt: str, started: float) -> dict:
    # One open per request: render_pages reports the page count from the
    # document it already holds, instead of opening the file a second time.
    total_pages, stream = render_pages(data, MAX_IMAGE_SIDE, PDF_MAX_PAGES)
    pages = []
    for index, image in stream:
        result = _recognise_page(image, prompt)
        image.close()
        result["page_index"] = index
        result["markdown"] = result.pop("text")
        pages.append(result)

    return {
        "markdown": "\n\n".join(page["markdown"] for page in pages),
        "pages": pages,
        "meta": _meta(
            started,
            len(pages),
            total_pages > len(pages),
            {"input_type": "pdf", "pdf_total_pages": total_pages, **_totals(pages)},
        ),
    }


def _handle_image(data: bytes, prompt: str, started: float) -> dict:
    with Image.open(io.BytesIO(data)) as image:
        result = _recognise_page(image.convert("RGB"), prompt)
    result["page_index"] = 0
    result["markdown"] = result.pop("text")
    return {
        "markdown": result["markdown"],
        "pages": [result],
        "meta": _meta(
            started,
            1,
            False,
            {
                "input_type": "image",
                "original_size": result["original_size"],
                "sent_size": result["sent_size"],
                **_totals([result]),
            },
        ),
    }


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

    prompt = job_input.get("prompt") or OCR_PROMPT
    try:
        data = resolve_payload(job_input, MAX_INPUT_BYTES)
    except ValueError as exc:
        return {"error": str(exc)}

    try:
        if is_pdf(data):
            return _handle_pdf(data, prompt, started)
        return _handle_image(data, prompt, started)
    except Exception as exc:  # surface the real cause instead of a bare timeout
        log.exception("job failed")
        return {"error": f"{type(exc).__name__}: {exc}"}


async def handler(event):
    """Run the blocking vLLM round-trip off the event loop so CONCURRENCY is real."""
    return await asyncio.to_thread(_handle_job, event)


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


if __name__ == "__main__":
    _process = start_vllm()
    _VLLM_READY_SECONDS = wait_for_vllm(_process, VLLM_START_TIMEOUT)
    log.info("handler concurrency=%s max_num_seqs=%s max_image_side=%s",
             CONCURRENCY, MAX_NUM_SEQS, MAX_IMAGE_SIDE)
    runpod.serverless.start({"handler": handler, "concurrency_modifier": adjust_concurrency})