# =============================================================================
# MinerU on ROCm 7.2.1 Docker Image
# 原生 Linux + Ubuntu 24.04 + ROCm 7.2.1 + PyTorch 2.11.0 + vllm + MinerU 3.2.0
#
# 构建前请根据你的 GPU 修改 ARCH 参数（默认 gfx1201 = RX 9070）
# =============================================================================

FROM ubuntu:24.04

# -- 构建参数 ---------------------------------------------------------------
# 改成你的 gfx 代号：gfx1201(RX 9070) gfx1200(RX 9060) gfx1100(RX 7900) gfx1101(RX 7800/7700) gfx1030(RX 6900/6800)
ARG ARCH=gfx1201
ARG PYTHON_VER=3.12
ARG VENV=/opt/mineru_venv
ARG TORCH_INDEX=https://download.pytorch.org/whl/rocm7.2

# -- 环境变量 ---------------------------------------------------------------
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 apt-get update && apt-get install -y --no-install-recommends \
    wget curl ca-certificates gnupg software-properties-common && \
    # 添加 AMD ROCm 仓库
    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 && \
    apt-get update && \
    # 安装 ROCm 基础组件
    apt-get install -y --no-install-recommends \
        rocminfo hip-dev miopen-hip && \
    # 修复 rocminfo / rocm-device-libs 版本（替换 Ubuntu 自带旧版）
    apt-get install -y --allow-downgrades \
        rocminfo=1.0.0.70201-38~24.04 \
        rocm-device-libs=1.0.0.70201-38~24.04 && \
    # 清理
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ===========================================================================
# 阶段 2：ROCm 头文件补丁（LLVM 22 兼容性修复）
# 这些是 ROCm 7.2.1 在 24.04 上的已知问题，每次 apt 升级 ROCm 后需重新应用
# ===========================================================================
RUN set -ex && \
    # 补丁 1: hipcc/clang 符号链接（hipcc.pl 硬编码 clang-17，实际是 clang-22）
    ln -sf /usr/bin/hipvars.pm /usr/share/perl5/hipvars.pm && \
    ln -sf /usr/bin/hipcc.pl /opt/rocm/bin/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 常量（__AMDGCN_WAVEFRONT_SIZE 在 LLVM 22 未定义）
    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()
    # 注意：只改 amd_warp_sync_functions.h，不要动 amd_warp_functions.h（那是定义本身）
    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 \
    # vllm 运行时依赖
    libgomp1 libopenblas0 && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ===========================================================================
# 阶段 4：CMake 4.0（vllm 要求 ≥ 4.0，Ubuntu 24.04 自带 3.28 不够）
# ===========================================================================
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} && \
    ${VENV}/bin/pip install --no-cache-dir -U pip setuptools wheel && \
    # 安装 PyTorch ROCm 版（锁定 2.11，≥ 2.12 在部分环境下有 rocprofiler 问题）
    ${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__); print('ROCm:', torch.version.hip); assert torch.version.hip is not None"

# ===========================================================================
# 阶段 6：ROCm 开发包（vllm 编译必需）
# ===========================================================================
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 && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

# ===========================================================================
# 阶段 7：amd-aiter + flash_attn
# ===========================================================================
RUN set -ex && \
    # aiter（AMD 优化的 attention 算子）
    cd /opt && git clone --recursive --depth 1 https://github.com/ROCm/aiter.git && \
    ${VENV}/bin/pip install --no-cache-dir -e /opt/aiter && \
    # flash_attn（Triton AMD 后端，锁定已验证的 commit）
    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 && \
    # 验证 PyTorch 没被覆盖
    ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten: {v}'; print('PyTorch OK:', v)"

