############################################################
#  Stage-0: 构建依赖轮子（PyTorch + SGLang + sgl_kernel）   #
############################################################
ARG CUDA_VERSION=12.8.1                     # 仅 12.6.1 / 12.8.1 测过
FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu22.04 AS builder

# ---- Python 环境 ----
RUN apt-get update && \
    apt-get install -y --no-install-recommends python3 python3-pip && \
    ln -sf /usr/bin/python3 /usr/bin/python && \
    python -m pip install --no-cache-dir --upgrade pip

# ---- PyTorch / torchvision / SGLang / sgl_kernel ----
ARG TORCH_VER=2.7.1                         # cu128 轮子从 2.7.1 起提供
ARG TV_VER=0.22.1                           # 与 torch-2.7.1 对齐
RUN case "$CUDA_VERSION" in \
        12.6.1) CUINDEX=126 ;; \
        12.8.1) CUINDEX=128 ;; \
        *) echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1 ;; \
    esac && \
    python -m pip install --no-cache-dir \
        torch==${TORCH_VER}+cu${CUINDEX} \
        torchvision==${TV_VER}+cu${CUINDEX} \
        --extra-index-url https://download.pytorch.org/whl/cu${CUINDEX} && \
    python -m pip install --no-cache-dir sglang==0.4.8.post1 && \
    # NCCL 2.27——CUDA 12.x 官方推荐
    python -m pip install --no-cache-dir nvidia-nccl-cu12==2.27.3 --force-reinstall --no-deps && \
    # ✅ 官方预编译，带 cp310 & cu128 标签；不再手写直链
    python -m pip install --no-cache-dir sgl-kernel==0.0.2.post17

############################################################
#  Stage-1: 生成最小运行镜像                                #
############################################################
FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1

# ---- Python runtime ----
RUN apt-get update && \
    apt-get install -y --no-install-recommends python3 python3-distutils && \
    ln -sf /usr/bin/python3 /usr/bin/python && \
    rm -rf /var/lib/apt/lists/*

# ---- 拷贝依赖 ----
COPY --from=builder /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages
COPY --from=builder /usr/local/bin /usr/local/bin

# ---- 拷贝模型（改成你的实际路径即可） ----
COPY ./Alibaba/Qwen3-8B /root/.cradle/Alibaba/Qwen3-8B

# ---- 端口 & 启动 ----
EXPOSE 30000
CMD ["python3", "-m", "sglang.launch_server", \
     "--host", "0.0.0.0", \
     "--port", "30000", \
     "--model-path", "/root/.cradle/Alibaba/Qwen3-8B/", \
     "--tp", "1", \
     "--api-key", "token-abc123"]