[backport] Update CUDA docker image and NCCL. (#8139) (#8162)

* Update CUDA docker image and NCCL. (#8139)

* Rest of the CI.

* CPU test dependencies.

tests/ci_build/Dockerfile.cpu

@@ -10,13 +10,13 @@ RUN \
     apt-get install -y software-properties-common && \
     add-apt-repository ppa:ubuntu-toolchain-r/test && \
     apt-get update && \
-    apt-get install -y tar unzip wget git build-essential doxygen graphviz llvm libasan2 libidn11 ninja-build gcc-8 g++-8 && \
+    apt-get install -y tar unzip wget git build-essential doxygen graphviz llvm libasan2 libidn11 ninja-build gcc-8 g++-8 openjdk-8-jdk-headless && \
     # CMake
     wget -nv -nc https://cmake.org/files/v3.14/cmake-3.14.0-Linux-x86_64.sh --no-check-certificate && \
     bash cmake-3.14.0-Linux-x86_64.sh --skip-license --prefix=/usr && \
     # Python
-    wget -O Miniconda3.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
-    bash Miniconda3.sh -b -p /opt/python
+    wget https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh && \
+    bash Mambaforge-Linux-x86_64.sh -b -p /opt/python
 
 ENV PATH=/opt/python/bin:$PATH
 ENV CC=gcc-8
@@ -24,10 +24,11 @@ ENV CXX=g++-8
 ENV CPP=cpp-8
 
 ENV GOSU_VERSION 1.10
+ENV JAVA_HOME /usr/lib/jvm/java-8-openjdk-amd64/
 
 # Create new Conda environment
 COPY conda_env/cpu_test.yml /scripts/
-RUN conda env create -n cpu_test --file=/scripts/cpu_test.yml
+RUN mamba env create -n cpu_test --file=/scripts/cpu_test.yml
 
 # Install lightweight sudo (not bound to TTY)
 RUN set -ex; \

tests/ci_build/Dockerfile.gpu

@@ -10,7 +10,7 @@ SHELL ["/bin/bash", "-c"]   # Use Bash as shell
 RUN \
     apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/3bf863cc.pub && \
     apt-get update && \
-    apt-get install -y wget unzip bzip2 libgomp1 build-essential && \
+    apt-get install -y wget unzip bzip2 libgomp1 build-essential openjdk-8-jdk-headless && \
     # Python
     wget -O Miniconda3.sh https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh && \
     bash Miniconda3.sh -b -p /opt/python
@@ -19,11 +19,14 @@ ENV PATH=/opt/python/bin:$PATH
 
 # Create new Conda environment with cuDF, Dask, and cuPy
 RUN \
-    conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.8 cudf=21.10* rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=21.10* dask-cudf=21.10* cupy=9.1* \
-        numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis
+    conda install -c conda-forge mamba && \
+    mamba create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
+        python=3.8 cudf=22.04* rmm=22.04* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=22.04* dask-cudf=22.04* cupy \
+        numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis \
+        pyspark cloudpickle cuda-python=11.7.0
 
 ENV GOSU_VERSION 1.10
+ENV JAVA_HOME /usr/lib/jvm/java-8-openjdk-amd64/
 
 # Install lightweight sudo (not bound to TTY)
 RUN set -ex; \

tests/ci_build/Dockerfile.gpu_build

@@ -24,7 +24,7 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.7.5-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
     apt-get update && \
     apt-get install -y --allow-downgrades --allow-change-held-packages libnccl2=${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-dev=${NCCL_VERSION}+cuda${CUDA_SHORT}
 

tests/ci_build/Dockerfile.gpu_build_centos7

@@ -4,7 +4,6 @@ ARG CUDA_VERSION_ARG
 
 # Install all basic requirements
 RUN \
-    rpm --erase gpg-pubkey-7fa2af80* && \
     curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
         > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
     yum install -y epel-release centos-release-scl && \
@@ -22,7 +21,7 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.7.3-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
     wget -nv -nc https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     rpm -i nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     yum -y update && \

tests/ci_build/Dockerfile.jvm_gpu_build

@@ -4,7 +4,6 @@ ARG CUDA_VERSION_ARG
 
 # Install all basic requirements
 RUN \
-    rpm --erase gpg-pubkey-7fa2af80* && \
     curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
         > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
     yum install -y epel-release centos-release-scl && \
@@ -25,12 +24,10 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.8.3-1 && \
-    wget -nv -nc https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
-    rpm -i nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
+    export NCCL_VERSION=2.13.4-1 && \
+    yum-config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo && \
     yum -y update && \
-    yum install -y libnccl-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-devel-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-static-${NCCL_VERSION}+cuda${CUDA_SHORT} && \
-    rm -f nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm;
+    yum install -y libnccl-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-devel-${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-static-${NCCL_VERSION}+cuda${CUDA_SHORT}
 
 ENV PATH=/opt/python/bin:/opt/maven/bin:$PATH
 ENV CC=/opt/rh/devtoolset-8/root/usr/bin/gcc

tests/ci_build/Dockerfile.rmm

@@ -18,7 +18,7 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION_ARG | grep -o -E '[0-9]+\.[0-9]'` && \
-    export NCCL_VERSION=2.7.5-1 && \
+    export NCCL_VERSION=2.13.4-1 && \
     apt-get update && \
     apt-get install -y --allow-downgrades --allow-change-held-packages libnccl2=${NCCL_VERSION}+cuda${CUDA_SHORT} libnccl-dev=${NCCL_VERSION}+cuda${CUDA_SHORT}
 
@@ -27,7 +27,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with RMM
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.9 rmm=22.06* cudatoolkit=$CUDA_VERSION_ARG cmake
+        python=3.9 rmm=22.04* cudatoolkit=$CUDA_VERSION_ARG cmake
 
 ENV GOSU_VERSION 1.10
 

tests/ci_build/build_python_wheels.sh

@@ -29,13 +29,15 @@ if [[ "$platform_id" == macosx_* ]]; then
         setup_env_var='CIBW_TARGET_OSX_ARM64=1'  # extra flag to be passed to setup.py
         export PYTHON_CROSSENV=1
         export MACOSX_DEPLOYMENT_TARGET=12.0
-        OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-arm64/llvm-openmp-11.1.0-hf3c4609_1.tar.bz2"
+        #OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-arm64/llvm-openmp-11.1.0-hf3c4609_1.tar.bz2"
+        OPENMP_URL="https://xgboost-ci-jenkins-artifacts.s3.us-west-2.amazonaws.com/llvm-openmp-11.1.0-hf3c4609_1-osx-arm64.tar.bz2"
     elif [[ "$platform_id" == macosx_x86_64 ]]; then
         # MacOS, Intel
         wheel_tag=macosx_10_15_x86_64.macosx_11_0_x86_64.macosx_12_0_x86_64
         cpython_ver=37
         export MACOSX_DEPLOYMENT_TARGET=10.13
-        OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-64/llvm-openmp-11.1.0-hda6cdc1_1.tar.bz2"
+        #OPENMP_URL="https://anaconda.org/conda-forge/llvm-openmp/11.1.0/download/osx-64/llvm-openmp-11.1.0-hda6cdc1_1.tar.bz2"
+        OPENMP_URL="https://xgboost-ci-jenkins-artifacts.s3.us-west-2.amazonaws.com/llvm-openmp-11.1.0-hda6cdc1_1-osx-64.tar.bz2"
     else
         echo "Platform not supported: $platform_id"
         exit 3

tests/ci_build/conda_env/cpu_test.yml

@@ -30,15 +30,13 @@ dependencies:
 - jsonschema
 - boto3
 - awscli
-- numba
-- llvmlite
 - py-ubjson
 - cffi
 - pyarrow
-- protobuf<=3.20
+- protobuf
+- pyspark>=3.3.0
+- cloudpickle
+- shap
+- modin
 - pip:
-  - shap
-  - ipython                     # required by shap at import time.
-  - sphinx_rtd_theme
   - datatable
-  - modin[all]

tests/ci_build/entrypoint.sh

@@ -20,9 +20,9 @@ else
 fi
 
 if [[ -n $CI_BUILD_UID ]] && [[ -n $CI_BUILD_GID ]]; then
-    groupadd -o -g "${CI_BUILD_GID}" "${CI_BUILD_GROUP}"
+    groupadd -o -g "${CI_BUILD_GID}" "${CI_BUILD_GROUP}" || true
     useradd -o -m -g "${CI_BUILD_GID}" -u "${CI_BUILD_UID}" \
-        "${CI_BUILD_USER}"
+        "${CI_BUILD_USER}" || true
     export HOME="/home/${CI_BUILD_USER}"
     shopt -s dotglob
     cp -r /root/* "$HOME/"

tests/ci_build/rename_whl.py

@@ -42,4 +42,4 @@ with cd(dirname):
 
     filesize = os.path.getsize(new_name) / 1024 / 1024  # MB
     msg = f"Limit of wheel size set by PyPI is exceeded. {new_name}: {filesize}"
-    assert filesize <= 200, msg
+    assert filesize <= 300, msg