Skip to content

Commit 98bed31

Browse files
hcho3hcho3
authored
Build VM image and containers targeting CUDA arm64 (#54)
* Install CUDA driver into VM image for ARM64 * Use rhel8 repo for CUDA * Build CUDA 13 image with ARM64
1 parent db07882 commit 98bed31

9 files changed

Lines changed: 65 additions & 13 deletions

‎.github/workflows/containers.yml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ jobs:
6161
runner: linux-arm64-cpu
6262
- image_repo: xgb-ci.manylinux_2_28_aarch64
6363
runner: linux-arm64-cpu
64+
- image_repo: xgb-ci.gpu_build_cuda13_rockylinux8_aarch64
65+
runner: linux-arm64-cpu
6466
steps:
6567
- name: Workflow trigger information
6668
run: |

‎containers/ci_container.yml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,14 @@ xgb-ci.gpu_build_cuda13_rockylinux8:
2222
build_args:
2323
CUDA_VERSION: "13.0.0"
2424
NCCL_VERSION: *nccl_version
25+
ARCH: x86_64
26+
27+
xgb-ci.gpu_build_cuda13_rockylinux8_aarch64:
28+
container_def: gpu_build_cuda13_rockylinux8
29+
build_args:
30+
CUDA_VERSION: "13.0.0"
31+
NCCL_VERSION: *nccl_version
32+
ARCH: aarch64
2533

2634
xgb-ci.gpu_build_r_rockylinux8:
2735
container_def: gpu_build_r_rockylinux8

‎containers/dockerfile/Dockerfile.gpu_build_cuda13_rockylinux8‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@ ARG CUDA_VERSION=notset
22
FROM nvcr.io/nvidia/cuda:$CUDA_VERSION-devel-rockylinux8
33
ARG CUDA_VERSION
44
ARG NCCL_VERSION
5-
ARG MINIFORGE_VERSION=25.3.1-0
5+
ARG ARCH=x86_64
6+
ARG MINIFORGE_VERSION=25.11.0-1
67
ARG CMAKE_VERSION=4.1.0
78

89
SHELL ["/bin/bash", "-c"]
@@ -16,30 +17,33 @@ ENV GOSU_VERSION=1.10
1617

1718
# Install all basic requirements
1819
RUN \
19-
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
20+
{ [ $ARCH = "aarch64" ] && export CUDA_REPO_ARCH="sbsa" || export CUDA_REPO_ARCH="x86_64"; } && \
21+
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/D42D0685.pub | sed '/^Version/d' \
2022
> /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
2123
dnf -y update && \
2224
dnf -y install dnf-plugins-core && \
2325
dnf config-manager --set-enabled powertools && \
2426
dnf install -y tar unzip wget xz git which ninja-build gcc-toolset-10-gcc gcc-toolset-10-binutils gcc-toolset-10-gcc-c++ && \
2527
# Miniforge
26-
wget -nv -O conda.sh https://github.com/conda-forge/miniforge/releases/download/$MINIFORGE_VERSION/Miniforge3-$MINIFORGE_VERSION-Linux-x86_64.sh && \
28+
wget -nv -O conda.sh https://github.com/conda-forge/miniforge/releases/download/$MINIFORGE_VERSION/Miniforge3-$MINIFORGE_VERSION-Linux-${ARCH}.sh && \
2729
bash conda.sh -b -p /opt/miniforge && \
2830
/opt/miniforge/bin/python -m pip install awscli && \
2931
# CMake
30-
wget -nv -O cmake.sh https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-x86_64.sh && \
32+
wget -nv -O cmake.sh https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-${ARCH}.sh && \
3133
bash cmake.sh --skip-license --prefix=/usr
3234

3335
# NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
3436
RUN \
37+
{ [ $ARCH = "aarch64" ] && export CUDA_REPO_ARCH="sbsa" || export CUDA_REPO_ARCH="x86_64"; } && \
3538
export NCCL_VERSION=$NCCL_VERSION && \
36-
dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && \
39+
dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo && \
3740
dnf -y update && \
3841
dnf install -y libnccl-${NCCL_VERSION}+cuda13.0 libnccl-devel-${NCCL_VERSION}+cuda13.0 libnccl-static-${NCCL_VERSION}+cuda13.0
3942

4043
# Install lightweight sudo (not bound to TTY)
4144
RUN set -ex; \
42-
wget -nv -nc -O /usr/local/bin/gosu "https://github.com/tianon/gosu/releases/download/$GOSU_VERSION/gosu-amd64" && \
45+
{ [ $ARCH = "aarch64" ] && export GOSU_ARCH="arm64" || export GOSU_ARCH="amd64"; } && \
46+
wget -nv -nc -O /usr/local/bin/gosu "https://github.com/tianon/gosu/releases/download/$GOSU_VERSION/gosu-${GOSU_ARCH}" && \
4347
chmod +x /usr/local/bin/gosu && \
4448
gosu nobody true
4549

‎containers/dockerfile/Dockerfile.gpu_build_r_rockylinux8‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ ENV GOSU_VERSION=1.10
1818

1919
# Install all basic requirements
2020
RUN \
21-
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
21+
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/D42D0685.pub | sed '/^Version/d' \
2222
> /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
2323
dnf -y update && \
2424
dnf -y install dnf-plugins-core && \

‎containers/dockerfile/Dockerfile.gpu_build_rockylinux8‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ ENV GOSU_VERSION=1.10
1717

1818
# Install all basic requirements
1919
RUN \
20-
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/D42D0685.pub | sed '/^Version/d' \
20+
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/D42D0685.pub | sed '/^Version/d' \
2121
> /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA && \
2222
dnf -y update && \
2323
dnf -y install dnf-plugins-core && \
@@ -34,7 +34,7 @@ RUN \
3434
# NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
3535
RUN \
3636
export NCCL_VERSION=$NCCL_VERSION && \
37-
dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo && \
37+
dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && \
3838
dnf -y update && \
3939
dnf install -y libnccl-${NCCL_VERSION}+cuda12.9 libnccl-devel-${NCCL_VERSION}+cuda12.9
4040

‎containers/dockerfile/Dockerfile.jvm_gpu_build‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ RUN \
3636
# NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
3737
RUN \
3838
export NCCL_VERSION=$NCCL_VERSION && \
39-
dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo && \
39+
dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && \
4040
dnf -y update && \
4141
dnf install -y libnccl-${NCCL_VERSION}+cuda12.9 libnccl-devel-${NCCL_VERSION}+cuda12.9 libnccl-static-${NCCL_VERSION}+cuda12.9
4242

‎vm_images/linux-arm64/bootstrap.sh‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,19 @@ sudo systemctl is-active --quiet docker.service || sudo systemctl start docker.s
3131
sudo systemctl is-enabled --quiet docker.service || sudo systemctl enable docker.service
3232
sleep 10 # Docker daemon takes time to come up after installing
3333
sudo docker info
34+
35+
## Install NVIDIA Container Toolkit
36+
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
37+
&& curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
38+
sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
39+
sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
40+
sudo apt-get update
41+
sudo apt-get install -y nvidia-container-toolkit
42+
sudo nvidia-ctk runtime configure --runtime=docker
43+
sudo systemctl restart docker
44+
45+
sleep 10
46+
sudo docker run --rm --gpus all ubuntu nvidia-smi
3447
sudo systemctl stop docker
3548

3649
## Install AWS CLI v2

‎vm_images/linux-arm64/install_drivers.sh‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
#!/bin/bash
2+
set -euo pipefail
3+
4+
## Install basic tools
5+
echo 'debconf debconf/frontend select Noninteractive' | sudo debconf-set-selections
6+
sudo apt-get update
7+
sudo apt-get install -y cmake git build-essential wget ca-certificates curl unzip
8+
9+
## Install CUDA Driver 580
10+
wget -nv https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbsa/cuda-keyring_1.1-1_all.deb
11+
sudo dpkg -i cuda-keyring_1.1-1_all.deb
12+
sudo apt-get update
13+
sudo apt-get -y install nvidia-open-580
14+
rm cuda-keyring_1.1-1_all.deb

‎vm_images/linux-arm64/linux-arm64.pkr.hcl‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ packer {
99

1010
locals {
1111
ami_name_prefix = "xgboost-ci"
12-
image_name = "RunsOn worker with Ubuntu 24.04 ARM64"
12+
image_name = "RunsOn worker with Ubuntu 24.04 ARM64 + CUDA driver 580"
1313
region = "us-west-2"
1414
timestamp = regex_replace(timestamp(), "[- TZ:]", "")
1515
volume_size = 40
@@ -33,7 +33,7 @@ source "amazon-ebs" "runs-on-linux-arm64" {
3333
ami_virtualization_type = "hvm"
3434
associate_public_ip_address = true
3535
communicator = "ssh"
36-
instance_type = "c6g.4xlarge"
36+
instance_type = "g5g.xlarge"
3737
region = "${local.region}"
3838
ssh_timeout = "10m"
3939
ssh_username = "ubuntu"
@@ -63,6 +63,17 @@ build {
6363
sources = ["source.amazon-ebs.runs-on-linux-arm64"]
6464

6565
provisioner "shell" {
66-
script = "bootstrap.sh"
66+
script = "install_drivers.sh"
67+
pause_after = "30s"
68+
}
69+
70+
provisioner "shell" {
71+
expect_disconnect = true
72+
inline = ["echo 'Reboot VM'", "sudo reboot"]
73+
}
74+
75+
provisioner "shell" {
76+
pause_before = "1m0s"
77+
script = "bootstrap.sh"
6778
}
6879
}

0 commit comments

Comments
 (0)