构建Tensorflow RDMA的Docker镜像

RDMA是一个远程通讯技术，它通过Kernel bypass等方式降低数据传输中的延迟和CPU消耗。
在分布式训练中，由于多个Worker之间或者Worker和Paramater Server 之间需要大量传输模型变量。当GPU到达一定数量后，受制于网络带宽以及TCP协议的延迟，通讯往往会成为计算性能的瓶颈，而在分布式训练中使用RDMA技术能够非常明显地提高训练速度。

Tensorflow

Tensorflow是谷歌开源的深度学习框架，它有丰富的平台支持和API，也可以非常轻松地构建分布式模型训练。
Tensorflow 在实现里支持RDMA作为其分布式场景的通讯协议，但是官方镜像默认没有支持RDMA。需要重新构建tensorflow，并开启RDMA相关的构建参数。 Tensorflow 对 RDMA的支持和实现协议参见文档： https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/verbs/README.md

如何build from source 可以见官网
https://www.tensorflow.org/install/source

开始构建

找一台GPU节点，运行develop容器，作为编译环境

docker run -it -w /tensorflow -v $PWD:/mnt -e HOST_PERMS="$(id -u):$(id -g)" tensorflow/tensorflow:nightly-devel-gpu-py3 bash

在dev容器中执行build

# 开启RDMA参数
export TF_NEED_GDR=1
export TF_NEED_VERBS=1
tensorflow/tools/ci_build/builds/configured GPU 
bazel build --config=opt --config=cuda //tensorflow/tools/pip_package:build_pip_package
./bazel-bin/tensorflow/tools/pip_package/build_pip_package /mnt  # create package
chown $HOST_PERMS /mnt/tensorflow-*.whl

可以尝试安装， 检查下编译结果：

pip uninstall tensorflow  # remove current version
pip install /mnt/tensorflow-*.whl
cd /tmp  # don't import from source directory
python -c "import tensorflow as tf; print(tf.contrib.eager.num_gpus())"

下载正式构建的dockerfile和所需文件：

curl -OL https://raw.githubusercontent.com/tensorflow/tensorflow/r1.12/tensorflow/tools/dockerfiles/dockerfiles/nvidia.Dockerfile
curl -OL https://raw.githubusercontent.com/tensorflow/tensorflow/r1.12/tensorflow/tools/dockerfiles/bashrc

对nvidia.Dockerfile 做一下修改，将之前编译的Tensorflow whl文件添加到镜像并执行安装

ARG TF_PACKAGE=tensorflow-gpu
RUN ${PIP} install ${TF_PACKAGE}

改为

ADD . /tensorflow-rdma-gpu 
ARG TF_PACKAGE=/tensorflow-rdma-gpu/*.whl
RUN ${PIP} install ${TF_PACKAGE}
RUN rm -rf /tensorflow-rdma-gpu 

构建正式镜像

docker build . -f nvidia.Dockerfile -t tensorflow-rdma

如果希望在镜像里加入一些mellanox的工具，可以在dockerfile最后再扩展一下。

RUN apt-get update && apt-get install -y --no-install-recommends \
        libibverbs-dev \
        librdmacm-dev \
        wget \
        && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*
RUN wget http://content.mellanox.com/ofed/MLNX_OFED-4.1-1.0.2.0/MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64.tgz && \
        tar -xzvf MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64.tgz && \
        MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64/mlnxofedinstall --user-space-only --without-fw-update --all -q && \
        cd .. && \
        rm -rf MLNX_OFED_LINUX-4.1-1.0.2.0-ubuntu16.04-x86_64 && \
        rm -rf *.tgz

这样你就可以体验RDMA给分布式机器学习带来的极致加速体验。