Update NVIDIA driver in dstack OS images (#3099)

* Update NVIDIA driver in dstack OS images

Update the driver to support NVIDIA B200.

- Update from the 535 to the 570 family.
- Update to Ubuntu 24.04, since Ubuntu 22.04 does
  not have the gcc version required for building
  the 570 driver.
- Switch from proprietary to open kernel modules.
- Since pre-Turing GPUs aren't supported by NVIDIA
  open kernel modules, conditionally choose
  between old and new dstack OS images based on
  the GPU name.
- Adjust handling `apt` race conditions - the
  existing hack did not work on OCI's Ubuntu
  24.04.
- Install `ufw` when building the image - it is
  missing in OCI's Ubuntu 24.04.

* Fix tests

* [Feature]: GCP A4 instances #3088

Bumped `base_image` to `0.11rc2`

* [Feature]: GCP A4 instances #3088

Updated tests

---------

Co-authored-by: Andrey Cheptsov <54148038+peterschmidt85@users.noreply.github.com>
Co-authored-by: peterschmidt85 <andrey.cheptsov@gmail.com>

Author: jvstme

scripts/packer/aws-image-cuda.json

@@ -4,7 +4,7 @@
     "aws_secret_key": "{{env `AWS_SECRET_ACCESS_KEY`}}",
     "region": "eu-west-1",
     "ssh_username": "ubuntu",
-    "base_ami": "ami-0cffefff2d52e0a23",
+    "base_ami": "ami-0bc691261a82b32bc",
     "instance_type": "c5.large",
     "subnet_id": "subnet-c39cb6a5",
     "docker_version": "",

scripts/packer/aws-image.json

@@ -4,7 +4,7 @@
     "aws_secret_key": "{{env `AWS_SECRET_ACCESS_KEY`}}",
     "region": "eu-west-1",
     "ssh_username": "ubuntu",
-    "base_ami": "ami-0cffefff2d52e0a23",
+    "base_ami": "ami-0bc691261a82b32bc",
     "instance_type": "c5.large",
     "subnet_id": "subnet-c39cb6a5",
     "docker_version": "",

scripts/packer/azure-image-cuda.json

@@ -23,8 +23,8 @@
     "managed_image_name": "{{user `build_prefix`}}dstack-cuda-{{user `image_version` | clean_resource_name}}",
     "os_type": "Linux",
     "image_publisher": "canonical",
-    "image_offer": "0001-com-ubuntu-server-jammy",
-    "image_sku": "22_04-lts-gen2",
+    "image_offer": "ubuntu-24_04-lts",
+    "image_sku": "server",
     "azure_tags": {
         "Name": "DSTACK-CUDA"
     },

scripts/packer/azure-image-grid.json

@@ -22,8 +22,8 @@
     "managed_image_name": "{{user `build_prefix`}}dstack-grid-{{user `image_version` | clean_resource_name}}",
     "os_type": "Linux",
     "image_publisher": "canonical",
-    "image_offer": "0001-com-ubuntu-server-jammy",
-    "image_sku": "22_04-lts-gen2",
+    "image_offer": "ubuntu-24_04-lts",
+    "image_sku": "server",
     "azure_tags": {
         "Name": "DSTACK-GRID"
     },

scripts/packer/azure-image.json

@@ -22,8 +22,8 @@
     "managed_image_name": "{{user `build_prefix`}}dstack-{{user `image_version` | clean_resource_name}}",
     "os_type": "Linux",
     "image_publisher": "canonical",
-    "image_offer": "0001-com-ubuntu-server-jammy",
-    "image_sku": "22_04-lts-gen2",
+    "image_offer": "ubuntu-24_04-lts",
+    "image_sku": "server",
     "azure_tags": {
         "Name": "DSTACK"
     },

scripts/packer/gcp-image-cuda.json

@@ -10,7 +10,7 @@
     {
       "type": "googlecompute",
       "project_id": "dstack",
-      "source_image": "ubuntu-2204-jammy-v20230714",
+      "source_image": "ubuntu-2404-noble-amd64-v20250828",
       "image_name": "{{user `build_prefix`}}dstack-cuda-{{user `image_version` | clean_resource_name}}",
       "instance_name": "{{user `build_prefix`}}dstack-cuda-{{user `image_version` | clean_resource_name}}",
       "ssh_username": "ubuntu",

scripts/packer/gcp-image.json

@@ -9,7 +9,7 @@
     {
       "type": "googlecompute",
       "project_id": "dstack",
-      "source_image": "ubuntu-2204-jammy-v20230714",
+      "source_image": "ubuntu-2404-noble-amd64-v20250828",
       "image_name": "{{user `build_prefix`}}dstack-{{user `image_version` | clean_resource_name}}",
       "instance_name": "{{user `build_prefix`}}dstack-{{user `image_version` | clean_resource_name}}",
       "ssh_username": "ubuntu",

scripts/packer/oci-image-cuda.json

@@ -16,7 +16,7 @@
       "compartment_ocid": "{{user `oci_compartment_ocid`}}",
       "subnet_ocid": "{{user `oci_subnet_ocid`}}",
       "shape": "VM.Standard2.1",
-      "base_image_ocid": "ocid1.image.oc1.eu-frankfurt-1.aaaaaaaaxroekfbow3kdrdjlwao6tsxxfcb23xmqrdjtjcay2ow52eijvzqa",
+      "base_image_ocid": "ocid1.image.oc1.eu-frankfurt-1.aaaaaaaahelib4o7g4fsjgck2lhxjmzonvbniwcmjjn2im4cxlksjgyzw5gq",
       "image_name": "{{user `build_prefix`}}dstack-cuda-{{user `image_version`}}",
       "instance_name": "packer-{{user `build_prefix`}}dstack-cuda-{{user `image_version`}}",
       "ssh_username": "ubuntu"
@@ -27,10 +27,6 @@
       "type": "shell",
       "inline": ["cloud-init status --long --wait"]
     },
-    {
-      "type": "shell",
-      "script": "provisioners/wait-for-dpkg-lock.sh"
-    },
     {
       "type": "shell",
       "scripts": [

scripts/packer/oci-image.json

@@ -15,7 +15,7 @@
       "compartment_ocid": "{{user `oci_compartment_ocid`}}",
       "subnet_ocid": "{{user `oci_subnet_ocid`}}",
       "shape": "VM.Standard2.1",
-      "base_image_ocid": "ocid1.image.oc1.eu-frankfurt-1.aaaaaaaaxroekfbow3kdrdjlwao6tsxxfcb23xmqrdjtjcay2ow52eijvzqa",
+      "base_image_ocid": "ocid1.image.oc1.eu-frankfurt-1.aaaaaaaahelib4o7g4fsjgck2lhxjmzonvbniwcmjjn2im4cxlksjgyzw5gq",
       "image_name": "{{user `build_prefix`}}dstack-{{user `image_version`}}",
       "instance_name": "packer-{{user `build_prefix`}}dstack-{{user `image_version`}}",
       "ssh_username": "ubuntu"
@@ -26,10 +26,6 @@
       "type": "shell",
       "inline": ["cloud-init status --long --wait"]
     },
-    {
-      "type": "shell",
-      "script": "provisioners/wait-for-dpkg-lock.sh"
-    },
     {
       "type": "shell",
       "scripts": [

scripts/packer/provisioners/cuda.sh

@@ -10,15 +10,13 @@ ARCH=$(uname -m)
 CUDA_DISTRO=$(. /etc/os-release;echo $ID$VERSION_ID | sed -e 's/\.//g')
 
 # based on https://docs.nvidia.com/datacenter/tesla/tesla-installation-notes/index.html#ubuntu-lts
-wget https://developer.download.nvidia.com/compute/cuda/repos/$CUDA_DISTRO/$ARCH/cuda-keyring_1.0-1_all.deb
-sudo dpkg -i cuda-keyring_1.0-1_all.deb
-rm cuda-keyring_1.0-1_all.deb
+wget https://developer.download.nvidia.com/compute/cuda/repos/$CUDA_DISTRO/$ARCH/cuda-keyring_1.1-1_all.deb
+sudo dpkg -i cuda-keyring_1.1-1_all.deb
+rm cuda-keyring_1.1-1_all.deb
 
 sudo apt-get update
-# Pinned dkms due to https://bugs.launchpad.net/ubuntu/+source/dkms/+bug/2112114
 sudo DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
-    dkms=2.8.7-2ubuntu2.2 \
-    cuda-drivers-$CUDA_DRIVERS_VERSION \
+    nvidia-driver-$CUDA_DRIVERS_VERSION-server-open \
     nvidia-fabricmanager-$CUDA_DRIVERS_VERSION \
     datacenter-gpu-manager-4-core datacenter-gpu-manager-4-proprietary datacenter-gpu-manager-exporter
 sudo systemctl enable nvidia-fabricmanager