a4x-vm.yaml

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
---

blueprint_name: a4x-jbvm
vars:
  deployment_name: # supply deployment name
  project_id: # supply project ID
  region: # supply region
  zone: # supply zone
  instance_count: # supply cluster size
  reservation_name: # supply reservation name
  machine_type: 'a4x-highgpu-4g'
  hostname_prefix: $(vars.deployment_name)
  source_image_project_id: ubuntu-os-accelerator-images
  source_image_family: ubuntu-accelerator-2404-arm64-with-nvidia-580
  instance_image:
    project: $(vars.project_id)
    family: jbvm-ubuntu2404-accelerator-arm64
  disk_size_gb: 100
  image_disk_size_gb: 100
  filestore_ip_range: 192.168.32.0/24
  net0_range: 192.168.0.0/19
  net1_range: 192.168.64.0/19
deployment_groups:
- group: primary
  modules:
  - id: a4x_net_0
    source: modules/network/vpc
    settings:
      network_name: $(vars.deployment_name)-a4x-net-0
      mtu: 8896
      subnetworks:
      - subnet_name: $(vars.deployment_name)-a4x-sub-0
        subnet_region: $(vars.region)
        subnet_ip: $(vars.net0_range)
      firewall_rules:
      - name: $(vars.deployment_name)-internal-0
        ranges: [$(vars.net0_range)]
        allow:
        - protocol: tcp
        - protocol: udp
        - protocol: icmp
  - id: a4x_net_1
    source: modules/network/vpc
    settings:
      network_name: $(vars.deployment_name)-a4x-net-1
      mtu: 8896
      subnetworks:
      - subnet_name: $(vars.deployment_name)-a4x-sub-1
        subnet_region: $(vars.region)
        subnet_ip: $(vars.net1_range)
      firewall_rules:
      - name: $(vars.deployment_name)-internal-1
        ranges: [$(vars.net1_range)]
        allow:
        - protocol: tcp
        - protocol: udp
        - protocol: icmp