# ===========================================================================
# 阶段 8：编译 vllm
# ===========================================================================
RUN set -ex && \
    # setuptools 升级（PEP 639 兼容）
    ${VENV}/bin/pip install --no-cache-dir -U \
        "setuptools>=77.0.3" setuptools_scm setuptools_rust wheel && \
    # 克隆 vllm main
    cd /opt && git clone --depth 1 https://github.com/vllm-project/vllm.git && \
    # 补丁 5：注释掉 vllm mamba 模块的 operator+ 定义（ROCm 7.2 头文件已自带）
    cd /opt/vllm && \
    sed -i '109,121s/^/\/\/ /' csrc/mamba/mamba_ssm/selective_scan.h && \
    echo "vllm mamba operator+ patch applied." && \
    # cmake 配置
    mkdir -p /opt/vllm_build && \
    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_PREFIX_PATH="${VENV}/lib/python${PYTHON_VER}/site-packages/torch/share/cmake" && \
    # ninja 编译（-j4 防 OOM，内存 > 32GB 可调高）
    cd /opt/vllm_build && ninja -j4 && \
    # 安装 .so 到 vllm 源码目录
    cp /opt/vllm_build/*.abi3.so /opt/vllm/vllm/ && \
    # pip install vllm（让 pip 解析运行时依赖：xgrammar, compressed_tensors 等）
    cd /opt/vllm && ${VENV}/bin/pip install --no-cache-dir -e . --no-build-isolation && \
    # 验证 PyTorch 没被 vllm 依赖覆盖
    ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten by vllm deps: {v}'; print('PyTorch OK:', v)" && \
    # 清理可能的 CUDA triton 残余
    ${VENV}/bin/pip uninstall -y triton triton-rocm 2>/dev/null; \
    ${VENV}/bin/pip install --no-cache-dir --force-reinstall \
        torch==2.11.0+rocm7.2 torchvision pytorch-triton-rocm \
        --index-url ${TORCH_INDEX} && \
    # 最终验证 vllm 平台检测
    ${VENV}/bin/python -c "
from vllm.platforms import current_platform
print('Platform:', type(current_platform).__name__)
print('is_rocm:', current_platform.is_rocm())
print('device_type:', current_platform.device_type)
assert current_platform.is_rocm(), 'vllm ROCm detection failed!'
print('vllm OK')
" && \
    # 清理构建目录（减小镜像体积，约 3-5GB）
    rm -rf /opt/vllm_build

# ===========================================================================
# 阶段 9：安装 MinerU + RDNA 适配补丁
# ===========================================================================
RUN set -ex && \
    ${VENV}/bin/pip install --no-cache-dir 'mineru[core]' && \
    # 验证 PyTorch 没被覆盖
    ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten: {v}'; print('PyTorch OK:', v)" && \
    # 定位 mineru infer 目录
    MINERU_INFER_DIR=$(${VENV}/bin/python -c "import mineru.model.utils.tools.infer; import os; print(os.path.dirname(mineru.model.utils.tools.infer.__file__))") && \
    echo "MinerU infer dir: ${MINERU_INFER_DIR}" && \
    # --- Patch A: predict_rec.py imgW 对齐到 32 ---
    ${VENV}/bin/python -c "
import re
f = '${MINERU_INFER_DIR}/predict_rec.py'
c = open(f).read()
# 在 imgW = max(min(... 之后插入 imgW = math.ceil(imgW / 32) * 32
old = '(imgW = max\(min\(imgW, self\.limited_max_width\), self\.limited_min_width\)\n)'
new = r'\1        imgW = math.ceil(imgW / 32) * 32\n'
c2 = re.sub(old, new, c)
if c2 == c:
    # 尝试找已经插入过的情况
    if 'math.ceil(imgW / 32)' not in c:
        raise RuntimeError('Patch A: cannot find imgW line in predict_rec.py')
    else:
        print('Patch A: already applied')
else:
    open(f, 'w').write(c2)
    print('Patch A: imgW 32-align inserted')
" && \
    # --- Patch B: predict_rec.py 批次填充 ---
    ${VENV}/bin/python -c "
f = '${MINERU_INFER_DIR}/predict_rec.py'
c = open(f).read()
# 在 norm_img_batch = np.concatenate(norm_img_batch) 前插入 padding 逻辑
old = '( {8}norm_img_batch = np\.concatenate\(norm_img_batch\))'
new = '''        actual_batch_size = len(norm_img_batch)
        if actual_batch_size < batch_num:
            pad_size = batch_num - actual_batch_size
            pad_img = np.zeros_like(norm_img_batch[0])
            for _ in range(pad_size):
                norm_img_batch.append(pad_img)
\\1'''
import re
c2 = re.sub(old, new, c)
if c2 == c:
    if 'actual_batch_size' not in c:
        raise RuntimeError('Patch B: cannot find norm_img_batch concatenation')
    else:
        print('Patch B: already applied')
else:
    open(f, 'w').write(c2)
    print('Patch B: batch padding inserted')
# 修改 range(len(rec_result)) → range(actual_batch_size)
c3 = open(f).read()
c4 = re.sub(r'for rno in range\(len\(rec_result\)\):', '                for rno in range(actual_batch_size):', c3)
open(f, 'w').write(c4)
" && \
    # --- Patch C: predict_det.py contiguous 检查 ---
    ${VENV}/bin/python -c "
f = '${MINERU_INFER_DIR}/predict_det.py'
c = open(f).read()
old = '( {8}inp = inp\.to\(self\.device\)\n)'
new = r'\1            if not inp.is_contiguous():\n                inp = inp.contiguous()\n'
import re
c2 = re.sub(old, new, c)
if c2 == c:
    if 'is_contiguous' not in c:
        raise RuntimeError('Patch C: cannot find inp.to(device) line')
    else:
        print('Patch C: already applied')
else:
    open(f, 'w').write(c2)
    print('Patch C: contiguous check inserted')
" && \
    echo "All MinerU RDNA patches applied."

# ===========================================================================
# 阶段 10：MIOpen 预热脚本
# ===========================================================================
COPY scripts/cache_warmer.py /opt/cache_warmer.py

# ===========================================================================
# 阶段 11：入口与最终设置
# ===========================================================================
RUN echo 'source /opt/mineru_venv/bin/activate' >> /etc/bash.bashrc && \
    echo "MinerU Docker image built successfully." && \
    ${VENV}/bin/python -c "
import torch, vllm, mineru
print('='*50)
print('MinerU ROCm Docker Image Ready')
print(f'  PyTorch : {torch.__version__}')
print(f'  ROCm    : {torch.version.hip}')
print(f'  vllm    : {vllm.__version__}')
print(f'  MinerU  : {mineru.__version__}')
print(f'  Arch    : ${ARCH}')
print('='*50)
"

# 容器入口：默认 bash，用户可 override
ENTRYPOINT ["/bin/bash", "-c"]
CMD ["bash"]
