pytorch
/
Dockerfile
108 строк · 4.1 Кб
1# syntax = docker/dockerfile:experimental
2#
3# NOTE: To build this you will need a docker version > 18.06 with
4# experimental enabled and DOCKER_BUILDKIT=1
5#
6# If you do not use buildkit you are not going to have a good time
7#
8# For reference:
9# https://docs.docker.com/develop/develop-images/build_enhancements/
10ARG BASE_IMAGE=ubuntu:20.04
11ARG PYTHON_VERSION=3.8
12
13FROM ${BASE_IMAGE} as dev-base
14RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
15build-essential \
16ca-certificates \
17ccache \
18cmake \
19curl \
20git \
21libjpeg-dev \
22libpng-dev && \
23rm -rf /var/lib/apt/lists/*
24RUN /usr/sbin/update-ccache-symlinks
25RUN mkdir /opt/ccache && ccache --set-config=cache_dir=/opt/ccache
26ENV PATH /opt/conda/bin:$PATH
27
28FROM dev-base as conda
29ARG PYTHON_VERSION=3.8
30# Automatically set by buildx
31ARG TARGETPLATFORM
32# translating Docker's TARGETPLATFORM into miniconda arches
33RUN case ${TARGETPLATFORM} in \
34"linux/arm64") MINICONDA_ARCH=aarch64 ;; \
35*) MINICONDA_ARCH=x86_64 ;; \
36esac && \
37curl -fsSL -v -o ~/miniconda.sh -O "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-${MINICONDA_ARCH}.sh"
38COPY requirements.txt .
39# Manually invoke bash on miniconda script per https://github.com/conda/conda/issues/10431
40RUN chmod +x ~/miniconda.sh && \
41bash ~/miniconda.sh -b -p /opt/conda && \
42rm ~/miniconda.sh && \
43/opt/conda/bin/conda install -y python=${PYTHON_VERSION} cmake conda-build pyyaml numpy ipython && \
44/opt/conda/bin/python -mpip install -r requirements.txt && \
45/opt/conda/bin/conda clean -ya
46
47FROM dev-base as submodule-update
48WORKDIR /opt/pytorch
49COPY . .
50RUN git submodule update --init --recursive
51
52FROM conda as build
53ARG CMAKE_VARS
54WORKDIR /opt/pytorch
55COPY --from=conda /opt/conda /opt/conda
56COPY --from=submodule-update /opt/pytorch /opt/pytorch
57RUN make triton
58RUN --mount=type=cache,target=/opt/ccache \
59export eval ${CMAKE_VARS} && \
60TORCH_CUDA_ARCH_LIST="3.5 5.2 6.0 6.1 7.0+PTX 8.0" TORCH_NVCC_FLAGS="-Xfatbin -compress-all" \
61CMAKE_PREFIX_PATH="$(dirname $(which conda))/../" \
62python setup.py install
63
64FROM conda as conda-installs
65ARG PYTHON_VERSION=3.8
66ARG CUDA_VERSION=12.1
67ARG CUDA_CHANNEL=nvidia
68ARG INSTALL_CHANNEL=pytorch-nightly
69# Automatically set by buildx
70# Note conda needs to be pinned to 23.5.2 see: https://github.com/pytorch/pytorch/issues/106470
71RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION} conda=23.5.2
72ARG TARGETPLATFORM
73
74# On arm64 we can only install wheel packages.
75RUN case ${TARGETPLATFORM} in \
76"linux/arm64") pip install --extra-index-url https://download.pytorch.org/whl/cpu/ torch torchvision torchaudio ;; \
77*) /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -c "${CUDA_CHANNEL}" -y "python=${PYTHON_VERSION}" pytorch torchvision torchaudio "pytorch-cuda=$(echo $CUDA_VERSION | cut -d'.' -f 1-2)" ;; \
78esac && \
79/opt/conda/bin/conda clean -ya
80RUN /opt/conda/bin/pip install torchelastic
81
82FROM ${BASE_IMAGE} as official
83ARG PYTORCH_VERSION
84ARG TRITON_VERSION
85ARG TARGETPLATFORM
86ARG CUDA_VERSION
87LABEL com.nvidia.volumes.needed="nvidia_driver"
88RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
89ca-certificates \
90libjpeg-dev \
91libpng-dev \
92&& rm -rf /var/lib/apt/lists/*
93COPY --from=conda-installs /opt/conda /opt/conda
94RUN if test -n "${TRITON_VERSION}" -a "${TARGETPLATFORM}" != "linux/arm64"; then \
95DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends gcc; \
96rm -rf /var/lib/apt/lists/*; \
97fi
98ENV PATH /opt/conda/bin:$PATH
99ENV NVIDIA_VISIBLE_DEVICES all
100ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
101ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
102ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
103ENV PYTORCH_VERSION ${PYTORCH_VERSION}
104WORKDIR /workspace
105
106FROM official as dev
107# Should override the already installed version from the official-image stage
108COPY --from=build /opt/conda /opt/conda
109