# 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: vLLM/b12x pins live in versions.env (single source of truth — the
# staged vllm-src.tar.gz / b12x-src.tar.gz on the build host are built from
# those pins, composition notes included). torch = 2.13.0 cu130 (the fork's
# requirements/cuda.txt pin); flashinfer = 0.6.17 (not a pip dep of the fork
# since 0.6.14, installed from the flashinfer.ai wheel index). b12x is
# pure-Python: CuTe-DSL / tvm-ffi JIT (nvidia-cutlass-dsl 4.6.2), no CUDA
# compile at install time.
#
# 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 deps-src
# Pre-fetched vllm build sub-deps (deps-src.tar.gz in the build context, staged
# on the build host): the node's direct github.com access is flaky (FetchContent
# git clones get "could not read Username" / hung-up) — these SRC_DIR env vars
# are each cmake's documented local-source escape hatch.
COPY deps-src.tar.gz /tmp/deps-src.tar.gz
RUN set -eux; \
    mkdir -p /deps2; \
    tar xzf /tmp/deps-src.tar.gz -C /deps2; \
    mv /deps2/triton/python/triton_kernels/triton_kernels /deps2/triton_kernels; \
    rm -rf /deps2/triton /tmp/deps-src.tar.gz

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
# Pre-staged r12 source tree (vllm-src.tar.gz in the build context, assembled
# on the build host: base VLLM_BASE_REF + VLLM_PR_HEADS merged in order, fresh
# lightweight .git). The node's github.com git access is WAF-blocked
# (401 on git-upload-pack; curl GET works) so no in-build git is possible.
COPY vllm-src.tar.gz /tmp/vllm-src.tar.gz
RUN set -eux; \
    mkdir -p /workspace/vllm; \
    tar xzf /tmp/vllm-src.tar.gz -C /workspace/vllm; \
    rm -f /tmp/vllm-src.tar.gz; \
    git config --global --add safe.directory /workspace/vllm; \
    cd /workspace/vllm; \
    git rev-parse HEAD
WORKDIR /workspace/vllm
# eugr/spark-vllm-docker build-time patches (main, 2026-09-02): the lab's
# Spark-tuned vllm fixes. All are source-aware/idempotent (skip when the
# vulnerable pattern is absent), so running the full set is safe. The two env
# gates are opt-in and ON for this build (sm_121a + b12x):
#   - preserve_sm12x_target: keep 12.1a in CUDA_SUPPORTED_ARCHS (their
#     standard CUDA-13 build collapses 12.1a -> 12.0)
#   - sm120_cooperative_topk: cooperative_topk fails on SM12.x at launch;
#     keep SM90-only (Spark fix)
#   - mrv2_speculator_cudagraph_pool: MTP speculator cudagraph pool fix
#   - spark_kv_cache_cleanup: DGX Spark UMA allocator cleanup
#   - b12x_c128a_topk_alignment / flashinfer_b12x_swigluoai: b12x kernel fixes
#   - routed_experts_weight_shape: modelopt MoE 2-element weight_shape load
#   - topk_softplus_sqrt_control_flow: dead CUDA kernel fix (upstream bug)
# Provenance: https://github.com/eugr/spark-vllm-docker (docker/patch_vllm_*.py)
COPY patches-eugr/patch_vllm_*.py /tmp/eugr-patches/
RUN VLLM_PRESERVE_SM12X_TARGET=1 \
    VLLM_PATCH_B12X_C128A_ALIGNMENT=1 \
    python3 /tmp/eugr-patches/patch_vllm_mrv2_speculator_cudagraph_pool.py . \
    && python3 /tmp/eugr-patches/patch_vllm_b12x_c128a_topk_alignment.py . \
    && python3 /tmp/eugr-patches/patch_vllm_flashinfer_b12x_swigluoai.py . \
    && python3 /tmp/eugr-patches/patch_vllm_topk_softplus_sqrt_control_flow.py . \
    && python3 /tmp/eugr-patches/patch_vllm_preserve_sm12x_target.py . \
    && python3 /tmp/eugr-patches/patch_vllm_diffusion_tensor_causal.py . \
    && python3 /tmp/eugr-patches/patch_vllm_gemma4_mtp_embedding_share.py . \
    && python3 /tmp/eugr-patches/patch_vllm_sm120_cooperative_topk.py . \
    && python3 /tmp/eugr-patches/patch_vllm_autogptq_symmetric_moe_qzeros.py . \
    && python3 /tmp/eugr-patches/patch_vllm_disable_minimax_qk_rmsnorm_ipc.py . \
    && python3 /tmp/eugr-patches/patch_vllm_routed_experts_weight_shape.py . \
    && python3 /tmp/eugr-patches/patch_vllm_spark_kv_cache_cleanup.py .
COPY --from=deps-src /deps2 /deps2
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    --mount=type=cache,id=ccache,target=/root/.ccache \
    --mount=type=cache,id=cargo-registry,target=/root/.cargo/registry \
    --mount=type=cache,id=rust-target,target=/workspace/vllm/vllm/target \
    git config --global protocol.version 0 \
    && uv pip install setuptools-rust setuptools-scm \
    && (command -v cargo >/dev/null 2>&1 || curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal --default-toolchain stable --no-modify-path) \
    && export PATH=/root/.cargo/bin:$PATH \
    && (command -v protoc >/dev/null 2>&1 || (apt-get update -qq && apt-get install -y --no-install-recommends protobuf-compiler libprotobuf-dev >/dev/null && rm -rf /var/lib/apt/lists/*)) \
    && _j="${BUILD_JOBS:-$(nproc)}" \
    && export GIT_PROTOCOL=0 MAX_JOBS="${_j}" CMAKE_BUILD_PARALLEL_LEVEL="${_j}" \
        VLLM_REQUIRE_RUST_FRONTEND=1 CARGO_BUILD_JOBS="${_j}" \
        TRITON_KERNELS_SRC_DIR=/deps2/triton_kernels \
        FMHA_SM100_SRC_DIR=/deps2/MSA \
        FLASH_MLA_SRC_DIR=/deps2/FlashMLA \
        FLASH_KDA_SRC_DIR=/deps2/FlashKDA \
        VLLM_FLASH_ATTN_SRC_DIR=/deps2/flash-attention \
        QUTLASS_SRC_DIR=/deps2/qutlass \
        TML_FA4_SRC_DIR=/deps2/tml-fa4 \
        VLLM_CUTLASS_SRC_DIR=/deps2/cutlass \
        CUTLASS_INCLUDE_DIR_PRE=/deps2/cutlass/include \
    && 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
# Pre-staged r12 b12x tree (same reason as vllm: node github WAF block);
# base B12X_BASE_REF + B12X_PRS merged on the build host.
COPY b12x-src.tar.gz /tmp/b12x-src.tar.gz
RUN set -eux; \
    mkdir -p /workspace/b12x; \
    tar xzf /tmp/b12x-src.tar.gz -C /workspace/b12x; \
    rm -f /tmp/b12x-src.tar.gz; \
    git config --global --add safe.directory /workspace/b12x; \
    cd /workspace/b12x; \
    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"]