# Two-stage, in-docker build of the local-inference-lab r12 stack for GB10
# (arm64 / SM121a), per the rtx6kpro source contract (models/glm-5.3-flash.md):
#
#   vLLM  = local-inference-lab/vllm composition commit
#           7d8a09c42c7ba743b9e936562aa9205f9d0fda9d
#           (dev/jovian-judgement @54f6e9826c20... + PR heads 550/552/553/554/555/557)
#   B12X  = local-inference-lab/b12x composition commit
#           d064ca4cc3aa25292f1a8756fa51b38134d1df84
#           (master @139e04048bc3b... + PR heads 260/267/268/269/270)
#   LMCache optional — NOT vendored (opt-in).
#
# The compile happens INSIDE docker (operator directive). Torch/flashinfer/
# CUDA come from the working day-0 arm64 cu130 base that already serves GLM
# 5.3 Flash; the fork sources install with --no-deps so the base toolchain
# stays pinned. sm_121f applies to block-scaled FP4 kernels (kk3 gencode
# rule); TORCH_CUDA_ARCH_LIST pins the extension build.
#
# Verification markers: pip lines for vllm==0.1.dev...; python -c
# "import vllm" in a `--gpus` container; quantprobe in modelopt.py.

ARG BASE_IMAGE
FROM ${BASE_IMAGE} AS build

ARG VLLM_REF
ARG B12X_REF
ARG VLLM_PRS="550 552 553 554 555 557"
ARG B12X_PRS="260 267 268 269 270"
ARG VLLM_REPO=https://github.com/local-inference-lab/vllm.git
ARG B12X_REPO=https://github.com/local-inference-lab/b12x.git

RUN set -eux; apt-get update && apt-get install -y --no-install-recommends git ca-certificates && rm -rf /var/lib/apt/lists/*

# The rtx6kpro "composition commit" (7d8a09c… / d064ca4c…) is NOT a public
# ref — reproduce it: base branch checkout + sequential merge of the open
# PR heads (refs/pull/*/head), ascending order.
RUN set -eux; \
    git config --global user.email build@sm121.lan; \
    git config --global user.name sm121-builder; \
    git clone ${VLLM_REPO} /tmp/vllm-src; \
    git -C /tmp/vllm-src checkout ${VLLM_REF}; \
    for pr in ${VLLM_PRS}; do \
        git -C /tmp/vllm-src fetch origin "refs/pull/$pr/head:prs/$pr"; \
        git -C /tmp/vllm-src merge --no-ff -m "PR $pr" "prs/$pr"; \
    done; \
    git clone ${B12X_REPO} /tmp/b12x-src; \
    git -C /tmp/b12x-src checkout ${B12X_REF}; \
    for pr in ${B12X_PRS}; do \
        git -C /tmp/b12x-src fetch origin "refs/pull/$pr/head:prs/$pr"; \
        git -C /tmp/b12x-src merge --no-ff -m "PR $pr" "prs/$pr"; \
    done

RUN set -eux; \
    cd /tmp/vllm-src; \
    pip install setuptools-rust; \
    apt-get update && apt-get install -y --no-install-recommends ninja-build && rm -rf /var/lib/apt/lists/*; \
    pip install cmake; \
    # cu13 wheels carry missing dev headers (cusparse.h etc.) outside the
    # default include path — symlink only what /usr/local/cuda/include lacks.
    INC=/usr/local/lib/python3.12/dist-packages/nvidia/cu13/include; \
    for h in "$INC"/*.h; do b=$(basename "$h"); [ -e "/usr/local/cuda/include/$b" ] || ln -s "$h" /usr/local/cuda/include/$b; done; \
    # The spinloop target needs an explicit nvrtc: the cu13 toolkit ships
    # under targets/sbsa-linux (ldconfig knows it, FindCUDA does not).
    NVRTC=$(ls /usr/local/cuda*/targets/sbsa-linux/lib/libnvrtc.so* | head -1); \
    ln -s "${NVRTC}" /usr/lib/aarch64-linux-gnu/libnvrtc.so; \
    CMAKE_ARGS="-DCUDAToolkit_ROOT=/usr/local/cuda -DCUDA_TOOLKIT_ROOT_DIR=/usr/local/cuda -DCUDA_nvrtc_LIBRARY=${NVRTC} -DCUDAToolkit_LIBRARY_DIR=/usr/local/cuda/targets/sbsa-linux/lib" \
    VLLM_TARGET_DEVICE=cuda TORCH_CUDA_ARCH_LIST="12.1a" \
        pip install --no-build-isolation --no-deps . 2>&1; \
    python3 -c "import vllm; print('vllm import ok', vllm.__version__)"

# B12X builds its own CUDA extension set; compile gate fails loudly on
# anchor drift. Keep base-provided deps untouched (--no-deps).
RUN set -eux; \
    cd /tmp/b12x-src; \
    TORCH_CUDA_ARCH_LIST="12.1f;12.1a" \
        pip install --no-build-isolation --no-deps . 2>&1; \
    python3 -c "import b12x; print('b12x import ok', b12x.__version__)" || \
    python3 -c "import importlib.util; \
        assert importlib.util.find_spec('b12x'), 'b12x module not found'; \
        print('b12x module present')"

RUN set -eux; python3 -c "import vllm; print('vllm', vllm.__version__)"

FROM ${BASE_IMAGE}

# Compiled artifacts only (site-packages), dropped onto a pristine base so
# the container keeps the day-0 entrypoint/env layout.
COPY --from=build /usr/local/lib/python3.12/dist-packages/ /usr/local/lib/python3.12/dist-packages/

RUN set -eux; \
    python3 -c "import vllm; print('final vllm', vllm.__version__)"