[CI] Build ARM CUDA container, part 3

.github/workflows/containers.yml

@@ -65,6 +65,8 @@ jobs:
             runner: linux-arm64-cpu
           - image_repo: xgb-ci.gpu_build_rockylinux8_aarch64
             runner: linux-arm64-cpu
+          - image_repo: xgb-ci.gpu_aarch64
+            runner: linux-arm64-cpu
     steps:
       - name: Workflow trigger information
         run: |

containers/ci_container.yml

@@ -55,6 +55,15 @@ xgb-ci.gpu:
     CUDA_VERSION: *cuda_version
     NCCL_VERSION: *nccl_version
     RAPIDS_VERSION: *rapids_version
+    ARCH: x86_64
+
+xgb-ci.gpu_aarch64:
+  container_def: gpu
+  build_args:
+    CUDA_VERSION: *cuda_version
+    NCCL_VERSION: *nccl_version
+    RAPIDS_VERSION: *rapids_version
+    ARCH: aarch64
 
 xgb-ci.clang_tidy:
   container_def: clang_tidy

containers/dockerfile/Dockerfile.gpu

@@ -4,7 +4,7 @@ ARG CUDA_VERSION
 ARG RAPIDS_VERSION
   # Should be first 4 digits (e.g. 24.06)
 ARG NCCL_VERSION
-ARG RAPIDSAI_CONDA_CHANNEL="rapidsai"
+ARG ARCH=x86_64
 ARG MINIFORGE_VERSION=24.9.2-0
 ARG PYTHON_VERSION=3.10
 
@@ -17,21 +17,22 @@ ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64/
 
 # Install all basic requirements
 RUN \
+    { [ $ARCH = "aarch64" ] && export CUDA_REPO_ARCH="sbsa" || export CUDA_REPO_ARCH="x86_64"; } && \
     export NCCL_VERSION=$NCCL_VERSION && \
-    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \
+    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${CUDA_REPO_ARCH}/3bf863cc.pub && \
     apt-get update && \
     apt-get install -y wget unzip bzip2 libgomp1 build-essential openjdk-8-jdk-headless && \
     apt-get install "libnccl2=${NCCL_VERSION}+cuda12.9" \
       "libnccl-dev=${NCCL_VERSION}+cuda12.9" -y --allow-change-held-packages && \
     # Miniforge
-    wget -nv -O conda.sh https://github.com/conda-forge/miniforge/releases/download/$MINIFORGE_VERSION/Miniforge3-$MINIFORGE_VERSION-Linux-x86_64.sh && \
+    wget -nv -O conda.sh https://github.com/conda-forge/miniforge/releases/download/$MINIFORGE_VERSION/Miniforge3-$MINIFORGE_VERSION-Linux-${ARCH}.sh && \
     bash conda.sh -b -p /opt/miniforge
 
 # Create new Conda environment with cuDF, Dask, and cuPy
 RUN \
     export NCCL_SHORT_VER=$(echo "$NCCL_VERSION" | cut -d "-" -f 1) && \
     export CUDA_SHORT_VER=$(echo "$CUDA_VERSION" | grep -o -E '[0-9]+\.[0-9]') && \
-    mamba create -y -n gpu_test -c ${RAPIDSAI_CONDA_CHANNEL} -c conda-forge -c nvidia \
+    mamba create -y -n gpu_test -c rapidsai -c conda-forge -c nvidia \
         python=$PYTHON_VERSION "cudf=$RAPIDS_VERSION.*" "rmm=$RAPIDS_VERSION.*" cuda-version=$CUDA_SHORT_VER \
         "nccl=${NCCL_SHORT_VER}" \
         dask \
@@ -44,7 +45,8 @@ RUN \
 
 # Install lightweight sudo (not bound to TTY)
 RUN set -ex; \
-    wget -nv -O /usr/local/bin/gosu "https://github.com/tianon/gosu/releases/download/$GOSU_VERSION/gosu-amd64" && \
+    { [ $ARCH = "aarch64" ] && export GOSU_ARCH="arm64" || export GOSU_ARCH="amd64"; } && \
+    wget -nv -O /usr/local/bin/gosu "https://github.com/tianon/gosu/releases/download/$GOSU_VERSION/gosu-${GOSU_ARCH}" && \
     chmod +x /usr/local/bin/gosu && \
     gosu nobody true
 

containers/update_rapids.sh

@@ -4,13 +4,9 @@ set -euo pipefail
 
 LATEST_RAPIDS_VERSION=$(gh api repos/rapidsai/cuml/releases/latest --jq '.name' | sed -e 's/^v\([[:digit:]]\+\.[[:digit:]]\+\).*/\1/')
 echo "LATEST_RAPIDS_VERSION = $LATEST_RAPIDS_VERSION"
-DEV_RAPIDS_VERSION=$(date +%Y-%m-%d -d "20${LATEST_RAPIDS_VERSION//./-}-01 + 2 month" | cut -c3-7 | tr - .)
-echo "DEV_RAPIDS_VERSION = $DEV_RAPIDS_VERSION"
 
 DIR_PATH=$( cd "$(dirname "${BASH_SOURCE[0]}")" ; pwd -P )
 CONTAINER_YAML="$DIR_PATH/ci_container.yml"
 
 sed -i "s/\&rapids_version \"[[:digit:]]\+\.[[:digit:]]\+\"/\&rapids_version \"${LATEST_RAPIDS_VERSION}\"/" \
   "$CONTAINER_YAML"
-sed -i "s/\&dev_rapids_version \"[[:digit:]]\+\.[[:digit:]]\+\"/\&dev_rapids_version \"${DEV_RAPIDS_VERSION}\"/" \
-  "$CONTAINER_YAML"