# =============================================================================
# HY-Motion Worker Image — ghcr.io/poly-hammer/ph-hy-motion
#
# Universal image for both self-hosted (Hatchet) and Modal workers.
# Weights are volume-mounted at runtime, NOT baked in.
#
# Self-hosted usage:
#   docker volume create hy-motion-weights
#   docker run --gpus all -v hy-motion-weights:/weights ph-hy-motion \
#       ph-worker download-weights --job hy-motion --all
#   docker run --gpus all -v hy-motion-weights:/weights ph-hy-motion \
#       ph-worker start --token <WORKER_TOKEN>
#
# Modal usage:
#   modal.Image.from_registry("ghcr.io/poly-hammer/ph-hy-motion:latest", ...)
#
# Build (from workers/self-hosted/ context):
#   docker compose build hy-motion
#
# Multi-CUDA build:
#   docker compose build hy-motion-cu121
#
# See workers/self-hosted/README.md for license notices.
# =============================================================================

# ---------- CUDA / PyTorch build arguments ----------
# One variant per row of the ph-hy-motion 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.
#
#   cu121  CUDA 12.1 / torch 2.5.1 — driver >=530, sm_75..sm_90
#   cu128  CUDA 12.8 / torch 2.7.1 — driver >=570, sm_75..sm_120  (default)
#   cu130  CUDA 13.0 / torch 2.9.0 — driver >=580, sm_75..sm_120
#
# The base tag and the wheel index must agree: a CUDA 13 base needlessly
# raises the host driver floor to 580 for wheels that only need 12.8.
ARG CUDA_TAG=12.8.1-cudnn-devel-ubuntu24.04
ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128
ARG TORCH_VERSION=2.7.1
ARG TORCHVISION_VERSION=0.22.1

FROM nvidia/cuda:${CUDA_TAG}

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

ENV DEBIAN_FRONTEND=noninteractive

# ---------- System packages ----------
RUN apt-get update && apt-get install -y --no-install-recommends \
        python3 python3-venv python3-dev python3-pip \
        git git-lfs curl ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# ---------- Python venv (PEP 668 — Ubuntu 24.04 blocks system-wide pip) ----------
RUN python3 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

RUN pip install --no-cache-dir --upgrade pip setuptools wheel

# 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 the interpreter is already
# provisioned, so only the venv creation step is skipped.
ARG TORCH_INDEX_URL
ARG TORCH_VERSION
ARG TORCHVISION_VERSION
ENV HY_MOTION_REPO_PATH=/opt/hy-motion
COPY . /opt/worker-src
RUN PH_IN_DOCKER=1 \
    PH_PYTHON=/opt/venv/bin/python \
    PH_GIT="$(command -v git)" \
    PH_WORKER_SRC=/opt/worker-src \
    PH_HY_MOTION_SRC=${HY_MOTION_REPO_PATH} \
    bash /opt/worker-src/poly_hammer_worker/jobs/hy_motion/recipes/install.sh \
        --profile "hy-motion/linux-x64/docker" --prefix /opt/venv \
        --wheelhouse /opt/worker-src/wheelhouse \
        --torch-channel "${TORCH_INDEX_URL##*/}" \
        --torch-version "${TORCH_VERSION}" \
        --torchvision-version "${TORCHVISION_VERSION}" \
    && rm -rf /opt/worker-src

# ---------- Runtime configuration ----------
ENV WEIGHTS_PATH=/weights
ENV HF_HUB_ENABLE_HF_TRANSFER=1
# High-performance HF Xet transfer backend — large weight files download
# noticeably faster on first run.
ENV HF_XET_HIGH_PERFORMANCE=1
ENV USE_HF_MODELS=0
ENV HF_HOME="${WEIGHTS_PATH}/huggingface"
ENV TRANSFORMERS_CACHE="${WEIGHTS_PATH}/huggingface"

ENV PH_JOB_CLASS="poly_hammer_worker.jobs.hy_motion.job.HyMotionAllJob"
ENV PYTHONPATH="${HY_MOTION_REPO_PATH}"

VOLUME /weights

# No ENTRYPOINT — Modal needs to run `python` directly inside the container.
# Self-hosted usage still works: `docker run ... ph-worker start --token X`
CMD ["ph-worker", "start"]
