構建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給分佈式機器學習帶來的極致加速體驗。