Update Azure GRID driver to 580.178.04

.github/workflows/vm-images.yml

@@ -19,6 +19,15 @@ on:
         description: "Build Azure images"
         type: boolean
         default: true
+      azure_variant:
+        description: "Azure image variants to build"
+        type: choice
+        options: ["all", "standard", "cuda", "grid"]
+        default: "all"
+      azure_new_gallery_version:
+        description: "Rebuild Azure images that already exist and publish them as a new gallery version"
+        type: boolean
+        default: false
       build_gcp:
         description: "Build GCP images"
         type: boolean
@@ -72,7 +81,7 @@ jobs:
       VERSION: ${{ github.run_number }}
     strategy:
       matrix:
-        variant: ["", "-cuda", "-grid"]
+        variant: ${{ fromJSON('{"all":["","-cuda","-grid"],"standard":[""],"cuda":["-cuda"],"grid":["-grid"]}')[inputs.azure_variant] }}
     steps:
       - uses: actions/checkout@v4
       - uses: Azure/login@v2
@@ -86,13 +95,13 @@ jobs:
           chmod +x packer
       - name: Run packer
         run: |
-          ./packer build -var-file=versions.json -var image_version=${{ inputs.image_version }} -var build_prefix=$VM_IMAGE_BUILD_PREFIX azure-image${{ matrix.variant }}.json
+          ./packer build ${{ inputs.azure_new_gallery_version && '-force' || '' }} -var-file=versions.json -var image_version=${{ inputs.image_version }} -var build_prefix=$VM_IMAGE_BUILD_PREFIX azure-image${{ matrix.variant }}.json
       - name: Publish azure image
         if: ${{ !inputs.staging }}
         run: |
           IMAGE_DEFINITION=${VM_IMAGE_BUILD_PREFIX}dstack${{ matrix.variant }}-${{ inputs.image_version }}
           IMAGE_NAME=${VM_IMAGE_BUILD_PREFIX}dstack${{ matrix.variant }}-${{ inputs.image_version }}
-          ../publish_azure_image.sh $IMAGE_DEFINITION $IMAGE_NAME
+          ../publish_azure_image.sh $IMAGE_DEFINITION $IMAGE_NAME ${{ inputs.azure_new_gallery_version }}
 
   build-gcp-images:
     if: inputs.build_gcp

scripts/packer/provisioners/install-nvidia-grid-driver-for-azure.sh

@@ -8,7 +8,7 @@ sudo apt-get update
 sudo DEBIAN_FRONTEND=noninteractive apt-get install build-essential linux-azure -y
 
 wget --no-verbose -O NVIDIA-Linux-x86_64-grid.run \
-    https://download.microsoft.com/download/2a04ca6a-9eec-40d9-9564-9cdea1ab795f/NVIDIA-Linux-x86_64-570.211.01-grid-azure.run
+    https://download.microsoft.com/download/4efc1974-7b6f-431f-ab2d-b8d99bb8a512/NVIDIA-Linux-x86_64-580.178.04-grid-azure.run
 chmod +x NVIDIA-Linux-x86_64-grid.run
 sudo ./NVIDIA-Linux-x86_64-grid.run --silent --disable-nouveau
 rm NVIDIA-Linux-x86_64-grid.run

scripts/publish_azure_image.sh

@@ -1,7 +1,11 @@
 #!/bin/bash
 
+set -o pipefail
+
 image_definition=$1
 image_name=$2
+# If true, publish a new gallery version when the image definition already has versions
+new_gallery_version=${3:-false}
 subscription_id=$AZURE_SUBSCRIPTION_ID
 
 if [ -z "$subscription_id" ]; then
@@ -37,17 +41,39 @@ function create_image_definition() {
         --features DiskControllerTypes=SCSI,NVMe
 }
 
+function get_latest_image_version() {
+    az sig image-version list \
+        --resource-group $resource_group \
+        --gallery-name $gallery_name \
+        --gallery-image-definition $image_definition \
+        --query "[].name" \
+        --output tsv |
+        sort -V |
+        tail -1
+}
+
 function create_image_version() {
-    echo Creating image version...
+    echo Creating image version $1...
     az sig image-version create \
         --resource-group $resource_group \
         --gallery-name $gallery_name \
         --gallery-image-definition $image_definition \
-        --gallery-image-version "0.0.1" \
+        --gallery-image-version "$1" \
         --target-regions "australiaeast" "brazilsouth" "canadacentral" "centralindia" "centralus" "eastasia" "eastus" "eastus2" "francecentral" "germanywestcentral" "japaneast" "koreacentral" "northeurope" "norwayeast" "qatarcentral" "southafricanorth" "southcentralus" "southeastasia" "swedencentral" "switzerlandnorth" "uaenorth" "uksouth" "westeurope" "westus2" "westus3" \
         --replica-count 1 \
         --managed-image "/subscriptions/${subscription_id}/resourceGroups/${resource_group}/providers/Microsoft.Compute/images/${image_name}"
 }
 
 get_image_definition > /dev/null || create_image_definition
-create_image_version
+
+latest_version=$(get_latest_image_version) || exit 1
+if [ -z "$latest_version" ]; then
+    image_version=0.0.1
+elif [ "$new_gallery_version" = true ]; then
+    image_version="${latest_version%.*}.$((${latest_version##*.} + 1))"
+else
+    echo "$image_definition already has gallery version $latest_version." \
+        "Enable azure_new_gallery_version to publish a new one."
+    exit 1
+fi
+create_image_version $image_version