- group: image-env
  modules:
  - id: image-build-script
    source: modules/scripts/startup-script
    settings:
     # Install ansible
      install_ansible: true
      # Allow ssh between nodes
      configure_ssh_host_patterns:
      - $(vars.hostname_prefix)-*
      runners:
      - type: ansible-local
        destination: setup_nvidia_cuda_environment.yml
        content: |
          # ---
          - name: Install NVIDIA repository packages
            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 }}/sbsa/cuda-keyring_1.1-1_all.deb
              cuda_repo_filename: /tmp/{{ cuda_repo_url | basename }}
              enable_nvidia_dcgm: false
              nvidia_packages:
              - nvidia-container-toolkit
              - cuda-toolkit-13
              - datacenter-gpu-manager-4-cuda13
            tasks:
            # Download the NVIDIA repo
            - name: Download NVIDIA CUDA repository keyring
              ansible.builtin.get_url:
                url: "{{ cuda_repo_url }}"
                dest: "{{ cuda_repo_filename }}"
            # Install the NVIDIA repo
            - name: Installing NVIDIA CUDA repository keyring
              ansible.builtin.apt:
                deb: "{{ cuda_repo_filename }}"
                state: present
            # Update apt cache
            - name: Update apt cache
              ansible.builtin.apt:
                update_cache: true
            # Install the package
            - name: Install NVIDIA fabric and CUDA
              ansible.builtin.apt:
                name: "{{ item }}"
                state: present
              loop: "{{ nvidia_packages }}"
            # Freeze NVIDIA fabric and CUDA
            - name: Freeze NVIDIA fabric and CUDA
              ansible.builtin.dpkg_selections:
                name: "{{ item }}"
                selection: hold
              loop: "{{ nvidia_packages }}"
            post_tasks:
            - name: Enable NVIDIA DCGM by default
              ansible.builtin.service:
                name: nvidia-dcgm.service
                state: started
                enabled: true
      # Install ibverbs-utils
      - type: ansible-local
        destination: install_ibverbs_packages.yml
        content: |
          ---
          - name: Install ibverbs-utils
            hosts: all
            become: true
            tasks:
            - name: Install ibverbs
              ansible.builtin.package:
                name:
                - ibverbs-utils
                state: present
     # Prerequisites for Mellanox/DOCA
      - type: ansible-local
        destination: install_mellanox_deps.yml
        content: |
          ---
          - name: Install Mellanox Prerequisites
            hosts: all
            become: true
            tasks:
            - name: Install build-essential, dkms, and kernel headers
              ansible.builtin.apt:
                name:
                - dkms
                - build-essential
                - "linux-headers-{{ ansible_kernel }}"
                state: present
                update_cache: true
      # Install DOCA-OFED
      - type: shell
        destination: install_mellanox.sh
        content: |
          #!/bin/bash
          set -e -o pipefail
          DEBIAN_FRONTEND=noninteractive
          mkdir -p /opt/src
          cd /opt/src
          wget https://www.mellanox.com/downloads/MFT/mft-4.34.0-145-aarch64-deb.tgz
          tar -xvf mft-4.34.0-145-aarch64-deb.tgz
          ./mft-4.34.0-145-aarch64-deb/install.sh
          rm -rf mft-4.34.0-145-aarch64-deb
          DOCA_URL="https://linux.mellanox.com/public/repo/doca/3.2.0/ubuntu24.04/arm64-sbsa/"
          BASE_URL=https://linux.mellanox.com/public/repo/doca
          DOCA_SUFFIX=3.2.0/ubuntu24.04/arm64-sbsa/
          DOCA_URL="$BASE_URL/$DOCA_SUFFIX"
          curl $BASE_URL/GPG-KEY-Mellanox.pub | gpg --dearmor > /etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub
          echo "deb [signed-by=/etc/apt/trusted.gpg.d/GPG-KEY-Mellanox.pub] $DOCA_URL ./" > /etc/apt/sources.list.d/doca.list
          sudo apt-get update
          sudo apt-get -y install doca-ofed

      # Install OpenMPI from source.
      - type: shell
        destination: install_openmpi.sh
        content: |
          #!/bin/bash
          set -e -o pipefail
          DEBIAN_FRONTEND=noninteractive

          mkdir -p /opt/src
          cd /opt/src

          wget https://download.open-mpi.org/release/open-mpi/v5.0/openmpi-5.0.8.tar.bz2
          tar -xvf openmpi-5.0.8.tar.bz2
          cd openmpi-5.0.8
          ./configure --with-cuda=/usr/local/cuda-13/targets/sbsa-linux
          make -j
          make install
          cd ../
          rm -rf openmpi-5.0.8.tar.bz2 openmpi-5.0.8

      # Install NCCL-GIB network plugin as per nccl-gib specification for a4x
      - type: ansible-local
        destination: install_nccl_gib.yml
        content: |
          ---
          - name: Install Google NCCL-GIB Plugin
            hosts: localhost
            become: true
            tasks:
            - name: Add artifact registry gpg key for us-apt.pkg.dev
              ansible.builtin.apt_key:
                url: https://us-apt.pkg.dev/doc/repo-signing-key.gpg
                state: present
            - name: Add artifact registry gpg key for packages.cloud.google.com
              ansible.builtin.apt_key:
                url: https://packages.cloud.google.com/apt/doc/apt-key.gpg
                state: present
            - name: Install Apt Transport AR Apt Repo
              apt_repository:
                repo: 'deb http://packages.cloud.google.com/apt apt-transport-artifact-registry-stable main'
                state: present
            - name: Install AR transport
              ansible.builtin.apt:
                name: "apt-transport-artifact-registry"
                update_cache: true
            - name: Install Google NCCL-GIB Plugin
              apt_repository:
                repo: "deb ar+https://us-apt.pkg.dev/projects/gce-ai-infra gpudirect-gib-apt main"
                state: present
            - name: Install NCCL-GIB Plugin
              ansible.builtin.apt:
                name: "nccl-gib"
                update_cache: true
            - name: Freeze NCCL GIB Plugin
              ansible.builtin.dpkg_selections:
                name: "nccl-gib"
                selection: hold
      - type: ansible-local
        destination: update_ldconfig.yml
        content: |
          ---
          - name: Update ldconfig
            hosts: localhost
            become: true
            tasks:
            - name: Configure ldconfig
              ansible.builtin.command: /usr/sbin/ldconfig

