# Universal Kimodo image for self-hosted Hatchet workers and Modal.
# Model and text-encoder weights are mounted at /weights, never baked in.

# ---------- CUDA build argument ----------
# One variant per row of the ph-kimodo matrix in
# backend/app/core/config/services.py — keep the two in sync, the CLI selects
# a variant from that catalog and pulls the tag CI publishes here.
#
#   cu128  NGC 25.02 / CUDA 12.8 / torch 2.7 — driver >=570, sm_75..sm_120  (default)
#   cu130  NGC 25.10 / CUDA 13.0 / torch 2.9 — driver >=580, sm_75..sm_120
#
# cu126 is not built here — see Dockerfile.cuda. NGC's CUDA 12.6 images ship a
# pre-release torch that Kimodo's pinned transformers cannot import.
#
# NGC 25.02 tops out at sm_120, so Blackwell (RTX 50xx) is native; older bases
# PTX-JIT every kernel from compute_90 — minutes of stall on a cold cache.
ARG NGC_TAG=25.02-py3

FROM nvcr.io/nvidia/pytorch:${NGC_TAG}

LABEL org.opencontainers.image.source="https://github.com/poly-hammer/poly-hammer-portal"

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1

RUN apt-get update && apt-get install -y --no-install-recommends \
        ca-certificates curl \
    && rm -rf /var/lib/apt/lists/*

# The install steps live in recipes/install.sh, which a native host runs
# verbatim. Building the image by executing that same script is what keeps the
# two from drifting; PH_IN_DOCKER tells it to use this image's tuned NGC torch
# instead of provisioning its own interpreter.
#
# Kimodo and motion_correction come from prebuilt wheels staged into
# wheelhouse/ by the build workflow, so this image needs no compiler and no
# git checkout of upstream.
COPY . /opt/worker-src
RUN --mount=type=cache,target=/root/.cache/pip \
    python -m pip install --upgrade pip setuptools wheel \
    && PH_IN_DOCKER=1 \
       PH_PYTHON="$(command -v python)" \
       PH_WORKER_SRC=/opt/worker-src \
       bash /opt/worker-src/poly_hammer_worker/jobs/kimodo/recipes/install.sh \
           --profile "kimodo/linux-x64/docker" --prefix /opt/ph-install \
           --wheelhouse /opt/worker-src/wheelhouse \
    && rm -rf /opt/worker-src

# CUDA_CACHE_* lives on the weights volume so PTX JIT'd for GPU architectures
# newer than this image's torch build (e.g. sm_120 / RTX 50xx) survives
# `docker run --rm` instead of costing minutes on every container start.
ENV WEIGHTS_PATH=/weights \
    CHECKPOINT_DIR=/weights/models \
    TEXT_ENCODERS_DIR=/weights/text-encoders \
    TEXT_ENCODER_MODE=local \
    TEXT_ENCODER_DEVICE=auto \
    HF_HOME=/weights/huggingface \
    HUGGINGFACE_HUB_CACHE=/weights/huggingface/hub \
    HF_XET_HIGH_PERFORMANCE=1 \
    CUDA_CACHE_PATH=/weights/.nv/ComputeCache \
    CUDA_CACHE_MAXSIZE=4294967296 \
    PH_JOB_CLASS=poly_hammer_worker.jobs.kimodo.job.KimodoJob

CMD ["ph-worker", "start"]
