feat: add CuBLAS support in Docker images (#403)

Signed-off-by: Sébastien Prud'homme <sebastien.prudhomme@gmail.com>

Dockerfile

@@ -1,15 +1,49 @@
 ARG GO_VERSION=1.20
-ARG BUILD_TYPE=
 FROM golang:$GO_VERSION
+
+ARG BUILD_TYPE=
+ARG GO_TAGS=
+ARG CUDA_MAJOR_VERSION=11
+ARG CUDA_MINOR_VERSION=7
+
+ENV BUILD_TYPE=${BUILD_TYPE}
+ENV GO_TAGS=${GO_TAGS}
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+ENV NVIDIA_REQUIRE_CUDA="cuda>=${CUDA_MAJOR_VERSION}.0"
+ENV NVIDIA_VISIBLE_DEVICES=all
+ENV HEALTHCHECK_ENDPOINT=http://localhost:8080/readyz
 ENV REBUILD=true
+
 WORKDIR /build
-RUN apt-get update && apt-get install -y cmake curl libgomp1 libopenblas-dev libopenblas-base libopencv-dev libopencv-core-dev libopencv-core4.5 ca-certificates
+
+RUN apt-get update && \
+    apt-get install -y ca-certificates cmake curl
+
+# CuBLAS requirements
+RUN if [ "${BUILD_TYPE}" = "cublas" ]; then \
+    apt-get install -y software-properties-common && \
+    apt-add-repository contrib && \
+    curl -O https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.0-1_all.deb && \
+    dpkg -i cuda-keyring_1.0-1_all.deb && \
+    rm -f cuda-keyring_1.0-1_all.deb && \
+    apt-get update && \
+    apt-get install -y cuda-nvcc-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libcublas-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
+    ; fi
+ENV PATH /usr/local/cuda/bin:${PATH}
+
+# OpenBLAS requirements
+RUN apt-get install -y libopenblas-dev
+
+# Stable Diffusion requirements
+RUN apt-get install -y libopencv-dev && \
+    ln -s /usr/include/opencv4/opencv2 /usr/include/opencv2
+
 COPY . .
-RUN ln -s /usr/include/opencv4/opencv2/ /usr/include/opencv2
 RUN make build
-ENV HEALTHCHECK_ENDPOINT=http://localhost:8080/readyz
+
 # Define the health check command
 HEALTHCHECK --interval=30s --timeout=360s --retries=10 \
   CMD curl -f $HEALTHCHECK_ENDPOINT || exit 1
+
 EXPOSE 8080
 ENTRYPOINT [ "/build/entrypoint.sh" ]

Dockerfile.dev

@@ -1,20 +1,79 @@
 ARG GO_VERSION=1.20
 ARG DEBIAN_VERSION=11
+FROM golang:$GO_VERSION as builder
+
 ARG BUILD_TYPE=
+ARG GO_TAGS=
+ARG CUDA_MAJOR_VERSION=11
+ARG CUDA_MINOR_VERSION=7
+
+ENV BUILD_TYPE=${BUILD_TYPE}
+ENV GO_TAGS=${GO_TAGS}
 
-FROM golang:$GO_VERSION as builder
 WORKDIR /build
-RUN apt-get update && apt-get install -y cmake libgomp1 libopenblas-dev libopenblas-base libopencv-dev libopencv-core-dev libopencv-core4.5 
-RUN ln -s /usr/include/opencv4/opencv2/ /usr/include/opencv2
+
+RUN apt-get update && \
+    apt-get install -y cmake
+
+# CuBLAS requirements
+RUN if [ "${BUILD_TYPE}" = "cublas" ]; then \
+    apt-get install -y software-properties-common && \
+    apt-add-repository contrib && \
+    curl -O https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.0-1_all.deb && \
+    dpkg -i cuda-keyring_1.0-1_all.deb && \
+    rm -f cuda-keyring_1.0-1_all.deb && \
+    apt-get update && \
+    apt-get install -y cuda-nvcc-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libcublas-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
+    ; fi
+ENV PATH /usr/local/cuda/bin:${PATH}
+
+# OpenBLAS requirements
+RUN apt-get install -y libopenblas-dev
+
+# Stable Diffusion requirements
+RUN apt-get install -y libopencv-dev && \
+    ln -s /usr/include/opencv4/opencv2 /usr/include/opencv2
+
 COPY . .
 RUN make build
 
 FROM debian:$DEBIAN_VERSION
-COPY --from=builder /build/local-ai /usr/bin/local-ai
-RUN apt-get update && apt-get install -y ca-certificates curl
+
+ARG BUILD_TYPE=
+ARG GO_TAGS=
+ARG CUDA_MAJOR_VERSION=11
+ARG CUDA_MINOR_VERSION=7
+
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+ENV NVIDIA_REQUIRE_CUDA="cuda>=${CUDA_MAJOR_VERSION}.0"
+ENV NVIDIA_VISIBLE_DEVICES=all
 ENV HEALTHCHECK_ENDPOINT=http://localhost:8080/readyz
+
+RUN apt-get update && \
+    apt-get install -y ca-certificates curl
+
+# CuBLAS requirements
+RUN if [ "${BUILD_TYPE}" = "cublas" ]; then \
+    apt-get install -y curl software-properties-common && \
+    apt-add-repository contrib && \
+    curl -O https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.0-1_all.deb && \
+    dpkg -i cuda-keyring_1.0-1_all.deb && \
+    rm -f cuda-keyring_1.0-1_all.deb && \
+    apt-get update && \
+    apt-get install -y cuda-cudart-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} libcublas-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
+    ; fi
+
+# OpenBLAS requirements
+RUN apt-get install -y libopenblas0
+
+# Stable Diffusion requirements
+RUN apt-get install -y libgomp1 libopencv-core4.5 libopencv-imgcodecs4.5
+
+COPY --from=builder /build/local-ai /usr/bin/local-ai
+
 # Define the health check command
 HEALTHCHECK --interval=30s --timeout=360s --retries=10 \
   CMD curl -f $HEALTHCHECK_ENDPOINT || exit 1
+
 EXPOSE 8080
-ENTRYPOINT [ "/usr/bin/local-ai" ]
+ENTRYPOINT [ "/usr/bin/local-ai" ]