ARG BUILD_BASE_IMAGE=registry.access.redhat.com/ubi9/ubi:9.6
ARG PYTHON_VERSION=3.12
ARG UV_EXTRA_INDEX_URL=https://repos.metax-tech.com/r/maca-pypi/simple
ARG UV_TRUSTED_HOST=repos.metax-tech.com

# may need passing a particular vllm version during build
ARG VLLM_VERSION
ARG MACA_VERSION
ARG CU_BRIDGE_VERSION=${MACA_VERSION}

#################### BASE BUILD IMAGE ####################
FROM ${BUILD_BASE_IMAGE} AS base
ARG UV_TRUSTED_HOST

# maca environment variables
ENV MACA_PATH=/opt/maca
ENV MACA_CLANG_PATH=/opt/maca/mxgpu_llvm/bin
ENV CUCC_PATH="${MACA_PATH}/tools/cu-bridge"
ENV CUDA_PATH=/root/cu-bridge/CUDA_DIR
ENV CUCC_CMAKE_ENTRY=2
ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
ENV PATH=/opt/mxdriver/bin:${MACA_PATH}/bin:${MACA_PATH}/mxgpu_llvm/bin:${MACA_PATH}/tools/cu-bridge/tools:${MACA_PATH}/tools/cu-bridge/bin:${PATH}
ENV LD_LIBRARY_PATH=/opt/mxdriver/lib:${MACA_PATH}/lib:${MACA_PATH}/mxgpu_llvm/lib:${MACA_PATH}/ompi/lib:${MACA_PATH}/ucx/lib:${LD_LIBRARY_PATH}

# uv environment variables
ENV VIRTUAL_ENV=/opt/venv
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_HTTP_TIMEOUT=6000
ENV UV_LINK_MODE=copy
ARG UV_EXTRA_INDEX_URL
ENV UV_EXTRA_INDEX_URL=${UV_EXTRA_INDEX_URL}
ARG UV_INDEX_URL
ENV UV_INDEX_URL=https://mirrors.aliyun.com/pypi/simple
ENV UV_TRUSTED_INDEX_HOST=mirrors.aliyun.com
ENV UV_OVERRIDE=/workspace/override.txt

# vllm compile option
ENV VLLM_INSTALL_PUNICA_KERNELS=1

# AI version arguments
ARG PYTHON_VERSION
ARG VLLM_VERSION
ARG VLLM_METAX_VERSION
ARG MACA_VERSION
ARG MEGATRON_VERSION
ARG SWIFT_VERSION
ARG CU_BRIDGE_VERSION
ARG TE_VERSION

WORKDIR /workspace
COPY override.txt /workspace/override.txt
COPY requirements_extra.txt /workspace/requirements_extra.txt

RUN printf "[metax-centos]\n\
name=Maca Driver Yum Repository\n\
baseurl=https://repos.metax-tech.com/r/metax-driver-centos-$(uname -m)/\n\
enabled=1\n\
gpgcheck=0" > /etc/yum.repos.d/metax-driver-centos.repo && \
    printf "[maca-sdk]\n\
name=Maca Sdk Yum Repository\n\
baseurl=https://repos.metax-tech.com/r/maca-sdk-rpm-$(uname -m)/\n\
enabled=1\n\
gpgcheck=0" > /etc/yum.repos.d/maca-sdk-rpm.repo

RUN yum install -y \
    python3-pip hostname \
    unzip vim git openblas-devel make cmake \
    ninja-build gcc g++ procps-ng \
    libibverbs librdmacm libibumad \
    binutils numactl-libs \
    metax-driver-${MACA_VERSION}* mxgvm \
    maca_sdk-${MACA_VERSION}* \
    && yum clean all && rm -rf /var/cache/yum /tmp/*

RUN python3 -m pip install uv -i $UV_INDEX_URL --trusted-host ${UV_TRUSTED_INDEX_HOST} && \
    uv venv /opt/venv --python=${PYTHON_VERSION}

RUN python3 --version && \
    uv self version


RUN git clone --depth 1 --branch ${SWIFT_VERSION} https://github.com/modelscope/ms-swift.git && \
    git clone --depth 1 --branch ${VLLM_METAX_VERSION} https://github.com/MetaX-MACA/vLLM-metax.git && \
    git clone --depth 1 --branch ${VLLM_VERSION} https://github.com/vllm-project/vllm.git && \
    git clone --depth 1 --branch ${MEGATRON_VERSION} https://github.com/NVIDIA/Megatron-LM.git

# Step 1: install MACA SDK, Metax-Driver and cu-bridge
# Metax-Driver mainly contains vbios and kmd files, which are not needed in a container.
# Here we keep the mx-smi management tool. Kernel version mismatch errors are ignored.
RUN cd /tmp/ && \
    export MACA_PATH=/opt/maca && \
    curl -o ${CU_BRIDGE_VERSION}.zip -LsSf https://gitee.com/metax-maca/cu-bridge/repository/archive/${CU_BRIDGE_VERSION}.zip && \
    unzip ${CU_BRIDGE_VERSION}.zip && \
    mv cu-bridge-${CU_BRIDGE_VERSION} cu-bridge && \
    chmod -R 755 cu-bridge && \
    cd cu-bridge && \
    mkdir build && cd build && \
    cmake -DCMAKE_INSTALL_PREFIX=/opt/maca/tools/cu-bridge ../ && \
    make && make install

# Step 2: trim unused MACA packages and install build prerequisites
RUN cd vLLM-metax && \
    uv pip install -r requirements/build.txt && \
    uv pip install build

# Step 3: install Metax python requirements
RUN cd vLLM-metax && \
    UV_HTTP_TIMEOUT=960 uv pip install -r requirements/maca.txt --trusted-host ${UV_TRUSTED_HOST}

# Step 4: build vLLM with empty device to avoid CUDA dependency
RUN cd vllm && \
    python3 use_existing_torch.py && \
    uv pip install -r requirements/build/cuda.txt && \
    VLLM_TARGET_DEVICE=empty uv pip install -v . --no-build-isolation

# Step 5: build vLLM-metax
RUN cd vLLM-metax && \
    uv pip install -r requirements/build.txt && \
    python3 -m build -w -n && \
    uv pip install dist/*.whl

# Step 6: install Megatron-LM
RUN cd /workspace/Megatron-LM && \
    if [ -f megatron/legacy/fused_kernels/__init__.py ]; then \
        sed -i 's/nvcc/cucc/g' megatron/legacy/fused_kernels/__init__.py; \
    fi && \
    uv pip install .

# Step 7: install transformer-engine
RUN uv pip install transformer_engine==${TE_VERSION} -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host ${UV_TRUSTED_HOST}

# Step 8: patch and install ms-swift v4.2.0 with Megatron extra dependencies
RUN sed -i '0,/^\(from \|import \)/{s//import vllm_metax.patch\n&/}' ms-swift/swift/__init__.py && \
    cd ms-swift && \
    uv pip install '.[megatron]'

# Step 9: install optional runtime dependencies used by swift 4.2.0
RUN uv pip install deepspeed -i https://repos.metax-tech.com/r/maca-pypi/simple --trusted-host ${UV_TRUSTED_HOST}
RUN uv pip install pip
RUN uv pip install -r requirements_extra.txt
RUN ln -sf ${CUDA_PATH}/bin/nvcc ${CUDA_PATH}/bin/cucc

# vllm installation may bring in incompatible CUDA-only wheels. Remove them here.
RUN uv pip uninstall flashinfer-python cupy-cuda12x
#################### FINAL IMAGE ####################