Dockerfile.msdp-apex

# Multi-stage build
# Dockerfile.msdp-apex

ARG PTH_VERSION

# 1/Building apex with pytorch:*-devel
FROM pytorch/pytorch:${PTH_VERSION}-devel AS apex-msdp-builder

ENV CUDA_HOME=/usr/local/cuda

# Install git
RUN apt-get update && apt-get install -y --no-install-recommends git && \
     rm -rf /var/lib/apt/lists/*

# Build apex
RUN echo "Setup NVIDIA Apex" && \
    tmp_apex_path="/tmp/apex" && \
    rm -rf $tmp_apex_path && \
    git clone https://github.com/NVIDIA/apex $tmp_apex_path && \
    cd $tmp_apex_path && \
    pip install packaging && \
    pip wheel -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" .

# For pip --use-feature option
RUN python -m pip install --upgrade pip

# MSDeepSpeed
ARG MSDP_VERSION

RUN conda install cmake llvmdev=9.0.1 -c conda-forge

RUN git clone https://github.com/microsoft/DeepSpeed.git -b ${MSDP_VERSION} /tmp/DeepSpeed && cd /tmp/DeepSpeed && \
    ./install.sh --allow_sudo

# Build using devel image
FROM pytorch/pytorch:${PTH_VERSION}-devel

# Apex
COPY --from=apex-msdp-builder /tmp/apex/apex-*.whl /tmp/apex/

RUN pip install --no-cache-dir /tmp/apex/apex-*.whl && \
    rm -fr /tmp/apex

# MSDeepSpeed
RUN conda install cmake llvmdev=9.0.1 -c conda-forge

COPY --from=apex-msdp-builder /tmp/DeepSpeed/dist/deepspeed-*.whl /msdp/

RUN cd /msdp && export CUDA_HOME=/usr/local/cuda && \
    pip install --no-cache-dir deepspeed-*.whl && \
    rm -fr /msdp

# Renew nvidia signing key
# https://developer.nvidia.com/blog/updating-the-cuda-linux-gpg-repository-key/
RUN apt-key del 7fa2af80 && \
    rm /etc/apt/sources.list.d/nvidia-ml.list && \
    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub

# Install tzdata / git
RUN apt-get update && \
    ln -fs /usr/share/zoneinfo/Europe/Paris /etc/localtime && \
    apt-get -y install --no-install-recommends tzdata git && \
    dpkg-reconfigure --frontend noninteractive tzdata && \
    apt-get autoremove -y && \
    apt-get clean -y && \
    rm -rf /var/lib/apt/lists/*

# Ignite main dependencies
RUN pip install --upgrade --no-cache-dir pytorch-ignite \
                                         tensorboard \
                                         tqdm \
                                         fire

# Replace pillow with pillow-simd
RUN pip uninstall -y pillow && \
    CC="cc -mavx2" pip install --upgrade --no-cache-dir --force-reinstall pillow-simd

# Checkout Ignite examples only
RUN mkdir -p pytorch-ignite-examples && \
    cd pytorch-ignite-examples && \
    git init && \
    git config core.sparsecheckout true && \
    echo examples >> .git/info/sparse-checkout && \
    git remote add -f origin https://github.com/pytorch/ignite.git && \
    git pull origin master && \
    # rm very large .git folder
    rm -rf .git