From a8a5c037cacd091955a5b39bca05b4f3ce75c105 Mon Sep 17 00:00:00 2001
From: kuangjux <18630816527@163.com>
Date: Sat, 26 Aug 2023 21:02:57 +0800
Subject: [PATCH] feat(env): add docker support.

- Added the necessary `Dockerfile` for running the project on CPU and CUDA environment.
- Added commands to the `Makefile` for convenient Docker startup.
- Added documentation in `docs/INSTALL_GUIDE_CN.md` explaining how to launch the Docker environment.

Co-authored-by: Xiaonan Song <songxiaonan@qiyuanlab.com>
Co-authored-by: Haojie Wang <haojie0429@gmail.com>
---
 Makefile                                      | 28 +++++++++++++++++++
 docs/INSTALL_GUIDE_CN.md                      | 23 +++++++++++++++
 .../infinitensor_ubuntu_22.04.dockerfile      | 19 +++++++++++++
 .../infinitensor_ubuntu_22.04_CUDA.dockerfile | 18 ++++++++++++
 4 files changed, 88 insertions(+)
 create mode 100644 scripts/dockerfile/infinitensor_ubuntu_22.04.dockerfile
 create mode 100644 scripts/dockerfile/infinitensor_ubuntu_22.04_CUDA.dockerfile

diff --git a/Makefile b/Makefile
index 0fcc9070..d6dfadd2 100644
--- a/Makefile
+++ b/Makefile
@@ -7,6 +7,20 @@ INTELCPU ?= off
 BACKTRACE ?= ON
 TEST ?= ON
 FORMAT_ORIGIN ?=
+# Docker build options
+DOCKER_NAME ?= infinitensor
+DOCKER_IMAGE_NAME ?= infinitensor
+DOCKER_FILE ?= infinitensor_ubuntu_22.04.dockerfile
+DOCKER_RUN_OPTION ?= 
+
+# CUDA option.
+ifeq ($(CUDA), ON)
+	DOCKER_IMAGE_NAME = infinitensor_cuda
+	DOCKER_NAME = infinitensor_cuda
+	DOCKER_FILE = infinitensor_ubuntu_22.04_CUDA.dockerfile
+	DOCKER_RUN_OPTION += --gpus all -it --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 -v `pwd`:`pwd` -w `pwd`
+endif
+
 
 CMAKE_OPT = -DCMAKE_BUILD_TYPE=$(TYPE)
 CMAKE_OPT += -DUSE_CUDA=$(CUDA)
@@ -39,3 +53,17 @@ test-cpp:
 test-onnx:
 	@echo
 	python3 pyinfinitensor/tests/test_onnx.py
+
+docker-build: 
+	docker build -f scripts/dockerfile/$(DOCKER_FILE) -t $(DOCKER_NAME) .
+
+docker-run:
+	docker run -t --name $(DOCKER_IMAGE_NAME) -d $(DOCKER_NAME) $(DOCKER_RUN_OPTION)
+
+docker-start:
+	docker start $(DOCKER_IMAGE_NAME)
+
+docker-exec:
+	docker exec -it $(DOCKER_IMAGE_NAME) bash
+
+
diff --git a/docs/INSTALL_GUIDE_CN.md b/docs/INSTALL_GUIDE_CN.md
index 053ae14e..285187ef 100644
--- a/docs/INSTALL_GUIDE_CN.md
+++ b/docs/INSTALL_GUIDE_CN.md
@@ -137,6 +137,29 @@
 
    安装成功后，您就可以使用本项目的 Python 接口进行编码并运行。具体使用方式可以参考项目样例代码 example/Resnet/resnet.py 以及用户使用手册
 
+## Docker
+
+本项目也提供了 Docker 的环境，您可以使用 `make docker-build` 或 `make docker-build CUDA=ON` 命令启动并编译 Dockerfile，您可以通过添加编译选项或者修改 Makefile 变量修改 docker image 名称或者所选的 Dockerfile 文件。
+ 
+由于在拉取 github repo 时需要将 ssh key 加入到 github profile 中，因此暂时注释掉拉取 repo 并编译项目的过程，由用户在进入 docker 后自己维护 ssh key（将 host 中的 ssh key 复制到 docker 中可能会遇到环境不一致的问题）。
+
+```shell
+# Build docker container.
+make docker-build
+# Run docker image.
+make docker-run
+# Execute docker image.
+make docker-exec
+```
+
+如果需要编译 CUDA 版，请使用如下命令：
+```shell
+# Build docker container.
+make docker-build CUDA=ON
+# Run docker image.
+make docker-run CUDA=ON
+```
+
 ## 技术支持
 
 如遇到问题，请联系我们技术支持团队
diff --git a/scripts/dockerfile/infinitensor_ubuntu_22.04.dockerfile b/scripts/dockerfile/infinitensor_ubuntu_22.04.dockerfile
new file mode 100644
index 00000000..ee8bb171
--- /dev/null
+++ b/scripts/dockerfile/infinitensor_ubuntu_22.04.dockerfile
@@ -0,0 +1,19 @@
+FROM ubuntu:22.04
+ENV DEBIAN_FRONTEND=noninteractive 
+
+# Install dependencies.
+RUN apt update && apt-get install -y git make cmake build-essential python-is-python3 python-dev-is-python3 python3-pip libdw-dev openssh-client
+# Generate ssh key.
+RUN ssh-keygen -t rsa -f /root/.ssh/id_rsa -q -P ""
+
+# Update pip and switch to Tsinghua source.
+RUN python -m pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade pip && pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+
+# TODO: Since SSH clone repo requires adding the SSH key to the GitHub profile, 
+# the process of pulling the project and compiling it has been temporarily commented.
+
+# Download InfiniTensor.
+# RUN git clone git@github.com:InfiniTensor/InfiniTensor.git /root/InfiniTensor --branch master --single-branch --recursive
+
+# Build and Install InfiniTensor
+# RUN cd /root/InfiniTensor && make install-python
diff --git a/scripts/dockerfile/infinitensor_ubuntu_22.04_CUDA.dockerfile b/scripts/dockerfile/infinitensor_ubuntu_22.04_CUDA.dockerfile
new file mode 100644
index 00000000..ee14fe5d
--- /dev/null
+++ b/scripts/dockerfile/infinitensor_ubuntu_22.04_CUDA.dockerfile
@@ -0,0 +1,18 @@
+FROM nvcr.io/nvidia/pytorch:23.07-py3
+
+# Install dependencies.
+RUN apt-get -o Acquire::Check-Valid-Until=false -o Acquire::Check-Date=false update && apt-get install -y git make cmake build-essential python-is-python3 python-dev-is-python3 python3-pip libdw-dev openssh-client
+# Generate ssh key.
+RUN ssh-keygen -t rsa -f /root/.ssh/id_rsa -q -P ""
+
+# Update pip and switch to Tsinghua source.
+RUN python -m pip install -i https://pypi.tuna.tsinghua.edu.cn/simple --upgrade pip && pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
+
+# TODO: Since SSH clone repo requires adding the SSH key to the GitHub profile, 
+# the process of pulling the project and compiling it has been temporarily commented.
+
+# Download InfiniTensor.
+# RUN git clone git@github.com:InfiniTensor/InfiniTensor.git /root?/InfiniTensor --branch master --single-branch --recursive
+
+# Build and Install InfiniTensor
+# RUN cd /root/InfiniTensor && make install-python CUDA=ON