Install CUDA driver into VM image for ARM64

Author: hcho3

vm_images/linux-arm64/bootstrap.sh

@@ -31,6 +31,19 @@ sudo systemctl is-active --quiet docker.service || sudo systemctl start docker.s
 sudo systemctl is-enabled --quiet docker.service || sudo systemctl enable docker.service
 sleep 10  # Docker daemon takes time to come up after installing
 sudo docker info
+
+## Install NVIDIA Container Toolkit
+curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
+  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
+    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
+    sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+sudo apt-get update
+sudo apt-get install -y nvidia-container-toolkit
+sudo nvidia-ctk runtime configure --runtime=docker
+sudo systemctl restart docker
+
+sleep 10
+sudo docker run --rm --gpus all ubuntu nvidia-smi
 sudo systemctl stop docker
 
 ## Install AWS CLI v2

vm_images/linux-arm64/install_drivers.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+set -euo pipefail
+
+## Install basic tools
+echo 'debconf debconf/frontend select Noninteractive' | sudo debconf-set-selections
+sudo apt-get update
+sudo apt-get install -y cmake git build-essential wget ca-certificates curl unzip
+
+## Install CUDA Driver 580
+wget -nv https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbsa/cuda-keyring_1.1-1_all.deb
+sudo dpkg -i cuda-keyring_1.1-1_all.deb
+sudo apt-get update
+sudo apt-get -y install nvidia-open-580
+rm cuda-keyring_1.1-1_all.deb

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

@@ -9,7 +9,7 @@ packer {
 
 locals {
   ami_name_prefix = "xgboost-ci"
-  image_name      = "RunsOn worker with Ubuntu 24.04 ARM64"
+  image_name      = "RunsOn worker with Ubuntu 24.04 ARM64 + CUDA driver 580"
   region          = "us-west-2"
   timestamp       = regex_replace(timestamp(), "[- TZ:]", "")
   volume_size     = 40
@@ -33,7 +33,7 @@ source "amazon-ebs" "runs-on-linux-arm64" {
   ami_virtualization_type     = "hvm"
   associate_public_ip_address = true
   communicator                = "ssh"
-  instance_type               = "c6g.4xlarge"
+  instance_type               = "g5g.xlarge"
   region                      = "${local.region}"
   ssh_timeout                 = "10m"
   ssh_username                = "ubuntu"
@@ -63,6 +63,17 @@ build {
   sources = ["source.amazon-ebs.runs-on-linux-arm64"]
 
   provisioner "shell" {
-    script = "bootstrap.sh"
+    script      = "install_drivers.sh"
+    pause_after = "30s"
+  }
+
+  provisioner "shell" {
+    expect_disconnect = true
+    inline            = ["echo 'Reboot VM'", "sudo reboot"]
+  }
+
+  provisioner "shell" {
+    pause_before = "1m0s"
+    script       = "bootstrap.sh"
   }
 }