From 7839eb14d3c75a589a3665e447e995745eee7f30 Mon Sep 17 00:00:00 2001
From: Andrei Betlen <abetlen@gmail.com>
Date: Wed, 3 May 2023 10:29:05 -0400
Subject: [PATCH] Add docker cuda image. Closes #143

---
 .github/workflows/build-docker.yaml | 34 ++++++++++++++++++++++++++++-
 Dockerfile.cuda                     | 15 +++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 Dockerfile.cuda

diff --git a/.github/workflows/build-docker.yaml b/.github/workflows/build-docker.yaml
index 16b00a2..2b340d8 100644
--- a/.github/workflows/build-docker.yaml
+++ b/.github/workflows/build-docker.yaml
@@ -36,4 +36,36 @@ jobs:
           push: true # push to registry
           pull: true # always fetch the latest base images
           platforms: linux/amd64,linux/arm64 # build for both amd64 and arm64
-          tags: ghcr.io/abetlen/llama-cpp-python:latest
\ No newline at end of file
+          tags: ghcr.io/abetlen/llama-cpp-python:latest
+
+  docker-cuda:
+    name: Build and push Docker image
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          submodules: "true"
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v2
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v2 
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build and push
+        uses: docker/build-push-action@v4
+        with:
+          file: Dockerfile.cuda
+          context: .
+          push: true # push to registry
+          pull: true # always fetch the latest base images
+          platforms: linux/amd64,linux/arm64 # build for both amd64 and arm64
+          tags: ghcr.io/abetlen/llama-cpp-python-cuda:latest
\ No newline at end of file
diff --git a/Dockerfile.cuda b/Dockerfile.cuda
new file mode 100644
index 0000000..a852f3c
--- /dev/null
+++ b/Dockerfile.cuda
@@ -0,0 +1,15 @@
+FROM nvidia/cuda:12.1.1-devel-ubuntu20.04
+
+# We need to set the host to 0.0.0.0 to allow outside access
+ENV HOST 0.0.0.0
+
+COPY . .
+
+# Install the package
+RUN apt update && apt install -y python3 python3-pip
+RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette
+
+RUN LLAMA_CUBLAS=1 python3 setup.py develop
+
+# Run the server
+CMD python3 -m llama_cpp.server
\ No newline at end of file