57 lines
1.8 KiB
Docker
57 lines
1.8 KiB
Docker
# ---------- Builder Base ----------
|
|
# Using NVIDIA NGC PyTorch container (26.03) with:
|
|
# - PyTorch 2.11.0a0 (bleeding edge)
|
|
# - CUDA 13.2.0
|
|
# - cuDNN 9.20, NCCL 2.29.7, TensorRT 10.16, TransformerEngine 2.13
|
|
# - Multi-arch: x86 + ARM SBSA (GH200 support)
|
|
FROM nvcr.io/nvidia/pytorch:26.03-py3 AS base
|
|
|
|
# Set arch lists for all targets
|
|
ARG TORCH_CUDA_ARCH_LIST="9.0a"
|
|
ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}
|
|
ENV DEBIAN_FRONTEND=noninteractive
|
|
|
|
# Install dependencies
|
|
RUN apt update && apt upgrade -y && \
|
|
apt install -y --no-install-recommends \
|
|
curl \
|
|
git \
|
|
libibverbs-dev \
|
|
zlib1g-dev && \
|
|
apt clean && \
|
|
rm -rf /var/lib/apt/lists/* /var/cache/apt/archives
|
|
|
|
# Set compiler paths
|
|
ENV CC=/usr/bin/gcc
|
|
ENV CXX=/usr/bin/g++
|
|
|
|
# Install uv
|
|
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh
|
|
|
|
# Set up workspace
|
|
WORKDIR /workspace
|
|
|
|
# Environment setup (PyTorch container already has CUDA paths set)
|
|
ENV CUDA_HOME=/usr/local/cuda
|
|
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
|
|
|
|
FROM base AS build-base
|
|
|
|
# Install build tools and dependencies
|
|
RUN pip install -U build cmake ninja pybind11 "setuptools>=77.0.3,<81.0.0"
|
|
|
|
# Clone the repo, apply the patch, and build
|
|
RUN git clone https://github.com/LMCache/LMCache.git && \
|
|
cd LMCache && \
|
|
git checkout dev && \
|
|
echo "\n\n========================================" && \
|
|
echo ">>> BUILDING LMCACHE FROM:" && \
|
|
echo ">>> BRANCH: $(git rev-parse --abbrev-ref HEAD)" && \
|
|
echo ">>> COMMIT: $(git rev-parse HEAD)" && \
|
|
echo ">>> DATE: $(git log -1 --format=%cd --date=short)" && \
|
|
echo "========================================\n\n" && \
|
|
pip install setuptools_scm && \
|
|
python -m build --wheel --no-isolation && \
|
|
cp dist/*.whl /workspace/
|
|
CMD ["/bin/bash"]
|