Bootstrap: docker
From: nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04

%post
    USE_CUDA=1
    CUDA_VERSION=11.3.1
    CUDNN_VERSION=8
    LINUX_DISTRO=ubuntu
    DISTRO_VERSION=20.04
    TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6"

    rm -f /etc/apt/apt.conf.d/docker-clean; \
    echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' \
    > /etc/apt/apt.conf.d/keep-cache
    # Downloads the latest package lists (important).
    apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
        ca-certificates \
        ccache \
        curl \
        git \
        libaio-dev && \
    apt-get install -y --no-install-recommends llvm-10 lldb-10 llvm-10-dev libllvm10 llvm-10-runtime
    rm -rf /var/lib/apt/lists/*
    # Install miniconda
    PYTHON_VERSION=3.8
    CONDA_URL=https://repo.anaconda.com/miniconda/Miniconda3-py38_4.10.3-Linux-x86_64.sh
    curl -fsSL -v -o ~/miniconda.sh -O ${CONDA_URL} && \
    chmod +x ~/miniconda.sh && \
    ~/miniconda.sh -b -p /opt/miniconda3 && \
    rm ~/miniconda.sh

    export PATH="/opt/miniconda3/bin:$PATH" && \
    conda install -y python=${PYTHON_VERSION} && \
    conda clean -ya
    which python
    echo /opt/miniconda3/lib >> /etc/ld.so.conf.d/conda.conf && ldconfig
    /usr/sbin/update-ccache-symlinks
    mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache && ccache --max-size 0
    conda install -y \
        pip \
        astunparse \
        cffi \
        cmake \
        future \
        mkl \
        mkl-include \
        ninja \
        numpy \
        pillow \
        pkgconfig \
        pyyaml \
        requests \
        setuptools \
        six \
        typing_extensions && \
    conda install -y -c conda-forge \
        libjpeg-turbo \
        libpng && \
    MAGMA_VERSION=113
    conda install -y -c pytorch \
        magma-cuda${MAGMA_VERSION}
    which conda

    cd /opt
    git clone --recursive --jobs 0 https://github.com/pytorch/pytorch.git /opt/pytorch
    git clone --recursive --jobs 0 https://github.com/pytorch/vision.git /opt/vision

    PYTORCH_VERSION_TAG="v1.10.2"
    TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
    cd /opt/pytorch
    if [ -n ${PYTORCH_VERSION_TAG} ]; then \
    git checkout ${PYTORCH_VERSION_TAG} && \
    git submodule sync && \
    git submodule update --init --recursive --jobs 0; \
    fi

    USE_CUDA=${USE_CUDA} USE_CUDNN=${USE_CUDA} USE_ROCM=0 \
    TORCH_NVCC_FLAGS=${TORCH_NVCC_FLAGS} \
    TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
    python setup.py install

    TORCHVISION_VERSION_TAG='v0.11.3'
    cd /opt/vision
    if [ -n ${TORCHVISION_VERSION_TAG} ]; then \
    git checkout ${TORCHVISION_VERSION_TAG} && \
    git submodule sync && \
    git submodule update --init --recursive --jobs 0; \
    fi

    FORCE_CUDA=${USE_CUDA} TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} \
    python setup.py install

    git clone https://github.com/NVIDIA/apex.git /opt/apex/
    cd /opt/apex/
    pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .

    pip install --extra-index-url https://developer.download.nvidia.com/compute/redist --upgrade nvidia-dali-cuda110
    # Install pipenv
    pip install pipenv
    pip install asyncio
    #pip install deepspeed
    #ds_report
    DS_BUILD_OPS=1 pip install deepspeed --global-option="build_ext" --global-option="-j8"
    ds_report
    pip install mlperf_compliance matplotlib Cython pycocotools

%environment
    export LANG=C.UTF-8
    export LC_ALL=C.UTF-8
    export PATH=/opt/miniconda3/bin:$PATH
    export PYTHONDONTWRITEBYTECODE=1
    export PYTHONUNBUFFERED=1
    export PYTHONIOENCODING=UTF-8
    export PIPENV_VENV_IN_PROJECT=1
