Dockerfile

FROM nvidia/cuda:11.3.0-devel-ubuntu20.04
ENV DEBIAN_FRONTEND=noninteractive

##############################################################################
# Configurable Version
##############################################################################
ARG PYTORCH_VERSION=1.10.0

##############################################################################
# Temporary Installation Directory
##############################################################################
ARG STAGE_DIR=/tmp
RUN mkdir -p ${STAGE_DIR}

##############################################################################
# Installation/Basic Utilities
##############################################################################
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential autotools-dev \
        git cmake g++ gcc autoconf libtool automake apt-utils \
        curl wget less unzip \
        htop iftop iotop ca-certificates openssh-client openssh-server iproute2 \
        iputils-ping net-tools sudo \
        libsndfile-dev libcupti-dev libjpeg-dev libpng-dev libaio-dev llvm-9-dev

##############################################################################
# Python
##############################################################################
ENV PYTHON_VERSION=3
RUN apt-get install -y python3 python3-dev && \
    rm -f /usr/bin/python && \
    ln -s /usr/bin/python3 /usr/bin/python && \
    curl -O https://bootstrap.pypa.io/get-pip.py && \
        python get-pip.py && \
        rm get-pip.py && \
    pip install --upgrade pip && \
    # Print python an pip version
    python -V && pip -V

##############################################################################
## Add trainer user
###############################################################################
RUN useradd --create-home --uid 1000 --shell /bin/bash trainer
RUN usermod -aG sudo trainer
RUN echo "trainer ALL=(ALL) NOPASSWD: ALL" >> /etc/sudoers

ENV USER trainer
ENV HOME /home/$USER
RUN echo $HOME

