# syntax=docker/dockerfile:1.7
# GLM-5.2 NF3-Hybrid for DGX Spark (aarch64/sm_121a)
# Adapted from timothystewart6/vllm-gb10

ARG CUDA_BASE_IMAGE

# ───── Stage 1: system-base ──────────────────────────────────────────
FROM ${CUDA_BASE_IMAGE} AS system-base

ARG UV_VERSION
ARG PYTORCH_INDEX_URL
ARG DEBIAN_FRONTEND=noninteractive

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    UV_SYSTEM_PYTHON=1 UV_LINK_MODE=copy

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential ninja-build cmake curl git wget \
    python3-dev python3-pip ccache \
    libibverbs-dev rdma-core libcudnn9-cuda-13 libcudnn9-dev-cuda-13 \
    && rm -rf /var/lib/apt/lists/* \
    && pip install --no-cache-dir uv==0.11.28

# ───── Stage 2: torch-base ───────────────────────────────────────────
FROM system-base AS torch-base

ARG TORCH_VERSION
ARG TORCHVISION_VERSION
ARG TRITON_VERSION
ARG PYTORCH_INDEX_URL

RUN pip install --no-cache-dir --no-deps \
    torch==2.12.0 torchvision==0.22.0 \
    --index-url "${PYTORCH_INDEX_URL}"

# ───── Stage 3: NCCL (nccl-canonical) ────────────────────────────────
FROM torch-base AS nccl-build

ARG NCCL_REPO
ARG NCCL_REF

RUN git clone --depth=1 -b "${NCCL_REF}" "${NCCL_REPO}" /opt/nccl
WORKDIR /opt/nccl
RUN make -j$(nproc) src.build \
    CUDA_HOME=/usr/local/cuda \
    NVCC_GENCODE="-gencode=arch=compute_121,code=sm_121" \
    && make install PREFIX=/usr/local && ldconfig

# ───── Stage 4: b12x ─────────────────────────────────────────────────
FROM torch-base AS b12x-build

ARG B12X_REPO
ARG B12X_REF
ARG TORCH_CUDA_ARCH_LIST=12.1a

ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    CUTE_DSL_ARCH=sm_121a

RUN git clone --recursive "${B12X_REPO}" /opt/b12x \
    && cd /opt/b12x && git checkout "${B12X_REF}"
WORKDIR /opt/b12x
RUN pip install --no-cache-dir -e . --no-build-isolation \
    -C setup.args="--cuda-arch=${TORCH_CUDA_ARCH_LIST}"

# ───── Stage 5: FlashInfer ───────────────────────────────────────────
FROM torch-base AS flashinfer-build

ARG FLASHINFER_REPO
ARG FLASHINFER_REF
ARG FLASHINFER_CUDA_ARCH_LIST=12.1a
ARG TORCH_CUDA_ARCH_LIST=12.1a

ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST} \
    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}

RUN git clone --recursive "${FLASHINFER_REPO}" /opt/flashinfer \
    && cd /opt/flashinfer && git checkout "${FLASHINFER_REF}" \
    && git submodule update --init --recursive
WORKDIR /opt/flashinfer
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
    && cd flashinfer-cubin \
    && uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
    && cd ../flashinfer-jit-cache \
    && uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
    && pip install --no-cache-dir /wheels/*.whl

# ───── Stage 6: vLLM (eldritch-enlightenment) ────────────────────────
FROM torch-base AS vllm-build

ARG VLLM_REPO
ARG VLLM_REF
ARG TORCH_CUDA_ARCH_LIST=12.1a

ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    CUTE_DSL_ARCH=sm_121a

# Clone vLLM at eldritch-enlightenment pin
RUN git clone --recursive "${VLLM_REPO}" /opt/vllm \
    && cd /opt/vllm && git checkout "${VLLM_REF}" \
    && git submodule update --init --recursive

WORKDIR /opt/vllm
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
    && pip install --no-cache-dir /wheels/vllm*.whl

# ───── Stage 7: runner (final image) ─────────────────────────────────
FROM torch-base AS runner

# Copy built artifacts from each stage
COPY --from=nccl-build     /usr/local/lib/libnccl*       /usr/local/lib/
COPY --from=nccl-build     /usr/local/include/nccl.h     /usr/local/include/
RUN ldconfig

COPY --from=b12x-build     /usr/local/lib/python3.12/dist-packages/b12x* \
    /usr/local/lib/python3.12/dist-packages/
COPY --from=flashinfer-build /usr/local/lib/python3.12/dist-packages/flashinfer* \
    /usr/local/lib/python3.12/dist-packages/
COPY --from=vllm-build     /usr/local/lib/python3.12/dist-packages/ \
    /usr/local/lib/python3.12/dist-packages/

# ── NF3 hybrid overlay ──────────────────────────────────────────────
COPY overlay/hybrid_loader.py  /usr/local/lib/python3.12/dist-packages/hybrid_loader.py
COPY overlay/nf3_kernel.py     /usr/local/lib/python3.12/dist-packages/nf3_kernel.py
COPY overlay/nvfp4_kernel.py   /usr/local/lib/python3.12/dist-packages/nvfp4_kernel.py
COPY overlay/nf3_replan.py     /usr/local/lib/python3.12/dist-packages/nf3_replan.py
COPY overlay/mxfp8_tier.json   /usr/local/lib/python3.12/dist-packages/mxfp8_tier.json

RUN python3 -c "import sys; open('/usr/local/lib/python3.12/dist-packages/hybrid_loader.pth','w').write('import hybrid_loader\n')"

# Verify imports
RUN python3 -c "import hybrid_loader; print('OK hybrid_loader'); import nf3_kernel; print('OK nf3_kernel'); import nvfp4_kernel; print('OK nvfp4_kernel'); import nf3_replan; print('OK nf3_replan'); import flashinfer; print(f'OK flashinfer {flashinfer.__version__}'); import torch; print(f'OK torch {torch.__version__} cuda {torch.version.cuda}'); import b12x; print(f'OK b12x {b12x.__version__}'); import vllm; print(f'OK vllm {vllm.__version__}')"

ENV TORCH_CUDA_ARCH_LIST=12.1a \
    CUTE_DSL_ARCH=sm_121a \
    FLASHINFER_CUDA_ARCH_LIST=12.1a

ENTRYPOINT ["vllm", "serve"]