1ARG BASE_TARGET=base 2ARG GPU_IMAGE=ubuntu:20.04 3FROM ${GPU_IMAGE} as base 4 5ENV DEBIAN_FRONTEND=noninteractive 6 7RUN apt-get clean && apt-get update 8RUN apt-get install -y curl locales g++ git-all autoconf automake make cmake wget unzip sudo 9# Just add everything as a safe.directory for git since these will be used in multiple places with git 10RUN git config --global --add safe.directory '*' 11 12RUN locale-gen en_US.UTF-8 13 14ENV LC_ALL en_US.UTF-8 15ENV LANG en_US.UTF-8 16ENV LANGUAGE en_US.UTF-8 17 18# Install openssl 19FROM base as openssl 20ADD ./common/install_openssl.sh install_openssl.sh 21RUN bash ./install_openssl.sh && rm install_openssl.sh 22 23# Install python 24FROM base as python 25ADD common/install_cpython.sh install_cpython.sh 26RUN apt-get update -y && \ 27 apt-get install build-essential gdb lcov libbz2-dev libffi-dev \ 28 libgdbm-dev liblzma-dev libncurses5-dev libreadline6-dev \ 29 libsqlite3-dev libssl-dev lzma lzma-dev tk-dev uuid-dev zlib1g-dev -y && \ 30 bash ./install_cpython.sh && \ 31 rm install_cpython.sh && \ 32 apt-get clean 33 34FROM base as conda 35ADD ./common/install_conda_docker.sh install_conda.sh 36RUN bash ./install_conda.sh && rm install_conda.sh 37 38FROM base as cpu 39# Install Anaconda 40COPY --from=conda /opt/conda /opt/conda 41# Install python 42COPY --from=python /opt/python /opt/python 43COPY --from=python /opt/_internal /opt/_internal 44ENV PATH=/opt/conda/bin:/usr/local/cuda/bin:$PATH 45# Install MKL 46ADD ./common/install_mkl.sh install_mkl.sh 47RUN bash ./install_mkl.sh && rm install_mkl.sh 48 49FROM cpu as cuda 50ADD ./common/install_cuda.sh install_cuda.sh 51ADD ./common/install_magma.sh install_magma.sh 52ENV CUDA_HOME /usr/local/cuda 53 54FROM cuda as cuda11.8 55RUN bash ./install_cuda.sh 11.8 56RUN bash ./install_magma.sh 11.8 57RUN ln -sf /usr/local/cuda-11.8 /usr/local/cuda 58 59FROM cuda as cuda12.1 60RUN bash ./install_cuda.sh 12.1 61RUN bash ./install_magma.sh 12.1 62RUN ln -sf /usr/local/cuda-12.1 /usr/local/cuda 63 64FROM cuda as cuda12.4 65RUN bash ./install_cuda.sh 12.4 66RUN bash ./install_magma.sh 12.4 67RUN ln -sf /usr/local/cuda-12.4 /usr/local/cuda 68 69FROM cpu as rocm 70ARG PYTORCH_ROCM_ARCH 71ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH} 72ENV MKLROOT /opt/intel 73# Adding ROCM_PATH env var so that LoadHip.cmake (even with logic updated for ROCm6.0) 74# find HIP works for ROCm5.7. Not needed for ROCm6.0 and above. 75# Remove below when ROCm5.7 is not in support matrix anymore. 76ENV ROCM_PATH /opt/rocm 77# No need to install ROCm as base docker image should have full ROCm install 78#ADD ./common/install_rocm.sh install_rocm.sh 79ADD ./common/install_rocm_drm.sh install_rocm_drm.sh 80ADD ./common/install_rocm_magma.sh install_rocm_magma.sh 81# gfortran and python needed for building magma from source for ROCm 82RUN apt-get update -y && \ 83 apt-get install gfortran -y && \ 84 apt-get install python -y && \ 85 apt-get clean 86 87RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh 88RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh 89 90# Install AOTriton 91COPY ./common/common_utils.sh common_utils.sh 92COPY ./aotriton_version.txt aotriton_version.txt 93COPY ./common/install_aotriton.sh install_aotriton.sh 94RUN bash ./install_aotriton.sh /opt/rocm && rm install_aotriton.sh aotriton_version.txt 95ENV AOTRITON_INSTALLED_PREFIX /opt/rocm/aotriton 96 97FROM ${BASE_TARGET} as final 98COPY --from=openssl /opt/openssl /opt/openssl 99# Install patchelf 100ADD ./common/install_patchelf.sh install_patchelf.sh 101RUN bash ./install_patchelf.sh && rm install_patchelf.sh 102# Install Anaconda 103COPY --from=conda /opt/conda /opt/conda 104# Install python 105COPY --from=python /opt/python /opt/python 106COPY --from=python /opt/_internal /opt/_internal 107ENV PATH=/opt/conda/bin:/usr/local/cuda/bin:$PATH 108