Dockerfile.pytorch-1.10.0-cuda11.3-cudnn8

FROM pytorch/pytorch:1.10.0-cuda11.3-cudnn8-devel

RUN apt-get update && apt-get install -y curl software-properties-common wget sudo
RUN add-apt-repository ppa:git-core/ppa -y
RUN sed -i 's/mozilla\/DST_Root_CA_X3.crt/!mozilla\/DST_Root_CA_X3.crt/g' /etc/ca-certificates.conf && update-ca-certificates
RUN curl -sSf https://apt.kitware.com/kitware-archive.sh | sh
RUN apt-get update && apt-get install -y git cmake
RUN curl https://sh.rustup.rs -sSf | sh -s -- --default-toolchain stable -y
ENV PATH=/root/.cargo/bin:${PATH}
RUN cargo install mdbook mdbook-linkcheck mdbook-katex mdbook-open-on-gh

RUN yes | python3 -m pip install -U setuptools wheel build pip

ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64/stubs/:/usr/local/lib64:/usr/local/lib"
ENV LIBRARY_PATH="/usr/local/cuda/lib64/stubs/:/usr/local/lib64:/usr/local/lib"
ENV PKG_CONFIG_PATH="/usr/local/cuda/pkgconfig/"
ENV CUDA_LIBRARY_PATH="/usr/local/cuda/lib64/"


# OpenMPI version 4.0.3
RUN apt-get update -y && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
        bzip2 \
        file \
        hwloc \
        libnuma-dev \
        make \
        openssh-client \
        perl \
        tar \
        wget && \
    rm -rf /var/lib/apt/lists/*
RUN mkdir -p /var/tmp && wget -q -nc --no-check-certificate -P /var/tmp https://www.open-mpi.org/software/ompi/v4.0/downloads/openmpi-4.0.3.tar.bz2 && \
    mkdir -p /var/tmp && tar -x -f /var/tmp/openmpi-4.0.3.tar.bz2 -C /var/tmp -j && \
    cd /var/tmp/openmpi-4.0.3 &&   ./configure --disable-getpwuid --disable-oshmem --enable-fortran --enable-mca-no-build=btl-uct --enable-orterun-prefix-by-default --with-cuda --without-verbs && \
    make -j$(nproc) && \
    make -j$(nproc) install && \
    rm -rf /var/tmp/openmpi-4.0.3 /var/tmp/openmpi-4.0.3.tar.bz2 && cd -

# hwloc
RUN mkdir -p /var/tmp && wget -q -nc --no-check-certificate -P /var/tmp https://download.open-mpi.org/release/hwloc/v2.5/hwloc-2.5.0.tar.bz2 && \
    mkdir -p /var/tmp && tar -x -f /var/tmp/hwloc-2.5.0.tar.bz2 -C /var/tmp -j && \
    cd /var/tmp/hwloc-2.5.0 &&   ./configure && \
    make -j$(nproc) && \
    make -j$(nproc) install && \
    rm -rf /var/tmp/hwloc* && cd -

# Redis
RUN add-apt-repository ppa:redislabs/redis
RUN apt-get update && apt-get install -y redis
RUN yes | python3 -m pip install -U redis

# pytorch-lightning
RUN pip3 install pytorch-lightning

RUN mkdir /bagua 
COPY examples/ /bagua/examples
COPY ./ /var/tmp/bagua
RUN cd /var/tmp/bagua && python3 -m pip install . && cd - && rm -rf /var/tmp/bagua