This commit is contained in:
hailin 2025-09-18 19:40:46 +08:00
parent cd4b493815
commit 0992c81029
1 changed files with 11 additions and 2 deletions

View File

@ -41,7 +41,7 @@ RUN echo "Building PyTorch with USE_DISTRIBUTED=$USE_DISTRIBUTED" && \
ARG CUDA_VERSION=12.6.1 ARG CUDA_VERSION=12.6.1
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS builder-extras FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04 AS builder-extras
ENV TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9;9.0" ENV TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;8.9"
ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8 ENV DEBIAN_FRONTEND=noninteractive PYTHONUNBUFFERED=1 LANG=C.UTF-8 LC_ALL=C.UTF-8
RUN apt-get update && apt-get install -y --no-install-recommends \ RUN apt-get update && apt-get install -y --no-install-recommends \
@ -73,7 +73,7 @@ WORKDIR /opt/flashinfer
# 覆盖你的目标算力3090=8.64090=8.9H100=9.0a;可按需增/减 # 覆盖你的目标算力3090=8.64090=8.9H100=9.0a;可按需增/减
ENV FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9 9.0a" ENV FLASHINFER_CUDA_ARCH_LIST="7.5 8.0 8.6 8.9"
# 先做 AOT 预编译,再直接打 wheel不隔离使用同一份自编 torch # 先做 AOT 预编译,再直接打 wheel不隔离使用同一份自编 torch
RUN python3 -m pip install --no-cache-dir numpy requests build "cuda-python>=12.0,<13" "nvidia-nvshmem-cu12" ninja pynvml && \ RUN python3 -m pip install --no-cache-dir numpy requests build "cuda-python>=12.0,<13" "nvidia-nvshmem-cu12" ninja pynvml && \
@ -134,11 +134,16 @@ RUN pip wheel "gradio==5.38.2" requests -w /wheels
RUN pip wheel pybase64==1.3.2 -w /wheels RUN pip wheel pybase64==1.3.2 -w /wheels
# 导出轮子的独立阶段 # 导出轮子的独立阶段
FROM scratch AS wheelhouse FROM scratch AS wheelhouse
COPY --from=builder-extras /wheels / COPY --from=builder-extras /wheels /
# 从宿主机目录 _wheelhouse/ 安装轮子的 runtime # 从宿主机目录 _wheelhouse/ 安装轮子的 runtime
ARG CUDA_VERSION=12.6.1 ARG CUDA_VERSION=12.6.1
FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu22.04 AS runtime-prebuilt FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu22.04 AS runtime-prebuilt
@ -179,6 +184,10 @@ RUN echo "/usr/local/cuda/extras/CUPTI/lib64" > /etc/ld.so.conf.d/cupti.conf &&
# 保险起见,再加一行环境变量(有些基础镜像不把 extras 加入 ld.so.conf # 保险起见,再加一行环境变量(有些基础镜像不把 extras 加入 ld.so.conf
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:${LD_LIBRARY_PATH} ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:${LD_LIBRARY_PATH}
############################################################################### ###############################################################################
# Stage 2 ─ runtime极简运行镜像仅离线安装 wheel # Stage 2 ─ runtime极简运行镜像仅离线安装 wheel
############################################################################### ###############################################################################