github.com/kubeflow/training-operator@v1.7.0/examples/pytorch/mnist/Dockerfile-mpi (about)

     1  FROM nvidia/cuda:9.0-cudnn7-devel-ubuntu16.04
     2  ARG PYTHON_VERSION=3.6
     3  
     4  RUN apt-get update && apt-get install -y --no-install-recommends \
     5           build-essential \
     6           cmake \
     7           git \
     8           curl \
     9           vim \
    10           wget \
    11           ca-certificates \
    12           openssh-client \
    13           libjpeg-dev \
    14           libpng-dev &&\
    15       rm -rf /var/lib/apt/lists/*
    16  
    17  RUN wget https://www.open-mpi.org/software/ompi/v3.0/downloads/openmpi-3.0.0.tar.gz && \
    18      gunzip -c openmpi-3.0.0.tar.gz | tar xf - && \
    19      cd openmpi-3.0.0 && \
    20      ./configure --prefix=/home/.openmpi --with-cuda && \
    21      make all install
    22  
    23  ENV PATH="$PATH:/home/.openmpi/bin"
    24  ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/home/.openmpi/lib/"
    25  
    26  RUN ompi_info --parsable --all | grep mpi_built_with_cuda_support:value
    27  RUN wget -O ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh  && \
    28       chmod +x ~/miniconda.sh && \
    29       ~/miniconda.sh -b -p /opt/conda && \
    30       rm ~/miniconda.sh && \
    31       /opt/conda/bin/conda update conda && \
    32       /opt/conda/bin/conda install -y python=$PYTHON_VERSION numpy pyyaml scipy ipython mkl mkl-include cython typing && \
    33       /opt/conda/bin/conda clean -ya
    34  ENV PATH /opt/conda/bin:$PATH
    35  # This must be done before pip so that requirements.txt is available
    36  WORKDIR /opt/pytorch
    37  
    38  RUN git clone --recursive https://github.com/pytorch/pytorch
    39  
    40  RUN TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
    41      CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
    42      cd pytorch/ && \
    43      pip install -v .
    44  
    45  RUN /opt/conda/bin/conda config --set ssl_verify False
    46  RUN pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org
    47  RUN pip install --trusted-host pypi.org --trusted-host files.pythonhosted.org torchvision tensorboardX==1.6.0
    48  
    49  WORKDIR /var
    50  ADD mnist.py /var
    51  
    52  ENTRYPOINT ["mpirun", "-n", "1", "--allow-run-as-root", "python", "/var/mnist.py"]