##############################################################################
# SSH Setup
##############################################################################
ARG SSHDIR=${HOME}/.ssh
RUN mkdir -p ${SSHDIR} \
&& touch ${SSHDIR}/sshd_config \
&& ssh-keygen -t rsa -f ${SSHDIR}/ssh_host_rsa_key -N '' \
&& cp ${SSHDIR}/ssh_host_rsa_key.pub ${SSHDIR}/authorized_keys \
&& cp ${SSHDIR}/ssh_host_rsa_key ${SSHDIR}/id_rsa \
&& echo "       IdentityFile ${SSHDIR}/id_rsa" >> ${SSHDIR}/config \
&& echo "       StrictHostKeyChecking no" >> ${SSHDIR}/config \
&& echo "       UserKnownHostsFile /dev/null" >> ${SSHDIR}/config \
&& echo "       Port 2022" >> ${SSHDIR}/config \
&& echo 'Port 2022' >> ${SSHDIR}/sshd_config \
&& echo 'UsePrivilegeSeparation no' >> ${SSHDIR}/sshd_config \
&& echo "HostKey ${SSHDIR}/ssh_host_rsa_key" >> ${SSHDIR}/sshd_config \ && echo "PidFile ${SSHDIR}/sshd.pid" >> ${SSHDIR}/sshd_config \
&& chmod -R 600 ${SSHDIR}/* \
&& chown -R ${USER}:${USER} ${SSHDIR}/
RUN eval `ssh-agent -s` && ssh-add ${SSHDIR}/id_rsa
EXPOSE 22

##############################################################################
# OPENMPI
##############################################################################
ARG OPENMPI_BASEVERSION=4.0
ARG OPENMPI_VERSION=${OPENMPI_BASEVERSION}.1
RUN cd ${STAGE_DIR} && \
    wget -q -O - https://download.open-mpi.org/release/open-mpi/v${OPENMPI_BASEVERSION}/openmpi-${OPENMPI_VERSION}.tar.gz | tar xzf - && \
    cd openmpi-${OPENMPI_VERSION} && \
    ./configure --prefix=/usr/local/openmpi-${OPENMPI_VERSION} && \
    make -j"$(nproc)" install && \
    ln -s /usr/local/openmpi-${OPENMPI_VERSION} /usr/local/mpi && \
    # Sanity check:
    test -f /usr/local/mpi/bin/mpic++ && \
    cd ${STAGE_DIR} && \
    rm -r ${STAGE_DIR}/openmpi-${OPENMPI_VERSION}
ENV PATH=/usr/local/mpi/bin:${PATH} \
    LD_LIBRARY_PATH=/usr/local/lib:/usr/local/mpi/lib:/usr/local/mpi/lib64:${LD_LIBRARY_PATH}
# Create a wrapper for OpenMPI to allow running as root by default
RUN mv /usr/local/mpi/bin/mpirun /usr/local/mpi/bin/mpirun.real && \
    echo '#!/bin/bash' > /usr/local/mpi/bin/mpirun && \
    echo 'mpirun.real --allow-run-as-root --prefix /usr/local/mpi "$@"' >> /usr/local/mpi/bin/mpirun && \
    chmod a+x /usr/local/mpi/bin/mpirun

##############################################################################
# EFA Setup
##############################################################################
RUN cd ${STAGE_DIR} && curl -O  https://efa-installer.amazonaws.com/aws-efa-installer-1.15.1.tar.gz && tar -xf aws-efa-installer-1.15.1.tar.gz && cd aws-efa-installer && ./efa_installer.sh -y -d -g  --skip-kmod --skip-limit-conf --no-verify && rm -rf ${STAGE_DIR}/aws-efa-installer-1.15.1.tar.gz && rm -rf ${STAGE_DIR}/aws-efa-installer-1.15.1

##############################################################################
# NCCL hacks
# Due to issue https://github.com/pytorch/pytorch/issues/69013, we have to remove nccl brought by nvidia docker
##############################################################################
RUN rm -rf /usr/include/nccl.h && rm -rf /usr/lib/x86_64-linux-gnu/libnccl.so.2

##############################################################################
# PyTorch
# setuptools set to 59.5.0 due to https://stackoverflow.com/questions/70520120/attributeerror-module-setuptools-distutils-has-no-attribute-version
##############################################################################
ARG TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0"
RUN pip3 install astunparse numpy ninja pyyaml mkl mkl-include setuptools==59.5.0 cmake cffi typing_extensions future six requests dataclasses
RUN cd ${STAGE_DIR} && git clone https://github.com/pytorch/pytorch.git && cd pytorch && git checkout v${PYTORCH_VERSION} && git submodule sync && git submodule update --init --recursive && TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} python3 setup.py install

##############################################################################
# Fairseq
##############################################################################
RUN git clone https://github.com/pytorch/fairseq ${STAGE_DIR}/fairseq
ARG FAIRSEQ_VERSION=0.11.0
ARG BUILD_CUDA_EXTENSIONS=1

# installing dependencies explicitly without torch and torchaudio
RUN pip3 install cffi cython dataclasses hydra-core==1.0.7 omegaconf==2.0.6 numpy regex sacrebleu tqdm bitarray
RUN cd ${STAGE_DIR}/fairseq && git checkout v${FAIRSEQ_VERSION} && BUILD_CUDA_EXTENSIONS=${BUILD_CUDA_EXTENSIONS} TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" pip3 install --no-deps -v . && rm -rf ${STAGE_DIR}/fairseq

##############################################################################
# Fairscale
##############################################################################
RUN git clone https://github.com/facebookresearch/fairscale ${STAGE_DIR}/fairscale
ARG FAIRSCALE_VERSION=0.4.5
## installing with no dependencies to make sure torch is not overridden
RUN cd ${STAGE_DIR}/fairscale && git checkout v${FAIRSCALE_VERSION} && sudo pip3 install --no-deps -v . && rm -rf ${STAGE_DIR}/fairscale

##############################################################################
## Install custom Apex
##############################################################################
RUN cd ${STAGE_DIR} && git clone https://github.com/szhengac/apex.git && cd apex && git checkout lans && TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" pip3 install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./ && rm -rf ${STAGE_DIR}/apex

##############################################################################
# Install AWS-OFI-NCCL plugin
##############################################################################
RUN git clone https://github.com/aws/aws-ofi-nccl.git ${STAGE_DIR}/aws-ofi-nccl \
    && cd ${STAGE_DIR}/aws-ofi-nccl \
    && git checkout aws  \
    && ./autogen.sh \
    && ./configure --prefix=$HOME/aws-ofi-nccl/install \
       --with-libfabric=/opt/amazon/efa/ \
       --with-cuda=/usr/local/cuda \
       --with-mpi=/opt/amazon/openmpi/ \
       --with-nccl=${STAGE_DIR}/pytorch/build/nccl \
    && make -j$(nproc) && make install && rm -rf ${STAGE_DIR}/aws-ofi-nccl

##############################################################################
## Install more python dependencies
##############################################################################
RUN pip3 install tensorboardX subword_nmt jellyfish

##############################################################################
# Set up runtime
##############################################################################
USER trainer
ENV LD_LIBRARY_PATH="${STAGE_DIR}/pytorch/build/nccl/lib:/usr/local/cuda/lib64:/opt/amazon/efa/lib:/opt/amazon/openmpi/lib:${HOME}/aws-ofi-nccl/install/lib:${LD_LIBRARY_PATH}"

ADD ./fairseq_train_wrapped /