Update NVIDIA driver in dstack OS images

jvstme · jvstme · commit 085feaecf4b5 · 2025-09-13T23:45:23.000+02:00
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.
diff --git a/scripts/packer/aws-image-cuda.json b/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": "",
diff --git a/scripts/packer/aws-image.json b/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": "",
diff --git a/scripts/packer/azure-image-cuda.json b/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"
     },
diff --git a/scripts/packer/azure-image-grid.json b/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"
     },
diff --git a/scripts/packer/azure-image.json b/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"
     },
diff --git a/scripts/packer/gcp-image-cuda.json b/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",
diff --git a/scripts/packer/gcp-image.json b/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",
diff --git a/scripts/packer/oci-image-cuda.json b/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": [
diff --git a/scripts/packer/oci-image.json b/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": [
diff --git a/scripts/packer/provisioners/cuda.sh b/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
diff --git a/scripts/packer/provisioners/kernel/apt-packages.sh b/scripts/packer/provisioners/kernel/apt-packages.sh
@@ -19,6 +19,7 @@ DEPS="
   tree
   jq
   gdb
+  ufw
   python3-pip
   python3-boto
   python3-boto3
diff --git a/scripts/packer/provisioners/kernel/apt-upgrade.sh b/scripts/packer/provisioners/kernel/apt-upgrade.sh
@@ -5,5 +5,39 @@
 
 set -e
 
-sudo apt-get update
+#
+# Run apt-get update, but retry in case the lock is held by another process.
+#
+# A better way of handling apt races is the `-o DPkg::Lock::Timeout=X` option,
+# but it does not work with `apt-get update`.
+#
+# This function was added specifically for the `oci` backend, where our build
+# process conflicts with OCI's instance agent.
+#
+apt_update_with_retry() {
+    local MAX_RETRIES=10
+    local RETRY_DELAY=3
+    local COUNT=0
+    local LOGFILE=$(mktemp)
+
+    while [ $COUNT -lt $MAX_RETRIES ]; do
+        set +e
+        sudo apt-get update 2>&1 | tee "$LOGFILE"
+        local EXIT_CODE=${PIPESTATUS[0]}
+        set -e
+
+        if grep -q "Could not get lock" "$LOGFILE"; then
+            echo "apt lock file is held by another process. Retrying in $RETRY_DELAY seconds..."
+            COUNT=$((COUNT + 1))
+            sleep $RETRY_DELAY
+        else
+            return $EXIT_CODE
+        fi
+    done
+
+    echo "apt-get update failed due to lock being held after $MAX_RETRIES attempts."
+    return 1
+}
+
+apt_update_with_retry
 sudo DEBIAN_FRONTEND=noninteractive apt-get -o DPkg::Lock::Timeout=60 dist-upgrade -y -q
diff --git a/scripts/packer/provisioners/wait-for-dpkg-lock.sh b/scripts/packer/provisioners/wait-for-dpkg-lock.sh
diff --git a/scripts/packer/versions.json b/scripts/packer/versions.json
@@ -1,4 +1,4 @@
 {
   "docker_version": "27.1.1",
-  "cuda_drivers_version": "535"
+  "cuda_drivers_version": "570"
 }
diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py
@@ -276,7 +276,11 @@ def create_instance(
                 image_id, username = self._get_image_id_and_username(
                     ec2_client=ec2_client,
                     region=instance_offer.region,
-                    cuda=len(instance_offer.instance.resources.gpus) > 0,
+                    gpu_name=(
+                        instance_offer.instance.resources.gpus[0].name
+                        if len(instance_offer.instance.resources.gpus) > 0
+                        else None
+                    ),
                     instance_type=instance_offer.instance.name,
                     image_config=self.config.os_images,
                 )
@@ -882,11 +886,13 @@ def _get_image_id_and_username_cache_key(
         self,
         ec2_client: botocore.client.BaseClient,
         region: str,
-        cuda: bool,
+        gpu_name: Optional[str],
         instance_type: str,
         image_config: Optional[AWSOSImageConfig] = None,
     ) -> tuple:
-        return hashkey(region, cuda, instance_type, image_config.json() if image_config else None)
+        return hashkey(
+            region, gpu_name, instance_type, image_config.json() if image_config else None
+        )
 
     @cachedmethod(
         cache=lambda self: self._get_image_id_and_username_cache,
@@ -897,13 +903,13 @@ def _get_image_id_and_username(
         self,
         ec2_client: botocore.client.BaseClient,
         region: str,
-        cuda: bool,
+        gpu_name: Optional[str],
         instance_type: str,
         image_config: Optional[AWSOSImageConfig] = None,
     ) -> tuple[str, str]:
         return aws_resources.get_image_id_and_username(
             ec2_client=ec2_client,
-            cuda=cuda,
+            gpu_name=gpu_name,
             instance_type=instance_type,
             image_config=image_config,
         )
diff --git a/src/dstack/_internal/core/backends/aws/resources.py b/src/dstack/_internal/core/backends/aws/resources.py
@@ -6,6 +6,8 @@
 
 import dstack.version as version
 from dstack._internal.core.backends.aws.models import AWSOSImageConfig
+from dstack._internal.core.backends.base.compute import requires_nvidia_proprietary_kernel_modules
+from dstack._internal.core.consts import DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES
 from dstack._internal.core.errors import BackendError, ComputeError, ComputeResourceNotFoundError
 from dstack._internal.utils.logging import get_logger
 
@@ -17,14 +19,14 @@
 
 def get_image_id_and_username(
     ec2_client: botocore.client.BaseClient,
-    cuda: bool,
+    gpu_name: Optional[str],
     instance_type: str,
     image_config: Optional[AWSOSImageConfig] = None,
 ) -> tuple[str, str]:
     if image_config is not None:
-        image = image_config.nvidia if cuda else image_config.cpu
+        image = image_config.nvidia if gpu_name else image_config.cpu
         if image is None:
-            logger.warning("%s image not configured", "nvidia" if cuda else "cpu")
+            logger.warning("%s image not configured", "nvidia" if gpu_name else "cpu")
             raise ComputeResourceNotFoundError()
         image_name = image.name
         image_owner = image.owner
@@ -35,9 +37,12 @@ def get_image_id_and_username(
         image_owner = DLAMI_OWNER_ACCOUNT_ID
         username = "ubuntu"
     else:
-        image_name = (
-            f"dstack-{version.base_image}" if not cuda else f"dstack-cuda-{version.base_image}"
-        )
+        if gpu_name is None:
+            image_name = f"dstack-{version.base_image}"
+        elif not requires_nvidia_proprietary_kernel_modules(gpu_name):
+            image_name = f"dstack-cuda-{version.base_image}"
+        else:
+            image_name = f"dstack-cuda-{DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES}"
         image_owner = DSTACK_ACCOUNT_ID
         username = "ubuntu"
     response = ec2_client.describe_images(
diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py
@@ -47,8 +47,10 @@
     get_gateway_user_data,
     get_user_data,
     merge_tags,
+    requires_nvidia_proprietary_kernel_modules,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
+from dstack._internal.core.consts import DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES
 from dstack._internal.core.errors import ComputeError, NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.gateways import (
@@ -369,25 +371,32 @@ def _parse_config_vpc_id(vpc_id: str) -> Tuple[str, str]:
 class VMImageVariant(enum.Enum):
     GRID = enum.auto()
     CUDA = enum.auto()
+    CUDA_WITH_PROPRIETARY_KERNEL_MODULES = enum.auto()
     STANDARD = enum.auto()
 
     @classmethod
     def from_instance_type(cls, instance: InstanceType) -> "VMImageVariant":
         if "_A10_v5" in instance.name:
             return cls.GRID
         elif len(instance.resources.gpus) > 0:
-            return cls.CUDA
+            if not requires_nvidia_proprietary_kernel_modules(instance.resources.gpus[0].name):
+                return cls.CUDA
+            else:
+                return cls.CUDA_WITH_PROPRIETARY_KERNEL_MODULES
         else:
             return cls.STANDARD
 
     def get_image_name(self) -> str:
-        name = "dstack-"
         if self is self.GRID:
-            name += "grid-"
+            return f"dstack-grid-{version.base_image}"
         elif self is self.CUDA:
-            name += "cuda-"
-        name += version.base_image
-        return name
+            return f"dstack-cuda-{version.base_image}"
+        elif self is self.CUDA_WITH_PROPRIETARY_KERNEL_MODULES:
+            return f"dstack-cuda-{DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES}"
+        elif self is self.STANDARD:
+            return f"dstack-{version.base_image}"
+        else:
+            raise ValueError(f"Unexpected image variant {self!r}")
 
 
 _SUPPORTED_VM_SERIES_PATTERNS = [
diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py
@@ -47,6 +47,22 @@
 DSTACK_SHIM_BINARY_NAME = "dstack-shim"
 DSTACK_RUNNER_BINARY_NAME = "dstack-runner"
 DEFAULT_PRIVATE_SUBNETS = ("10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16")
+NVIDIA_GPUS_REQUIRING_PROPRIETARY_KERNEL_MODULES = frozenset(
+    # All NVIDIA architectures prior to Turing do not support Open Kernel Modules and require
+    # proprietary modules. This list is incomplete, update when necessary.
+    [
+        "v100",
+        "p100",
+        "p40",
+        "p4",
+        "m60",
+        "m40",
+        "m4",
+        "k80",
+        "k40",
+        "k20",
+    ]
+)
 
 GoArchType = Literal["amd64", "arm64"]
 
@@ -887,3 +903,12 @@ def merge_tags(
         for k, v in resource_tags.items():
             res.setdefault(k, v)
     return res
+
+
+def requires_nvidia_proprietary_kernel_modules(gpu_name: str) -> bool:
+    """
+    Returns:
+        Whether this NVIDIA GPU requires NVIDIA proprietary kernel modules
+        instead of open kernel modules.
+    """
+    return gpu_name.lower() in NVIDIA_GPUS_REQUIRING_PROPRIETARY_KERNEL_MODULES
diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py
diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py
diff --git a/src/dstack/_internal/core/backends/oci/resources.py b/src/dstack/_internal/core/backends/oci/resources.py
diff --git a/src/dstack/_internal/core/consts.py b/src/dstack/_internal/core/consts.py

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@`
`10`	`10`	`{`
`11`	`11`	`"type": "googlecompute",`
`12`	`12`	`"project_id": "dstack",`
`13`		`- "source_image": "ubuntu-2204-jammy-v20230714",`
	`13`	`+ "source_image": "ubuntu-2404-noble-amd64-v20250828",`
`14`	`14`	"image_name": "{{user `build_prefix`}}dstack-cuda-{{user `image_version` \| clean_resource_name}}",
`15`	`15`	"instance_name": "{{user `build_prefix`}}dstack-cuda-{{user `image_version` \| clean_resource_name}}",
`16`	`16`	`"ssh_username": "ubuntu",`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,7 @@`
`9`	`9`	`{`
`10`	`10`	`"type": "googlecompute",`
`11`	`11`	`"project_id": "dstack",`
`12`		`- "source_image": "ubuntu-2204-jammy-v20230714",`
	`12`	`+ "source_image": "ubuntu-2404-noble-amd64-v20250828",`
`13`	`13`	"image_name": "{{user `build_prefix`}}dstack-{{user `image_version` \| clean_resource_name}}",
`14`	`14`	"instance_name": "{{user `build_prefix`}}dstack-{{user `image_version` \| clean_resource_name}}",
`15`	`15`	`"ssh_username": "ubuntu",`