mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 12:54:11 +08:00
ab8fbd26f8
Fixes following: https://github.com/pytorch/pytorch/actions/runs/3242695506/jobs/5316334351 crash in Docker builds introduced by: #82682 The PR seems to introduce some changes not compatible with cuda 11.3 which is used by our Docker builds This is a reland of original pr: https://github.com/pytorch/pytorch/pull/86941 (Created this new PR to start fresh) Which was reverted because conda install, installed wrong version of pytorch. It installed pytorch for cuda 11.3 still rather then 11.6 This should be fixed now with Release 1.13 Pull Request resolved: https://github.com/pytorch/pytorch/pull/87858 Approved by: https://github.com/seemethere, https://github.com/malfet, https://github.com/izaitsevfb
97 lines
3.5 KiB
Docker
97 lines
3.5 KiB
Docker
# syntax = docker/dockerfile:experimental
|
|
#
|
|
# NOTE: To build this you will need a docker version > 18.06 with
|
|
# experimental enabled and DOCKER_BUILDKIT=1
|
|
#
|
|
# If you do not use buildkit you are not going to have a good time
|
|
#
|
|
# For reference:
|
|
# https://docs.docker.com/develop/develop-images/build_enhancements/
|
|
ARG BASE_IMAGE=ubuntu:18.04
|
|
ARG PYTHON_VERSION=3.8
|
|
|
|
FROM ${BASE_IMAGE} as dev-base
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
build-essential \
|
|
ca-certificates \
|
|
ccache \
|
|
cmake \
|
|
curl \
|
|
git \
|
|
libjpeg-dev \
|
|
libpng-dev && \
|
|
rm -rf /var/lib/apt/lists/*
|
|
RUN /usr/sbin/update-ccache-symlinks
|
|
RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache
|
|
ENV PATH /opt/conda/bin:$PATH
|
|
|
|
FROM dev-base as conda
|
|
ARG PYTHON_VERSION=3.8
|
|
# Automatically set by buildx
|
|
ARG TARGETPLATFORM
|
|
# translating Docker's TARGETPLATFORM into miniconda arches
|
|
RUN case ${TARGETPLATFORM} in \
|
|
"linux/arm64") MINICONDA_ARCH=aarch64 ;; \
|
|
*) MINICONDA_ARCH=x86_64 ;; \
|
|
esac && \
|
|
curl -fsSL -v -o ~/miniconda.sh -O "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh"
|
|
COPY requirements.txt .
|
|
RUN chmod +x ~/miniconda.sh && \
|
|
~/miniconda.sh -b -p /opt/conda && \
|
|
rm ~/miniconda.sh && \
|
|
/opt/conda/bin/conda install -y python=${PYTHON_VERSION} cmake conda-build pyyaml numpy ipython && \
|
|
/opt/conda/bin/python -mpip install -r requirements.txt && \
|
|
/opt/conda/bin/conda clean -ya
|
|
|
|
FROM dev-base as submodule-update
|
|
WORKDIR /opt/pytorch
|
|
COPY . .
|
|
RUN git submodule update --init --recursive --jobs 0
|
|
|
|
FROM conda as build
|
|
WORKDIR /opt/pytorch
|
|
COPY --from=conda /opt/conda /opt/conda
|
|
COPY --from=submodule-update /opt/pytorch /opt/pytorch
|
|
RUN --mount=type=cache,target=/opt/ccache \
|
|
TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
|
|
CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
|
|
python setup.py install
|
|
|
|
FROM conda as conda-installs
|
|
ARG PYTHON_VERSION=3.8
|
|
ARG CUDA_VERSION=11.6
|
|
ARG CUDA_CHANNEL=nvidia
|
|
ARG INSTALL_CHANNEL=pytorch-nightly
|
|
# Automatically set by buildx
|
|
RUN /opt/conda/bin/conda update -y conda
|
|
RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION}
|
|
ARG TARGETPLATFORM
|
|
|
|
# On arm64 we can only install wheel packages
|
|
RUN case ${TARGETPLATFORM} in \
|
|
"linux/arm64") pip install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchtext ;; \
|
|
*) /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -c "${CUDA_CHANNEL}" -y "python=${PYTHON_VERSION}" pytorch torchvision torchtext "pytorch-cuda=$(echo $CUDA_VERSION | cut -d'.' -f 1-2)" ;; \
|
|
esac && \
|
|
/opt/conda/bin/conda clean -ya
|
|
RUN /opt/conda/bin/pip install torchelastic
|
|
|
|
FROM ${BASE_IMAGE} as official
|
|
ARG PYTORCH_VERSION
|
|
LABEL com.nvidia.volumes.needed="nvidia_driver"
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
ca-certificates \
|
|
libjpeg-dev \
|
|
libpng-dev && \
|
|
rm -rf /var/lib/apt/lists/*
|
|
COPY --from=conda-installs /opt/conda /opt/conda
|
|
ENV PATH /opt/conda/bin:$PATH
|
|
ENV NVIDIA_VISIBLE_DEVICES all
|
|
ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
|
|
ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
|
|
ENV PYTORCH_VERSION ${PYTORCH_VERSION}
|
|
WORKDIR /workspace
|
|
|
|
FROM official as dev
|
|
# Should override the already installed version from the official-image stage
|
|
COPY --from=build /opt/conda /opt/conda
|