[Docker] Update the CUDA version in the default Docker image to 12.8 (from 12.1) (#3166)

* [Docker] Update the CUDA version in the default Docker image to 12.8 (from 12.1) #3163

* Pin base image version for Azure GRID image

* [Azure] Downgrade Linux Kernel to 6.8 as a workaround to install Grid driver

* Bugfix

* [Docker] Update the CUDA version in the default Docker image to 12.8 (from 12.1) #3163

Updated base_image to 0.11

---------

Co-authored-by: Jvst Me <git@jvst.me>

Author: peterschmidt85

docker/base/Dockerfile

@@ -2,7 +2,7 @@
 ARG UBUNTU_VERSION
 
 # Build stage
-FROM nvidia/cuda:12.1.1-base-ubuntu${UBUNTU_VERSION}.04 AS builder
+FROM nvidia/cuda:12.8.1-base-ubuntu${UBUNTU_VERSION}.04 AS builder
 
 ENV NCCL_HOME=/opt/nccl
 ENV CUDA_HOME=/usr/local/cuda

docker/base/Dockerfile.common

@@ -1,6 +1,6 @@
 ARG UBUNTU_VERSION
 
-FROM nvidia/cuda:12.1.1-base-ubuntu${UBUNTU_VERSION}.04
+FROM nvidia/cuda:12.8.1-base-ubuntu${UBUNTU_VERSION}.04
 
 ARG _UV_HOME="/opt/uv"
 

scripts/packer/azure-image-grid.json

@@ -63,6 +63,15 @@
         "./install-docker.sh --version {{user `docker_version`}}"
       ]
     },
+    {
+      "type": "shell",
+      "script": "provisioners/downgrade-azure-kernel.sh"
+    },
+    {
+      "type": "shell",
+      "inline": ["sudo reboot"],
+      "expect_disconnect": true
+    },
     {
       "type": "shell",
       "script": "provisioners/install-nvidia-grid-driver-for-azure.sh"

scripts/packer/provisioners/downgrade-azure-kernel.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+# based on https://learn.microsoft.com/en-us/azure/virtual-machines/extensions/hpccompute-gpu-linux#known-issues
+# this is a temporary solution only required until the issue is fixed
+
+set -e
+
+# Install 6.8 kernel
+sudo apt-get update
+sudo DEBIAN_FRONTEND=noninteractive apt install linux-image-6.8.0-1015-azure linux-headers-6.8.0-1015-azure -y
+
+# Update the Grub entry name
+grub_entry_name="$(sudo grep -Po "menuentry '\KUbuntu, with Linux 6\.8[^(']+" /boot/grub/grub.cfg | sort -V | head -1)"
+sudo sed -i "s/^\s*GRUB_DEFAULT=.*$/GRUB_DEFAULT='Advanced options for Ubuntu>$grub_entry_name'/" /etc/default/grub
+sudo update-grub
+
+# Disable the kernel package upgrade
+sudo apt-mark hold $(dpkg --get-selections | grep -Po "^linux[^\t]+${grub_entry_name##* }")

src/dstack/_internal/core/backends/vastai/compute.py

@@ -47,7 +47,7 @@ def __init__(self, config: VastAIConfig):
                     "reliability2": {"gte": 0.9},
                     "inet_down": {"gt": 128},
                     "verified": {"eq": True},
-                    "cuda_max_good": {"gte": 12.1},
+                    "cuda_max_good": {"gte": 12.8},
                     "compute_cap": {"gte": 600},
                 }
             )

src/dstack/version.py

@@ -5,5 +5,5 @@
 
 __version__ = "0.0.0"
 __is_release__ = False
-base_image = "0.11rc2"
+base_image = "0.11"
 base_image_ubuntu_version = "22.04"

src/tests/_internal/server/routers/test_runs.py

@@ -333,7 +333,7 @@ def get_dev_env_run_dict(
                 " && tail -f /dev/null"
             ),
         ]
-        image_name = "dstackai/base:0.11rc2-base-ubuntu22.04"
+        image_name = "dstackai/base:0.11-base-ubuntu22.04"
 
     return {
         "id": run_id,