# syntax=docker/dockerfile:1.7
# ==========================================================================
# GLM-5.2 NF3-Hybrid — DGX Spark (aarch64/sm_121a)
#
# Builds vLLM (local-inference-lab/eldritch-enlightenment) + b12x
# + FlashInfer from source, then overlays the NF3 Triton kernels.
#
# Usage:
#   docker build \
#     --build-arg VLLM_SRC=/path/to/vllm \
#     --build-arg B12X_SRC=/path/to/b12x \
#     --build-arg NCCL_SRC=/path/to/nccl \
#     -t vllm-nf3-hybrid:probe .
#
# On dgx4.lan:
#   VLLM_SRC=/home/kilork/projects/github.com/vllm-eldritch
#   B12X_SRC=/home/kilork/projects/github.com/b12x-nf3
#   NCCL_SRC=/home/kilork/projects/github.com/nccl-canonical
# ==========================================================================

ARG CUDA_VERSION=13.2.0
ARG UBUNTU_VERSION=24.04
ARG GPU_ARCH=12.1a

FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}

ARG GPU_ARCH
ARG DEBIAN_FRONTEND=noninteractive

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    TORCH_CUDA_ARCH_LIST=${GPU_ARCH} \
    CUTE_DSL_ARCH=sm_121a \
    FLASHINFER_CUDA_ARCH_LIST=${GPU_ARCH} \
    MAX_JOBS=$(nproc) \
    CMAKE_BUILD_PARALLEL_LEVEL=$(nproc)

# ── System packages ─────────────────────────────────────────────────
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential ninja-build cmake curl git wget \
    python3-dev python3-pip python3-venv ccache \
    libibverbs-dev libibverbs1 rdma-core \
    libcudnn9-cuda-13 libcudnn9-dev-cuda-13 \
    rapidjson-dev libxml2-dev \
    && rm -rf /var/lib/apt/lists/*

# ── Python build essentials ─────────────────────────────────────────
RUN pip install --no-cache-dir uv pybind11 setuptools-scm wheel

# ── PyTorch 2.12.0+cu132 (aarch64) ──────────────────────────────────
RUN pip install --no-cache-dir --no-deps torch==2.12.0 \
    && pip install --no-cache-dir --no-deps torchvision==0.22.0

# ── NCCL ────────────────────────────────────────────────────────────
COPY nccl /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

# ── b12x ────────────────────────────────────────────────────────────
COPY b12x /opt/b12x
WORKDIR /opt/b12x
RUN pip install --no-cache-dir -e . --no-build-isolation \
    -C setup.args="--cuda-arch=${GPU_ARCH}"

# ── FlashInfer ──────────────────────────────────────────────────────
RUN git clone --recursive https://github.com/flashinfer-ai/flashinfer.git /opt/flashinfer \
    && cd /opt/flashinfer \
    && git checkout 25dd814e03791e370f96c3148242f0dc8de504ac \
    && git submodule update --init --recursive
WORKDIR /opt/flashinfer
RUN 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

# ── vLLM (eldritch-enlightenment) ───────────────────────────────────
COPY vllm /opt/vllm
WORKDIR /opt/vllm
RUN git checkout 45c1582e9b80ba83e71c3a6458e71da4736fbdc4 2>/dev/null || true \
    && git submodule update --init --recursive 2>/dev/null || true

# Build vLLM wheel
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

# ── 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

# .pth hook: hybrid_loader fires when any Python process imports from vllm
RUN python3 -c "import sys; open('/usr/local/lib/python3.12/dist-packages/hybrid_loader.pth','w').write('import hybrid_loader\n')"

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

ENTRYPOINT ["vllm", "serve"]