# This deployment group builds a custom image using Packer + Ansible + Bash
- group: image
  modules:
  - id: custom-image
    source: modules/packer/custom-image
    kind: packer
    use:
    - a4x_net_0
    - image-build-script  # From the above image-build-script
    settings:
      machine_type: t2a-standard-16
      disk_size: $(vars.image_disk_size_gb)
      state_timeout: 45m  # Give the build plenty of time to finish
      # Start from this source image
      source_image_family: $(vars.source_image_family)
      source_image_project_id: [$(vars.source_image_project_id)]
      image_family: $(vars.instance_image.family)
      omit_external_ip: false
      metadata:
        user-data: |
          #cloud-config
          create_hostname_file: true
          write_files:
          - path: /etc/apt/apt.conf.d/20auto-upgrades
            permissions: '0644'
            owner: root
            content: |
              APT::Periodic::Update-Package-Lists "0";
              APT::Periodic::Unattended-Upgrade "0";
- group: vm-env
  modules:
  - id: a4x_rdma_net
    source: modules/network/gpu-rdma-vpc
    settings:
      mtu: 8896
      network_name: $(vars.deployment_name)-a4x-rdma-net
      network_profile: https://www.googleapis.com/compute/beta/projects/$(vars.project_id)/global/networkProfiles/$(vars.zone)-vpc-roce
      network_routing_mode: REGIONAL
      subnetworks_template:
        name_prefix: a4x-mrdma-sub
        count: 4
        ip_range: 192.168.128.0/24
        region: $(vars.region)
  - id: homefs
    source: modules/file-system/filestore
    use: [a4x_net_0]
    settings:
      filestore_tier: HIGH_SCALE_SSD
      size_gb: 10240
      local_mount: /home
      reserved_ip_range: $(vars.filestore_ip_range)

# This deployment group  creates the VMs.
- group: vms
  modules:
  - id: a4x_vms
    source: modules/compute/vm-instance
    use: [homefs]
    settings:
      machine_type: $(vars.machine_type)
      instance_count: $(vars.instance_count)
      reservation_name: $(vars.reservation_name)
      name_prefix: $(vars.deployment_name)
      placement_policy:
        collocation: COLLOCATED
      disk_type: hyperdisk-balanced
      disk_size_gb: $(vars.disk_size_gb)
      automatic_restart: true
      on_host_maintenance: TERMINATE
      instance_image: $(vars.instance_image)
      network_interfaces:
        $(concat(
          [{
            network=null,
            subnetwork=a4x_net_0.subnetwork_self_link,
            subnetwork_project=vars.project_id,
            nic_type="GVNIC",
            queue_count=null,
            network_ip=null,
            stack_type=null,
            access_config=[{nat_ip=null, public_ptr_domain_name=null, network_tier=null}],
            ipv6_access_config=[],
            alias_ip_range=[]
          },
          {
            network=null,
            subnetwork=a4x_net_1.subnetwork_self_link,
            subnetwork_project=vars.project_id,
            nic_type="GVNIC",
            queue_count=null,
            network_ip=null,
            stack_type=null,
            access_config=[{nat_ip=null, public_ptr_domain_name=null, network_tier=null}],
            ipv6_access_config=[],
            alias_ip_range=[]
          }],
          a4x_rdma_net.subnetwork_interfaces
        ))
  - id: wait-for-vms
    source: community/modules/scripts/wait-for-startup
    settings:
      instance_names: $(a4x_vms.name)
      timeout: 7200