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

ARG ARG_TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0"
ENV TORCH_CUDA_ARCH_LIST=$ARG_TORCH_CUDA_ARCH_LIST

# 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 wheel --no-cache-dir --global-option="--cpp_ext" --global-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

# 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