FutureCudaDockerfile

ARG FROM_IMAGE_NAME=pytorch/pytorch:2.9.1-cuda13.0-cudnn9-runtime
FROM ${FROM_IMAGE_NAME}

# Install build dependencies and llama-cpp-python with CUDA support
ENV DEBIAN_FRONTEND=noninteractive

# Install Python and pip
RUN apt-get update && apt-get install -y python3 python3-pip build-essential cmake ninja-build wget && \
apt-get clean && rm -rf /var/lib/apt/lists/*

ENV PYTHONPATH=/code
WORKDIR /code

RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/cuda-keyring_1.1-1_all.deb && dpkg -i cuda-keyring_1.1-1_all.deb

RUN apt-get update && apt-get -y install cuda-toolkit-13-0 && apt autoremove -y

ENV GGML_CUDA=on
ENV GGML_CUDA_ARCH="89;90;100;120"
ENV CUDA_DOCKER_ARCH="89;90;100;120"
ENV TORCH_CUDA_ARCH_LIST="8.9;9.0;10.0;12.0"

ENV LD_LIBRARY_PATH=/usr/local/cuda-13.0/lib64:$LD_LIBRARY_PATH
ENV PATH=/usr/local/cuda-13.0/bin:$PATH
ENV CUDA_HOME=/usr/local/cuda-13.0

# Install llama-cpp-python with CUDA support for GGUF models
RUN CUDACXX=/usr/local/cuda-13.0/bin/nvcc CMAKE_ARGS="-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=${GGML_CUDA_ARCH}" FORCE_CMAKE=1 CMAKE_BUILD_PARALLEL_LEVEL=2 pip install llama-cpp-python --no-cache-dir --force-reinstall --upgrade

# Install Python dependencies
RUN pip install "fastapi[standard]" "uvicorn[standard]" httpx --break-system-packages

COPY ./app /code

EXPOSE 8080

CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8080"]