Docker + Miniconda + VSCode Remote SSH 在 Ubuntu 搭建多版本共存的 TensorFlow-GPU 远程开发环境

工作原因，需要用到 TensorFlow-GPU，之前一直是机器装好了 Nvidia 驱动、CUDA、cudnn 后紧着一个 TensorFlow 版本一直用，有创建虚拟环境的需求就用 virtualenv 新建一个环境，再装上 TensorFlow-GPU 和一堆依赖。

这个过程虽然操作不难，但重复工作着实无聊，而且 virtualenv 管理虚拟环境并不方便，而更重要的是虚拟环境并不能解决同一 OS 上 TensorFlow 和 CUDA 的多版本共存问题！

最近组内的服务器重装系统，刚好最近也没那么忙了，环境重建也是由我处理，就顺手多学点容器技术，自己琢磨了一套感觉用着比较舒服的解决方案。

背景

为什么需要多版本共存

安装的 TensorFlow 版本不同，所需的 CUDA 版本也有可能不一样，所以不同项目依赖了不同版本的 TensorFlow，就可能需要安装对应版本的 CUDA。

但是因为 CUDA 的安装会修改软链接 /usr/local/cuda/ 指向的路径，也就会改变了当前系统的默认 CUDA。所以在没有人为干涉的情况下，系统同时只能识别到一个版本的 CUDA，也就是最后安装的 CUDA。

因此依赖其他版本 CUDA 的 TensorFlow 就无法正常使用 GPU 了。

时下的解决方案

此前遇到 TensorFlow-GPU 共存的问题通常有这几个解决方案：

按照正常流程安装需要的 TensorFLow1/TensorFlow2 以及所有对应版本的 CUDA，通过脚本在切换虚拟环境时指定要使用的 CUDA。
不论是 TensorFLow1 还是 TensorFlow2，大版本几乎没有不向下兼容的情况，所以尽量安装高版本的 TensorFlow1，然后根据已安装的 TensorFlow1 依赖的 CUDA 版本，在其他虚拟环境安装同样依赖该 CUDA 版本的 TensorFlow2。
使用 TensorFlow 官方提供的 Docker 镜像。

现有方案的缺陷

需要安装多个CUDA、写脚本……具体我没有成功实施过，单纯是安装 CUDA 就因为网络等原因失败多次，具体有什么缺陷，不太好说。
一旦遇到需要更高版本的 TensorFlow2 或者更低版本的 TensorFlow1 的项目，这个方案就无法应对。
这里我认为最简单、最合理的方式是第 3 个。但是仍然存在局限性：
- 每个环境独立的项目都要新建一个容器，也可能需要下载其他版本 TensorFlow 的镜像，而每一个镜像都有 3~4GB 那么大。 (btw，重装系统前，我们的服务器是双系统的，而 Ubuntu 的主分区只有 200GB 不到，别说镜像文件了，平时训练个 BERT 就该清理旧的 checkpoint 了，所以尽量少的镜像和容器，让容器能够复用对我来说相当重要)
- 如果有远程开发的需求，则需要管理多个容器的端口映射，即使是使用 Docker-Compose 我也不觉得有多方便。更何况还要专门记下每个项目使用的容器对应的端口号。

亿点废话

曾经在网络上搜过多版本共存的问题，答案不多，能数得过来的基本都是第一种修改环境变量的方式。当时网络上并没见到有人提到第二种方式，不过第二种也并不是我本人直接想到的。只是一次巧合，我发现两个 conda 环境分别装了 tensorflow-gpu-1.14 和 tensorflow-gpu-2.1.0，而且两个都能正常使用 GPU，于是专门查了下才确认这两个版本都是支持我们服务器上装着的 CUDA-10.0。

虽说是巧合，但也算是临时的解决方案吧，确实能一定程度上解决问题，也那么用了很长时间，直到研究依赖了 TensorFlow-GPU>=2.4.0 的 text-to-text-transfer-transformer。当时由于时间关系，就简单了解了 TensorFlow 官方文档上通过 Docker 使用 TensorFlow 的方式临时搭了个环境。

有一说一，虽然因为公司严格的网络规则导致各种各样麻烦的问题，以至于环境搭建非常困难，但在环境独立的容器根本不用考虑太多环境、库冲突等问题，还是很舒服的。

事前准备

首先需要确认你的显卡是支持 CUDA 的 NVIDIA 显卡，在支持 CUDA® 的 GPU 卡可以查看
安装 NVIDIA® GPU 驱动程序

