108 lines
4.0 KiB
Docker
108 lines
4.0 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:20.04
|
|
ARG PYTHON_VERSION=3.8
|
|
|
|
FROM ${BASE_IMAGE} as dev-base
|
|
RUN apt-get update && DEBIAN_FRONTEND=noninteractive 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 .
|
|
# Manually invoke bash on miniconda script per https://github.com/conda/conda/issues/10431
|
|
RUN chmod +x ~/miniconda.sh && \
|
|
bash ~/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
|
|
|
|
FROM conda as build
|
|
ARG CMAKE_VARS
|
|
WORKDIR /opt/pytorch
|
|
COPY --from=conda /opt/conda /opt/conda
|
|
COPY --from=submodule-update /opt/pytorch /opt/pytorch
|
|
RUN make triton
|
|
RUN --mount=type=cache,target=/opt/ccache \
|
|
export eval ${CMAKE_VARS} && \
|
|
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.7
|
|
ARG CUDA_CHANNEL=nvidia
|
|
ARG INSTALL_CHANNEL=pytorch-nightly
|
|
# Automatically set by buildx
|
|
# Note conda needs to be pinned to 23.5.2 see: https://github.com/pytorch/pytorch/issues/106470
|
|
RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION} conda=23.5.2
|
|
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 torchaudio ;; \
|
|
*) /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)" ;; \
|
|
esac && \
|
|
/opt/conda/bin/conda clean -ya
|
|
RUN /opt/conda/bin/pip install torchelastic
|
|
|
|
FROM ${BASE_IMAGE} as official
|
|
ARG PYTORCH_VERSION
|
|
ARG TRITON_VERSION
|
|
ARG TARGETPLATFORM
|
|
ARG CUDA_VERSION
|
|
LABEL com.nvidia.volumes.needed="nvidia_driver"
|
|
RUN apt-get update && DEBIAN_FRONTEND=noninteractive 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
|
|
RUN if test -n "${TRITON_VERSION}" -a "${TARGETPLATFORM}" != "linux/arm64"; then \
|
|
DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends gcc; \
|
|
rm -rf /var/lib/apt/lists/*; \
|
|
fi
|
|
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
|