update

.ci/docker/Dockerfile

@@ -1,43 +1,45 @@
-FROM ubuntu:20.04
+ARG BASE_IMAGE
+ARG CUDA_VERSION
+ARG UBUNTU_VERSIOn
+FROM ${BASE_IMAGE}
 
 ENV DEBIAN_FRONTEND noninteractive
 
 # Install common dependencies (so that this step can be cached separately)
-COPY ./install_base.sh install_base.sh
+COPY ./common/install_base.sh install_base.sh
 RUN bash ./install_base.sh && rm install_base.sh
 
 # Install user
-COPY ./install_user.sh install_user.sh
+COPY ./common/install_user.sh install_user.sh
 RUN bash ./install_user.sh && rm install_user.sh
 
-COPY ./install_docs_reqs.sh install_docs_reqs.sh
+COPY ./common/install_docs_reqs.sh install_docs_reqs.sh
 RUN bash ./install_docs_reqs.sh && rm install_docs_reqs.sh
 
-# Install conda and other packages (e.g., numpy, pytest)
+# Install conda and other packages
 ENV ANACONDA_PYTHON_VERSION=3.10
 ENV CONDA_CMAKE yes
 ENV DOCS yes
 ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
-ENV CUDA_VERSION=11.8
 COPY ./requirements.txt /opt/conda/
-COPY ./install_conda.sh install_conda.sh
-COPY ./common_utils.sh common_utils.sh
+COPY ./common/install_conda.sh install_conda.sh
+COPY ./common/common_utils.sh common_utils.sh
 RUN bash ./install_conda.sh && rm install_conda.sh common_utils.sh /opt/conda/requirements.txt
 
 # Install gcc
 # ENV GCC_VERSION 9
 # COPY ./install_gcc.sh install_gcc.sh
 # RUN bash ./install_gcc.sh && rm install_gcc.sh
 
-# Install cuda and cudnn
-RUN wget -q https://raw.githubusercontent.com/pytorch/builder/main/install_cuda.sh -O install_cuda.sh
+# Install cuda
+RUN wget -q https://raw.githubusercontent.com/pytorch/builder/main/common/install_cuda.sh -O install_cuda.sh
 RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh
 ENV DESIRED_CUDA ${CUDA_VERSION}
 ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
 
 # Install CUDNN
 ENV CUDNN_VERSION=8
-COPY ./install_cudnn.sh install_cudnn.sh
+COPY ./common/install_cudnn.sh install_cudnn.sh
 RUN if [ "${CUDNN_VERSION}" -eq 8 ]; then bash install_cudnn.sh; fi
 RUN rm install_cudnn.sh
 

.ci/docker/build.sh

@@ -10,6 +10,10 @@ set -exu
 IMAGE_NAME="$1"
 shift
 
+export CUDA_VERSION="12.1.1"
+export UBUNTU_VERSION="20.04"
+
+export BASE_IMAGE="nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04"
 echo "Building ${IMAGE_NAME} Docker image"
 
 docker build \

.ci/docker/install_conda.sh → .ci/docker/common/install_conda.sh

@@ -52,7 +52,7 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
   # and focal cannot find conda mkl with stock cmake, so we need a cmake from conda
   conda_install cmake
 
-  conda_install magma-cuda$(TMP=${CUDA_VERSION/./};echo ${TMP%.*[0-9]}) -c pytorch
+  conda_install "magma-cuda$(TMP=${CUDA_VERSION/./};echo ${TMP%.*[0-9]})" -c pytorch
 
   # Install some other packages, including those needed for Python test reporting
   pip_install -r /opt/conda/requirements.txt