From 9ddc08be17bf0d3cea16aafe44191905a6ed4931 Mon Sep 17 00:00:00 2001
From: One <imone@tuta.io>
Date: Sat, 25 May 2024 02:25:07 +0800
Subject: [PATCH] [docker] serving and ci workflow

---
 .github/workflows/docker_build_and_push.yaml | 34 +++++++++++++++++++
 docker/serving/Dockerfile                    | 35 ++++++++++++++++++++
 docker/serving/start.sh                      | 21 ++++++++++++
 3 files changed, 90 insertions(+)
 create mode 100644 .github/workflows/docker_build_and_push.yaml
 create mode 100644 docker/serving/Dockerfile
 create mode 100644 docker/serving/start.sh

diff --git a/.github/workflows/docker_build_and_push.yaml b/.github/workflows/docker_build_and_push.yaml
new file mode 100644
index 0000000..b69e794
--- /dev/null
+++ b/.github/workflows/docker_build_and_push.yaml
@@ -0,0 +1,34 @@
+name: Docker Build and Push
+
+on:
+  workflow_run:
+    workflows: ["Publish to PyPI.org"]
+    types:
+      - completed
+
+jobs:
+  docker:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+      
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+      
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: ./docker/serving
+          platforms: linux/amd64
+          push: true
+          tags: ${{ secrets.DOCKERHUB_TAG }}
diff --git a/docker/serving/Dockerfile b/docker/serving/Dockerfile
new file mode 100644
index 0000000..4e950f6
--- /dev/null
+++ b/docker/serving/Dockerfile
@@ -0,0 +1,35 @@
+FROM python:3.11-slim-bookworm
+
+######### Setup system
+
+RUN mkdir /workspace && mkdir /workspace/transformers_cache
+WORKDIR /workspace
+
+ENV HF_HOME /workspace/transformers_cache
+
+######### Install system dependencies
+RUN apt update && apt install -y git bash curl wget libxml2
+
+# Install ssh server, remove all pre-generated ssh host keys, and disable password auth
+RUN apt install -y openssh-server && \
+    rm -f /etc/ssh/ssh_host_* && \
+    sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/g' /etc/ssh/sshd_config
+
+# Install CUDA (for FlashAttention 2)
+RUN wget -q --show-progress --progress=bar:force:noscroll -O cuda_installer https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run && \
+    chmod +x cuda_installer && \
+    ./cuda_installer --silent --toolkit --override && \
+    rm -f cuda_installer
+
+######### Install OpenChat
+# Install OpenChat
+RUN pip3 install ninja packaging torch
+RUN pip3 install ochat
+
+######### Install Cloudflared
+RUN wget -q --show-progress --progress=bar:force:noscroll -O /cloudflared https://github.com/cloudflare/cloudflared/releases/latest/download/cloudflared-linux-amd64 && chmod +x /cloudflared
+
+######### Startup script
+
+COPY start.sh /start.sh
+ENTRYPOINT ["/start.sh"]
diff --git a/docker/serving/start.sh b/docker/serving/start.sh
new file mode 100644
index 0000000..677944d
--- /dev/null
+++ b/docker/serving/start.sh
@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# start ssh server
+if [ -n "$PUBLIC_KEY" ]; then
+    mkdir -p ~/.ssh
+    echo "$PUBLIC_KEY" >> ~/.ssh/authorized_keys
+    chmod 700 -R ~/.ssh
+
+    dpkg-reconfigure openssh-server  # generate ssh keys
+    service ssh start
+fi
+
+# start cloudflare tunnel
+if [ -n "$CLOUDFLARED_TUNNEL_ARGS" ]; then
+    /cloudflared $CLOUDFLARED_TUNNEL_ARGS &
+fi
+
+# start openchat server
+python3 -m ochat.serving.openai_api_server --model $MODEL --host 127.0.0.1 --port 18888 --engine-use-ray --worker-use-ray --disable-log-requests --disable-log-stats $ARGS &
+
+wait