Add build-conda-images.yml in pytorch/pytorch (#128563)

* Add build-conda-images.yml in pytorch/pytorch * Add build_docker.sh to run .github/workflows/build-conda-images.yml * Create build-conda-images.yml * Delete .github/build-conda-images.yml Add build-conda-images.yml in pytorch/pytorch (#128563) (#128700) * Small adjustment for docker build * build_docker * Add docker file * Add common Add build-conda-images.yml in pytorch/pytorch (#128563) (#128962) * Add jni.h path * test * test * fixpath Update script path moving conda builds from builder to pytorch (#129167) * Add jni.h path * test * test * fixpath * test * test * test * test * test * test * test * test * test * test * tets * test * test Cleanup build docker images (#129273) * Cleanup docker move * add docker file Add calculate-docker-image when WITH_PUSH is true Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml test Update build-conda-images.yml Update build-conda-images.yml test test test test Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build-conda-images.yml Update build_docker_conda.sh test test test test test test test test test test test test test test test test test keep_both_tag test_conda_builds test test test test test test test test test
pytorch · Jul 20, 2024 · 45af303 · 45af303
1 parent 637ab85
commit 45af303
Show file tree

Hide file tree

Showing 7 changed files with 520 additions and 0 deletions.
diff --git a/.ci/docker/common/install_conda_docker.sh b/.ci/docker/common/install_conda_docker.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+set -ex
+
+# Anaconda
+# Latest anaconda is using openssl-3 which is incompatible with all currently published versions of git
+# Which are using openssl-1.1.1, see https://anaconda.org/anaconda/git/files?version=2.40.1 for example
+MINICONDA_URL=https://repo.anaconda.com/miniconda/Miniconda3-py311_23.5.2-0-Linux-x86_64.sh
+wget -q $MINICONDA_URL
+# NB: Manually invoke bash per https://github.com/conda/conda/issues/10431
+bash $(basename "$MINICONDA_URL") -b -p /opt/conda
+rm $(basename "$MINICONDA_URL")
+export PATH=/opt/conda/bin:$PATH
+# See https://github.com/pytorch/builder/issues/1473
+# Pin conda to 23.5.2 as it's the last one compatible with openssl-1.1.1
+conda install -y conda=23.5.2 conda-build anaconda-client git ninja
+# The cmake version here needs to match with the minimum version of cmake
+# supported by PyTorch (3.18). There is only 3.18.2 on anaconda
+/opt/conda/bin/pip3 install cmake==3.18.2
+conda remove -y --force patchelf
diff --git a/.ci/docker/common/install_cuda.sh b/.ci/docker/common/install_cuda.sh
@@ -0,0 +1,235 @@
+#!/bin/bash
+
+set -ex
+
+CUDNN_VERSION=9.1.0.70
+
+function install_cusparselt_040 {
+ # cuSparseLt license: https://docs.nvidia.com/cuda/cusparselt/license.html
+ mkdir tmp_cusparselt && pushd tmp_cusparselt
+ wget -q https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/libcusparse_lt-linux-x86_64-0.4.0.7-archive.tar.xz
+ tar xf libcusparse_lt-linux-x86_64-0.4.0.7-archive.tar.xz
+ cp -a libcusparse_lt-linux-x86_64-0.4.0.7-archive/include/* /usr/local/cuda/include/
+ cp -a libcusparse_lt-linux-x86_64-0.4.0.7-archive/lib/* /usr/local/cuda/lib64/
+ popd
+ rm -rf tmp_cusparselt
+}
+
+function install_cusparselt_052 {
+ # cuSparseLt license: https://docs.nvidia.com/cuda/cusparselt/license.html
+ mkdir tmp_cusparselt && pushd tmp_cusparselt
+ wget -q https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/libcusparse_lt-linux-x86_64-0.5.2.1-archive.tar.xz
+ tar xf libcusparse_lt-linux-x86_64-0.5.2.1-archive.tar.xz
+ cp -a libcusparse_lt-linux-x86_64-0.5.2.1-archive/include/* /usr/local/cuda/include/
+ cp -a libcusparse_lt-linux-x86_64-0.5.2.1-archive/lib/* /usr/local/cuda/lib64/
+ popd
+ rm -rf tmp_cusparselt
+}
+
+function install_118 {
+ echo "Installing CUDA 11.8 and cuDNN ${CUDNN_VERSION} and NCCL 2.15 and cuSparseLt-0.4.0"
+ rm -rf /usr/local/cuda-11.8 /usr/local/cuda
+ # install CUDA 11.8.0 in the same container
+ wget -q https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
+ chmod +x cuda_11.8.0_520.61.05_linux.run
+ ./cuda_11.8.0_520.61.05_linux.run --toolkit --silent
+ rm -f cuda_11.8.0_520.61.05_linux.run
+ rm -f /usr/local/cuda && ln -s /usr/local/cuda-11.8 /usr/local/cuda
+
+ # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+ mkdir tmp_cudnn && cd tmp_cudnn
+ wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-${CUDNN_VERSION}_cuda11-archive.tar.xz -O cudnn-linux-x86_64-${CUDNN_VERSION}_cuda11-archive.tar.xz
+ tar xf cudnn-linux-x86_64-${CUDNN_VERSION}_cuda11-archive.tar.xz
+ cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda11-archive/include/* /usr/local/cuda/include/
+ cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda11-archive/lib/* /usr/local/cuda/lib64/
+ cd ..
+ rm -rf tmp_cudnn
+
+ # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
+ # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
+ git clone -b v2.20.5-1 --depth 1 https://github.com/NVIDIA/nccl.git
+ cd nccl && make -j src.build
+ cp -a build/include/* /usr/local/cuda/include/
+ cp -a build/lib/* /usr/local/cuda/lib64/
+ cd ..
+ rm -rf nccl
+
+ install_cusparselt_040
+
+ ldconfig
+}
+
+function install_121 {
+ echo "Installing CUDA 12.1 and cuDNN ${CUDNN_VERSION} and NCCL 2.20.5 and cuSparseLt-0.5.2"
+ rm -rf /usr/local/cuda-12.1 /usr/local/cuda
+ # install CUDA 12.1.0 in the same container
+ wget -q https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run
+ chmod +x cuda_12.1.1_530.30.02_linux.run
+ ./cuda_12.1.1_530.30.02_linux.run --toolkit --silent
+ rm -f cuda_12.1.1_530.30.02_linux.run
+ rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.1 /usr/local/cuda
+
+ # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+ mkdir tmp_cudnn && cd tmp_cudnn
+ wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz -O cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
+ tar xf cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
+ cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/include/* /usr/local/cuda/include/
+ cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/lib/* /usr/local/cuda/lib64/
+ cd ..
+ rm -rf tmp_cudnn
+
+ # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
+ # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
+ git clone -b v2.20.5-1 --depth 1 https://github.com/NVIDIA/nccl.git
+ cd nccl && make -j src.build
+ cp -a build/include/* /usr/local/cuda/include/
+ cp -a build/lib/* /usr/local/cuda/lib64/
+ cd ..
+ rm -rf nccl
+
+ install_cusparselt_052
+
+ ldconfig
+}
+
+function install_124 {
+ echo "Installing CUDA 12.4 and cuDNN ${CUDNN_VERSION} and NCCL 2.20.5 and cuSparseLt-0.5.2"
+ rm -rf /usr/local/cuda-12.4 /usr/local/cuda
+ # install CUDA 12.4.0 in the same container
+ wget -q https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
+ chmod +x cuda_12.4.0_550.54.14_linux.run
+ ./cuda_12.4.0_550.54.14_linux.run --toolkit --silent
+ rm -f cuda_12.4.0_550.54.14_linux.run
+ rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.4 /usr/local/cuda
+
+ # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+ mkdir tmp_cudnn && cd tmp_cudnn
+ wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz -O cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
+ tar xf cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
+ cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/include/* /usr/local/cuda/include/
+ cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/lib/* /usr/local/cuda/lib64/
+ cd ..
+ rm -rf tmp_cudnn
+
+ # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
+ # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
+ git clone -b v2.20.5-1 --depth 1 https://github.com/NVIDIA/nccl.git
+ cd nccl && make -j src.build
+ cp -a build/include/* /usr/local/cuda/include/
+ cp -a build/lib/* /usr/local/cuda/lib64/
+ cd ..
+ rm -rf nccl
+
+ install_cusparselt_052
+
+ ldconfig
+}
+
+function prune_118 {
+ echo "Pruning CUDA 11.8 and cuDNN"
+ #####################################################################################
+ # CUDA 11.8 prune static libs
+ #####################################################################################
+ export NVPRUNE="/usr/local/cuda-11.8/bin/nvprune"
+ export CUDA_LIB_DIR="/usr/local/cuda-11.8/lib64"
+
+ export GENCODE="-gencode arch=compute_35,code=sm_35 -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+ export GENCODE_CUDNN="-gencode arch=compute_35,code=sm_35 -gencode arch=compute_37,code=sm_37 -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+
+ if [[ -n "$OVERRIDE_GENCODE" ]]; then
+ export GENCODE=$OVERRIDE_GENCODE
+ fi
+
+ # all CUDA libs except CuDNN and CuBLAS (cudnn and cublas need arch 3.7 included)
+ ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis" \
+ | xargs -I {} bash -c \
+ "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
+
+ # prune CuDNN and CuBLAS
+ $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
+ $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
+
+ #####################################################################################
+ # CUDA 11.8 prune visual tools
+ #####################################################################################
+ export CUDA_BASE="/usr/local/cuda-11.8/"
+ rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2022.3.0 $CUDA_BASE/nsight-systems-2022.4.2/
+}
+
+function prune_121 {
+ echo "Pruning CUDA 12.1"
+ #####################################################################################
+ # CUDA 12.1 prune static libs
+ #####################################################################################
+ export NVPRUNE="/usr/local/cuda-12.1/bin/nvprune"
+ export CUDA_LIB_DIR="/usr/local/cuda-12.1/lib64"
+
+ export GENCODE="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+ export GENCODE_CUDNN="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+
+ if [[ -n "$OVERRIDE_GENCODE" ]]; then
+ export GENCODE=$OVERRIDE_GENCODE
+ fi
+
+ # all CUDA libs except CuDNN and CuBLAS
+ ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis" \
+ | xargs -I {} bash -c \
+ "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
+
+ # prune CuDNN and CuBLAS
+ $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
+ $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
+
+ #####################################################################################
+ # CUDA 12.1 prune visual tools
+ #####################################################################################
+ export CUDA_BASE="/usr/local/cuda-12.1/"
+ rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2023.1.0 $CUDA_BASE/nsight-systems-2023.1.2/
+}
+
+function prune_124 {
+ echo "Pruning CUDA 12.4"
+ #####################################################################################
+ # CUDA 12.4 prune static libs
+ #####################################################################################
+ export NVPRUNE="/usr/local/cuda-12.4/bin/nvprune"
+ export CUDA_LIB_DIR="/usr/local/cuda-12.4/lib64"
+
+ export GENCODE="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+ export GENCODE_CUDNN="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
+
+ if [[ -n "$OVERRIDE_GENCODE" ]]; then
+ export GENCODE=$OVERRIDE_GENCODE
+ fi
+
+ # all CUDA libs except CuDNN and CuBLAS
+ ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis" \
+ | xargs -I {} bash -c \
+ "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
+
+ # prune CuDNN and CuBLAS
+ $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
+ $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
+
+ #####################################################################################
+ # CUDA 12.1 prune visual tools
+ #####################################################################################
+ export CUDA_BASE="/usr/local/cuda-12.4/"
+ rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2024.1.0 $CUDA_BASE/nsight-systems-2023.4.4/
+}
+
+# idiomatic parameter and option handling in sh
+while test $# -gt 0
+do
+ case "$1" in
+ 11.8) install_118; prune_118
+ ;;
+ 12.1) install_121; prune_121
+ ;;
+ 12.4) install_124; prune_124
+ ;;
+ *) echo "bad argument $1"; exit 1
+ ;;
+ esac
+ shift
+done
diff --git a/.ci/docker/common/install_mnist.sh b/.ci/docker/common/install_mnist.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+set -ex
+
+mkdir -p /usr/local/mnist/
+
+cd /usr/local/mnist
+
+for img in train-images-idx3-ubyte.gz train-labels-idx1-ubyte.gz t10k-images-idx3-ubyte.gz t10k-labels-idx1-ubyte.gz; do
+ wget -q https://ossci-datasets.s3.amazonaws.com/mnist/$img
+ gzip -d $img
+done
diff --git a/.ci/docker/common/install_patchelf.sh b/.ci/docker/common/install_patchelf.sh
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+set -ex
+
+# Pin the version to latest release 0.17.2, building newer commit starts
+# to fail on the current image
+git clone -b 0.17.2 --single-branch https://github.com/NixOS/patchelf
+cd patchelf
+sed -i 's/serial/parallel/g' configure.ac
+./bootstrap.sh
+./configure
+make
+make install
+cd ..
+rm -rf patchelf
diff --git a/.ci/docker/conda/Dockerfile b/.ci/docker/conda/Dockerfile
@@ -0,0 +1,101 @@
+ARG CUDA_VERSION=10.2
+ARG BASE_TARGET=cuda${CUDA_VERSION}
+FROM centos:7 as base
+
+ENV LC_ALL en_US.UTF-8
+ENV LANG en_US.UTF-8
+ENV LANGUAGE en_US.UTF-8
+
+ARG DEVTOOLSET_VERSION=9
+RUN sed -i s/mirror.centos.org/vault.centos.org/g /etc/yum.repos.d/*.repo
+RUN sed -i s/^#.*baseurl=http/baseurl=http/g /etc/yum.repos.d/*.repo
+RUN sed -i s/^mirrorlist=http/#mirrorlist=http/g /etc/yum.repos.d/*.repo
+RUN yum update -y
+RUN yum install -y wget curl perl util-linux xz bzip2 git patch which unzip
+# Just add everything as a safe.directory for git since these will be used in multiple places with git
+RUN git config --global --add safe.directory '*'
+RUN yum install -y yum-utils centos-release-scl
+RUN yum-config-manager --enable rhel-server-rhscl-7-rpms
+RUN sed -i s/mirror.centos.org/vault.centos.org/g /etc/yum.repos.d/*.repo
+RUN sed -i s/^#.*baseurl=http/baseurl=http/g /etc/yum.repos.d/*.repo
+RUN sed -i s/^mirrorlist=http/#mirrorlist=http/g /etc/yum.repos.d/*.repo
+RUN yum install -y devtoolset-${DEVTOOLSET_VERSION}-gcc devtoolset-${DEVTOOLSET_VERSION}-gcc-c++ devtoolset-${DEVTOOLSET_VERSION}-gcc-gfortran devtoolset-${DEVTOOLSET_VERSION}-binutils
+# EPEL for cmake
+RUN wget http://dl.fedoraproject.org/pub/epel/epel-release-latest-7.noarch.rpm && \
+ rpm -ivh epel-release-latest-7.noarch.rpm && \
+ rm -f epel-release-latest-7.noarch.rpm
+# cmake
+RUN yum install -y cmake3 && \
+ ln -s /usr/bin/cmake3 /usr/bin/cmake
+ENV PATH=/opt/rh/devtoolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
+ENV LD_LIBRARY_PATH=/opt/rh/devtoolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/devtoolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH
+
+RUN yum install -y autoconf aclocal automake make sudo
+RUN rm -rf /usr/local/cuda-*
+
+FROM base as patchelf
+# Install patchelf
+ADD ./common/install_patchelf.sh install_patchelf.sh
+RUN bash ./install_patchelf.sh && rm install_patchelf.sh && cp $(which patchelf) /patchelf
+
+FROM base as openssl
+# Install openssl
+ADD ./common/install_openssl.sh install_openssl.sh
+RUN bash ./install_openssl.sh && rm install_openssl.sh
+
+FROM base as conda
+# Install Anaconda
+ADD ./common/install_conda_docker.sh install_conda.sh
+RUN bash ./install_conda.sh && rm install_conda.sh
+
+# Install CUDA
+FROM base as cuda
+ARG CUDA_VERSION=10.2
+RUN rm -rf /usr/local/cuda-*
+ADD ./common/install_cuda.sh install_cuda.sh
+ENV CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
+# Preserve CUDA_VERSION for the builds
+ENV CUDA_VERSION=${CUDA_VERSION}
+# Make things in our path by default
+ENV PATH=/usr/local/cuda-${CUDA_VERSION}/bin:$PATH
+
+FROM cuda as cuda11.8
+RUN bash ./install_cuda.sh 11.8
+ENV DESIRED_CUDA=11.8
+
+FROM cuda as cuda12.1
+RUN bash ./install_cuda.sh 12.1
+ENV DESIRED_CUDA=12.1
+
+FROM cuda as cuda12.4
+RUN bash ./install_cuda.sh 12.4
+ENV DESIRED_CUDA=12.4
+
+# Install MNIST test data
+FROM base as mnist
+ADD ./common/install_mnist.sh install_mnist.sh
+RUN bash ./install_mnist.sh
+
+FROM base as all_cuda
+COPY --from=cuda11.8 /usr/local/cuda-11.8 /usr/local/cuda-11.8
+COPY --from=cuda12.1 /usr/local/cuda-12.1 /usr/local/cuda-12.1
+COPY --from=cuda12.4 /usr/local/cuda-12.4 /usr/local/cuda-12.4
+
+# Final step
+FROM ${BASE_TARGET} as final
+COPY --from=openssl /opt/openssl /opt/openssl
+COPY --from=patchelf /patchelf /usr/local/bin/patchelf
+COPY --from=conda /opt/conda /opt/conda
+
+# Add jni.h for java host build.
+COPY ./common/install_jni.sh install_jni.sh
+COPY ./java/jni.h jni.h
+RUN bash ./install_jni.sh && rm install_jni.sh
+
+ENV PATH /opt/conda/bin:$PATH
+COPY --from=mnist /usr/local/mnist /usr/local/mnist
+RUN rm -rf /usr/local/cuda
+RUN chmod o+rw /usr/local
+RUN touch /.condarc && \
+ chmod o+rw /.condarc && \
+ chmod -R o+rw /opt/conda