Update base image driver to 580 in a3u and a4h slurm blueprints (#5456)

Author: saara-tyagi27

examples/machine-learning/a3-ultragpu-8g/a3ultra-slurm-blueprint.yaml

@@ -27,7 +27,7 @@ vars:
   # Image settings
   base_image:
     project: ubuntu-os-accelerator-images
-    image: ubuntu-accelerator-2204-amd64-with-nvidia-570-v20260210
+    image: ubuntu-accelerator-2404-amd64-with-nvidia-580-v20260522
   image_build_machine_type: n2-standard-16
   build_slurm_from_git_ref: 6.12.1
   # Cluster env settings
@@ -38,7 +38,7 @@ vars:
   local_ssd_mountpoint: /mnt/localssd
   instance_image:
     project: $(vars.project_id)
-    family: $(vars.deployment_name)-u22
+    family: $(vars.deployment_name)-u24
   disk_size_gb: 100
   nccl_gib_version: 1.1.0
   libnccl_version: 2.27.5-1+cuda12.9
@@ -51,6 +51,7 @@ vars:
 
   # Managed Lustre is only supported in specific regions and zones
   # Please refer https://cloud.google.com/managed-lustre/docs/locations
+
   # Managed-Lustre instance name. This should be unique for each deployment.
   lustre_instance_id: $(vars.deployment_name)-lustre
 
@@ -125,7 +126,6 @@ deployment_groups:
             "install_gcsfuse": true,
             "install_lustre": false,
             "install_managed_lustre": true,
-            "install_nvidia_repo": true,
             "install_ompi": true,
             "allow_kernel_upgrades": false,
             "monitoring_agent": "cloud-ops",
@@ -150,75 +150,79 @@ deployment_groups:
           * - nofile 1048576
           * - cpu unlimited
           * - rtprio unlimited
-      - type: shell
-        destination: install-cuda-toolkit.sh
-        content: |
-          #!/bin/bash
-          set -e -o pipefail
-          add-nvidia-repositories -y
-          apt install -y \
-              cuda-toolkit-12-8 \
-              datacenter-gpu-manager-4-cuda12=1:4.5.3-1 \
-              datacenter-gpu-manager-4-dev=1:4.5.3-1 \
-              datacenter-gpu-manager-4-core=1:4.5.3-1
-
-          # 2. Hold ALL components
-          apt-mark hold \
-              datacenter-gpu-manager-4-cuda12 \
-              datacenter-gpu-manager-4-dev \
-              datacenter-gpu-manager-4-core
 
       - type: ansible-local
-        destination: settings_nvidia_dcgm.yml
+        destination: install_cuda_and_dcgm.yml
         content: |
           ---
-          - name: Nvidia and DGMA settings
+          - name: Install and Configure CUDA 13 / DCGM
             hosts: all
             become: true
+            vars:
+              distribution: "{{ ansible_distribution | lower }}{{ ansible_distribution_version | replace('.','') }}"
+              cuda_repo_url: "https://developer.download.nvidia.com/compute/cuda/repos/{{ distribution }}/x86_64/cuda-keyring_1.1-1_all.deb"
+              nvidia_packages:
+                - cuda-toolkit-13-0
+                - datacenter-gpu-manager-4-cuda13
+                - datacenter-gpu-manager-4-dev
+                - datacenter-gpu-manager-4-core
             tasks:
-            - name: Reduce NVIDIA repository priority
-              ansible.builtin.copy:
-                dest: /etc/apt/preferences.d/cuda-repository-pin-600
-                mode: 0o0644
-                owner: root
-                group: root
-                content: |
-                  Package: *
-                  Pin: release l=NVIDIA CUDA
-                  Pin-Priority: 400
-            - name: Create nvidia-persistenced override directory
-              ansible.builtin.file:
-                path: /etc/systemd/system/nvidia-persistenced.service.d
-                state: directory
-                owner: root
-                group: root
-                mode: 0o755
-            - name: Configure nvidia-persistenced override
-              ansible.builtin.copy:
-                dest: /etc/systemd/system/nvidia-persistenced.service.d/persistence_mode.conf
-                owner: root
-                group: root
-                mode: 0o644
-                content: |
-                  [Service]
-                  ExecStart=
-                  ExecStart=/usr/bin/nvidia-persistenced --user nvidia-persistenced --verbose
-              notify: Reload SystemD
+              - name: Install NVIDIA repository keyring
+                ansible.builtin.apt:
+                  deb: "{{ cuda_repo_url }}"
+                  state: present
+
+              - name: Update apt cache
+                ansible.builtin.apt:
+                  update_cache: yes
+
+              - name: Install NVIDIA packages
+                ansible.builtin.apt:
+                  name: "{{ item }}"
+                  state: present
+                  allow_downgrade: yes
+                loop: "{{ nvidia_packages }}"
+
+              - name: Freeze NVIDIA packages
+                ansible.builtin.dpkg_selections:
+                  name: "{{ item }}"
+                  selection: hold
+                loop: "{{ nvidia_packages }}"
+
+              - name: Create nvidia-persistenced override directory
+                ansible.builtin.file:
+                  path: /etc/systemd/system/nvidia-persistenced.service.d
+                  state: directory
+                  owner: root
+                  group: root
+                  mode: 0o755
+
+              - name: Configure nvidia-persistenced override
+                ansible.builtin.copy:
+                  dest: /etc/systemd/system/nvidia-persistenced.service.d/persistence_mode.conf
+                  owner: root
+                  group: root
+                  mode: 0o644
+                  content: |
+                    [Service]
+                    ExecStart=
+                    ExecStart=/usr/bin/nvidia-persistenced --user nvidia-persistenced --verbose
+                notify: Reload SystemD
+
             handlers:
             - name: Reload SystemD
               ansible.builtin.systemd:
                 daemon_reload: true
+
             post_tasks:
-            - name: Disable NVIDIA DCGM by default (enable during boot on GPU nodes)
-              ansible.builtin.service:
-                name: nvidia-dcgm.service
-                state: stopped
-                enabled: false
-            - name: Disable nvidia-persistenced SystemD unit (enable during boot on GPU nodes)
-              ansible.builtin.service:
-                name: nvidia-persistenced.service
-                state: stopped
-                enabled: false
+              - name: Disable NVIDIA services by default (Slurm starts them on boot)
+                ansible.builtin.service:
+                  name: "{{ item }}"
+                  state: stopped
+                  enabled: false
+                loop:
+                  - nvidia-dcgm.service
+                  - nvidia-persistenced.service
 
       - type: ansible-local
         destination: install_ibverbs_utils.yml
@@ -263,6 +267,7 @@ deployment_groups:
       metadata:
         user-data: |
           #cloud-config
+          create_hostname_file: true
           write_files:
           - path: /etc/apt/apt.conf.d/20auto-upgrades
             permissions: '0644'
@@ -453,6 +458,21 @@ deployment_groups:
           ENROOT_CACHE_PATH      $(vars.local_ssd_mountpoint)/${UID}/enroot/cache
           ENROOT_DATA_PATH       $(vars.local_ssd_mountpoint)/${UID}/enroot/data
           ENROOT_TEMP_PATH       $(vars.local_ssd_mountpoint)/${UID}/enroot
+      - type: shell
+        destination: fix_enroot_env.sh
+        content: |
+          #!/bin/bash
+          # 1. Disable Ubuntu 24.04 restriction on unprivileged user namespaces
+          # Required for enroot/pyxis to work
+          if [ -f /proc/sys/kernel/apparmor_restrict_unprivileged_userns ]; then
+            sysctl -w kernel.apparmor_restrict_unprivileged_userns=0
+            echo "kernel.apparmor_restrict_unprivileged_userns = 0" > /etc/sysctl.d/90-enroot-namespaces.conf
+          fi
+
+          # 2. Ensure /etc/hostname exists (needed by enroot/pyxis)
+          if [ ! -f /etc/hostname ]; then
+            hostname > /etc/hostname
+          fi
       # Install NCCL
       - type: ansible-local
         destination: install_nccl.yml
@@ -618,7 +638,18 @@ deployment_groups:
         ThreadsPerCore: 2
       additional_networks:
         $(concat(
-          a3ultra-slurm-net-1.instance_additional_networks,
+          [{
+            network=null,
+            subnetwork=a3ultra-slurm-net-1.subnetwork_self_link,
+            subnetwork_project=vars.project_id,
+            nic_type="GVNIC",
+            queue_count=null,
+            network_ip="",
+            stack_type=null,
+            access_config=[],
+            ipv6_access_config=[],
+            alias_ip_range=[]
+          }],
           a3ultra-slurm-rdma-net.subnetwork_interfaces
         ))