# The LMCache Standalone Dockerfile is used to build a LMCache image
# without vLLM integration.

ARG CUDA_VERSION=13.0
ARG UBUNTU_VERSION=24.04
ARG NGC_VERSION=25.09
ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:${NGC_VERSION}-cuda${CUDA_VERSION}-devel-ubuntu${UBUNTU_VERSION}
# Override to a stable version for releases.
ARG VLLM_VERSION=nightly

#################### BASE BUILD IMAGE ####################
# Prepare basic build environment

FROM ${BASE_IMAGE} AS base

ARG CUDA_VERSION
ARG PYTHON_VERSION=3.12
ARG UBUNTU_VERSION
ENV DEBIAN_FRONTEND=noninteractive
ENV PATH="/opt/venv/bin:${PATH}"

# Install Python and other dependencies
RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \
    && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \
    && apt-get update -y \
    && apt-get install -y --no-install-recommends \
        ccache software-properties-common git curl sudo \
        python3 python3-dev python3-venv tzdata \
    && rm -rf /var/lib/apt/lists/* \
    && ldconfig /usr/local/cuda-$(echo $CUDA_VERSION | cut -d. -f1,2)/compat/ \
    && curl -LsSf https://astral.sh/uv/install.sh | sh \
    && mv ~/.local/bin/uv /usr/local/bin/ \
    && mv ~/.local/bin/uvx /usr/local/bin/ \
    && uv venv /opt/venv \
    && . /opt/venv/bin/activate \
    && python3 --version

WORKDIR /workspace

# Install small non-torch runtime dependencies; torch is installed in lmcache-build
# against the exact CUDA version so the compiled C extensions match at runtime.
# The CLI requirements (e.g. ``openai`` for the bench client) are also installed
# here so ``lmcache`` subcommand discovery in ``lmcache.cli.commands`` -- which
# eagerly imports every command module -- does not crash with
# ``ModuleNotFoundError: No module named 'openai'`` when the standalone image
# boots ``lmcache server``. See LMCache#3353.
RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,source=requirements/cli.txt,target=/tmp/cli.txt \
    . /opt/venv/bin/activate && \
    uv pip install ray nvidia-ml-py -r /tmp/cli.txt

# CUDA arch list used by torch
ARG torch_cuda_arch_list='7.5 8.0 8.6 8.9 9.0 10.0 12.0+PTX'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}

#################### TORCH RESOLVER ##########################
# Resolve the torch version vLLM uses for this CUDA tag into /torch.pin.
# vLLM is not installed; only the pin is consumed by downstream stages.

FROM base AS torch-resolver

ARG CUDA_VERSION
ARG VLLM_VERSION

RUN --mount=type=cache,target=/root/.cache/uv,id=uv-resolver \
    . /opt/venv/bin/activate && \
    CUDA_TAG=cu$(echo ${CUDA_VERSION} | tr -d '.') && \
    if [ "$VLLM_VERSION" = "nightly" ]; then \
        echo "vllm" > /tmp/req.in && \
        uv pip compile /tmp/req.in --quiet --prerelease=allow \
            --extra-index-url https://wheels.vllm.ai/nightly/${CUDA_TAG} \
            --extra-index-url https://download.pytorch.org/whl/${CUDA_TAG} \
            --index-strategy unsafe-best-match \
            > /tmp/resolved.txt ; \
    else \
        echo "vllm==${VLLM_VERSION}" > /tmp/req.in && \
        uv pip compile /tmp/req.in --quiet --prerelease=allow \
            --extra-index-url https://download.pytorch.org/whl/${CUDA_TAG} \
            > /tmp/resolved.txt ; \
    fi && \
    grep -E '^torch==' /tmp/resolved.txt > /torch.pin && \
    echo "Resolved torch pin (matches vllm ${VLLM_VERSION} on ${CUDA_TAG}):" && \
    cat /torch.pin

#################### LMCache Build ##########################
# Build LMCache wheel

FROM base AS lmcache-build

ARG CUDA_VERSION

# install build dependencies
COPY ./requirements/build.txt build.txt
COPY --from=torch-resolver /torch.pin /torch.pin

# Max jobs used by Ninja to build extensions
ARG max_jobs=2
ENV MAX_JOBS=${max_jobs}

# Number of threads used by nvcc
ARG nvcc_threads=8
ENV NVCC_THREADS=$nvcc_threads

# --index-url (singular) keeps PyPI out of torch resolution so uv can't
# pick PyPI's default cu13 wheel on a cu12.9 host.
RUN --mount=type=cache,target=/root/.cache/uv \
    . /opt/venv/bin/activate && \
    CUDA_TAG=cu$(echo ${CUDA_VERSION} | tr -d '.') && \
    uv pip install -r /torch.pin \
        --index-url https://download.pytorch.org/whl/${CUDA_TAG} && \
    uv pip install -r build.txt

ARG LMCACHE_COMMIT_ID=1

COPY . /workspace/LMCache
WORKDIR /workspace/LMCache

# Build LMCache wheel (don't install yet)
RUN --mount=type=cache,target=/root/.cache/ccache,id=ccache \
    --mount=type=cache,target=/root/.cache/uv,id=uv-cache,sharing=locked \
    . /opt/venv/bin/activate && \
    export LMCACHE_CUDA_MAJOR=$(echo ${CUDA_VERSION} | cut -d. -f1) && \
    python3 -c 'import torch; print("TORCH=", torch.__version__)' && \
    python3 setup.py bdist_wheel --dist-dir=dist_lmcache

#################### LMCache Final Image ##########################
# Clean install of LMCache from wheel

FROM base AS lmcache-final

ARG CUDA_VERSION

# Copy the built wheel from build stage
COPY --from=lmcache-build /workspace/LMCache/dist_lmcache/*.whl /tmp/
COPY --from=torch-resolver /torch.pin /torch.pin

# Install torch from the CUDA index first; the wheel's torch dep is then
# already satisfied, so uv won't re-resolve it against PyPI default cu13.
RUN --mount=type=cache,target=/root/.cache/uv \
    . /opt/venv/bin/activate && \
    CUDA_TAG=cu$(echo ${CUDA_VERSION} | tr -d '.') && \
    uv pip install -r /torch.pin \
        --index-url https://download.pytorch.org/whl/${CUDA_TAG} && \
    uv pip install /tmp/*.whl --verbose && \
    rm -rf /tmp/*.whl /torch.pin

WORKDIR /workspace

# Default shell entrypoint (no vLLM server)
CMD ["/bin/bash"]