# syntax=docker/dockerfile:1.7
# DeepSeek-V4-Flash DSpark — from-source build for DGX Spark (aarch64/sm_121a)
# EXPERIMENTAL. Builds vLLM + FlashInfer + NCCL + b12x from the
# local-inference-lab "eldritch" b12x fork at the newest DS4/DSpark branch, so we
# are not limited to the pre-built aidendle94 image.
#
# Derived from docker-images/glm52-nf3-hybrid (same generic from-source harness:
# torch/deps from the cu130 wheel index via a hash-locked requirements file,
# FlashInfer/vLLM built from source wheels, b12x editable install). GLM-specific
# overlay/patches are dropped; only the branch-pinned versions.env differs.
#
# All build inputs pinned in versions.env and locks/. We do NOT reuse the GLM
# python-build/python-runtime locks (they pin torch 2.11 / flashinfer 0.6.12).
# Regenerate them for THIS stack first:
#     bash gen-locks.sh     # -> writes locks/python-build.txt & python-runtime.txt
#
# Build: cd docker-images/deepseek-v4-flash-dspark-src && 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 NCCL_REPO NCCL_REF
# NCCL_REF may be a tag, a branch, or a commit SHA (the DS4 fork pins a SHA), so
# check out the ref rather than `git clone -b` (which only accepts branch/tag).
RUN git clone ${NCCL_REPO} /opt/nccl \
 && cd /opt/nccl \
 && git checkout ${NCCL_REF} \
 && 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
# Force DGX Spark arch (sm_121a). The fork's CMake otherwise falls back to the
# torch wheel's sm_120, but our FP4/`cvt.e2m1x2.f32` paths require sm_121a.
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 ${VLLM_REPO} /workspace/vllm \
 && cd /workspace/vllm \
 && git checkout ${VLLM_COMMIT} \
 && git submodule update --init --recursive
COPY patches/ /patches/
COPY apply-arch.py /apply-arch.py
RUN cd /workspace/vllm \
 && for p in /patches/*.patch; do [ -f "$p" ] && patch -p1 --forward < "$p" || true; done \
 && python3 /apply-arch.py
WORKDIR /workspace/vllm
RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
    --mount=type=cache,id=ccache,target=/root/.ccache \
    echo "DIAG patch=$(command -v patch || echo MISSING)" \
 && echo "DIAG forced_arch_lines=$(grep -c _forced_arch CMakeLists.txt || echo 0)" \
 && echo "DIAG supported:" && grep -n 'CUDA_SUPPORTED_ARCHS \"7.5' CMakeLists.txt \
 && echo "DIAG TORCH_CUDA_ARCH_LIST=[${TORCH_CUDA_ARCH_LIST} ]" \
 && _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 --no-deps /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

ENTRYPOINT []
CMD ["bash"]