# syntax=docker/dockerfile:1.7
# GLM-5.2 NF3-Hybrid — self-contained build for DGX Spark (aarch64/sm_121a)
#
# All build inputs pinned in versions.env and locks/. No external repo needed
# beyond the package registries and git repos each stage clones.
#
# Build: cd docker-images/glm52-nf3-hybrid && bash build.sh

ARG CUDA_BASE_IMAGE
ARG CUDA_BASE_DIGEST

FROM ${CUDA_BASE_IMAGE}@${CUDA_BASE_DIGEST} AS apt-base
ARG UV_VERSION
ARG PYTORCH_INDEX_URL
ARG PYPI_INDEX_URL
ARG TORCH_CUDA_ARCH_LIST=12.1a
ENV DEBIAN_FRONTEND=noninteractive \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    UV_SYSTEM_PYTHON=1 UV_BREAK_SYSTEM_PACKAGES=1 UV_LINK_MODE=copy \
    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
COPY locks/apt-sources.list /tmp/apt-sources.list
COPY locks/apt-packages.txt /tmp/apt-packages.txt
COPY locks/python-bootstrap.txt /tmp/python-bootstrap.txt
COPY locks/python-build.txt /tmp/python-build.txt
RUN rm -f /etc/apt/sources.list.d/*.list /etc/apt/sources.list.d/*.sources \
 && cp /tmp/apt-sources.list /etc/apt/sources.list \
 && apt-get update \
 && grep -vE '^\s*(#|$)' /tmp/apt-packages.txt \
      | xargs apt-get install -y --no-install-recommends \
 && rm -rf /var/lib/apt/lists/* \
 && python3 -m pip install --no-cache-dir --require-hashes --no-deps \
      -r /tmp/python-bootstrap.txt --index-url ${PYPI_INDEX_URL}

FROM apt-base AS torch-base
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv pip install --require-hashes -r /tmp/python-build.txt \
      --index-url ${PYPI_INDEX_URL} --extra-index-url ${PYTORCH_INDEX_URL} \
      --index-strategy unsafe-best-match

FROM torch-base AS base
ARG TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION TRITON_VERSION
ARG NCCL_REPO NCCL_REF
RUN git clone --depth 1 -b ${NCCL_REF} ${NCCL_REPO} /opt/nccl \
 && cd /opt/nccl \
 && make -j $(nproc) src.build \
      NVCC_GENCODE="-gencode=arch=compute_121,code=sm_121" \
 && make install PREFIX=/usr/local && ldconfig

FROM base AS flashinfer-builder
ARG FLASHINFER_REPO FLASHINFER_REF
ARG FLASHINFER_CUDA_ARCH_LIST=12.1a
ARG SOURCE_DATE_EPOCH BUILD_JOBS
ENV FLASHINFER_CUDA_ARCH_LIST=${FLASHINFER_CUDA_ARCH_LIST}
RUN git clone --recursive ${FLASHINFER_REPO} /workspace/flashinfer \
 && cd /workspace/flashinfer \
 && git checkout ${FLASHINFER_REF} \
 && git submodule update --init --recursive
WORKDIR /workspace/flashinfer
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    --mount=type=cache,id=ccache,target=/root/.ccache \
    _j="${BUILD_JOBS:-$(nproc)}" \
 && export MAX_JOBS="${_j}" CMAKE_BUILD_PARALLEL_LEVEL="${_j}" \
 && sed -i '/^license-files/d' pyproject.toml \
 && sed -i 's|^license = .*|license = {text = "Apache-2.0"}|' pyproject.toml \
 && uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
 && cd flashinfer-cubin \
 && sed -i '/^license-files/d' pyproject.toml \
 && sed -i 's|^license = .*|license = {text = "Apache-2.0"}|' pyproject.toml \
 && uv build --no-build-isolation --wheel . --out-dir=/wheels -v \
 && cd ../flashinfer-jit-cache \
 && sed -i '/^license-files/d' pyproject.toml \
 && sed -i 's|^license = .*|license = {text = "Apache-2.0"}|' pyproject.toml \
 && uv build --no-build-isolation --wheel . --out-dir=/wheels -v

FROM base AS vllm-builder
ARG VLLM_REPO VLLM_REF VLLM_COMMIT
ARG SOURCE_DATE_EPOCH BUILD_JOBS
RUN git clone --recursive ${VLLM_REPO} /workspace/vllm \
 && cd /workspace/vllm \
 && git checkout ${VLLM_REF} \
 && git submodule update --init --recursive
COPY patches/ /patches/
RUN cd /workspace/vllm \
 && for p in /patches/*.patch; do [ -f "$p" ] && patch -p1 --forward < "$p" || true; done
WORKDIR /workspace/vllm
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    --mount=type=cache,id=ccache,target=/root/.ccache \
    _j="${BUILD_JOBS:-$(nproc)}" \
 && export MAX_JOBS="${_j}" CMAKE_BUILD_PARALLEL_LEVEL="${_j}" \
 && python3 use_existing_torch.py \
 && sed -i '/^license-files/d' pyproject.toml \
 && sed -i 's|^license = .*|license = {text = "Apache-2.0"}|' pyproject.toml \
 && uv build --no-build-isolation --wheel . --out-dir=/wheels -v

FROM base AS b12x-builder
ARG B12X_REPO B12X_REF B12X_COMMIT
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} \
 && pip install --no-cache-dir -e . --no-build-isolation \
    -C setup.args="--cuda-arch=${TORCH_CUDA_ARCH_LIST}"

FROM base AS runner
ARG TORCH_CUDA_ARCH_LIST=12.1a
ARG PYTORCH_INDEX_URL PYPI_INDEX_URL
ARG CUTLASS_DSL_VERSION
COPY locks/python-runtime.txt /tmp/python-runtime.txt
WORKDIR /workspace
COPY checksums/tiktoken.sha256 /tmp/tiktoken.sha256
RUN mkdir -p tiktoken_encodings \
 && wget -O tiktoken_encodings/o200k_base.tiktoken https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken \
 && wget -O tiktoken_encodings/cl100k_base.tiktoken https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken \
 && cd tiktoken_encodings && sha256sum -c /tmp/tiktoken.sha256
ENV TIKTOKEN_ENCODINGS_BASE=/workspace/tiktoken_encodings
RUN apt-get update && apt-get install -y --no-install-recommends libboost-dev libaio-dev \
 && rm -rf /var/lib/apt/lists/*

# Install all wheels
RUN --mount=type=bind,from=flashinfer-builder,source=/wheels,target=/fi-wheels \
    --mount=type=bind,from=vllm-builder,source=/wheels,target=/vllm-wheels \
    --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv pip install --require-hashes -r /tmp/python-runtime.txt \
      --index-url ${PYPI_INDEX_URL} --extra-index-url ${PYTORCH_INDEX_URL} \
      --index-strategy unsafe-best-match \
 && pip install --no-cache-dir /fi-wheels/*.whl /vllm-wheels/*.whl

# b12x: symlink full source tree (must include integration/ attention/ etc.)
COPY --from=b12x-builder /opt/b12x/ /opt/b12x/
RUN rm -rf /usr/local/lib/python3.12/dist-packages/b12x* \
 && ln -sf /opt/b12x/b12x /usr/local/lib/python3.12/dist-packages/b12x \
 && pip install --no-cache-dir -e /opt/b12x --no-build-isolation --no-deps \
    -C setup.args="--cuda-arch=${TORCH_CUDA_ARCH_LIST}"
RUN pip install --no-cache-dir \
    nvidia-cutlass-dsl==${CUTLASS_DSL_VERSION} \
    nvidia-cutlass-dsl-libs-cu13==${CUTLASS_DSL_VERSION}

# NCCL symlink for torch
RUN rm -f /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2 \
 && ln -s /usr/local/lib/libnccl.so.2 \
      /usr/local/lib/python3.12/dist-packages/nvidia/nccl/lib/libnccl.so.2

# NF3 hybrid overlay
COPY overlay/hybrid_loader.py              /usr/local/lib/python3.12/dist-packages/hybrid_loader.py
COPY overlay/hybrid.pth                    /usr/local/lib/python3.12/dist-packages/hybrid.pth
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
COPY overlay/b12x/moe/fused/w4a16/prepare.py /usr/local/lib/python3.12/dist-packages/b12x/moe/fused/w4a16/prepare.py
COPY overlay/b12x/moe/fused/w4a16/host.py    /usr/local/lib/python3.12/dist-packages/b12x/moe/fused/w4a16/host.py
COPY overlay/b12x/cute/fp4.py                 /usr/local/lib/python3.12/dist-packages/b12x/cute/fp4.py
COPY overlay/vllm/v1/attention/backends/mla/indexer.py \
      /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/indexer.py

ENTRYPOINT []
CMD ["bash"]