stoke-gpu-mpi.Dockerfile

############
# BASE IMAGE
############
# Pull from the base ubuntu image w/ the correct ARCH and CUDA
FROM nvidia/cuda:11.1-base-ubuntu18.04 as base

##########
# ENV VARS
##########
ENV LANG C.UTF-8
ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1

######################
# INSTALL DEPENDENCIES
######################
# Use bash instead of sh
# Needed for string substitution
SHELL ["/bin/bash", "-c"]
# All CUDA dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        cuda-command-line-tools-11-1 \
        cuda-libraries-11-1 \
        libcudnn8=8.1.0.77-1+cuda11.2 \
        libnccl2=2.8.4-1+cuda11.1 \
        && apt-get clean && \
        rm -rf /var/lib/apt/lists/*

# All CUDA-dev dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
        cuda-libraries-dev-11-1 \
        cuda-compiler-11-1 \
        libcudnn8-dev=8.1.0.77-1+cuda11.2 \
        libnccl-dev=2.8.4-1+cuda11.1 \
        && apt-get clean && \
        rm -rf /var/lib/apt/lists/*

# Pickup additional dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
        ca-certificates \
        cmake \
        curl \
        g++-7 \
        git \
        ibverbs-providers \
        libfreetype6-dev \
        libhdf5-serial-dev \
        libibverbs1 \
        libjpeg-dev \
        libpng-dev \
        librdmacm1 \
        libsndfile1 \
        libzmq3-dev \
        openssh-client \
        pkg-config \
        software-properties-common \
        unzip \
        vim \
        wget \
        && apt-get clean && \
        rm -rf /var/lib/apt/lists/*

# For CUDA profiling, TensorFlow requires CUPTI.
ENV LD_LIBRARY_PATH /usr/local/cuda/extras/CUPTI/lib64:/usr/local/cuda/lib64:$LD_LIBRARY_PATH

# Link the libcuda stub and reconfigure dynamic linker run-time bindings
RUN ln -s /usr/local/cuda/lib64/stubs/libcuda.so /usr/local/cuda/lib64/stubs/libcuda.so.1 \
    && echo "/usr/local/cuda/lib64/stubs" > /etc/ld.so.conf.d/z-cuda-stubs.conf \
    && ldconfig

# Download public key for github.com
RUN mkdir -p -m 0600 ~/.ssh && ssh-keyscan github.com >> ~/.ssh/known_hosts

############
# PIP/Python
############
# Add deadsnakes which has all Python versions
RUN add-apt-repository ppa:deadsnakes/ppa
# Install python X.X-dev and pip via apt-get
RUN apt-get update && apt-get install -y python3.8-dev python3.8-distutils python3-pip python3-setuptools && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*
# Create the symlink to the correct version of python X.X
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 2
# Upgrade pip -- this will link the pip cmd to the correct version of python X.X
RUN python3.8 -m pip --no-cache-dir install --upgrade pip
# Upgrade setuptools
RUN pip --no-cache-dir install --upgrade setuptools
# Some tools expect a "python" binary
RUN ln -s $(which python3.8) /usr/local/bin/python

################
# PYTORCH -- GPU
################
RUN pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html

##########
# OPEN-MPI
##########
WORKDIR /usr/src/code/
RUN wget https://download.open-mpi.org/release/open-mpi/v4.1/openmpi-4.1.1.tar.gz && \
    tar zxf openmpi-4.1.1.tar.gz
WORKDIR /usr/src/code/openmpi-4.1.1
RUN ./configure --enable-orterun-prefix-by-default && make -j $(nproc) all
RUN make install && ldconfig

################
# HOROVOD -- GPU
################
RUN ldconfig /usr/local/cuda/targets/x86_64-linux/lib/stubs
RUN HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_WITH_PYTORCH=1 pip install --no-cache-dir horovod[pytorch] && ldconfig

##########
# OPEN SSH
##########
# Install OpenSSH for MPI to communicate between containers
RUN apt-get update && apt-get install -y --no-install-recommends openssh-client openssh-server && \
    mkdir -p /var/run/sshd
# Allow OpenSSH to talk to containers without asking for confirmation
RUN cat /etc/ssh/ssh_config | grep -v StrictHostKeyChecking > /etc/ssh/ssh_config.new && \
    echo "    StrictHostKeyChecking no" >> /etc/ssh/ssh_config.new && \
    mv /etc/ssh/ssh_config.new /etc/ssh/ssh_config

#############
# NVIDIA APEX
#############
WORKDIR /usr/src/code/
RUN git clone https://github.com/NVIDIA/apex
WORKDIR /usr/src/code/apex
RUN pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

# Install stoke w/ MPI support
RUN pip install --no-cache-dir --trusted-host pypi.python.org stoke[mpi]