# LLaMA-Factory + Megatron Bridge (CUDA) runtime
# Mirrors the verified host env: Python 3.12, PyTorch 2.12.1+cu126,
# TransformerEngine 2.17, megatron-core 0.18, megatron-bridge 0.5.
#
# CUDA user-mode libs come from the PyTorch cu126 wheels (same as the host venv).
# Layers are intentionally few (helps vfs / nested-docker disk usage).
#
# Build from repo root:
#   bash docker/docker-cuda/build-megatron.sh

ARG BASE_IMAGE=ubuntu:24.04
FROM ${BASE_IMAGE}

ARG PIP_INDEX=https://mirrors.aliyun.com/pypi/simple
ARG PYPI_TRUSTED_HOST=mirrors.aliyun.com
ARG TORCH_INDEX=https://download.pytorch.org/whl/cu126
ARG APT_MIRROR=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_ROOT_USER_ACTION=ignore \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    PIP_CONSTRAINT="" \
    MAX_JOBS=8 \
    PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    VLLM_WORKER_MULTIPROC_METHOD=spawn \
    DISABLE_VERSION_CHECK=1 \
    USE_MEGATRON_BRIDGE=1 \
    GRADIO_SERVER_PORT=7860 \
    API_PORT=8000 \
    http_proxy= \
    https_proxy=

SHELL ["/bin/bash", "-c"]
WORKDIR /app

# System deps + Ubuntu Python 3.12
RUN if [ -f /etc/apt/sources.list.d/ubuntu.sources ]; then \
      sed -i "s|http://archive.ubuntu.com/ubuntu|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu|${APT_MIRROR}|g" /etc/apt/sources.list.d/ubuntu.sources; \
    elif [ -f /etc/apt/sources.list ]; then \
      sed -i "s|http://archive.ubuntu.com/ubuntu/|${APT_MIRROR}|g; s|http://security.ubuntu.com/ubuntu/|${APT_MIRROR}|g" /etc/apt/sources.list; \
    fi && \
    apt-get update && apt-get install -y --no-install-recommends \
        ca-certificates curl git vim wget \
        build-essential ninja-build cmake libgomp1 zip unzip \
        python3 python3-pip python3-dev python3-venv && \
    ln -sf /usr/bin/python3 /usr/bin/python && \
    rm -rf /var/lib/apt/lists/* && \
    pip install --no-cache-dir --upgrade pip setuptools wheel packaging ninja pybind11 \
        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX}

COPY . /app

# PyTorch + TE + Megatron Bridge + LLaMA-Factory in one layer
RUN pip install --no-cache-dir \
        torch==2.12.1 torchvision==0.27.1 torchaudio==2.11.0 \
        --index-url ${TORCH_INDEX} && \
    pip install --no-cache-dir --no-build-isolation \
        "transformer-engine[pytorch]==2.17.0" \
        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
    pip install --no-cache-dir --no-build-isolation \
        "megatron-bridge==0.5.0" \
        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
    pip install --no-cache-dir --no-build-isolation -e . \
        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
    pip install --no-cache-dir --no-build-isolation \
        -r requirements/metrics.txt \
        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
    pip install --no-cache-dir --no-build-isolation \
        "megatron-bridge==0.5.0" \
        --trusted-host ${PYPI_TRUSTED_HOST} --index-url ${PIP_INDEX} && \
    python - <<'PY'
import torch
import megatron.core
import transformer_engine
from megatron.bridge import AutoBridge  # noqa: F401
import llamafactory
print("torch", torch.__version__, "cuda", torch.version.cuda)
print("te", transformer_engine.__version__)
print("megatron-bridge import ok")
PY

EXPOSE 7860 8000
CMD ["bash"]