# =============================================================================
# MinerU on ROCm 7.2.1 Docker Image（HK 云服务器 / 海外版）
# 原生 Linux + Ubuntu 24.04 + ROCm 7.2.1 + PyTorch 2.11.0 + vllm + MinerU 3.2.0
# =============================================================================

FROM ubuntu:24.04

# -- 构建参数 ---------------------------------------------------------------
ARG ARCH=gfx1201
ARG PYTHON_VER=3.12
ARG VENV=/opt/mineru_venv
ARG TORCH_INDEX=https://download.pytorch.org/whl/rocm7.2

# -- GitHub 访问（国内无镜像，需代理）----------------------------------------
ARG GIT_PROXY=http://127.0.0.1:8118

# -- 国内镜像 ----------------------------------------------------------------
ARG PIP_INDEX=https://pypi.tuna.tsinghua.edu.cn/simple

# -- 环境变量 ---------------------------------------------------------------
ENV DEBIAN_FRONTEND=noninteractive \
    PATH=/opt/rocm/bin:/opt/rocm/llvm/bin:${VENV}/bin:${PATH} \
    PYTORCH_ROCM_ARCH=${ARCH} \
    FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE \
    MINERU_MODEL_SOURCE=huggingface \
    TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 \
    HSA_ENABLE_SDMA=1 \
    VLLM_TARGET_DEVICE=rocm

WORKDIR /opt

# ===========================================================================
# 阶段 1：换国内源 + 安装 ROCm 7.2.1
# ===========================================================================
RUN sed -i 's|http://.*archive.ubuntu.com|http://mirrors.tuna.tsinghua.edu.cn|g' /etc/apt/sources.list.d/ubuntu.sources && \
    sed -i 's|http://.*security.ubuntu.com|http://mirrors.tuna.tsinghua.edu.cn|g' /etc/apt/sources.list.d/ubuntu.sources && \
    apt-get update && apt-get install -y --no-install-recommends \
    wget curl ca-certificates gnupg software-properties-common && \
    wget -q https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
        gpg --dearmor | tee /etc/apt/trusted.gpg.d/rocm.gpg > /dev/null && \
    echo 'deb [arch=amd64] https://repo.radeon.com/rocm/apt/7.2.1 noble main' \
        > /etc/apt/sources.list.d/rocm.list && \
    printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' \
        > /etc/apt/preferences.d/rocm-pin-600 && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
        rocminfo rocm-device-libs hip-dev miopen-hip && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ===========================================================================
# 阶段 2：ROCm 头文件补丁
# ===========================================================================
RUN set -ex && \
    # 补丁 1: hipcc/clang 符号链接
    ln -sf /usr/bin/hipvars.pm /usr/share/perl5/hipvars.pm && \
    ([ -f /usr/bin/hipcc.pl ] && ln -sf /usr/bin/hipcc.pl /opt/rocm/bin/hipcc) || echo "hipcc.pl not found, keeping apt-installed hipcc" && \
    ln -sf /opt/rocm/llvm/bin/clang-22 /opt/rocm/llvm/bin/clang-17 && \
    ln -sf /opt/rocm/llvm/bin/clang++   /opt/rocm/llvm/bin/clang++-17 && \
    # 补丁 2: __hip_internal::conditional → std::conditional
    find /opt/rocm/include/hip -name "*.h" \
        -exec sed -i 's/__hip_internal::conditional/std::conditional/g' {} + && \
    # 补丁 3: warpSize 常量
    find /opt/rocm/include/hip -name "amd_warp_functions.h" \
        -exec sed -i 's/static constexpr int warpSize = __AMDGCN_WAVEFRONT_SIZE;/constexpr int warpSize = 32;/g' {} + && \
    # 补丁 4: __activemask() → __builtin_amdgcn_read_exec()
    sed -i 's/__activemask()/__builtin_amdgcn_read_exec()/g' \
        /opt/rocm/include/hip/amd_detail/amd_warp_sync_functions.h && \
    echo "ROCm 7.2.1 header patches applied."

# ===========================================================================
# 阶段 3：系统依赖 + Python 3.12
# ===========================================================================
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential git ninja-build pkg-config \
    python${PYTHON_VER} python${PYTHON_VER}-venv python${PYTHON_VER}-dev \
    libnuma-dev libdrm2 libhwloc-dev libgl1 \
    libgomp1 libopenblas0 && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ===========================================================================
