# syntax=docker/dockerfile:1.7
# GLM-5.3-Flash r27 (jovian-judgement, qualified 2026-09-06) — GB10/Spark build.
#
# Clone-at-pin build, NO staged tarballs: the Dockerfile clones every source
# at the pins in versions.env and gates each one (HEAD commit + tree hash).
# A fresh clone of this repo + `bash build.sh` reproduces this image end to
# end. Source contract: local-inference-lab/rtx6kpro master
# models/glm-5.3-flash/build/glm53-jovian-judgement-community-20260906-r27.source.lock
# (vLLM 63a82f8d / tree f54cd9ca, B12X e8ad299b / f3cd8a9e [includes
# comm.roce RoCEnante, PR #295 + follow-ups — see state doc], LMCache
# 822330db / 008ac3e0, FlashKDA 3b225bf2; merge order documented in vLLM
# issue 651).
#
# Runtime: CUDA 13.2 devel aarch64 base + torch 2.13.0 cu130 + flashinfer
# 0.6.17 — the same pins the r27 tree's requirements/cuda.txt asks for,
# proven on GB10 (deviation from their RTX-PRO-6000 cu133 runtime documented
# in versions.env). GB10 sm_121 fixes baked in (patch_cuda_sm12x.py, eugr
# patch set). LMCache wheel + cuMem interposer built from the locked source
# in-container (engine-driven cache transfers sidecar for vram/native/lmcache
# cache modes).
#
# Build: cd docker-images/jovian-judgement && TAG=... bash build.sh

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
ARG APT_SOURCE
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-snapshot.list
COPY locks/apt-packages.txt /tmp/apt-packages.txt
COPY locks/python-bootstrap.txt /tmp/python-bootstrap.txt
COPY locks/deps.lock /tmp/deps.lock
# APT_SOURCE (versions.env):
#   snapshot (default, historical behavior): the pinned point-in-time mirror
#       + exact package versions — strictly reproducible, but
#       snapshot.ubuntu.com is slow (~30 kB/s) and intermittently 503s
#       (12-min apt stall, 2026-09-08).
#   current: the live archive.ubuntu.com — fast and reliable; installs the
#       pinned package names unversioned. Acceptable: the engine stack is
#       pinned by the vLLM/B12X/LMCache source pins, not by apt.
RUN set -eux; \
    if [ "${APT_SOURCE:-snapshot}" = "current" ]; then \
        printf '%s\n' \
            "deb http://ports.ubuntu.com/ubuntu-ports/ noble main restricted universe multiverse" \
            "deb http://ports.ubuntu.com/ubuntu-ports/ noble-updates main restricted universe multiverse" \
            "deb http://ports.ubuntu.com/ubuntu-ports/ noble-security main restricted universe multiverse" \
            > /etc/apt/sources.list; \
        PKGS="$(grep -vE '^\s*(#|$)' /tmp/apt-packages.txt | cut -d= -f1)"; \
    else \
        cp /tmp/apt-sources-snapshot.list /etc/apt/sources.list; \
        PKGS="$(grep -vE '^\s*(#|$)' /tmp/apt-packages.txt)"; \
    fi; \
    rm -f /etc/apt/sources.list.d/*.list /etc/apt/sources.list.d/*.sources; \
    apt-get update; \
    apt-get install -y --no-install-recommends ${PKGS}; \
    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}"

# DeepGEMM: cloned in-build (proven to work on the nodes via git; pins are
# gated again in fetch-deps.sh against the vLLM tree's own declaration).
FROM torch-base AS deepgemm-builder
ARG DEEPGEMM_REPO DEEPGEMM_BRANCH DEEPGEMM_REF
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

# ---- vLLM @ r27 integration branch, gated by commit + tree hash ----
FROM torch-base AS vllm-clone
ARG VLLM_REPO VLLM_PIN VLLM_TREE
RUN set -eux; \
    git config --global protocol.version 0; \
    git init -q /workspace/vllm; \
    git -C /workspace/vllm fetch -q --depth 1 ${VLLM_REPO} ${VLLM_PIN}; \
    git -C /workspace/vllm checkout -q FETCH_HEAD; \
    cd /workspace/vllm; \
    test "$(git rev-parse HEAD)" = "${VLLM_PIN}"; \
    test "$(git rev-parse HEAD^{tree})" = "${VLLM_TREE}"; \
    git config --global --add safe.directory /workspace/vllm

# ---- All build sub-deps, cloned at pins the r27 tree itself declares ----
# REGEN_LOCK=1: fetch-deps.sh rewrites /tmp/deps.lock from the staged state
# (pin cuts) instead of verifying against it — export the stage to harvest
# the new lock, then commit it.
FROM vllm-clone AS deps-fetched
ARG REGEN_LOCK=""
COPY fetch-deps.sh /tmp/fetch-deps.sh
RUN if [ -n "$REGEN_LOCK" ]; then bash /tmp/fetch-deps.sh --regen-lock /tmp/deps.lock; else bash /tmp/fetch-deps.sh; fi

FROM deps-fetched AS vllm-builder
WORKDIR /workspace/vllm
# Hard-declared (do NOT rely on stage-inherited ENVs): the buildkit build
# has no GPU, so torch's device-capability autodetect is unavailable and
# the wheel must be told the target arch explicitly (the r27 slip: kernels
# came out sm120-only -> cudaErrorNoKernelImageForDevice at boot).
ENV TORCH_CUDA_ARCH_LIST=12.1a GIT_PROTOCOL=0
ARG DEEPGEMM_REF
ARG VLLM_PRETEND_VERSION
ARG SOURCE_DATE_EPOCH
ARG BUILD_JOBS
# Empty default overrides (profile `FLASH_KDA_SRC_DIR=`) make the vLLM tree
# fall back to its own pinned FlashKDA + in-tree packed-checkpoints patch
# (required by the r28.1+ flashkda.cmake API gate); the default keeps the
# r27-style shared /deps2 checkout for the older trees.
ARG FLASH_KDA_SRC_DIR=/deps2/FlashKDA
# eugr/spark-vllm-docker build-time patches (main, 2026-09-02): all are
# source-aware/idempotent (skip when the vulnerable pattern is absent), so
# running the full set is safe on the r27 tree. 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 build collapses 12.1a -> 12.0)
#   - sm120_cooperative_topk: keep SM90-only (Spark fix)
#   - mrv2_speculator_cudagraph_pool / spark_kv_cache_cleanup (UMA snapshot
#     flow) / b12x c128a alignment / flashinfer b12x swigluoai / routed
#     expert weight shape / topk_softplus_sqrt / others: see each script.
COPY patches-eugr/patch_vllm_*.py /tmp/eugr-patches/
COPY --from=deepgemm-builder /deps/deepgemm /deps/deepgemm
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 .
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 --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}" \
        SETUPTOOLS_SCM_PRETEND_VERSION="${VLLM_PRETEND_VERSION}" \
        TRITON_KERNELS_SRC_DIR=/deps2/triton_kernels \
        FMHA_SM100_SRC_DIR=/deps2/MSA \
        FLASH_MLA_SRC_DIR=/deps2/FlashMLA \
        FLASH_KDA_SRC_DIR="${FLASH_KDA_SRC_DIR}" \
        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

# ---- B12X @ r27 integration branch, gated; pure-python CuTe-DSL wheel ----
FROM torch-base AS b12x-clone
ARG B12X_REPO B12X_PIN B12X_TREE
RUN set -eux; \
    git config --global protocol.version 0; \
    git init -q /workspace/b12x; \
    git -C /workspace/b12x fetch -q --depth 1 ${B12X_REPO} ${B12X_PIN}; \
    git -C /workspace/b12x checkout -q FETCH_HEAD; \
    cd /workspace/b12x; \
    test "$(git rev-parse HEAD)" = "${B12X_PIN}"; \
    test "$(git rev-parse HEAD^{tree})" = "${B12X_TREE}"; \
    git config --global --add safe.directory /workspace/b12x

FROM b12x-clone AS b12x-builder
WORKDIR /workspace/b12x
ARG B12X_PRETEND_VERSION
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    SETUPTOOLS_SCM_PRETEND_VERSION_FOR_B12X="${B12X_PRETEND_VERSION}" \
    uv build --no-build-isolation --wheel . --out-dir=/wheels -v

# ---- LMCache @ r27 source lock: wheel + cuMem interposer from source ----
FROM torch-base AS lmcache-builder
ARG LMCACHE_REPO LMCACHE_PIN LMCACHE_TREE LMCACHE_VERSION
ARG BUILD_JOBS
RUN set -eux; \
    git config --global protocol.version 0; \
    git init -q /workspace/lmcache; \
    git -C /workspace/lmcache fetch -q --depth 1 ${LMCACHE_REPO} ${LMCACHE_PIN}; \
    git -C /workspace/lmcache checkout -q FETCH_HEAD; \
    cd /workspace/lmcache; \
    test "$(git rev-parse HEAD^{tree})" = "${LMCACHE_TREE}"; \
    git config --global --add safe.directory /workspace/lmcache
WORKDIR /workspace/lmcache
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    uv pip install "setuptools>=77.0.3,<81.0.0" setuptools_scm wheel packaging ninja \
    && SETUPTOOLS_SCM_PRETEND_VERSION_FOR_LMCACHE="${LMCACHE_VERSION}" \
       BUILD_WITH_CUDA=1 \
       LMCACHE_CUDA_MAJOR=13 \
       TORCH_CUDA_ARCH_LIST=12.1a \
       ENABLE_CXX11_ABI=1 \
       MAX_JOBS="${BUILD_JOBS:-4}" \
       NVCC_THREADS=auto \
       python3 -m pip wheel --no-build-isolation --no-deps --wheel-dir /lmcache-artifacts . \
    && make -C csrc/cumem_ipc_interposer CUDA_HOME=/usr/local/cuda \
    && install -Dm755 csrc/cumem_ipc_interposer/liblmcache_cumem_shareable.so \
       /lmcache-artifacts/liblmcache_cumem_shareable.so \
    && cd /lmcache-artifacts && sha256sum *.whl liblmcache_cumem_shareable.so > SHA256SUMS

# ---- Runner: flattened GB10 serving image ----
FROM torch-base AS runner
ARG PYPI_INDEX_URL
ARG LMCACHE_VERSION
COPY checksums/tiktoken.sha256 /tmp/tiktoken.sha256
COPY patch_cuda_sm12x.py /tmp/patch_cuda_sm12x.py
COPY lmcache-runtime-requirements.txt /tmp/lmcache-runtime-requirements.txt
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=bind,from=lmcache-builder,source=/lmcache-artifacts,target=/lmcache-artifacts \
    --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 --no-deps -r /tmp/lmcache-runtime-requirements.txt \
 && pip install --no-cache-dir --no-deps --force-reinstall /b12x-wheels/*.whl \
 && pip install --no-cache-dir --no-deps --force-reinstall /lmcache-artifacts/*.whl \
 && install -Dm755 /lmcache-artifacts/liblmcache_cumem_shareable.so \
      /opt/lmcache/lib/liblmcache_cumem_shareable.so \
 && install -m 644 /tmp/lmcache-runtime-requirements.txt /opt/lmcache/ \
 && python3 /tmp/patch_cuda_sm12x.py \
 && find /usr/local/lib/python3.12/dist-packages/vllm /opt/lmcache \
      -name "__pycache__" -type d -prune -exec rm -rf {} + 2>/dev/null || true
# b12x loader allocation patch (discord, 2026-09-07): pinned_wc -> managed
# in b12x/integration/vllm/loader.py — two occurrences (the weight_pool call
# and the matching audit label). Reported prefill 1350 -> ~2000 tok/s. No
# recompiles; hard-gated: any pinned_wc left behind fails the build.
#
# Conditional since the r38 cut (2026-09-19): b12x master 0f3a8cbf replaced
# the hardcoded allocation with weight_allocation() in b12x/loader/_pool.py
# (auto-selects "managed" on GB10: pageable_memory_access + host_page_tables
# caps), so the legacy pinned_wc strings are gone and the sed is a no-op.
# Legacy pins (r27..r37 b12x) still carry the pattern and take the gated path.
RUN set -eux; \
    F=/usr/local/lib/python3.12/dist-packages/b12x/integration/vllm/loader.py; \
    if grep -q 'pinned_wc' "$F" 2>/dev/null; then \
      sed -i 's/allocation="pinned_wc"/allocation="managed"/; \
              s/"allocation": "pinned_wc"/"allocation": "managed"/' "$F"; \
      test "$(grep -c 'allocation="managed"' "$F")" = "1"; \
      grep -Fq '"allocation": "managed"' "$F"; \
      ! grep -q 'pinned_wc' "$F"; \
    else \
      echo "b12x loader carries no pinned_wc (new auto-allocation b12x) — patch skipped"; \
      test -f "$F"; \
    fi
# Build gate: the r27 stack must import with the GB10 patches in place and
# carry the qualified package trees (marker spot-checks from the source lock).
RUN set -eux; \
    python3 -c "import vllm; print('vllm', vllm.__version__)"; \
    python3 -c "from importlib.metadata import version; print('b12x', version('b12x'))"; \
    python3 -c "from importlib.metadata import version as v; \
      assert v('lmcache') == '${LMCACHE_VERSION}', v('lmcache'); print('lmcache', v('lmcache'))"; \
    python3 -c "import lmcache, lmcache.cuda_ops, lmcache.lmcache_native, lmcache.lmcache_fs; print('lmcache modules 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; \
    grep -Fq 'def apply_top_k_top_p_probs' /usr/local/lib/python3.12/dist-packages/vllm/v1/sample/ops/topk_topp_sampler.py; \
    grep -Fq '_PREPARED_ONESHOT_LAUNCHERS' /usr/local/lib/python3.12/dist-packages/b12x/comm/pcie/_oneshot_cute.py; \
    grep -Fq 'self._transfer_workspace_slot_count = self._commit_workers' /usr/local/lib/python3.12/dist-packages/lmcache/v1/multiprocess/transfer_context/async_engine_driven.py; \
    test -s /opt/lmcache/lib/liblmcache_cumem_shareable.so

ENV LMCACHE_EXPECTED_SOURCE_ROOT=/usr/local/lib/python3.12/site-packages/lmcache
ENTRYPOINT []
CMD ["bash"]