rdma_pytorch.dockerfile

# SwitchML Project
# @file rdma_pytorch.dockerfile
# @brief Generates an image that can be used to run pytorch with switchml rdma integrated.
#
# IMPORTANT: For RDMA to work, you must pass `--cap-add=IPC_LOCK --device=/dev/infiniband/<NIC name>` arguments when starting
# You also still need to manually disable ICRC checking on the host machine. Take a look at our disable_icrc.sh script in the scripts folder.
#
# A typical command to start the container is:
# docker run -it --gpus all --net=host --cap-add=IPC_LOCK --device=/dev/infiniband/uverbs0 --name switchml-rdma-pytorch <name_of_the_image_created_from_this_file>
#

FROM nvidia/cuda:11.2.2-devel-ubuntu18.04

# Set default shell to /bin/bash
SHELL ["/bin/bash", "-cu"]

ARG TIMEOUTS=1
ARG VCL=1
ARG DEBUG=0

# General packages needed
RUN apt-get update && \
    apt install -y \
    git \
    wget \
    vim \
    nano \
    build-essential \
    net-tools \
    sudo \
    gpg \ 
    lsb-release \
    software-properties-common

# Add kitware's APT repository for cmake
RUN wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | \
    gpg --dearmor - | \
    tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null && \
    apt-add-repository "deb https://apt.kitware.com/ubuntu/ $(lsb_release -cs) main" && \
    apt update

# General and RDMA client library requirements
RUN apt install -y \
    gcc \
    make \
    libboost-program-options-dev \
    libgoogle-glog-dev \
    autoconf \
    libtool \
    pkg-config \
    libibverbs-dev \
    libhugetlbfs-dev \
    cmake \
    libssl-dev

# Clone the switchml repo and compile the client library with the benchmarks and examples.
ARG SWITCHML_UPDATED
RUN git clone https://github.com/p4lang/p4app-switchML.git /home/switchml && \
    cd /home/switchml/dev_root && \
    git submodule update --init --recursive -- third_party/vcl && \
    git submodule update --init --recursive -- third_party/grpc && \
    make RDMA=1 TIMEOUTS=${TIMEOUTS} VCL=${VCL} DEBUG=${DEBUG}

# Register the compiled GRPC
# You can skip this step however you would need to use the LD_LIBRARY_PATH variable each time you run 
# any application with switchml
RUN echo /home/switchml/dev_root/third_party/grpc/build/lib > /etc/ld.so.conf.d/000_grpc.conf && \
    ldconfig

# At this point the microbenchmark can be run with the rdma backend.

# Install miniconda
RUN cd /usr/local && \
    wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ./miniconda.sh && \
    bash miniconda.sh -b -p /usr/local/conda && \
    rm miniconda.sh && \
    /usr/local/conda/bin/conda init bash && \
    /usr/local/conda/bin/conda install -y numpy jupyter matplotlib
ENV PATH $PATH:/usr/local/conda/bin

# Install pytorch requirments
RUN conda env update --name base --file /home/switchml/dev_root/frameworks_integration/pytorch_patch/environment.yml

# Download PyTorch and checkout to 1.7.1 
RUN git clone https://github.com/pytorch/pytorch.git /home/pytorch && \
    cd /home/pytorch && \
    git checkout 57bffc3 && \
    git submodule sync && \
    git submodule update --init --recursive

# Make and apply the PyTorch patch
RUN cd /home/switchml/dev_root/ && \
    make pytorch_patch RDMA=1 && \
    git apply /home/switchml/dev_root/build/switchml_pytorch.patch

# Build PyTorch
RUN cd /home/pytorch && \
    CUDA_HOME=/usr/local/conda/ BUILD_TEST=0 /usr/local/conda/bin/python setup.py install --prefix=/usr/local/conda/ 2>&1 | tee build.log

# At this point you have the base conda environment with switchml integrated into pytorch.
# You can run training scripts and benchmarks and switchml will kick in for all reduce operations
# that it can handle. Just make sure to use the gloo pytorch backend as that's
# the pytorch backend we integrate into.