# syntax=docker/dockerfile:1.6
# ============================================================
# Stage 1: Build CTranslate2 from source with custom HIP arches
# ============================================================
FROM rocm/dev-ubuntu-22.04:7.2 AS builder

ARG CT2_GIT_REF=master
ARG HIP_ARCHITECTURES="gfx1010;gfx1011;gfx1012;gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
ARG CXX_FLAGS="-msse4.1 -O3 -Wno-deprecated-literal-operator"
ARG HIP_FLAGS="-O3 -Wno-deprecated-literal-operator"

RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get update && \
    apt-get install -y --no-install-recommends \
        rocm-hip-runtime-dev hipblas-common-dev hipblas-dev hipcub-dev \
        hiprand-dev rocprim-dev rocrand-dev rocthrust-dev \
        python3-dev python3-pip wget git ca-certificates && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

WORKDIR /root

# Intel oneAPI MKL
ENV ONEAPI_VERSION=2025.3
RUN wget -q https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB && \
    apt-key add *.PUB && rm *.PUB && \
    echo "deb https://apt.repos.intel.com/oneapi all main" > /etc/apt/sources.list.d/oneAPI.list && \
    apt-get update && \
    apt-get install -y --no-install-recommends intel-oneapi-mkl-devel-$ONEAPI_VERSION && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

RUN python3 -m pip --no-cache-dir install cmake==3.22.*

# oneDNN
ENV ONEDNN_VERSION=3.10.2
RUN wget -q https://github.com/oneapi-src/oneDNN/archive/refs/tags/v${ONEDNN_VERSION}.tar.gz && \
    tar xf *.tar.gz && rm *.tar.gz && cd oneDNN-* && \
    cmake -DCMAKE_BUILD_TYPE=Release -DONEDNN_LIBRARY_TYPE=STATIC \
          -DONEDNN_BUILD_EXAMPLES=OFF -DONEDNN_BUILD_TESTS=OFF \
          -DONEDNN_ENABLE_WORKLOAD=INFERENCE \
          -DONEDNN_ENABLE_PRIMITIVE="CONVOLUTION;REORDER" \
          -DONEDNN_BUILD_GRAPH=OFF . && \
    make -j$(nproc) install && cd .. && rm -r oneDNN-*

# OpenMPI
ENV OPENMPI_VERSION=4.1.6
RUN wget -q https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-${OPENMPI_VERSION}.tar.bz2 && \
    tar xf *.tar.bz2 && rm *.tar.bz2 && cd openmpi-* && \
    ./configure && make -j$(nproc) install && cd .. && rm -r openmpi-*

# Clone CTranslate2 source instead of needing it in the build context
RUN git clone --recursive --branch ${CT2_GIT_REF} \
        https://github.com/OpenNMT/CTranslate2.git /root/ct2

WORKDIR /root/ct2

ENV CTRANSLATE2_ROOT=/opt/ctranslate2
ENV LD_LIBRARY_PATH=/usr/local/lib/:${LD_LIBRARY_PATH}

RUN mkdir build_tmp && cd build_tmp && \
    cmake -DCMAKE_INSTALL_PREFIX=${CTRANSLATE2_ROOT} \
          -DCMAKE_C_COMPILER=amdclang -DCMAKE_CXX_COMPILER=amdclang++ \
          -DWITH_HIP=ON -DWITH_MKL=ON -DWITH_DNNL=ON -DOPENMP_RUNTIME=COMP \
          -DCMAKE_HIP_ARCHITECTURES="${HIP_ARCHITECTURES}" \
          -DGPU_TARGETS="${HIP_ARCHITECTURES}" \
          -DCMAKE_BUILD_TYPE=Release \
          -DCMAKE_CXX_FLAGS="${CXX_FLAGS}" \
          -DCMAKE_HIP_FLAGS="${HIP_FLAGS}" \
          .. && \
    VERBOSE=1 make -j$(nproc) install

ENV LANG=en_US.UTF-8

RUN cd python && \
    python3 -m pip --no-cache-dir install -r install_requirements.txt && \
    python3 setup.py bdist_wheel --dist-dir $CTRANSLATE2_ROOT

# ============================================================
# Stage 2: Final runtime image (CT2 base + your app)
# ============================================================
FROM rocm/dev-ubuntu-22.04:7.2

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
    --mount=type=cache,target=/var/lib/apt,sharing=locked \
    apt-get update && apt-get install -y --no-install-recommends \
        rocm-hip-libraries openmpi-bin libgomp1 python3-pip \
        ffmpeg python3 curl gosu tzdata libomp5 libdw1 \
    && ln -s /usr/lib/x86_64-linux-gnu/libomp.so.5 /usr/lib/x86_64-linux-gnu/libomp.so \
    && ldconfig \
    && apt-get clean && rm -rf /var/lib/apt/lists/*

ENV CTRANSLATE2_ROOT=/opt/ctranslate2
ENV ROCM_ROOT=/opt/rocm

# Bring in the freshly built CTranslate2
COPY --from=builder $CTRANSLATE2_ROOT $CTRANSLATE2_ROOT
RUN --mount=type=cache,target=/root/.cache/pip \
    python3 -m pip install $CTRANSLATE2_ROOT/*.whl && \
    rm $CTRANSLATE2_ROOT/*.whl

# torch nightly with rocm7.2 — own layer so app dep changes don't re-download torch
RUN python3 -m pip install --no-cache-dir -U --pre \
        torch torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.2

# Application Python deps — merged into one layer
RUN --mount=type=cache,target=/root/.cache/pip \
    python3 -m pip install -U \
        numpy fastapi requests uvicorn python-multipart ffmpeg-python \
        watchdog tqdm more-itertools transformers tokenizers huggingface-hub av scipy \
        onnxruntime \
    && python3 -m pip install -U --no-deps faster-whisper stable-ts-whisperless

# Application files
COPY entrypoint.sh /
WORKDIR /subgen
COPY launcher.py subgen.py language_code.py /subgen/
RUN mkdir -p /cache && chmod 777 /cache

# Caches
ENV XDG_CACHE_HOME=/cache \
    HF_HOME=/cache/huggingface \
    MPLCONFIGDIR=/cache/matplotlib

# Library paths and HSA preload (matches your previous image)
ENV LD_LIBRARY_PATH="/opt/ctranslate2/lib:/opt/rocm-7.2.0/lib:/opt/rocm-7.2.0/lib64:${LD_LIBRARY_PATH}"
ENV LD_PRELOAD="/opt/rocm-7.2.0/lib/libhsa-runtime64.so.1"
ENV CT2_CUDA_ALLOCATOR="cub_caching"

ENTRYPOINT ["/entrypoint.sh"]
CMD ["python3", "-u", "launcher.py"]