From b2a4fb6648055330b1d3c4bd6469c83c07cc161d Mon Sep 17 00:00:00 2001
From: qingfengfenga <yanxiaole_vip@163.com>
Date: Tue, 19 Sep 2023 17:18:30 +0800
Subject: [PATCH] Add Docker start

---
 .dockerignore      |  2 ++
 .env               |  7 +++++++
 .env.template      |  7 +++++++
 .gitignore         |  2 ++
 Dockerfile         | 14 ++++++++++++++
 README.md          | 15 +++++++++++++++
 docker-compose.yml | 43 +++++++++++++++++++++++++++++++++++++++++++
 web_demo.py        |  2 +-
 8 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 .dockerignore
 create mode 100644 .env
 create mode 100644 .env.template
 create mode 100644 .gitignore
 create mode 100644 Dockerfile
 create mode 100644 docker-compose.yml

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..68c284b
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,2 @@
+.env
+/THUDM
\ No newline at end of file
diff --git a/.env b/.env
new file mode 100644
index 0000000..4d9ca94
--- /dev/null
+++ b/.env
@@ -0,0 +1,7 @@
+# Start 
+
+# 默认启动webui
+START_MODE=web_demo.py
+
+# 启动Api
+# START_MODE=api.py
diff --git a/.env.template b/.env.template
new file mode 100644
index 0000000..89b59c8
--- /dev/null
+++ b/.env.template
@@ -0,0 +1,7 @@
+# Start 
+
+# 默认启动webui
+START_MODE=web_demo.py
+
+# 启动Api
+# START_MODE=api.py
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..68c284b
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+.env
+/THUDM
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..bb735b9
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,14 @@
+FROM nvcr.io/nvidia/pytorch:23.06-py3
+
+WORKDIR /app
+
+# RUN git clone --depth=1 https://github.com/THUDM/ChatGLM2-6B.git /app
+
+COPY . .
+
+RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && \
+    pip install -r requirements.txt
+
+EXPOSE 7860 8000
+
+CMD python ${CLI_ARGS}
\ No newline at end of file
diff --git a/README.md b/README.md
index 27ce17f..3a7bcb7 100644
--- a/README.md
+++ b/README.md
@@ -279,7 +279,22 @@ if __name__ == "__main__":
         if hasattr(chunk.choices[0].delta, "content"):
             print(chunk.choices[0].delta.content, end="", flush=True)
 ```
+### Docker 部署
 
+复制一份配置文件 `.env.template` 到`.env`，默认启动模式为`webui`
+
+如前文所述，在`THUDM\chatglm2-6b`文件夹放置好对应的模型文件
+
+编译镜像，启动程序并加载`int4`模型
+> 显存小于`8G`可以使用此启动方式
+```
+docker-compose --profile int4 up -d
+```
+编译镜像，启动程序并加载完整模型
+> 大约需要`13G`显存
+```
+docker-compose --profile int16 up -d
+```
 
 ## 低成本部署
 
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 0000000..48f9b99
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,43 @@
+version: "3.3"
+services:
+  chatglm2-6b:
+    profiles: ["int16"]
+    build:
+      context: .
+    ports:
+      - 7860:7860
+      - 8000:8000
+    stdin_open: true
+    tty: true
+    volumes:
+      - ./THUDM/chatglm2-6b:/app/THUDM/chatglm2-6b
+    environment:
+      - CLI_ARGS=$START_MODE
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              device_ids: ['0']
+              capabilities: [gpu]
+
+  chatglm2-6b-int4:
+    profiles: ["int4"]
+    build:
+      context: .
+    ports:
+      - 7860:7860
+      - 8000:8000
+    stdin_open: true
+    tty: true
+    volumes:
+      - ./THUDM/chatglm2-6b-int4:/app/THUDM/chatglm2-6b
+    environment:
+      - CLI_ARGS=$START_MODE
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              device_ids: ['0']
+              capabilities: [gpu]
\ No newline at end of file
diff --git a/web_demo.py b/web_demo.py
index 1af24c9..4b5363c 100644
--- a/web_demo.py
+++ b/web_demo.py
@@ -105,4 +105,4 @@ def reset_state():
 
     emptyBtn.click(reset_state, outputs=[chatbot, history, past_key_values], show_progress=True)
 
-demo.queue().launch(share=False, inbrowser=True)
+demo.queue().launch(server_name="0.0.0.0",share=False, inbrowser=True)