From 48679e8b74936ad65c3800ac39c536a1997e946e Mon Sep 17 00:00:00 2001
From: Yee Kit <yeek3063@gmail.com>
Date: Fri, 26 Jan 2024 23:06:50 +0800
Subject: [PATCH] Update Dockerfile with Cuda support by default

---
 Dockerfile | 36 ++++++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 3f318ce..bef2e7f 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,31 +1,52 @@
-# Use the official Python Image
-FROM python:3.11.4
+ARG CUDA_IMAGE="12.1.1-devel-ubuntu22.04"
+# Build with Cuda 12.1.1 and Ubuntu 22.04
+FROM nvidia/cuda:${CUDA_IMAGE}
 
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user
+
 # Install the dependencies
-RUN apt-get update && apt-get -y install cmake protobuf-compiler libopenblas-dev liblapack-dev liblapacke-dev libeigen3-dev libboost-all-dev
+RUN apt-get update && apt-get upgrade -y \
+    && apt-get install -y git build-essential \
+    python3.11 gcc wget \
+    ocl-icd-opencl-dev opencl-headers clinfo \
+    cmake protobuf-compiler pkg-config \
+    libclblast-dev libopenblas-dev \
+    liblapack-dev liblapacke-dev libeigen3-dev libboost-all-dev \
+    && mkdir -p /etc/OpenCL/vendors && echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd
+
+# Install pip for python 3.11
+RUN wget https://bootstrap.pypa.io/get-pip.py && \
+    python3.11 get-pip.py && \
+    rm get-pip.py
 
 # Switch to the user 'user'
 USER user
 
+# Setting build related env vars
+ENV CUDA_DOCKER_ARCH=all
+ENV LLAMA_CUBLAS=1
+
 # Set home to the user's home directory and Poetry's environment variables
 ENV HOME=/home/user \
 	PATH=/home/user/.local/bin:$PATH \
+    PYTHONUNBUFFERED=1 \
     POETRY_NO_INTERACTION=1 \
     POETRY_VIRTUALENVS_IN_PROJECT=1 \
     POETRY_VIRTUALENVS_CREATE=1 \
     POETRY_CACHE_DIR=/tmp/poetry_cache \
-    CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS"
+    # Build llama-cpp-python with default cuda support
+    CMAKE_ARGS="-DLLAMA_CUBLAS=on"
+    # CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS"
 
 # Set the working directory to /app
 WORKDIR $HOME/app
 
 # Update pip and wheel
-RUN pip install --upgrade pip setuptools wheel
+RUN python3.11 -m pip install --upgrade pip setuptools wheel
 
 # Install poetry
-RUN pip install poetry
+RUN python3.11 -m pip install poetry
 
 # Copy the poetry files
 COPY --chown=user ./backend/pyproject.toml ./backend/poetry.lock $HOME/app/
@@ -34,8 +55,7 @@ COPY --chown=user ./backend/pyproject.toml ./backend/poetry.lock $HOME/app/
 COPY --chown=user ./backend $HOME/app
 
 # Install the dependencies
-RUN poetry lock --no-update
-RUN poetry install --without dev && \
+RUN poetry install --without dev,torch-cpu && \
     rm -rf /tmp/poetry_cache
 
 # Change to the package directory