# syntax=docker/dockerfile:1.7
# GLM-5.3-Flash r12 (jovian-judgement) — CLEAN build for DGX Spark (aarch64/sm_121a)
#
# Per operator direction: start from a raw CUDA image (solar-open2-nvfp4
# pattern), NOT the day-0 vllm image. No baked-vllm contamination: torch,
# flashinfer, vLLM and b12x all come from pinned sources in this file.
#
# Stack (rtx6kpro models/glm-5.3-flash.md source contract):
#   vLLM  = local-inference-lab/vllm @ 54f6e9826c (dev/jovian-judgement base)
#           + sequential merges of PR heads 550,552,553,554,555,557
#           (the doc's "composition commit" 7d8a09c… is not a public ref)
#   b12x  = local-inference-lab/b12x @ 139e04048b (master) + PR heads
#           260,267,268,269,270. b12x is pure-Python: the kernels are
#           CuTe-DSL / tvm-ffi JIT (nvidia-cutlass-dsl 4.6.2), no CUDA
#           compile at install time.
#   torch = 2.13.0 cu130 (the fork's requirements/cuda.txt pin)
#   flashinfer = 0.6.17 (day-0-era; not a pip dep of the fork since 0.6.14,
#           so installed from the flashinfer.ai wheel index)
#
# GB10 sm_121 fix baked in (patch_cuda_sm12x.py): the sm12x auto-priority
# list in vllm/platforms/cuda.py offers only TRITON_MLA (no sparse) and
# FLASHINFER_MLA_SPARSE_SM120 (head 576 only) — neither fits GLM head-512.
# B12X (supports (12,0)/(12,1)) is inserted first.
#
# Build: cd docker-images/glm53-jovian-r12 && bash build.sh   (on a Spark node)

ARG CUDA_BASE_IMAGE
ARG CUDA_BASE_DIGEST
ARG UV_VERSION

FROM ${CUDA_BASE_IMAGE}@${CUDA_BASE_DIGEST} AS apt-base
ARG PYPI_INDEX_URL
ENV DEBIAN_FRONTEND=noninteractive \
    PIP_BREAK_SYSTEM_PACKAGES=1 \
    UV_SYSTEM_PYTHON=1 UV_BREAK_SYSTEM_PACKAGES=1 UV_LINK_MODE=copy
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
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 --no-deps \
      -r /tmp/python-bootstrap.txt --index-url ${PYPI_INDEX_URL}

FROM apt-base AS torch-base
ARG PYPI_INDEX_URL
ARG PYTORCH_INDEX_URL
ARG TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION
ARG FLASHINFER_VERSION
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv pip install \
      --index-url ${PYPI_INDEX_URL} \
      --extra-index-url ${PYTORCH_INDEX_URL} \
      --index-strategy unsafe-best-match \
      "torch==${TORCH_VERSION}" \
      "torchvision==${TORCHVISION_VERSION}" \
      "torchaudio==${TORCHAUDIO_VERSION}" \
      "flashinfer-python==${FLASHINFER_VERSION}"

# The fork's r12 composition: base branch + sequential merge of the open
# PR heads (ascending). The PR head pseudo-refs are GitHub's stable refs.
FROM torch-base AS deepgemm-builder
ARG DEEPGEMM_REPO DEEPGEMM_REF DEEPGEMM_BRANCH
RUN set -eux; \
    git config --global protocol.version 0; \
    git clone --branch ${DEEPGEMM_BRANCH} ${DEEPGEMM_REPO} /deps/deepgemm; \
    cd /deps/deepgemm; \
    git checkout ${DEEPGEMM_REF}; \
    git submodule update --init --depth 1 -- third-party/cutlass third-party/fmt

FROM deepgemm-builder AS vllm-builder
ARG VLLM_REPO
ARG VLLM_BASE_REF
ARG VLLM_PR_HEADS
ARG SOURCE_DATE_EPOCH
ARG BUILD_JOBS
ENV TORCH_CUDA_ARCH_LIST=12.1a GIT_PROTOCOL=0 DEEPGEMM_SRC_DIR=/deps/deepgemm
RUN set -eux; \
    git config --global user.email build@sm121.lan; \
    git config --global user.name sm121-builder; \
    git config --global protocol.version 0; \
    git clone --recursive ${VLLM_REPO} /workspace/vllm; \
    cd /workspace/vllm; \
    git checkout ${VLLM_BASE_REF}; \
    # PR head commits fetched by SHA (refs/pull/N/head vanishes when a PR
    # closes — 550 closed unmerged; raw-SHA fetches stay available).
    for pair in ${VLLM_PR_HEADS}; do \
        pr="${pair%%:*}"; sha="${pair#*:}"; \
        git fetch origin "${sha}:refs/prs/${pr}"; \
        git merge --no-ff -m "PR ${pr}" "refs/prs/${pr}"; \
    done; \
    git rev-parse HEAD
WORKDIR /workspace/vllm
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    --mount=type=cache,id=ccache,target=/root/.ccache \
    uv pip install setuptools-rust setuptools-scm \
    && _j="${BUILD_JOBS:-$(nproc)}" \
    && export MAX_JOBS="${_j}" CMAKE_BUILD_PARALLEL_LEVEL="${_j}" \
    && python3 use_existing_torch.py \
    && uv build --no-build-isolation --wheel . --out-dir=/wheels -v

FROM torch-base AS b12x-builder
ARG B12X_REPO
ARG B12X_BASE_REF
ARG B12X_PRS
ARG SOURCE_DATE_EPOCH
RUN set -eux; \
    git config --global user.email build@sm121.lan; \
    git config --global user.name sm121-builder; \
    git config --global protocol.version 0; \
    git clone ${B12X_REPO} /workspace/b12x; \
    cd /workspace/b12x; \
    git checkout ${B12X_BASE_REF}; \
    for pr in ${B12X_PRS}; do \
        git fetch origin "refs/pull/$pr/head:prs/$pr"; \
        git merge --no-ff -m "PR $pr" "prs/$pr"; \
    done; \
    git rev-parse HEAD
WORKDIR /workspace/b12x
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv build --no-build-isolation --wheel . --out-dir=/wheels -v

FROM torch-base AS runner
ARG PYPI_INDEX_URL
COPY checksums/tiktoken.sha256 /tmp/tiktoken.sha256
COPY patch_cuda_sm12x.py /tmp/patch_cuda_sm12x.py
WORKDIR /workspace
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 --mount=type=bind,from=vllm-builder,source=/wheels,target=/vllm-wheels \
    --mount=type=bind,from=b12x-builder,source=/wheels,target=/b12x-wheels \
    --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    pip install --no-cache-dir --no-build-isolation /vllm-wheels/*.whl \
 && pip install --no-cache-dir /b12x-wheels/*.whl \
 && python3 /tmp/patch_cuda_sm12x.py \
 && find /usr/local/lib/python3.12/dist-packages/vllm -name "__pycache__" \
      -type d -prune -exec rm -rf {} + 2>/dev/null || true
# Build gate: the r12 stack must import with the GB10 patches in place.
RUN set -eux; \
    python3 -c "import vllm; print('vllm', vllm.__version__)"; \
    python3 -c "import b12x; print('b12x ok')"; \
    python3 -c "import vllm._flashkda_C; print('flashkda_C ok')"; \
    python3 -c "from vllm.config.attention import AttentionBackendEnum as E; \
assert E.B12X.value and 'b12x' in E.B12X.value.lower(); print('B12X enum ok')"; \
    grep -q "AttentionBackendEnum.B12X," /usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py

ENTRYPOINT []
CMD ["bash"]