- 当我们落魄到需要去借别人的算力跑实验的时候,如何寄人篱下却依然衣冠楚楚。
使用 Docker 打包深度学习环境
当我们落魄到需要去借别人的算力跑实验的时候,我们一定要尽可能高效率的完成自己的实验来少给别人添麻烦,那么配置环境就是一个相当耗时的工作,而且如果你有需要 root 权限,事情也会变得非常尴尬。这个时候,
Docker
就可以帮我们做好环境的迁移工作。毕竟,Docker 已经愈发变成一种基础设施了,尤其是对于那些卡特别多的团队,不像我们这种小实验室,人均 root 权限。而 nvidia-docker 为 Docker 提供了很好的显卡支持,非常适合我们来打包一个炼丹炉镜像。因此,我们需要在本地打包一个开箱即用的炼丹炉镜像,然后就可以拷贝到对方服务器上直接进入工作状态。
那么,为了让师弟师妹能够更加方便的“提镜像跑路”,我先把实验室的基础设施建好,从 Docker 的安装干起。今天你不努力,明天就有人替你努力。
—— 佚名
NVIDIA-Docker 的安装部署
NVIDIA-Docker 首先依赖 Docker 本身,故而第一步需要安装 Docker,其步骤如下:
安装必要依赖;
1
2sudo apt update
sudo apt install ca-certificates curl gnupg lsb-release添加 Docker 官方 GPG 公钥;
1
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
添加 Docker 镜像源;
1
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
安装 Docker-CE;
1
2sudo apt update
sudo apt install docker-ce docker-ce-cli containerd.io(Optional) 验证 Docker 安装情况。
1
sudo docker run hello-world
在完成了 Docker 的安装之后,即可安装 nvidia-docker,其步骤如下:
配置 nvidia-docker 的存储库和 GPG 公钥;
1
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
(Optional) 添加 Experimental 功能支持;
1
curl -s -L https://nvidia.github.io/nvidia-container-runtime/experimental/$distribution/nvidia-container-runtime.list | sudo tee /etc/apt/sources.list.d/nvidia-container-runtime.list
安装 nvidia-docker;
1
2sudo apt update
sudo apt install -y nvidia-docker2重新启动 Docker 守护进程以完成安装。
1
sudo systemctl restart docker
Docker 本地工作环境的补充工作
既然都安装了 Docker 环境,那么我们这个小团队就也可以使用 Docker 来工作了,
毕竟,给一个未来的萌新师弟师妹 root 权限可能带来的影响实在是令人难以接受,君不见劳总连续 rm -rf / 了两次。一种安全的方式就是给 Docker 用户组权限,这样可以让每一个用户以自己想要的方式在容器里折腾。Docker 命令通过
UNIX Socket
建立通讯,默认情况下其只支持root
用户或docker
用户组中的用户访问。因此我们需要先创建 docker 用户组,可以使用sudo groupadd docker
来完成,随后,可以通过sudo usermod -aG docker [username]
的命令将某一用户加入 docker 用户组。如果上述修改未生效,可以重启 Docker 服务,重新登陆用户,使配置修改生效。在实际工作中可以使用 Docker 命令启动容器,并指定端口映射以及文件挂载等操作,但我个人更喜欢通过
docker-compose
来启动,毕竟,不会有人每次都重头敲一遍吧,反复按上键找命令也很不方便吧。Docker-Compose 的安装较为简单,有两个步骤:
下载 docker-compose 的 release 文件(以最新版为例);
1
sudo curl -L "https://github.com/docker/compose/releases/download/1.29.2/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
授予其可执行权限。
1
sudo chmod +x /usr/local/bin/docker-compose
为了让
docker-compose
命令自动调用nvidia-docker
来执行命令,修改/etc/docker/daemon.json
文件如下,此时,Docker 命令的执行就会由nvidia-docker
来完成。1
2
3
4
5
6
7
8
9{
"default-runtime": "nvidia",
"runtimes": {
"nvidia": {
"path": "nvidia-container-runtime",
"runtimeArgs": []
}
}
}
创建一个基于 Anaconda 的 Docker 镜像
为了更便捷的做一个开箱即用的炼丹炉镜像,我们可以充分利用一些现成的 Docker 镜像。比如,Anaconda 有自己官方的 Docker 镜像,我们可以以其作为基镜像快速构建一个面向我们自己具体 Task 的镜像文件。
Anaconda 在 continuumio’s Profile | Docker Hub 上提供了其官方 Docker 镜像,可以使用其 miniconda3 镜像作为基镜像创建一个自己的镜像文件,通过
conda env export > env.yml
可以将本地的工作环境的全部依赖加以记录,将该配置文件COPY
到Dockerfile
中,从而在创建包含我们所需依赖的一个炼丹炉镜像。一个典型的
Dockerfile
如下所示,包含了基础的 miniconda 环境、GPU支持的 PyTorch 框架以及 env.yml 中所要求的其他依赖。其中,libgl1-mesa-glx
是解决 miniconda3 基镜像中,找不到libGL.so.1
文件的问题。1
2
3
4
5
6
7
8
9
10
11
12
13
14FROM continuumio/miniconda3
LABEL author="Li Yingping"
RUN apt update -q && \
apt install -q -y libgl1-mesa-glx
COPY env.yml /build/env.yml
RUN conda install pip
RUN pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio==0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
RUN conda env update -f /build/env.yml
CMD [ "/bin/bash" ]完成 Dockerfile 的编写后,即可使用
docker build -t [image_name]:[tag_name] .
命令从当前文件夹构建一个镜像。构建的过程需要大量下载文件,耗时且需要消耗大量流量,
至少,PyTorch 下载就需要 2G 多,为了不给对方带来不必要的流量负担,而且,万一对方为了避免挖矿切断了公网连接呢,可以将本地构建好的镜像导出成文件,拷贝到对方的服务器中加载。- Docker 导出镜像的命令:
docker save [image_name]:[tag_name] -o [filename].tar
; - Docker 载入镜像的命令:
docker load -i [image_name]:[tag_name]
。
- Docker 导出镜像的命令:
创建一个启动工程的 Docker-Compose 配置
一个典型的
docker-compose.yml
如下所示,其中image
字段指定使用哪个镜像运行程序;volumes
字段指定需要将本地的哪些文件挂载到容器中,一般我们要挂载的包括代码文件夹、数据集和日志文件夹等;ports
字段定义端口映射关系,一般我们要考虑映射Jupyter Notebook
或Tensor Board
等应用的端口;environment
字段定义虚拟机中的环境变量,NVIDIA_VISIBLE_DEVICES
定义物理机中哪些 GPU 设备对虚拟机可见,CUDA_VISIBLE_DEVICES
定义虚拟机中哪些 GPU 设备对 PyTorch 可见;shm_size
字段定义共享给虚拟机的内存空间;working_dir
定义默认的工作目录,一般是容器中的代码目录;command
定义启动命令,一般是执行某一shell
脚本或运行某一.py
文件等。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18version: "3.8"
services:
train:
stdin_open: true
tty: true
build: .
image: [image_name]:[tag_name]
volumes:
- /path/on/physical/machine:/path/on/virtual/machine
ports:
- [local_port]:[container_port]
environment:
NVIDIA_VISIBLE_DEVICES: all
CUDA_VISIBLE_DEVICES: [0, 1, 2, ...]
shm_size: 32G
working_dir: /path/on/virtual/machine
command: "/bin/bash run.sh"当完成了 Docker-Compose 配置文件的编写之后,我们就可以使用
docker-compose up
的方式启动服务,需要关闭服务时使用docker-compose down
,当文件夹下存在多个配置文件时-f
选项可以用于指定将使用的配置文件,-d
选项可以将所启动服务指定在后台静默运行。当我们需要进入容器内进行 debug 时,可以使用
docker ps
命令获取当前运行中的容器的CONTAINER ID
然后使用docker exec -it [CONTAINER ID] /bin/bash
进入容器内 shell 开展调试。By the way,我们也可以把容器的启动命令简单的设成 bash 本身,这样我们可以启动一个拥有完整环境的炼丹炉,不使用 volume 映射的方式挂载文件,这样挂载的文件夹与物理机是同步的,而使用
docker cp
命令独立的将文件拷贝进容器,这样可以避免本地修改了代码而影响一个正在运行的训练过程。
使用 rsync 传输数据集
当我们落魄到需要去借别人的算力跑实验的时候,Docker 可以帮我们实现运行环境的友好迁移,但数据集的迁移也是一个很大的问题。
- 对于小规模的数据集可以在
Dockerfile
中用COPY
命令将数据集复制到镜像内,制作一个更加开箱即用的炼丹炉镜像。 - 但是对于大规模数据集,这样会导致导出的镜像文件体积过大,显得笨重,因此,我们考虑使用网络传输大规模数据集,
毕竟,局域网又不消耗流量,也不会让资金窘迫的我们我们变得更加落魄。
- 对于小规模的数据集可以在
rsync
是一款基于scp
的文件同步工具,支持增量备份,可以配合其他工具实现多种触发类型的计划任务。这里,我们把它退化成一个普通的文件传输工具,用于将本地的数据集推送到远端。rsync
的安装可以通过包管理工具朴素地安装,即sudo apt install rsync
。rsync
的基本语法如下,其支持三种工作模式,本地文件同步、通过 Remote Shell 的远程文件同步和通过 Rsync Deamon 的远程文件同步,三种命令均符合rsync [OPTION...] SRC... [DEST]
的基本结构。1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16Local:
rsync [OPTION...] SRC... [DEST]
Access via remote shell:
# Pull:
rsync [OPTION...] [USER@]HOST:SRC... [DEST]
# Push:
rsync [OPTION...] SRC... [USER@]HOST:DEST
Access via rsync daemon:
# Pull:
rsync [OPTION...] [USER@]HOST::SRC... [DEST]
rsync [OPTION...] rsync://[USER@]HOST[:PORT]/SRC... [DEST]
# Push:
rsync [OPTION...] SRC... [USER@]HOST::DEST
rsync [OPTION...] SRC... rsync://[USER@]HOST[:PORT]/DEST就我们的任务而言,我们需要的是其中的第二种方式,可以使用下列命令完成数据集的传输,过程中需要输入相应服务器的密码以建立连接。
1
2
3
4如果在本地将数据推送到远程服务器
rsync -r /path/to/local/dataset [remote_username]@[remote_ip]:/path/to/remote/dataset
如果使用SSH连接了远程服务器希望拉取本地的数据
rsync -r [local_username]@[local_ip]:/path/to/local/dataset /path/to/remote/dataset需要注意的是,源路径末尾带斜线,传输的是目录中所有的文件而不包括目录本身,而不带斜线的话,会在目的路径创建一个同名目录以及所有文件。具体的,上面的示例会在远程服务器中呈现
/path/to/remote/dataset/dataset/***(files)
。rsync
的其他用法可以参考下列链接,或者简单地通过rsync -h
或tldr rsync
了解。rsync 官方网站:rsync (samba.org)
阮一峰 rsync 教程:rsync 用法教程 - 阮一峰的网络日志 (ruanyifeng.com)
使用脚本批量运行程序
当我们落魄到需要去借别人的算力跑实验的时候,我们一定要尽可能高效率的完成自己的实验来少给别人添麻烦。众所周知,
炼丹就是黑心资本家对于显卡这一劳动力的无情剥削,因此,充分利用显卡的空余时间能帮助我们尽早完成炼丹任务。与其熬夜等结果,不如写个脚本自动值守。一个朴素的想法是,构建训练任务队列,轮询显卡的使用情况,监测到有空闲显卡,则从队列中选择一个工程启动。
人与猴子最大的区别在于创造并使用工具。
—— 佚名(忘了谁说的,只能请出佚名兄了,要不,鲁迅?)
使用 shutil 批量创建实验工程
- 在消融实验中,难免存在大量基本代码类似,仅仅参数或设置不同的情况,为了高效的进行实验,可以率先将同样的工程模板复制若干份,每份对应一个实验,尔后,集中修改每一个实验中不同的部分,这个复制的过程可以使用标准库
shutil
中的copytree()
函数来完成,对于不需要复制的部分,可以在参数中传入 glob 形式的匹配模板来指定。1
shutil.copytree([base_exper_dir], [target_exper_dir], ignore=shutil.ignore_patterns('apex', 'vis'))
使用 pynvml 监控显卡使用情况
pynvml 是英伟达显卡管理库 NVML 的一个 Python 绑定,可以通过该模块获取到类似
nvidia-smi
命令的查询结果,使用pip install nvidia-ml-py
命令可以完成模块的安装。在我们的任务需求中,我们需要监控显卡使用情况,因此,需要获取主机上的显卡数目,并查询每一张显卡上的显存占用,从而判断是否可以支持工程的运行。其他功能接口可以参考下述仓库:pynvml 官方 Repo:https://github.com/gpuopenanalytics/pynvml
NVML API 文档:https://docs.nvidia.com/deploy/nvml-api/
pynvml 使用前后需要先后调用
nvmlInit()
和nvmlShutdown()
两个函数来管理上下文,nvmlDeviceGetCount()
函数返回设备上总的 GPU 的数目,随后可以通过循环的方式查询每一张显卡的使用情况。给定 gpu_id 可以通过nvmlDeviceGetHandleByIndex()
返回一个对应该设备的句柄,nvmlDeviceGetMemoryInfo()
接受一个句柄作为输入,返回其显存使用情况,其返回值的used
属性即是已使用的显存,为方便计算,将结果右移 20 位,将单位转换为 MB。此时,即可通过一个预设的阈值判断是否有程序在占用显存,从而判定其是否在使用。核心代码如下:1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16def query_devices_mems():
pynvml.nvmlInit()
devices = pynvml.nvmlDeviceGetCount()
used_mems = []
for id_ in range(devices):
device = pynvml.nvmlDeviceGetHandleByIndex(id_)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(device)
mem_used = mem_info.used >> 20
used_mems.append(mem_used)
pynvml.nvmlShutdown()
return used_mems
def query_free_gpus(threshold=2000):
used_mems = query_devices_mems()
free_gpus = [idx for idx, mem in enumerate(used_mems) if mem < threshold]
return free_gpus
使用 subprocess 启动实验工程
- 当监测到空闲 GPU 之后,即可取出队列中的下一个工程并运行,一种简单的方式是通过标准库中的 subprocess 模块来启动。其优势在于,
subprocess.run()
接口提供了丰富的自定义参数,env
参数允许可以传入子进程运行的环境变量,我们可以通过复制当前环境变量并修改CUDA_VISIBLE_DEVICES
的值来为工程指定到空闲的 GPU 上。同时,cwd
字段允许指定工作目录,可以在此处传入对应工程路径以保证程序中的相对路径可以正确解析到需要的位置。1
2
3
4
5
6
7
8
9
10
11
12
13def run_script(script_path, gpu_id):
command = f"/bin/bash {script_path}"
subenv = os.environ.copy()
subenv["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
ret = subprocess.run(
command, shell=True,
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
encoding="utf-8", env=subenv,
cwd=op.dirname(script_path))
if ret.returncode == 0:
print(f"[SUBPROCESS] Experiment {script_path.split('/')[-1]} run successfully ...")
else:
print(f"[SUBPROCESS] Experiment {script_path.split('/')[-1]} failed with {ret}")
一个简单的调度例程
- 通过第三方库
fire
可以为程序创建命令行接口,下面的例程可以通过python xxxx.py create_experiments
来创建工程,修改好每一个工程之后可以通过python xxxx.py dispatch
来启动调度。种子撒出去了之后等几天再来收菜就好了。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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90import fire
import itertools
import os
import os.path as op
import pynvml
import shutil
import subprocess
import threading
import time
import tqdm
exper_names = [
"Exper1", "Exper2"
]
path = "[PATH]"
base_exper_dir = "[SOURCE_PATH]"
def create_experiments():
if not op.exists(path):
os.mkdir(path)
for name in tqdm.tqdm(exper_names):
abs_path = op.join(path, name)
shutil.copytree(base_exper_dir, abs_path,
ignore=shutil.ignore_patterns('apex', 'vis'))
def create_experiments():
if not op.exists(path):
os.mkdir(path)
for name in tqdm.tqdm(exper_names):
abs_path = op.join(path, name)
shutil.copytree(base_exper_dir, abs_path,
ignore=shutil.ignore_patterns('apex', 'vis'))
def query_devices_mems():
pynvml.nvmlInit()
devices = pynvml.nvmlDeviceGetCount()
used_mems = []
for id_ in range(devices):
device = pynvml.nvmlDeviceGetHandleByIndex(id_)
mem_info = pynvml.nvmlDeviceGetMemoryInfo(device)
mem_used = mem_info.used >> 20
used_mems.append(mem_used)
pynvml.nvmlShutdown()
return used_mems
def query_free_gpus(threshold = 2000):
used_mems = query_devices_mems()
free_gpus = [idx for idx, mem in enumerate(used_mems) if mem < threshold]
return free_gpus
def run_script(script_path, gpu_id):
command = f"/bin/bash {script_path}"
subenv = os.environ.copy()
subenv["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
ret = subprocess.run(
command, shell=True,
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
encoding="utf-8", env=subenv,
cwd=op.dirname(script_path))
if ret.returncode == 0:
print(f"[SUBPROCESS] Experiment {script_path.split('/')[-1]} run successfully ...")
else:
print(f"[SUBPROCESS] Experiment {script_path.split('/')[-1]} failed with {ret}")
def wait_until_run(project, interval=2):
free_gpus = query_free_gpus()
iter_ = itertools.cycle(["\\", "|", "/", "-"])
while len(free_gpus) == 0:
print(f"\r[DISPATCHER] Waiting for free GPU {next(iter_)}", end="")
time.sleep(interval)
free_gpus = query_free_gpus()
else:
print(f"\r[DISPATCHER] Waiting for free GPU {next(iter_)}")
free_gpu = free_gpus[0]
print(f"[DISPATCHER] GPU-{free_gpu} is free ...")
script_path = op.join(project, "run.sh")
print(f"[DISPATCHER] Start running {project} ...")
t = threading.Thread(target=run_script, kwargs={"script_path": script_path, "gpu_id": free_gpu})
t.start()
def dispatch():
for exper_name in exper_names:
project_path = op.join(path, exper_name)
print(f"[DISPATCHER] {exper_name} is waiting ...")
wait_until_run(project_path)
time.sleep(30)
if __name__ == "__main__":
fire.Fire()