# GPU worker image for the docling parse path, deployed as a RunPod
# serverless endpoint. Counterpart to deploy/Dockerfile (the CPU set-of-mark
# cloud job): this one carries the `gpu` extra and the docling model weights.
#
# Everything the worker needs is baked in at build time — dependencies and
# model weights — so a cold start never downloads anything. Build context is
# the repo root:
#   docker build -f runpod/Dockerfile -t quber-runpod .
#
# No CUDA base image: the pip `nvidia-*` wheels carry the CUDA runtime and
# only the host driver is needed. The Python version is pinned by the base
# image itself; 3.14 wheel sets fail at parse time with cuDNN sublibrary
# loading errors.
FROM python:3.13-slim AS base

COPY --from=ghcr.io/astral-sh/uv:0.5.11 /uv /uvx /bin/

ENV UV_LINK_MODE=copy \
    UV_COMPILE_BYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    PATH="/app/.venv/bin:$PATH" \
    # NVIDIA_VISIBLE_DEVICES is deliberately NOT baked: RunPod injects the
    # per-worker device assignment at container create, and a baked value
    # can shadow it — observed live as every worker reporting no visible
    # CUDA device. Local runs use --gpus (or an nvidia default runtime),
    # which injects the same variables.
    NVIDIA_DRIVER_CAPABILITIES=compute,utility

# OpenCV runtime (RapidOCR imports cv2) plus the font essentials so pages
# rasterized for OCR resolve missing fonts the same way everywhere.
# gcc/g++/libc6-dev: torch inductor JIT-compiles kernels at runtime on some
# documents — the CUDA path needs a C compiler and the CPU path needs a C++
# compiler. Without them the parse dies with "Failed to find C compiler" /
# "InvalidCxxCompiler: No working C++ compiler found ... g++" (both seen
# live on 28-page filings; small documents never trigger the compile path).
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgl1 \
    libglib2.0-0 \
    gcc \
    g++ \
    libc6-dev \
    fontconfig \
    fonts-urw-base35 \
    fonts-liberation \
    fonts-dejavu-core \
    fonts-noto-core \
    && fc-cache -f \
    && rm -rf /var/lib/apt/lists/*

# Drop privileges before anything lands in /app: syncing as the runtime user
# means nothing ever needs a `chown -R`, which in overlayfs would copy the
# whole venv into a duplicate layer (measured: +9.2GB on this image).
RUN useradd --create-home --uid 10001 quber \
    && mkdir /app && chown quber:quber /app
USER quber
ENV HF_HOME=/home/quber/.cache/huggingface

WORKDIR /app

# Core + gpu (docling, onnxruntime-gpu, nvidia wheels) + runpod (the SDK the
# handler runs on). This is the ~9GB layer; it changes only with uv.lock.
# --no-cache keeps uv's wheel archive out of the layer — with it, the same
# layer measured 18.6GB and overflowed the CI runner's disk.
COPY --chown=quber:quber pyproject.toml uv.lock README.md ./
COPY --chown=quber:quber src/ ./src/
RUN uv sync --frozen --no-dev --no-cache --extra gpu --extra runpod

# Bake the model weights by running the real parse once on CPU. This pulls
# exactly what the tuned-financial preset loads — layout, TableFormer,
# picture classifier, RapidOCR — into the image (~550MB) and fails the build
# if the pipeline cannot run. The warmup document is a 4KB synthetic page
# with tables so TableFormer executes.
COPY --chown=quber:quber documents/synthetic/sbs2.pdf /app/warmup.pdf
RUN quber parse /app/warmup.pdf --device cpu -o /tmp/warmup --format json \
    && rm -rf /tmp/warmup

# The embedding model the playground's ingestion sends here. Baked for the
# same reason as the parse models: a worker must not reach HuggingFace.
RUN python -c "from sentence_transformers import SentenceTransformer; SentenceTransformer('BAAI/bge-large-en-v1.5')"

# After the warmup has populated the cache, cut the runtime loose from the
# HuggingFace API entirely: with weights cached, docling still pings HF for a
# revision check at pipeline build, and a worker must not depend on HF being
# reachable. Set only after warmup — the warmup itself downloads.
ENV HF_HUB_OFFLINE=1

COPY --chown=quber:quber runpod/handler.py runpod/start.sh /app/runpod/

ENTRYPOINT ["/app/runpod/start.sh"]