Dockerfile.train

# Copyright (c) 2021, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

ARG IMAGE_VERSION=21.06-tf2-py3
FROM nvcr.io/nvidia/tensorflow:${IMAGE_VERSION} AS devel

ARG UBUNTU_VERSION
ARG CMAKE_VERSION=3.19.6
ARG PYTHON_VERSION=3.8
ARG RAPIDS_VERSION=21.06
ARG HWLOC_VERSION=2.4.1
ARG MPI_VERSION=4.1.0
ARG MPI4PY_VERSION=3.0.3

ARG RELEASE=false
ARG SM="60;61;70;75;80"
ARG VAL_MODE=OFF
ARG ENABLE_MULTINODES=ON

RUN apt-get remove --auto-remove -y python3 cmake && \
    rm -rfv /usr/local/mpi /usr/local/ucx

RUN apt-get update -y && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
        vim \
        gdb \
        git \
        wget \
        tar \
        unzip \
        curl \
        clang-format \
        libboost-all-dev \
        libboost-regex-dev \
        libboost-system-dev \
        libboost-filesystem-dev \
        libtbb-dev \
        aptitude \
        libtool \
        numactl \
        libaio-dev \
        libnuma-dev \
        libibverbs-dev && \
    rm -rf /var/lib/apt/lists/*

RUN mkdir -p /var/tmp && wget -q -nc --no-check-certificate -P /var/tmp http://repo.anaconda.com/miniconda/Miniconda3-4.7.12-Linux-x86_64.sh && \
    bash /var/tmp/Miniconda3-4.7.12-Linux-x86_64.sh -b -p /opt/conda && \
    /opt/conda/bin/conda init && \
    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
    /opt/conda/bin/conda update -n base -c defaults conda && \
    /opt/conda/bin/conda clean -afy && \
    rm -rf /var/tmp/Miniconda3-4.7.12-Linux-x86_64.sh && \
    rm -rf /opt/conda/lib/libtinfo.so.6 && \
    ln -s /usr/lib/x86_64-linux-gnu/libtinfo.so.6.2 /opt/conda/lib/libtinfo.so.6
ENV CPATH=/opt/conda/include:$CPATH \
    LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH \
    LIBRARY_PATH=/opt/conda/lib:$LIBRARY_PATH \
    PATH=/opt/conda/bin:$PATH \
    CONDA_PREFIX=/opt/conda

RUN NCCL_VERSION=`find / -name "libnccl.so.?.*" | cut -d'.' -f 3-` && \
    CUDNN_VERSION=`find / -name "libcudnn.so.?.*" | cut -d'.' -f 3-` && \
    conda install -c rapidsai -c nvidia -c conda-forge -c anaconda \
        cudf=${RAPIDS_VERSION} \
        python=${PYTHON_VERSION} \
        cmake=${CMAKE_VERSION} \
        nccl=$NCCL_VERSION \
        cudnn=$CUDNN_VERSION \
        ucx \
        openmpi=${MPI_VERSION} \
        openmpi-mpicc=${MPI_VERSION} \
        mpi4py=${MPI4PY_VERSION} && \
    conda clean -afy
ENV OMPI_MCA_plm_rsh_agent=sh \
    OMPI_MCA_opal_cuda_support=true

ENV NCCL_LAUNCH_MODE=PARALLEL \
    NCCL_COLLNET_ENABLE=0

ENV SHARP_COLL_NUM_COLL_GROUP_RESOURCE_ALLOC_THRESHOLD=0 \
    SHARP_COLL_LOCK_ON_COMM_INIT=1 \
    SHARP_COLL_LOG_LEVEL=3

COPY clean_conda.sh .
RUN chmod +x ./clean_conda.sh && \
    ./clean_conda.sh google && \
    rm -rfv ./clean_conda.sh

# CUDA-Aware hwloc
RUN mkdir -p /var/tmp && wget -q -nc --no-check-certificate -P /var/tmp https://download.open-mpi.org/release/hwloc/v2.4/hwloc-${HWLOC_VERSION}.tar.gz && \
    mkdir -p /var/tmp && tar -x -f /var/tmp/hwloc-${HWLOC_VERSION}.tar.gz -C /var/tmp && \
    cd /var/tmp/hwloc-${HWLOC_VERSION} && \
    ./configure CPPFLAGS="-I/usr/local/cuda/include/ -L/usr/local/cuda/lib64/" LDFLAGS="-L/usr/local/cuda/lib64" --enable-cuda && \
    make -j$(nproc) && make install && \
    rm -rf /var/tmp/hwloc-${HWLOC_VERSION} /var/tmp/hwloc-${HWLOC_VERSION}.tar.gz

RUN pip3 install ortools sklearn && \
    pip3 cache purge

RUN ln -s /usr/lib/x86_64-linux-gnu/libibverbs.so.1.11.32.1 /usr/lib/x86_64-linux-gnu/libibverbs.so

# HugeCTR
RUN if [ "$RELEASE" = "true" ]; \
    then \
      mkdir -p /var/tmp && cd /var/tmp && git clone --depth=1 --branch master https://github.com/NVIDIA/HugeCTR.git HugeCTR && cd - && \
      cd /var/tmp/HugeCTR && \
      git submodule update --init --recursive && \
      mkdir build && cd build && \
      cmake -DCMAKE_BUILD_TYPE=Release -DSM=$SM \
            -DVAL_MODE=$VAL_MODE -DENABLE_MULTINODES=$ENABLE_MULTINODES .. && \
      make -j$(nproc) && make install && \
      chmod +x /usr/local/hugectr/bin/* && \
      chmod +x /usr/local/hugectr/lib/* && \
      rm -rf /var/tmp/HugeCTR; \
    else \
      echo "Build container for development successfully"; \
    fi
ENV PATH=/usr/local/hugectr/bin:$PATH \
    LIBRARY_PATH=/usr/local/hugectr/lib:$LIBRARY_PATH \
    LD_LIBRARY_PATH=/usr/local/hugectr/lib:$LD_LIBRARY_PATH \
    PYTHONPATH=/usr/local/hugectr/lib:$PYTHONPATH

HEALTHCHECK NONE
ENTRYPOINT []
CMD ["/bin/bash"]