# syntax=docker/dockerfile:1
#
# Training image for exp-002 (UnifiedRenderer distillation, continued from exp-001).
# Build context must be the repo root, e.g.:
#   docker build -f experiments/exp-002/Dockerfile -t pseudopros-exp002:latest .
#
# Targets NVIDIA H100 (Hopper, compute capability 9.0) via CUDA 12.1 wheels.
# The cluster's CUDA 13.0 / driver 580.x combination is newer than the toolkit
# this image links against; NVIDIA's driver forward-compatibility covers that,
# so there's no need to chase the exact on-cluster CUDA version here.
FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

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

# Python 3.12 (matches pyproject.toml requires-python) via deadsnakes - Ubuntu
# 22.04's CUDA base images ship Python 3.10 by default. Plus runtime libs
# needed by opencv-python/mediapipe (libgl1/libglib2.0-0) and ffmpeg for
# HDTF clip trimming.
RUN apt-get update && apt-get install -y --no-install-recommends \
      software-properties-common \
    && add-apt-repository -y ppa:deadsnakes/ppa \
    && apt-get update && apt-get install -y --no-install-recommends \
      python3.12 python3.12-venv python3.12-dev \
      ffmpeg libgl1 libglib2.0-0 git curl \
    && curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12 \
    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 \
    && update-alternatives --install /usr/bin/python  python  /usr/bin/python3.12 1 \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

# CUDA-enabled torch/torchvision, pinned to the exact version pyproject.toml's
# ">=2.2,<2.3" range allows. Installed before `pip install -e .` so that step
# sees the requirement already satisfied and doesn't pull the CPU wheel over it.
RUN pip install \
      torch==2.2.2+cu121 torchvision==0.17.2+cu121 \
      --index-url https://download.pytorch.org/whl/cu121

# pseudopros itself (ArcFace embedder, FOMM, LDM synthesizer) - needed by the
# precompute phase, which still runs the current production pipeline to
# generate distillation targets.
COPY pyproject.toml README.md ./
COPY src ./src
RUN pip install -e .

# Training-only deps, not part of the installable pseudopros package.
RUN pip install \
      mediapipe==0.10.* \
      tqdm>=4.66 \
      yt-dlp>=2024.1

COPY experiments/exp-002/train.py experiments/exp-002/eval.py \
     experiments/exp-002/download_hdtf.py experiments/exp-002/entrypoint.sh \
     ./experiments/exp-002/
RUN chmod +x ./experiments/exp-002/entrypoint.sh

# OpenShift's restricted SCC runs containers as an arbitrary non-root UID with
# GID 0 - group-own and group-write everything under /app so that works
# regardless of which UID the platform assigns.
RUN useradd -u 1001 -g 0 -M -s /sbin/nologin appuser \
    && chgrp -R 0 /app && chmod -R g=u /app
USER 1001

ENTRYPOINT ["/app/experiments/exp-002/entrypoint.sh"]
CMD ["--help"]