安装 Docker Engine

参照 Install Docker Engine | Docker Documentation 操作即可

安装 Docker-Compose

Docker-Compose 可以方便我们管理多个容器，非常推荐使用！

下载 Docker-Compose 的二进制文件

1
sudo curl -L "https://github.com/docker/compose/releases/download/1.29.1/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose

添加可执行权限

1
sudo chmod +x /usr/local/bin/docker-compose

安装 nvidia-container-toolkit

要让 TensorFlow 能够在 Docker 容器中使用 GPU，需要在启动容器时指定 --gpus 参数, 而这个参数需要有 nvidia-container-toolkit 支持才行。

设置 Nvidia-Docker 安装源

1
2
3
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
    && wget --no-check-certificate -qO - https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
    && wget --no-check-certificate -qO - https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list

更新软件源并安装 nvidia-container-toolkit

1
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit

重启 Docker

1
sudo systemctl restart docker 

这一步单看 TensorFlow 的文档是看不出这一步的，由于没有经验，再加上公司网络安全的原因，折腾了好久，最终答案可以参考这个 Issue。虽然指令和 issue 里的有点区别，其实安装思路基本是固定的。

烦人的坑

这里也算是为了自己吧，记录一下因为网络踩到的坑：公司由于 Linux 系统无法安装公司要求的证书，https:// 链接均无法正常访问，参考以下几种对策：

使用 wget --no-check-certificate 下载证书
在 windows 下好证书再拷贝到 linux 机器上执行 sudo apt-key
设置代理，本质其实和第二个一样

下载 TensorFlow 的 Docker 镜像

这里以 tensorflow-gpu-1.15.5 和 tensorflow-gpu-latest 为例：

1
2
docker pull tensorflow/tensorflow:1.15.5-gpu
docker pull tensorflow/tensorflow:latest-gpu

骚操作的正式开始

编写 Dockerfile

自己重做镜像主要是为了以后重建容器方便，TensorFlow1 和 TensorFlow2 都做镜像，内容除了作成的镜像名和端口映射外完全一致。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
FROM tensorflow/tensorflow:1.15.5-gpu-py3

# 这个脚本主要是安装软件，详见下文
COPY ./init_env1.sh ./sources.list /
RUN chmod u+x /init_env*.sh 
RUN /init_env1.sh

# 这个脚本主要是添加用户、设置 ssh
COPY ./init_env2.sh /
RUN chmod u+x /init_env*.sh 
RUN /init_env2.sh

EXPOSE 22
# 设置一个挂载目录，将所有数据挂载到宿主机
VOLUME /AI
WORKDIR /AI
# 将默认用户修改为新建的用户
USER ai

Note: Dockerfile 的编写应尽量避免过多的使用 COPY, RUN 等指令，能合并就合并。因为每一行指令都会创建一个 layer，layer 越多，作成的镜像越大 我这里分开这么多写是因为编写该镜像过程中，init_env1.sh 更新源、安装软件太耗时间。且 init_env2.sh 有多次修改，而 init_env1.sh 没怎么修改过，为了节约镜像打包的时间专门分开处理。

Dockerfile 需要用到的脚本

init_env1.sh

用于安装之后可能用到的相关软件，因为 TensorFlow 提供的镜像并没有这些。可根据需要修改、添加其他软件。

