Dockerfile.gpu

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

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

# Add default users for Ubuntu and Amazon Linux for ease of use
RUN apt-get update && apt-get install -y --no-install-recommends sudo
RUN groupadd --gid 1000 ubuntu && \
    useradd --uid 1000 --gid ubuntu -G sudo ubuntu && \
    echo "ubuntu ALL=(ALL) NOPASSWD:ALL" >/etc/sudoers.d/ubuntu && \
    mkdir -p /home/ubuntu && \
    chown ubuntu:ubuntu /home/ubuntu
RUN groupadd --gid 500 ec2-user && \
    useradd --uid 500 --gid ec2-user -G sudo ec2-user && \
    echo "ec2-user ALL=(ALL) NOPASSWD:ALL" >/etc/sudoers.d/ec2-user && \
    mkdir -p /home/ec2-user && \
    chown ec2-user:ec2-user /home/ec2-user

ENV CUDNN_VERSION=7.6.5.32-1+cuda10.1
ENV NCCL_VERSION=2.4.8-1+cuda10.1
ENV PYTHON_VERSION=3.6

RUN apt-get update && apt-get install -y --allow-downgrades --allow-change-held-packages --no-install-recommends \
    build-essential \
    cmake \
    g++-4.8 \
    git \
    curl \
    vim \
    wget \
    ca-certificates \
    libcudnn7=${CUDNN_VERSION} \
    libnccl2=${NCCL_VERSION} \
    libnccl-dev=${NCCL_VERSION} \
    libjpeg-dev \
    libpng-dev \
    python${PYTHON_VERSION} \
    python${PYTHON_VERSION}-dev \
    python${PYTHON_VERSION}-distutils \
    librdmacm1 \
    libibverbs1 \
    ibverbs-providers

RUN ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python

RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
    python get-pip.py && \
    rm get-pip.py

# Install Open MPI
RUN mkdir /tmp/openmpi && \
    cd /tmp/openmpi && \
    wget https://download.open-mpi.org/release/open-mpi/v4.0/openmpi-4.0.3.tar.gz && \
    tar zxf openmpi-4.0.3.tar.gz && \
    cd openmpi-4.0.3 && \
    ./configure --enable-orterun-prefix-by-default && \
    make -j $(nproc) all && \
    make install && \
    ldconfig && \
    rm -rf /tmp/openmpi

# Install OpenSSH for MPI to communicate between containers
RUN 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

# Install MXNet
ENV MXNET_VERSION=1.6.0
RUN pip install mxnet-cu101==${MXNET_VERSION}

# TODO: Remove this section when the issue of missing MKLDNN headers is fixed in
#       MXNet
#RUN mkdir /tmp/mkldnn && \
#    cd /tmp/mkldnn && \
#    wget https://github.com/oneapi-src/oneDNN/archive/v0.21.5.tar.gz && \
#    tar zxf v0.21.5.tar.gz && \
#    cd oneDNN-0.21.5 && \
#    cp -r include /usr/local/lib/python3.6/dist-packages/mxnet/include/mkldnn && \
#    mkdir build && \
#    cd build && \
#    cmake .. && \
#    cp include/mkldnn_version.h /usr/local/lib/python3.6/dist-packages/mxnet/include/mkldnn && \
#    rm -rf /tmp/mkldnn

# Install Horovod and the MPI Python library, temporarily using CUDA stubs
ARG REQS_HOROVOD
RUN ldconfig /usr/local/cuda/targets/x86_64-linux/lib/stubs && \
    HOROVOD_GPU_ALLREDUCE=NCCL HOROVOD_GPU_BROADCAST=NCCL HOROVOD_WITH_MXNET=1 \
        pip install --no-cache-dir ${REQS_HOROVOD} && \
    ldconfig

# Sockeye Python dependencies
ARG REQS_BASE
RUN pip install --no-cache-dir ${REQS_BASE}

# Everything below this ARG command re-runs if the local commit has changed
ARG SOCKEYE_COMMIT

# Install Sockeye, including Docker entry point script
COPY . /opt/sockeye
RUN cd /opt/sockeye && \
    pip install --no-cache-dir --no-deps --editable . && \
    cp /opt/sockeye/sockeye_contrib/docker/entrypoint.sh /usr/local/bin/ && \
    chmod +x /usr/local/bin/entrypoint.sh

# Set entry point to use CUDA stubs when needed
ENTRYPOINT ["entrypoint.sh"]