# 阶段 4：CMake 4.0
# ===========================================================================
RUN cd /tmp && \
    wget -q https://github.com/Kitware/CMake/releases/download/v4.0.0/cmake-4.0.0-linux-x86_64.tar.gz && \
    tar -xzf cmake-4.0.0-linux-x86_64.tar.gz && \
    cp -r cmake-4.0.0-linux-x86_64/bin/* /usr/local/bin/ && \
    cp -r cmake-4.0.0-linux-x86_64/share/* /usr/local/share/ && \
    rm -rf cmake-4.0.0-linux-x86_64* && \
    cmake --version

# ===========================================================================
# 阶段 5：Python 虚拟环境 + PyTorch ROCm
# ===========================================================================
RUN python${PYTHON_VER} -m venv ${VENV} && \
    mkdir -p /root/.pip && \
    echo "[global]" > /root/.pip/pip.conf && \
    echo "index-url = ${PIP_INDEX}" >> /root/.pip/pip.conf && \
    ${VENV}/bin/pip install --no-cache-dir -U pip setuptools wheel && \
    # PyTorch ROCm wheels 在海外 CDN，通过代理下载
    export http_proxy=${GIT_PROXY} https_proxy=${GIT_PROXY} && \
    ${VENV}/bin/pip install --no-cache-dir --pre \
        torch==2.11.0+rocm7.2 \
        torchvision \
        pytorch-triton-rocm \
        --index-url ${TORCH_INDEX} && \
    ${VENV}/bin/python -c "import torch; print('PyTorch:', torch.__version__); assert torch.version.hip is not None"

# ===========================================================================
# 阶段 6：ROCm 开发包
# ===========================================================================
RUN apt-get update && apt-get install -y --no-install-recommends \
    hipblas-dev hiprand-dev hipsparse-dev hipsparselt-dev \
    hipsolver-dev hipcub-dev rocprim-dev rocthrust-dev \
    rocblas-dev rocrand-dev hipfft-dev hipblaslt-dev \
    rocsolver-dev rocm-cmake rocm-core && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ===========================================================================
# 阶段 6.5：安装 amdsmi
# ===========================================================================
RUN if [ -d /opt/rocm/share/amd_smi ]; then \
        cp -r /opt/rocm/share/amd_smi /opt/amd_smi && \
        cd /opt/amd_smi && ${VENV}/bin/pip install --no-cache-dir . --no-build-isolation && \
        echo "amdsmi installed."; \
    else \
        echo "amdsmi source not found, skipping."; \
    fi

# ===========================================================================
# 阶段 7：amd-aiter + flash_attn
# ===========================================================================
RUN set -ex && \
    git config --global http.proxy ${GIT_PROXY} && \
    git config --global https.proxy ${GIT_PROXY} && \
    cd /opt && git clone --recursive --depth 1 https://github.com/ROCm/aiter.git && \
    ${VENV}/bin/pip install --no-cache-dir -e /opt/aiter && \
    cd /opt && git clone --recursive https://github.com/Dao-AILab/flash-attention.git && \
    cd flash-attention && git checkout bba578d43974c1d3ba157ab597124dd0fe2ccdb4 && \
    ${VENV}/bin/pip install --no-cache-dir --no-build-isolation -e /opt/flash-attention && \
    ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten: {v}'; print('PyTorch OK:', v)"

# ===========================================================================
# 阶段 8a：编译 vllm（cmake + ninja，最耗时，单独缓存）
# ===========================================================================
RUN set -ex && \
    git config --global http.proxy ${GIT_PROXY} && \
    git config --global https.proxy ${GIT_PROXY} && \
    ${VENV}/bin/pip install --no-cache-dir -U \
        "setuptools>=77.0.3" setuptools_scm setuptools_rust wheel && \
    cd /opt && git clone --depth 1 https://github.com/vllm-project/vllm.git && \
    cd /opt/vllm && \
    if [ -f csrc/mamba/mamba_ssm/selective_scan.h ]; then \
        sed -i '109,121s/^/\/\/ /' csrc/mamba/mamba_ssm/selective_scan.h && \
        echo "vllm mamba operator+ patch applied."; \
    else \
        echo "vllm mamba selective_scan.h not found (upstream removed), skipping patch."; \
    fi && \
    mkdir -p /opt/rocm/lib/cmake/hiprand && \
    printf 'include(/opt/rocm/lib/cmake/rocrand/rocrand-config.cmake)\nif(TARGET roc::rocrand AND NOT TARGET hip::hiprand)\n  add_library(hip::hiprand ALIAS roc::rocrand)\nendif()\n' \
        > /opt/rocm/lib/cmake/hiprand/hiprand-config.cmake && \
    mkdir -p /opt/rocm/lib/cmake/hipblas && \
    printf 'include(/opt/rocm/lib/cmake/rocblas/rocblas-config.cmake)\nif(TARGET roc::rocblas AND NOT TARGET hip::hipblas)\n  add_library(hip::hipblas ALIAS roc::rocblas)\nendif()\n' \
        > /opt/rocm/lib/cmake/hipblas/hipblas-config.cmake && \
    # 创建 HIP cmake config（apt 安装的 hip-dev 缺少 find_package 需要的配置）
    mkdir -p /opt/rocm/lib/cmake/hip && \
    printf 'set(HIP_FOUND TRUE)\nset(HIP_VERSION "7.2.1")\nset(HIP_PLATFORM amd)\nset(HIP_COMPILER /opt/rocm/llvm/bin/clang++)\nset(HIP_RUNTIME rocm)\nset(HIP_INCLUDE_DIRS /opt/rocm/include)\nset(HIP_LIBRARIES /opt/rocm/lib/libamdhip64.so)\nif(NOT TARGET hip::host)\n  add_library(hip::host INTERFACE IMPORTED)\n  set_target_properties(hip::host PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${HIP_INCLUDE_DIRS}")\nendif()\nif(NOT TARGET hip::device)\n  add_library(hip::device INTERFACE IMPORTED)\n  set_target_properties(hip::device PROPERTIES INTERFACE_COMPILE_OPTIONS "${HIP_COMPILER}")\nendif()\n' \
        > /opt/rocm/lib/cmake/hip/hip-config.cmake && \
    # 补丁: ROCm 7.2.1 中 hip_version.h 已移除, 创建兼容头文件供 PyTorch LoadHIP.cmake 使用
    if [ ! -f /opt/rocm/include/hip/hip_version.h ]; then \
        mkdir -p /opt/rocm/include/hip && \
        printf '#pragma once\n#define HIP_VERSION_MAJOR 7\n#define HIP_VERSION_MINOR 2\n#define HIP_VERSION_PATCH 53211\n#define HIP_VERSION_GITDATE 0\n#define HIP_VERSION (HIP_VERSION_MAJOR * 10000000 + HIP_VERSION_MINOR * 100000 + HIP_VERSION_PATCH)\n' \
            > /opt/rocm/include/hip/hip_version.h && \
        echo "hip_version.h created."; \
    fi && \
    mkdir -p /opt/vllm_build && \
    export HIP_PLATFORM=amd HIP_PATH=/opt/rocm ROCM_PATH=/opt/rocm && \
    cmake -S /opt/vllm -B /opt/vllm_build -G Ninja \
        -DCMAKE_BUILD_TYPE=RelWithDebInfo \
        -DVLLM_TARGET_DEVICE=rocm \
        -DVLLM_PYTHON_EXECUTABLE=${VENV}/bin/python \
        -DHIP_ROOT_DIR=/opt/rocm \
        -DROCM_PATH=/opt/rocm \
        -DCMAKE_HIP_ARCHITECTURES=${ARCH} \
        -DCMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \
        -DHIP_COMPILER=/opt/rocm/llvm/bin/clang++ \
        -DHIP_PATH=/opt/rocm \
        -DHIP_INCLUDE_DIR=/opt/rocm/include \
        -DCMAKE_PREFIX_PATH="/opt/rocm;${VENV}/lib/python${PYTHON_VER}/site-packages/torch/share/cmake" && \
    cd /opt/vllm_build && ninja -j8 && \
    cp /opt/vllm_build/*.abi3.so /opt/vllm/vllm/ && \
    echo "vllm C++ build complete (layer cached)."

# ===========================================================================
# 阶段 8b：安装 vllm + 平台补丁 + 验证
# ===========================================================================
RUN set -ex && \
    cd /opt/vllm && ${VENV}/bin/pip install --no-cache-dir -e . --no-build-isolation && \
    ${VENV}/bin/pip install --no-cache-dir regex && \
    ${VENV}/bin/python /opt/patch_vllm_platform.py && \
    ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten: {v}'; print('PyTorch OK:', v)" && \
    ${VENV}/bin/pip install --no-cache-dir --force-reinstall \
        torch==2.11.0+rocm7.2 torchvision pytorch-triton-rocm \
        --index-url ${TORCH_INDEX} && \
    ${VENV}/bin/pip uninstall -y triton triton-rocm 2>/dev/null; \
    ${VENV}/bin/python -c "import vllm; print('vllm import OK:', vllm.__version__)" && \
    rm -rf /opt/vllm_build

# ===========================================================================
# 阶段 8.5：复制辅助脚本
# ===========================================================================
COPY scripts/apply_mineru_patches.py /opt/apply_mineru_patches.py
COPY scripts/cache_warmer.py /opt/cache_warmer.py
COPY scripts/patch_vllm_platform.py /opt/patch_vllm_platform.py

# ===========================================================================
# 阶段 9：安装 MinerU + RDNA 适配补丁
# ===========================================================================
RUN set -ex && \
    ${VENV}/bin/pip install --no-cache-dir 'mineru[core]' && \
    ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten: {v}'; print('PyTorch OK:', v)" && \
    ${VENV}/bin/python /opt/apply_mineru_patches.py

# ===========================================================================
# 阶段 10：入口与最终验证
# ===========================================================================
RUN echo 'source /opt/mineru_venv/bin/activate' >> /etc/bash.bashrc && \
    ${VENV}/bin/python -c "import torch, vllm, mineru; print('='*50); print('MinerU ROCm Docker Image Ready'); print(f'  PyTorch: {torch.__version__}'); print(f'  vllm:   {vllm.__version__}'); print(f'  MinerU:  {mineru.__version__}'); print('='*50)"

ENTRYPOINT ["/bin/bash", "-c"]
CMD ["bash"]