Build ARM CUDA container, part 2 (#56)

Author: hcho3

.github/workflows/containers.yml

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

containers/ci_container.yml

@@ -16,6 +16,15 @@ xgb-ci.gpu_build_rockylinux8:
     CUDA_VERSION: *cuda_version
     NCCL_VERSION: *nccl_version
     RAPIDS_VERSION: *rapids_version
+    ARCH: x86_64
+
+xgb-ci.gpu_build_rockylinux8_aarch64:
+  container_def: gpu_build_rockylinux8
+  build_args:
+    CUDA_VERSION: *cuda_version
+    NCCL_VERSION: *nccl_version
+    RAPIDS_VERSION: *rapids_version
+    ARCH: aarch64
 
 xgb-ci.gpu_build_cuda13_rockylinux8:
   container_def: gpu_build_cuda13_rockylinux8

containers/dockerfile/Dockerfile.gpu_build_rockylinux8

@@ -3,6 +3,7 @@ FROM nvcr.io/nvidia/cuda:$CUDA_VERSION-devel-rockylinux8
 ARG CUDA_VERSION
 ARG NCCL_VERSION
 ARG RAPIDS_VERSION
+ARG ARCH=x86_64
 ARG MINIFORGE_VERSION=24.9.2-0
 ARG CMAKE_VERSION=3.31.2
 
@@ -17,24 +18,26 @@ ENV GOSU_VERSION=1.10
 
 # Install all basic requirements
 RUN \
-    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/D42D0685.pub | sed '/^Version/d' \
+    { [ $ARCH = "aarch64" ] && export CUDA_REPO_ARCH="sbsa" || export CUDA_REPO_ARCH="x86_64"; } && \
+    curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/D42D0685.pub | sed '/^Version/d' \
         > /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
     dnf -y update && \
     dnf -y install dnf-plugins-core && \
     dnf config-manager --set-enabled powertools && \
     dnf install -y tar unzip wget xz git which ninja-build gcc-toolset-10-gcc gcc-toolset-10-binutils gcc-toolset-10-gcc-c++ && \
     # 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 && \
     /opt/miniforge/bin/python -m pip install awscli && \
     # CMake
-    wget -nv -O cmake.sh https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-x86_64.sh && \
+    wget -nv -O cmake.sh https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-${ARCH}.sh && \
     bash cmake.sh --skip-license --prefix=/usr
 
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
+    { [ $ARCH = "aarch64" ] && export CUDA_REPO_ARCH="sbsa" || export CUDA_REPO_ARCH="x86_64"; } && \
     export NCCL_VERSION=$NCCL_VERSION && \
-    dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && \
+    dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo && \
     dnf -y update && \
     dnf install -y libnccl-${NCCL_VERSION}+cuda12.9 libnccl-devel-${NCCL_VERSION}+cuda12.9
 
@@ -72,7 +75,8 @@ RUN git clone -b branch-${RAPIDS_VERSION} https://github.com/rapidsai/rmm.git --
 
 # Install lightweight sudo (not bound to TTY)
 RUN set -ex; \
-    wget -nv -nc -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 -nc -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