# syntax=docker/dockerfile:1.7
# Solar Open2 250B NVFP4 (Nota) — self-contained build for DGX Spark (aarch64/sm_121a)
#
# Stack mirrors glm52-nf3-hybrid (cu130, pinned torch/NCCL/flashinfer/cutlass-dsl)
# minus the b12x / NF3-hybrid suite: Solar is a stock vLLM serve using the
# UpstageAI solar_open2 fork (solar_open2 reasoning/tool parsers + logits
# processor). NVFP4 (compressed-tensors) is auto-detected from the model config.
#
# All build inputs pinned in versions.env and locks/. No external repo needed
# beyond the package registries and git repos each stage clones.
#
# Build: cd docker-images/solar-open2-nvfp4 && 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 TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION TRITON_VERSION
ARG NCCL_REPO NCCL_REF
RUN git clone --depth 1 -b ${NCCL_REF} ${NCCL_REPO} /opt/nccl \
 && cd /opt/nccl \
 && 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
RUN git clone --recursive ${VLLM_REPO} /workspace/vllm \
 && cd /workspace/vllm \
 && git checkout ${VLLM_REF} \
 && git submodule update --init --recursive
WORKDIR /workspace/vllm
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}" \
 && 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 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 /fi-wheels/*.whl /vllm-wheels/*.whl

# cutlass-dsl (used by flashinfer/vllm CuTeDSL paths at runtime)
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"]