1
2
3
4
5
6
7
8
cp /sources.list /etc/apt/sources.list
rm -Rf /var/lib/apt/lists/* 
# 更新源和软件之前，一定要禁止 CUDA 的更新（不然要这 TensorFlow 的镜像有何用）
apt-mark hold cuda*
apt-mark hold machine-learning
apt-get update 
# 安装所需软件
apt-get -y install openssh-server openssh-sftp-server vim sudo net-tools

init_env2.sh

用于设置 SSH 以及登录用户

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
config_file="/etc/ssh/sshd_config"
sftp_server="/usr/lib/openssh/sftp-server"

# modify configs start
# 禁止 root 用户通过 ssh 登录
if grep -q "^PermitRootLogin" $config_file;then
   sed -i '/^PermitRootLogin/s/yes/no/' $config_file
else
   sed -i '$a PermitRootLogin no' $config_file
fi

# 这里设置是为了让 WinSCP 等客户端可以通过 sftp 连接容器
if grep -q "^Subsystem" $config_file;then
   sed -i '/^Subsystem/s/$sftp_server/internal-sftp/' $config_file
else
   sed -i '$a Subsystem  sftp  internal-sftp' $config_file
fi
# modify configs end

# 添加用户:ai，设置用户目录:/AI，修改密码:123456
# 注意了解 `useradd` 和 `adduser` 的区别
useradd -d /AI -U -m ai && echo "ai:123456" | chpasswd
# 为 ai 用户设置 sudo 权限
sed -i '$a ai   ALL=(ALL) NOPASSWD: ALL' /etc/sudoers

Note: 为什么不使用 root 用户？因为即使是在 Docker 容器里，直接使用 root 用户也是存在安全隐患的，如果容器挂载了宿主机的目录，那么在容器内使用 root 用户操作挂载目录实际上和在宿主机使用 root 用户操作该目录是一样。就连 TensorFlow 的容器也会在你进入容器时判断当前用户的身份，如果是 root 则会建议你使用非 root 用户

sources.list

软件源设置，这里用的阿里云，当然也可以替换成别的源。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
deb http://mirrors.aliyun.com/ubuntu/ bionic main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ bionic-security main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ bionic-updates main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ bionic-proposed main restricted universe multiverse
deb http://mirrors.aliyun.com/ubuntu/ bionic-backports main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ bionic main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ bionic-security main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ bionic-updates main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ bionic-proposed main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ bionic-backports main restricted universe multiverse

打包镜像

在 Dockerfile 所在目录执行：

1
docker build . -t tf:1.15.5

新建容器并启动

通过以下 Docker 命令就可以启动支持 GPU 的容器

1
docker run -it --name tf1 --gpus all -p 10022:22 -v /home/ai/tf_project:/AI tf:1.15.5

如果只需要一个虚拟环境/一个版本的 TensorFlow，到这里就已经足够了，但这并不是本文的重点。接下来继续讲骚操作，技巧其实很简单：活用 VOLUME 挂载目录。

安装 Miniconda/Anaconda

conda 安装很简单，从官网下载对应脚本，添加可执行权限后执行，按照交互提示操作就行，这里不再赘述。

使用 Docker-Compose 管理容器

※ 假设装好的 Miniconda 所在根目录为 /home/ai/miniconda

创建任意目录作为项目目录，比如 /home/ai/ai_docker
目录下新建文件 docker-compose.yml，和 Dockerfile 一样，文件名固定

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
# 由于大多数都是 Docker 指令中能看到的，功能相同，下面只会对一些特殊的指令做解释。
# 注意这是本文件使用的 Docker-Compose 版本，决定你能在这个文件使用哪些指令
version: "3.9"

services:
  # 服务名，随意
  tensorflow1:
    # 指定镜像
    image: tf:1.15.5
    container_name: tf1
    ports:
      - "10022:22"
    restart: always
    # 指定用户，在下一节环境变量配置
    user: ${CURRENT_UID}
    # 使容器保持后台运行
    tty: true
    networks:
      - ai-bridge
    # 使容器能使用 GPU
    deploy:
      resources:
        reservations:
          devices:
          - driver: nvidia
            capabilities: [gpu]
    # 目录挂载，骚的根源
    volumes:
      # 可以直接将原来的用户目录先挂载过来，授权 ro
      # 目的：如果有需要和容器环境一起使用但不依赖 TensorFlow 版本的软件，如：个别分词器，
      # 直接装在宿主机即可，容器内只需要配置环境变量
      - /home/ai:/home/ai:ro
      # 将宿主机的 Miniconda 目录挂载到容器的同等目录下，授权 rw
      # 目的：共享宿主机的 conda 环境，不同项目只要依赖的 TensorFlow 版本相同，
      # 就只需要创建 conda 环境克隆已有环境的 TensorFlow，在新环境装项目依赖即可。
      # 
      # 此外，远程开发甚至可以不为任何容器分配端口映射，配合下一条挂载项，直接连接宿主机的项目
      # 和 conda 环境即可。当然，如果没什么意外，前面安装 SSH 等软件的步骤都可以省去。
      - /home/ai/miniconda3:/home/ai/miniconda3:rw
      # 将项目目录挂载到容器的默认工作目录，同时这里也是前面指定的、容器里的用户目录。
      - /home/ai/tensorflow_volume:/AI

  # 与上一节仅服务名、镜像、端口映射不同，其他一致。添加新容器时也一样
  tensorflow2:
    image: tf:latest
    container_name: tf2
    ports:
      - "20022:22"
    restart: always
    user: ${CURRENT_UID}
    tty: true
    networks:
      - ai-bridge
    deploy:
      resources:
        reservations:
          devices:
          - driver: nvidia
            capabilities: [gpu]
    volumes:
      - /home/ai:/home/ai:ro
      - /home/ai/miniconda3:/home/ai/miniconda3:rw
      - /home/ai/tensorflow_volume:/AI

networks:
  ai-bridge:
    driver: bridge

1
2
写入当前用户的 uid 和 gid
CURRENT_UID=1000:1000

Note: 如何查看当前用户的 id？执行 id -u 和 id -g 就好了

启动容器

1
2
cd /home/ai/ai_docker
docker-compose up -d

使用 docker-compose ps 查看容器状态，是 Up 就启动成功了。

骚操作的后续补充

为容器初始化 conda

虽然完全可以进入容器后再激活对应的 conda 环境，但用得久了着实感觉繁琐。可以通过用户登录会自动调用 .bashrc 或 .profile 等脚本的机制来实现。

关闭 conda 的自动初始化环境因为上面配置了容器的 conda 环境实际上是和宿主机共用的，为了避免登录脚本激活容器后 conda 又切换回默认环境，首先关闭该设置

1
2
3
cat <<EOF >>~/.condarc
auto_activate_base: false
EOF

修改宿主机 .bashrc，添加环境切换函数安装 conda (并启用自动初始化)后，在宿主机的 .bashrc 找到类似下面的代码

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
# >>> conda initialize >>>
# !! Contents within this block are managed by 'conda init' !!
__conda_setup="$('/home/ai/miniconda3/bin/conda' 'shell.bash' 'hook' 2> /dev/null)"
if [ $? -eq 0 ]; then
    eval "$__conda_setup"
else
    if [ -f "/home/ai/miniconda3/etc/profile.d/conda.sh" ]; then
        . "/home/ai/miniconda3/etc/profile.d/conda.sh"
    else
        export PATH="/home/ai/miniconda3/bin:$PATH"
    fi
fi
unset __conda_setup
# <<< conda initialize <<<

在上述代码之后添加

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
tf1(){
  if [ -z "$1" ]; then
    env=tf1
  else
    env=$1
  fi
  docker exec -it -e env=$1 -e TF_VERSION=TensorFlow1 tf1 bash
}
tf2(){
  if [ -z "$1" ]; then
    env=tf2
  else
    env=$1
  fi
  docker exec -it -e env=$1 -e TF_VERSION=TensorFlow2 tf2 bash
}

配置宿主机默认 conda 环境在宿主机 .bashrc 最后添加

1
conda activate base

在容器 .bashrc 添加

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
# >>> auto activate env >>>
if [ -n "$env" ]; then
  conda activate $env  # 进入容器自动激活指定的 conda 环境 
else
  # 如果没有指定 conda 环境，则自动进入容器自身 tf 版本对应的 conda 环境
  tf1=`pip freeze | grep tensorflow-gpu==1.15.5`  # 查询当前容器默认环境的 tf1 版本，1.15.5 修改成自己安装的版本
  tf2=`pip freeze | grep tensorflow-gpu==2.1.0`  # 查询当前容器默认环境的 tf1 版本，2.1.0 修改成自己安装的版本
  if [ -n "$tf1" ]; then
    conda activate tf1
  elif [ -n "$tf2" ]; then
    conda activate tf2
  fi
fi
# <<< auto activate env <<<

至此，通过 tf1 和 tf2 就能直接进入容器并激活相对的 conda 环境了。举个例子：

进入 TensorFlow1 容器：tf1
进入 TensorFlow2 容器：tf2
进入 TensorFlow1 容器，并激活名为 BERT 的虚拟环境：tf1 BERT
进入 TensorFlow2 容器，并激活名为 T5 的虚拟环境：tf1 T5

简单的使用说明

TODO 2021/09/14 目前还是没觉得哪些操作需要记下来

使用 VSCode 远程开发

安装插件 Remote - SSH 连接宿主机即可，即使不查资料也不难操作。由于写本文的时间已经太晚了而且这一节不算重点，暂且告一段落了。。。

问题小记

使用 bash 作为容器的默认 shell 工具时，sftp 登录时可能会遇到 “收到了过大的数据包” 之类的报错。可以参考以下网址，在 .bashrc 文件中打印语句前通过 [[ $- == *i* ]] 判断当前终端是否是交互终端。

https://my.oschina.net/leejun2005/blog/342865

一些已知存在但仍不知如何解决的小问题

TODO 2021/09/14 本来是有一两个问题来着，其中最影响效率的问题在问题小记记下来了。别的暂时没什么了

背景