From a9489cd9d99d9a7bd7f24db8c16379bd6da41ca9 Mon Sep 17 00:00:00 2001
From: chenjw28 <792430652@qq.com>
Date: Thu, 4 Jun 2026 13:52:04 +0800
Subject: [PATCH] x

---
 docker/Dockerfile | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 0a4d00b..98027b9 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -148,7 +148,7 @@ RUN set -ex && \
     ${VENV}/bin/python -c "import torch; v=torch.__version__; assert 'rocm' in v, f'PyTorch overwritten: {v}'; print('PyTorch OK:', v)"
 
 # ===========================================================================
-# 阶段 8：编译 vllm
+# 阶段 8a：编译 vllm（git clone + cmake + ninja，最耗时，单独一层缓存）
 # ===========================================================================
 RUN set -ex && \
     # git 通过代理访问 GitHub
@@ -159,11 +159,11 @@ RUN set -ex && \
         "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 头文件已自带）
+    # 补丁 5：注释掉 vllm mamba 模块的 operator+ 定义
     cd /opt/vllm && \
     sed -i '109,121s/^/\/\/ /' csrc/mamba/mamba_ssm/selective_scan.h && \
     echo "vllm mamba operator+ patch applied." && \
-    # cmake 别名兜底：ROCm 7.2 可能缺少 hiprand/hipblas cmake target
+    # cmake 别名兜底
     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 && \
@@ -184,18 +184,23 @@ RUN set -ex && \
     cd /opt/vllm_build && ninja -j8 && \
     # 安装 .so 到 vllm 源码目录
     cp /opt/vllm_build/*.abi3.so /opt/vllm/vllm/ && \
-    # pip install vllm（让 pip 解析运行时依赖）
+    echo "vllm C++ build complete (layer cached)."
+
+# ===========================================================================
+# 阶段 8b：安装 vllm + 验证 PyTorch（独立层，失败不影响 8a 缓存）
+# ===========================================================================
+RUN set -ex && \
+    # 安装 vllm（让 pip 解析运行时依赖：xgrammar, regex, 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)" && \
     # 先重装 ROCm PyTorch 覆盖可能的 CUDA 版，再清理 CUDA triton 元数据
-    # 顺序重要：pytorch-triton-rocm 和 triton 共享 triton/ 物理目录，必须先重装后卸载
     ${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; \
-    # vllm 平台模块导入验证（GPU 检测只能在运行时，容器构建时无 GPU 设备）
-    ${VENV}/bin/python -c "from vllm.platforms import current_platform; print('Platform module:', type(current_platform).__name__); print('vllm import OK')" && \
+    # 确保 vllm 基本可导入（构建时无 GPU，仅验证 import）
+    ${VENV}/bin/python -c "import vllm; print('vllm import OK:', vllm.__version__)" && \
     # 清理构建目录（减小镜像体积，约 3-5GB）
     rm -rf /opt/vllm_build